Make models more consistent with existing models

Author: oschwald

geoip2/database.py

@@ -99,10 +99,7 @@ def connection_type(self, ip_address):
         :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'))
+        return self._flat_model_for(geoip2.models.ConnectionType, ip_address)
 
     def domain(self, ip_address):
         """Get the Domain object for the IP address
@@ -112,9 +109,7 @@ def domain(self, ip_address):
         :returns: :py:class:`geoip2.models.Domain` object
 
         """
-        record = self._get(ip_address)
-        return geoip2.models.Domain(ip_address=ip_address,
-                                    domain=record.get('domain'))
+        return self._flat_model_for(geoip2.models.Domain, ip_address)
 
     def isp_org(self, ip_address):
         """Get the ISPOrg object for the IP address
@@ -124,14 +119,7 @@ def isp_org(self, ip_address):
         :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'))
+        return self._flat_model_for(geoip2.models.ISPOrg, ip_address)
 
     def _get(self, ip_address):
         record = self._db_reader.get(ip_address)
@@ -145,6 +133,11 @@ def _model_for(self, model_class, ip_address):
         record.setdefault('traits', {})['ip_address'] = ip_address
         return model_class(record, locales=self._locales)
 
+    def _flat_model_for(self, model_class, ip_address):
+        record = self._get(ip_address)
+        record['ip_address'] = ip_address
+        return model_class(record)
+
     def close(self):
         """Closes the GeoIP2 database"""
 

geoip2/models.py

@@ -326,9 +326,10 @@ class ConnectionType(object):
       :type: unicode
     """
 
-    def __init__(self, ip_address=None, connection_type=None):
-        self.connection_type = connection_type
-        self.ip_address = ip_address
+    def __init__(self, raw):
+        self.connection_type = raw.get('connection_type')
+        self.ip_address = raw.get('ip_address')
+        self.raw = raw
 
 
 class Domain(object):
@@ -351,9 +352,10 @@ class Domain(object):
 
     """
 
-    def __init__(self, ip_address=None, domain=None):
-        self.domain = domain
-        self.ip_address = ip_address
+    def __init__(self, raw):
+        self.domain = raw.get('domain')
+        self.ip_address = raw.get('ip_address')
+        self.raw = raw
 
 
 class ISPOrg(object):
@@ -395,13 +397,11 @@ class ISPOrg(object):
     """
 
     # pylint:disable=too-many-arguments
-    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
+    def __init__(self, raw):
+        self.autonomous_system_number = raw.get('autonomous_system_number')
+        self.autonomous_system_organization = raw.get(
+            'autonomous_system_organization')
+        self.isp = raw.get('isp')
+        self.organization = raw.get('organization')
+        self.ip_address = raw.get('ip_address')
+        self.raw = raw