Replace raw attribute with to_dict method

Author: oschwald

HISTORY.rst

@@ -4,9 +4,12 @@
 History
 -------
 
-4.9.0
+5.0.0
 ++++++++++++++++++
 
+* BREAKING: The ``raw`` attribute on the model classes has been replaced
+  with a ``to_dict()`` method. This can be used to get a representation of
+  the object that is suitable for serialization.
 * IMPORTANT: Python 3.9 or greater is required. If you are using an older
   version, please use an earlier release.
 * ``metro_code`` on ``geoip2.record.Location`` has been deprecated. The

geoip2/mixins.py

@@ -1,15 +1,50 @@
-"""This package contains utility mixins"""
+"""This package contains internal utilities"""
 
 # pylint: disable=too-few-public-methods
 from abc import ABCMeta
 from typing import Any
 
 
-class SimpleEquality(metaclass=ABCMeta):
-    """Naive __dict__ equality mixin"""
+class Model(metaclass=ABCMeta):
+    """Shared methods for MaxMind model classes"""
 
     def __eq__(self, other: Any) -> bool:
         return isinstance(other, self.__class__) and self.__dict__ == other.__dict__
 
     def __ne__(self, other):
         return not self.__eq__(other)
+
+    # pylint: disable=too-many-branches
+    def to_dict(self):
+        """Returns a dict of the object suitable for serialization"""
+        result = {}
+        for key, value in self.__dict__.items():
+            if key.startswith("_"):
+                continue
+            if hasattr(value, "to_dict") and callable(value.to_dict):
+                if d := value.to_dict():
+                    result[key] = d
+            elif isinstance(value, (list, tuple)):
+                ls = []
+                for e in value:
+                    if hasattr(e, "to_dict") and callable(e.to_dict):
+                        if e := e.to_dict():
+                            ls.append(e)
+                    elif e is not None:
+                        ls.append(e)
+                if ls:
+                    result[key] = ls
+            # We only have dicts of strings currently. Do not bother with
+            # the general case.
+            elif isinstance(value, dict):
+                if value:
+                    result[key] = value
+            elif value is not None and value is not False:
+                result[key] = value
+
+        # network is a property for performance reasons
+        # pylint: disable=no-member
+        if hasattr(self, "network") and self.network is not None:
+            result["network"] = str(self.network)
+
+        return result

geoip2/models.py

@@ -17,10 +17,10 @@
 from typing import Any, cast, Dict, Optional, Sequence, Union
 
 import geoip2.records
-from geoip2.mixins import SimpleEquality
+from geoip2.mixins import Model
 
 
-class Country(SimpleEquality):
+class Country(Model):
     """Model for the Country web service and Country database.
 
     This class provides the following attributes:
@@ -97,12 +97,9 @@ def __init__(
         self.maxmind = geoip2.records.MaxMind(**raw_response.get("maxmind", {}))
 
         self.traits = geoip2.records.Traits(**raw_response.get("traits", {}))
-        self.raw = raw_response
 
     def __repr__(self) -> str:
-        return (
-            f"{self.__module__}.{self.__class__.__name__}({self.raw}, {self._locales})"
-        )
+        return f"{self.__module__}.{self.__class__.__name__}({self.to_dict()}, {self._locales})"
 
 
 class City(Country):
@@ -321,22 +318,27 @@ class Enterprise(City):
     """
 
 
-class SimpleModel(SimpleEquality, metaclass=ABCMeta):
+class SimpleModel(Model, metaclass=ABCMeta):
     """Provides basic methods for non-location models"""
 
-    raw: Dict[str, Union[bool, str, int]]
     ip_address: str
     _network: Optional[Union[ipaddress.IPv4Network, ipaddress.IPv6Network]]
     _prefix_len: int
 
     def __init__(self, raw: Dict[str, Union[bool, str, int]]) -> None:
-        self.raw = raw
-        self._network = None
-        self._prefix_len = cast(int, raw.get("prefix_len"))
+        if network := raw.get("network"):
+            self._network = ipaddress.ip_network(network, False)
+            self._prefix_len = self._network.prefixlen
+        else:
+            # This case is for MMDB lookups where performance is paramount.
+            # This is why we don't generate the network unless .network is
+            # used.
+            self._network = None
+            self._prefix_len = cast(int, raw.get("prefix_len"))
         self.ip_address = cast(str, raw.get("ip_address"))
 
     def __repr__(self) -> str:
-        return f"{self.__module__}.{self.__class__.__name__}({self.raw})"
+        return f"{self.__module__}.{self.__class__.__name__}({self.to_dict()})"
 
     @property
     def network(self) -> Optional[Union[ipaddress.IPv4Network, ipaddress.IPv6Network]]:

geoip2/records.py

@@ -13,14 +13,14 @@
 from abc import ABCMeta
 from typing import Dict, Optional, Type, Sequence, Union
 
-from geoip2.mixins import SimpleEquality
+from geoip2.mixins import Model
 
 
-class Record(SimpleEquality, metaclass=ABCMeta):
+class Record(Model, metaclass=ABCMeta):
     """All records are subclasses of the abstract class ``Record``."""
 
     def __repr__(self) -> str:
-        args = ", ".join(f"{k}={v!r}" for k, v in self.__dict__.items())
+        args = ", ".join(f"{k}={v!r}" for k, v in self.to_dict().items())
         return f"{self.__module__}.{self.__class__.__name__}({args})"
 
 

tests/models_test.py

@@ -13,6 +13,9 @@
 
 
 class TestModels(unittest.TestCase):
+    def setUp(self):
+        self.maxDiff = 20_000
+
     def test_insights_full(self) -> None:
         raw = {
             "city": {
@@ -85,7 +88,6 @@ def test_insights_full(self) -> None:
                 "is_satellite_provider": True,
                 "is_tor_exit_node": True,
                 "isp": "Comcast",
-                "network_speed": "cable/DSL",
                 "organization": "Blorg",
                 "static_ip_score": 1.3,
                 "user_count": 2,
@@ -131,7 +133,7 @@ def test_insights_full(self) -> None:
         self.assertEqual(
             type(model.traits), geoip2.records.Traits, "geoip2.records.Traits object"
         )
-        self.assertEqual(model.raw, raw, "raw method returns raw input")
+        self.assertEqual(model.to_dict(), raw, "to_dict() method matches raw input")
         self.assertEqual(
             model.subdivisions[0].iso_code, "MN", "div 1 has correct iso_code"
         )
@@ -290,7 +292,9 @@ def test_city_full(self) -> None:
         self.assertEqual(
             type(model.traits), geoip2.records.Traits, "geoip2.records.Traits object"
         )
-        self.assertEqual(model.raw, raw, "raw method returns raw input")
+        self.assertEqual(
+            model.to_dict(), raw, "to_dict method output matches raw input"
+        )
         self.assertEqual(model.continent.geoname_id, 42, "continent geoname_id is 42")
         self.assertEqual(model.continent.code, "NA", "continent code is NA")
         self.assertEqual(
@@ -338,7 +342,7 @@ def test_city_full(self) -> None:
             True,
             "traits is_setellite_provider is True",
         )
-        self.assertEqual(model.raw, raw, "raw method produces raw output")
+        self.assertEqual(model.to_dict(), raw, "to_dict method matches raw input")
 
         self.assertRegex(
             str(model), r"^geoip2.models.City\(\{.*geoname_id.*\}, \[.*en.*\]\)"

tests/webservice_test.py

@@ -110,7 +110,7 @@ def test_country_ok(self):
             country.traits.network, ipaddress.ip_network("1.2.3.0/24"), "network"
         )
         self.assertTrue(country.traits.is_anycast)
-        self.assertEqual(country.raw, self.country, "raw response is correct")
+        self.assertEqual(country.to_dict(), self.country, "raw response is correct")
 
     def test_me(self):
         self.httpserver.expect_request(
@@ -357,6 +357,7 @@ def setUp(self):
         self.client_class = Client
         self.client = Client(42, "abcdef123456")
         self.client._base_uri = self.httpserver.url_for("/geoip/v2.1")
+        self.maxDiff = 20_000
 
     def run_client(self, v):
         return v
@@ -368,6 +369,7 @@ def setUp(self):
         self.client_class = AsyncClient
         self.client = AsyncClient(42, "abcdef123456")
         self.client._base_uri = self.httpserver.url_for("/geoip/v2.1")
+        self.maxDiff = 20_000
 
     def tearDown(self):
         self._loop.run_until_complete(self.client.close())