Allow defining a custom geolocation function (fixes #70)

docs/index.rst

@@ -14,6 +14,7 @@ Basic metrics tracking for your `Flask`_ application. The core of library is ver
 The following is optional:
 
 * `freegeoip.net <http://freegeoip.net/>`_ integration for storing geography of the visitor.
+* Alternatively you can define a custom geolocation function, for use with the geoip2 package for example.
 * Unique visitor tracking if you are wanting to use Flask's cookie storage.
 * Summation hooks for live count of common web analysis statistics such as hit counts.
 
@@ -55,6 +56,8 @@ Usage
     app.config['TRACK_USAGE_USE_FREEGEOIP'] = False
     # You can use a different instance of freegeoip like so
     # app.config['TRACK_USAGE_FREEGEOIP_ENDPOINT'] = 'https://example.org/api/'
+    # (You can also define a custom geolocation function when initializing the extension)
+
     app.config['TRACK_USAGE_INCLUDE_OR_EXCLUDE_VIEWS'] = 'include'
 
     # We will just print out the data for the example
@@ -68,7 +71,9 @@ Usage
     t = TrackUsage(app, [
         PrintWriter(),
         OutputWriter(transform=lambda s: "OUTPUT: " + str(s))
-    ])
+        ],
+        custom_geolocation_function # (Optional)
+    )
 
     # Include the view in the metrics
     @t.include
@@ -156,7 +161,8 @@ TRACK_USAGE_USE_FREEGEOIP
 
 **Default**: False
 
-Turn FreeGeoIP integration on or off. If set to true, then geography information is also stored in the usage logs.
+Turn FreeGeoIP integration on or off. If set to true, then the geography information is also stored in the usage logs.
+Alternatively you can define a custom geolocation lookup function when initializing the extension.
 
 .. versionchanged:: 1.1.
    The default server for using geoip integration changed to extreme-ip-lookup.com

src/flask_track_usage/__init__.py

@@ -56,7 +56,7 @@ class TrackUsage(object):
     Tracks basic usage of Flask applications.
     """
 
-    def __init__(self, app=None, storage=None, _fake_time=None):
+    def __init__(self, app=None, storage=None, ip_lookup_func=None, _fake_time=None):
         """
         Create the instance.
 
@@ -78,9 +78,9 @@ def __init__(self, app=None, storage=None, _fake_time=None):
         self._fake_time = _fake_time
 
         if app is not None and storage is not None:
-            self.init_app(app, storage)
+            self.init_app(app, storage, ip_lookup_func)
 
-    def init_app(self, app, storage):
+    def init_app(self, app, storage, ip_lookup_func):
         """
         Initialize the instance with the app.
 
@@ -90,6 +90,7 @@ def init_app(self, app, storage):
         """
         self.app = app
         self._storages = storage
+        self.ip_lookup_func = ip_lookup_func
         self._use_freegeoip = app.config.get(
             'TRACK_USAGE_USE_FREEGEOIP', False)
         self._freegeoip_endpoint = app.config.get(
@@ -190,7 +191,10 @@ def after_request(self, response):
             data['username'] = str(ctx.request.authorization.username)
         elif getattr(self.app, 'login_manager', None) and current_user and not current_user.is_anonymous:
             data['username'] = str(current_user)
-        if self._use_freegeoip:
+        if self.ip_lookup_func:
+            clean_ip = quote_plus(str(ctx.request.remote_addr))
+            data['ip_info'] = self.ip_lookup_func(clean_ip)
+        elif self._use_freegeoip:
             clean_ip = quote_plus(str(ctx.request.remote_addr))
             if '{ip}' in self._freegeoip_endpoint:
                 url = self._freegeoip_endpoint.format(ip=clean_ip)