Added new database models and methods

Author: oschwald

geoip2/database.py

@@ -51,9 +51,7 @@ def __init__(self, filename, locales=None):
     def country(self, ip_address):
         """Get the Country record object for the IP address
 
-        :param ip_address: IPv4 or IPv6 address as a string. If no address
-          is provided, the address that the web service is called from will
-          be used.
+        :param ip_address: IPv4 or IPv6 address as a string.
 
         :returns: :py:class:`geoip2.models.Country` object
 
@@ -64,9 +62,7 @@ def country(self, ip_address):
     def city(self, ip_address):
         """Get the City record object for the IP address
 
-        :param ip_address: IPv4 or IPv6 address as a string. If no address
-          is provided, the address that the web service is called from will
-          be used.
+        :param ip_address: IPv4 or IPv6 address as a string.
 
         :returns: :py:class:`geoip2.models.City` object
 
@@ -77,8 +73,7 @@ def city_isp_org(self, ip_address):
         """Get the CityISPOrg record object for the IP address
 
         :param ip_address: IPv4 or IPv6 address as a string. If no address
-          is provided, the address that the web service is called from will
-          be used.
+          is provided.
 
         :returns: :py:class:`geoip2.models.CityISPOrg` object
 
@@ -89,20 +84,64 @@ def city_isp_org(self, ip_address):
     def omni(self, ip_address):
         """Get the Omni record object for the IP address
 
-        :param ip_address: IPv4 or IPv6 address as a string. If no address
-          is provided, the address that the web service is called from will
-          be used.
+        :param ip_address: IPv4 or IPv6 address as a string.
 
         :returns: :py:class:`geoip2.models.Omni` object
 
         """
         return self._model_for(geoip2.models.Omni, ip_address)
 
-    def _model_for(self, model_class, ip_address):
+    def connection_type(self, ip_address):
+        """Get the ConnectionType object for the IP address
+
+        :param ip_address: IPv4 or IPv6 address as a string.
+
+        :returns: :py:class:`geoip2.models.ConnectionType` object
+
+        """
+        record = self._get(ip_address)
+        return geoip2.models.ConnectionType(ip_address=ip_address,
+                                            connection_type=record.get(
+                                                'connection_type'))
+
+    def domain(self, ip_address):
+        """Get the Domain object for the IP address
+
+        :param ip_address: IPv4 or IPv6 address as a string.
+
+        :returns: :py:class:`geoip2.models.Domain` object
+
+        """
+        record = self._get(ip_address)
+        return geoip2.models.Domain(ip_address=ip_address,
+                                    domain=record.get('domain'))
+
+    def isp_org(self, ip_address):
+        """Get the ISPOrg object for the IP address
+
+        :param ip_address: IPv4 or IPv6 address as a string.
+
+        :returns: :py:class:`geoip2.models.ISPOrg` object
+
+        """
+        record = self._get(ip_address)
+        return geoip2.models.ISPOrg(ip_address=ip_address,
+                                    autonomous_system_number=record.get(
+                                        'autonomous_system_number'),
+                                    autonomous_system_organization=record.get(
+                                        'autonomous_system_organization'),
+                                    isp=record.get('isp'),
+                                    organization=record.get('organization'))
+
+    def _get(self, ip_address):
         record = self._db_reader.get(ip_address)
         if record is None:
             raise geoip2.errors.AddressNotFoundError(
                 "The address %s is not in the database." % ip_address)
+        return record
+
+    def _model_for(self, model_class, ip_address):
+        record = self._get(ip_address)
         record.setdefault('traits', {})['ip_address'] = ip_address
         return model_class(record, locales=self._locales)
 

geoip2/models.py

@@ -16,10 +16,9 @@
 
 class Country(object):
 
-    """Model class for the GeoIP2 Country end point
+    """Model class for the GeoIP2 Country Web Service or Database
 
-    This class provides the following methods, each of which returns a record
-    object.
+    This class provides the following attributes:
 
     .. attribute:: continent
 
@@ -93,7 +92,7 @@ def __init__(self, raw_response, locales=None):
 
 class City(Country):
 
-    """Model class for the GeoIP2 Precision City end point
+    """Model class for the GeoIP2 Precision City Web Service or City Database
 
     .. attribute:: city
 
@@ -299,3 +298,109 @@ class Omni(CityISPOrg):
       :type: :py:class:`geoip2.records.Traits`
 
     """
+
+
+class ConnectionType(object):
+
+    """Model class for the GeoIP2 Connection-Type Database
+
+    This class provides the following attribute:
+
+    .. attribute:: connection_type
+
+      The connection type may take the following values:
+
+      - Dialup
+      - Cable/DSL
+      - Corporate
+      - Cellular
+
+      Additional values may be added in the future.
+
+      :type: unicode
+
+    .. attribute:: ip_address
+
+      The IP address used in the lookup.
+
+      :type: unicode
+    """
+
+    def __init__(self, ip_address=None, connection_type=None):
+        self.connection_type = connection_type
+        self.ip_address = ip_address
+
+
+class Domain(object):
+
+    """Model class for the GeoIP2 Domain Database
+
+    This class provides the following attribute:
+
+    .. attribute:: domain
+
+      The domain associated with the IP address.
+
+      :type: unicode
+
+    .. attribute:: ip_address
+
+      The IP address used in the lookup.
+
+      :type: unicode
+
+    """
+
+    def __init__(self, ip_address=None, domain=None):
+        self.domain = domain
+        self.ip_address = ip_address
+
+
+class ISPOrg(object):
+
+    """Model class for the GeoIP2 ISP-Org Database
+
+    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:: isp
+
+      The name of the ISP associated with the IP address.
+
+      :type: unicode
+
+    .. attribute:: organization
+
+      The name of the organization associated with the IP address.
+
+      :type: unicode
+
+    .. attribute:: ip_address
+
+      The IP address used in the lookup.
+
+      :type: unicode
+    """
+
+    def __init__(self, ip_address=None,
+                 autonomous_system_number=None,
+                 autonomous_system_organization=None,
+                 isp=None,
+                 organization=None):
+        self.autonomous_system_number = autonomous_system_number
+        self.autonomous_system_organization = autonomous_system_organization
+        self.isp = isp
+        self.organization = organization
+        self.ip_address = ip_address

tests/data

@@ -40,7 +40,6 @@ def test_has_ip_address(self):
         reader = geoip2.database.Reader(
             'tests/data/test-data/GeoIP2-City-Test.mmdb')
         record = reader.country('81.2.69.160')
-        record = reader.omni('81.2.69.160')
         self.assertEqual(record.traits.ip_address, '81.2.69.160')
         reader.close()
 
@@ -61,3 +60,38 @@ def test_invalid_address(self):
                                     "IPv4 or IPv6 address"):
             reader.city('invalid')
         reader.close()
+
+    def test_connection_type(self):
+        reader = geoip2.database.Reader(
+            'tests/data/test-data/GeoIP2-Connection-Type-Test.mmdb')
+        ip_address = '1.0.1.0'
+
+        record = reader.connection_type(ip_address)
+        self.assertEqual(record.connection_type, 'Cable/DSL')
+        self.assertEqual(record.ip_address, ip_address)
+        reader.close()
+
+    def test_domain(self):
+        reader = geoip2.database.Reader(
+            'tests/data/test-data/GeoIP2-Domain-Test.mmdb')
+
+        ip_address = '1.2.0.0'
+        record = reader.domain(ip_address)
+        self.assertEqual(record.domain, 'maxmind.com')
+        self.assertEqual(record.ip_address, ip_address)
+
+        reader.close()
+
+    def test_isp_org(self):
+        reader = geoip2.database.Reader(
+            'tests/data/test-data/GeoIP2-ISP-Org-Test.mmdb')
+
+        ip_address = '2001:1700::'
+        record = reader.isp_org(ip_address)
+        self.assertEqual(record.autonomous_system_number, 6730)
+        self.assertEqual(record.autonomous_system_organization,
+                         'Sunrise Communications AG')
+        # XXX - Add org/isp when available
+        self.assertEqual(record.ip_address, ip_address)
+
+        reader.close()