Merge pull request #43 from maxmind/greg/asn

Add GeoLite2 ASN database + run yapf

Author: oalders

HISTORY.rst

@@ -3,6 +3,11 @@
 History
 -------
 
+2.5.0
+++++++++++++++++++
+
+* Added support for GeoLite2 ASN database.
+
 2.4.2 (2016-12-08)
 ++++++++++++++++++
 

README.rst

@@ -185,6 +185,22 @@ Anonymous IP Database
     '128.101.101.101'
     >>> reader.close()
 
+ASN Database
+^^^^^^^^^^^^
+
+.. code-block:: pycon
+
+    >>> import geoip2.database
+    >>>
+    >>> # This creates a Reader object. You should use the same object
+    >>> # across multiple requests as creation of it is expensive.
+    >>> with geoip2.database.Reader('/path/to/GeoLite2-ASN.mmdb') as reader:
+    >>>     response = reader.asn('1.128.0.0')
+    >>>     response.autonomous_system_number
+    1221
+    >>>     response.autonomous_system_organization
+    'Telstra Pty Ltd'
+
 Connection-Type Database
 ^^^^^^^^^^^^^^^^^^^^^^^^
 

geoip2/compat.py

@@ -6,12 +6,14 @@
 # pylint: skip-file
 
 if sys.version_info[0] == 2:
+
     def compat_ip_address(address):
         """Intended for internal use only."""
         if isinstance(address, bytes):
             address = address.decode()
         return ipaddress.ip_address(address)
 else:
+
     def compat_ip_address(address):
         """Intended for internal use only."""
         return ipaddress.ip_address(address)

geoip2/database.py

@@ -120,6 +120,17 @@ def anonymous_ip(self, ip_address):
         return self._flat_model_for(geoip2.models.AnonymousIP,
                                     'GeoIP2-Anonymous-IP', ip_address)
 
+    def asn(self, ip_address):
+        """Get the ASN object for the IP address.
+
+        :param ip_address: IPv4 or IPv6 address as a string.
+
+        :returns: :py:class:`geoip2.models.ASN` object
+
+        """
+        return self._flat_model_for(geoip2.models.ASN, 'GeoLite2-ASN',
+                                    ip_address)
+
     def connection_type(self, ip_address):
         """Get the ConnectionType object for the IP address.
 
@@ -168,8 +179,8 @@ def _get(self, database_type, ip_address):
         if database_type not in self.metadata().database_type:
             caller = inspect.stack()[2][3]
             raise TypeError("The %s method cannot be used with the "
-                            "%s database" %
-                            (caller, self.metadata().database_type))
+                            "%s database" % (caller,
+                                             self.metadata().database_type))
         record = self._db_reader.get(ip_address)
         if record is None:
             raise geoip2.errors.AddressNotFoundError(

geoip2/models.py

@@ -368,6 +368,40 @@ def __init__(self, raw):
         self.raw = raw
 
 
+class ASN(SimpleModel):
+    """Model class for the GeoLite2 ASN.
+
+    This class provides the following attribute:
+
+    .. attribute:: autonomous_system_number
+
+      The autonomous system number associated with the IP address.
+
+      :type: int
+
+    .. attribute:: autonomous_system_organization
+
+      The organization associated with the registered autonomous system number
+      for the IP address.
+
+      :type: unicode
+
+    .. attribute:: ip_address
+
+      The IP address used in the lookup.
+
+      :type: unicode
+    """
+
+    # pylint:disable=too-many-arguments
+    def __init__(self, raw):
+        self.autonomous_system_number = raw.get('autonomous_system_number')
+        self.autonomous_system_organization = raw.get(
+            'autonomous_system_organization')
+        self.ip_address = raw.get('ip_address')
+        self.raw = raw
+
+
 class ConnectionType(SimpleModel):
     """Model class for the GeoIP2 Connection-Type.
 
@@ -424,7 +458,7 @@ def __init__(self, raw):
         self.raw = raw
 
 
-class ISP(SimpleModel):
+class ISP(ASN):
     """Model class for the GeoIP2 ISP.
 
     This class provides the following attribute:
@@ -463,10 +497,6 @@ class ISP(SimpleModel):
 
     # pylint:disable=too-many-arguments
     def __init__(self, raw):
-        self.autonomous_system_number = raw.get('autonomous_system_number')
-        self.autonomous_system_organization = raw.get(
-            'autonomous_system_organization')
+        super(ISP, self).__init__(raw)
         self.isp = raw.get('isp')
         self.organization = raw.get('organization')
-        self.ip_address = raw.get('ip_address')
-        self.raw = raw

geoip2/records.py

@@ -50,9 +50,8 @@ def __init__(self, locales=None, **kwargs):
     def name(self):
         """Dict with locale codes as keys and localized name as value."""
         # pylint:disable=E1101
-        return next(
-            (self.names.get(x) for x in self._locales
-             if x in self.names), None)
+        return next((self.names.get(x) for x in self._locales
+                     if x in self.names), None)
 
 
 class City(PlaceRecord):
@@ -243,8 +242,8 @@ class RepresentedCountry(Country):
 
     """
 
-    _valid_attributes = set(['confidence', 'geoname_id', 'iso_code', 'names',
-                             'type'])
+    _valid_attributes = set(
+        ['confidence', 'geoname_id', 'iso_code', 'names', 'type'])
 
 
 class Location(Record):
@@ -314,9 +313,11 @@ class Location(Record):
 
     """
 
-    _valid_attributes = set(['average_income', 'accuracy_radius', 'latitude',
-                             'longitude', 'metro_code', 'population_density',
-                             'postal_code', 'postal_confidence', 'time_zone'])
+    _valid_attributes = set([
+        'average_income', 'accuracy_radius', 'latitude', 'longitude',
+        'metro_code', 'population_density', 'postal_code', 'postal_confidence',
+        'time_zone'
+    ])
 
 
 class MaxMind(Record):
@@ -594,14 +595,17 @@ class Traits(Record):
 
     """
 
-    _valid_attributes = set(
-        ['autonomous_system_number', 'autonomous_system_organization',
-         'connection_type', 'domain', 'is_anonymous_proxy',
-         'is_legitimate_proxy', 'is_satellite_provider', 'isp', 'ip_address',
-         'organization', 'user_type'])
+    _valid_attributes = set([
+        'autonomous_system_number', 'autonomous_system_organization',
+        'connection_type', 'domain', 'is_anonymous_proxy',
+        'is_legitimate_proxy', 'is_satellite_provider', 'isp', 'ip_address',
+        'organization', 'user_type'
+    ])
 
     def __init__(self, **kwargs):
-        for k in ['is_anonymous_proxy', 'is_legitimate_proxy',
-                  'is_satellite_provider']:
+        for k in [
+                'is_anonymous_proxy', 'is_legitimate_proxy',
+                'is_satellite_provider'
+        ]:
             kwargs[k] = bool(kwargs.get(k, False))
         super(Traits, self).__init__(**kwargs)

geoip2/webservice.py

@@ -141,13 +141,14 @@ def _response_for(self, path, model_class, ip_address):
         if ip_address != 'me':
             ip_address = str(compat_ip_address(ip_address))
         uri = '/'.join([self._base_uri, path, ip_address])
-        response = requests.get(uri,
-                                auth=(self._user_id, self._license_key),
-                                headers={
-                                    'Accept': 'application/json',
-                                    'User-Agent': self._user_agent()
-                                },
-                                timeout=self._timeout)
+        response = requests.get(
+            uri,
+            auth=(self._user_id, self._license_key),
+            headers={
+                'Accept': 'application/json',
+                'User-Agent': self._user_agent()
+            },
+            timeout=self._timeout)
         if response.status_code == 200:
             body = self._handle_success(response, uri)
             return model_class(body, locales=self._locales)

tests/data

@@ -25,7 +25,6 @@
 
 
 class BaseTestReader(object):
-
     def test_language_list(self):
         reader = geoip2.database.Reader(
             'tests/data/test-data/GeoIP2-Country-Test.mmdb',
@@ -76,6 +75,26 @@ def test_anonymous_ip(self):
         self.assertEqual(record.ip_address, ip_address)
         reader.close()
 
+    def test_asn(self):
+        reader = geoip2.database.Reader(
+            'tests/data/test-data/GeoLite2-ASN-Test.mmdb')
+
+        ip_address = '1.128.0.0'
+        record = reader.asn(ip_address)
+        self.assertEqual(record.autonomous_system_number, 1221)
+        self.assertEqual(record.autonomous_system_organization,
+                         'Telstra Pty Ltd')
+        self.assertEqual(record.ip_address, ip_address)
+
+        self.assertRegex(
+            str(record),
+            r'geoip2.models.ASN\(.*1\.128\.0\.0.*\)',
+            'str representation is correct')
+
+        self.assertEqual(record, eval(repr(record)), "ASN repr can be eval'd")
+
+        reader.close()
+
     def test_city(self):
         reader = geoip2.database.Reader(
             'tests/data/test-data/GeoIP2-City-Test.mmdb')
@@ -101,7 +120,8 @@ def test_connection_type(self):
             str(record), r'ConnectionType\(\{.*Cable/DSL.*\}\)',
             'ConnectionType str representation is reasonable')
 
-        self.assertEqual(record, eval(repr(record)),
+        self.assertEqual(record,
+                         eval(repr(record)),
                          "ConnectionType repr can be eval'd")
 
         reader.close()
@@ -127,8 +147,8 @@ def test_domain(self):
             str(record), r'Domain\(\{.*maxmind.com.*\}\)',
             'Domain str representation is reasonable')
 
-        self.assertEqual(record, eval(repr(record)),
-                         "Domain repr can be eval'd")
+        self.assertEqual(record,
+                         eval(repr(record)), "Domain repr can be eval'd")
 
         reader.close()