Add __repr__ and __eq__ methods to model/record classes

Author: oschwald

geoip2/models.py

@@ -10,7 +10,9 @@
 http://dev.maxmind.com/geoip/geoip2/web-services for more details.
 
 """
-# pylint:disable=R0903
+# pylint: disable=too-many-instance-attributes,too-few-public-methods
+from abc import ABCMeta
+
 import geoip2.records
 
 
@@ -66,6 +68,7 @@ class Country(object):
     def __init__(self, raw_response, locales=None):
         if locales is None:
             locales = ['en']
+        self._locales = locales
         self.continent = \
             geoip2.records.Continent(locales,
                                      **raw_response.get('continent', {}))
@@ -76,18 +79,27 @@ def __init__(self, raw_response, locales=None):
             geoip2.records.Country(locales,
                                    **raw_response.get('registered_country',
                                                       {}))
-        # pylint:disable=bad-continuation
         self.represented_country \
             = geoip2.records.RepresentedCountry(locales,
                                                 **raw_response.get(
-                                                'represented_country', {}))
+                                                    'represented_country', {}))
 
         self.maxmind = \
             geoip2.records.MaxMind(**raw_response.get('maxmind', {}))
 
         self.traits = geoip2.records.Traits(**raw_response.get('traits', {}))
         self.raw = raw_response
 
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __repr__(self):
+        return '{module}.{class_name}({data}, {locales})'.format(
+            module=self.__module__,
+            class_name=self.__class__.__name__,
+            data=self.raw,
+            locales=self._locales)
+
 
 class City(Country):
 
@@ -230,7 +242,24 @@ class Insights(City):
     """
 
 
-class ConnectionType(object):
+class SimpleModel(object):
+
+    """Provides basic methods for non-location models"""
+
+    __metaclass__ = ABCMeta
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __repr__(self):
+        # pylint: disable=no-member
+        return '{module}.{class_name}({data})'.format(
+            module=self.__module__,
+            class_name=self.__class__.__name__,
+            data=str(self.raw))
+
+
+class ConnectionType(SimpleModel):
 
     """Model class for the GeoIP2 Connection-Type
 
@@ -262,7 +291,7 @@ def __init__(self, raw):
         self.raw = raw
 
 
-class Domain(object):
+class Domain(SimpleModel):
 
     """Model class for the GeoIP2 Domain
 
@@ -288,7 +317,7 @@ def __init__(self, raw):
         self.raw = raw
 
 
-class ISP(object):
+class ISP(SimpleModel):
 
     """Model class for the GeoIP2 ISP
 

geoip2/records.py

@@ -22,6 +22,16 @@ def __init__(self, **kwargs):
     def __setattr__(self, name, value):
         raise AttributeError("can't set attribute")
 
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __repr__(self):
+        args = ', '.join('%s=%r' % x for x in self.__dict__.items())
+        return '{module}.{class_name}({data})'.format(
+            module=self.__module__,
+            class_name=self.__class__.__name__,
+            data=args)
+
 
 class PlaceRecord(Record):
 

tests/database_test.py

@@ -15,6 +15,7 @@
 
 if sys.version_info[0] == 2:
     unittest.TestCase.assertRaisesRegex = unittest.TestCase.assertRaisesRegexp
+    unittest.TestCase.assertRegex = unittest.TestCase.assertRegexpMatches
 
 
 class TestReader(unittest.TestCase):
@@ -78,6 +79,14 @@ def test_connection_type(self):
         record = reader.connection_type(ip_address)
         self.assertEqual(record.connection_type, 'Cable/DSL')
         self.assertEqual(record.ip_address, ip_address)
+
+        self.assertRegex(
+            str(record), r'ConnectionType\(\{.*Cable/DSL.*\}\)',
+            'ConnectionType str representation is reasonable')
+
+        self.assertEqual(record, eval(repr(record)),
+                         "ConnectionType repr can be eval'd")
+
         reader.close()
 
     def test_domain(self):
@@ -89,6 +98,13 @@ def test_domain(self):
         self.assertEqual(record.domain, 'maxmind.com')
         self.assertEqual(record.ip_address, ip_address)
 
+        self.assertRegex(
+            str(record), r'Domain\(\{.*maxmind.com.*\}\)',
+            'Domain str representation is reasonable')
+
+        self.assertEqual(record, eval(repr(record)),
+                         "Domain repr can be eval'd")
+
         reader.close()
 
     def test_isp(self):
@@ -104,4 +120,11 @@ def test_isp(self):
         self.assertEqual(record.organization, 'Telstra Internet')
         self.assertEqual(record.ip_address, ip_address)
 
+        self.assertRegex(
+            str(record), r'ISP\(\{.*Telstra.*\}\)',
+            'ISP str representation is reasonable')
+
+        self.assertEqual(record, eval(repr(record)),
+                         "ISP repr can be eval'd")
+
         reader.close()

tests/models_test.py

@@ -15,6 +15,7 @@
 
 if sys.version_info[0] == 2:
     unittest.TestCase.assertRaisesRegex = unittest.TestCase.assertRaisesRegexp
+    unittest.TestCase.assertRegex = unittest.TestCase.assertRegexpMatches
 
 
 class TestModels(unittest.TestCase):
@@ -58,7 +59,7 @@ def test_insights_full(self):
                     'geoname_id': 123,
                     'iso_code': 'HP',
                     'names': {'en': 'Hennepin'},
-                }
+            }
             ],
             'registered_country': {
                 'geoname_id': 2,
@@ -133,6 +134,21 @@ def test_insights_full(self):
         self.assertEqual(model.location.metro_code, 765,
                          'correct metro_code')
 
+        self.assertRegex(
+            str(
+                model), r'^geoip2.models.Insights\(\{.*geoname_id.*\}, \[.*en.*\]\)',
+            'Insights str representation looks reasonable')
+
+        self.assertEqual(
+            model, eval(repr(model)), "Insights repr can be eval'd")
+
+        self.assertRegex(
+            str(model.location), r'^geoip2.records.Location\(.*longitude=.*\)',
+            'Location str representation is reasonable')
+
+        self.assertEqual(model.location, eval(repr(model.location)),
+                         "Location repr can be eval'd")
+
     def test_insights_min(self):
         model = geoip2.models.Insights({'traits': {'ip_address': '5.6.7.8'}})
         self.assertEqual(type(model), geoip2.models.Insights,
@@ -229,6 +245,9 @@ def test_city_full(self):
                          'traits is_setellite_provider is True')
         self.assertEqual(model.raw, raw, 'raw method produces raw output')
 
+        self.assertRegex(
+            str(model), r'^geoip2.models.City\(\{.*geoname_id.*\}, \[.*en.*\]\)')
+
     def test_unknown_keys(self):
         model = geoip2.models.City({'traits': {'ip_address': '1.2.3.4',
                                                'invalid': 'blah'},