Add type annotations throughout code

Author: oschwald

geoip2/database.py

@@ -5,11 +5,12 @@
 
 """
 import inspect
+from typing import Any, cast, List, Optional, Type, Union
 
 import maxminddb
 
 # pylint: disable=unused-import
-from maxminddb import (
+from maxminddb import (  # type: ignore
     MODE_AUTO,
     MODE_MMAP,
     MODE_MMAP_EXT,
@@ -21,6 +22,17 @@
 import geoip2
 import geoip2.models
 import geoip2.errors
+from geoip2.types import IPAddress
+from geoip2.models import (
+    ASN,
+    AnonymousIP,
+    City,
+    ConnectionType,
+    Country,
+    Domain,
+    Enterprise,
+    ISP,
+)
 
 
 class Reader(object):
@@ -47,7 +59,9 @@ class Reader(object):
 
 """
 
-    def __init__(self, fileish, locales=None, mode=MODE_AUTO):
+    def __init__(
+        self, fileish: str, locales: Optional[List[str]] = None, mode: int = MODE_AUTO
+    ) -> None:
         """Create GeoIP2 Reader.
 
         :param fileish: The string path to the GeoIP2 database, or an existing
@@ -94,13 +108,13 @@ def __init__(self, fileish, locales=None, mode=MODE_AUTO):
         self._db_type = self._db_reader.metadata().database_type
         self._locales = locales
 
-    def __enter__(self):
+    def __enter__(self) -> "Reader":
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, exc_type: None, exc_value: None, traceback: None) -> None:
         self.close()
 
-    def country(self, ip_address):
+    def country(self, ip_address: IPAddress) -> Country:
         """Get the Country object for the IP address.
 
         :param ip_address: IPv4 or IPv6 address as a string.
@@ -109,83 +123,101 @@ def country(self, ip_address):
 
         """
 
-        return self._model_for(geoip2.models.Country, "Country", ip_address)
+        return cast(
+            Country, self._model_for(geoip2.models.Country, "Country", ip_address)
+        )
 
-    def city(self, ip_address):
+    def city(self, ip_address: IPAddress) -> City:
         """Get the City object for the IP address.
 
         :param ip_address: IPv4 or IPv6 address as a string.
 
         :returns: :py:class:`geoip2.models.City` object
 
         """
-        return self._model_for(geoip2.models.City, "City", ip_address)
+        return cast(City, self._model_for(geoip2.models.City, "City", ip_address))
 
-    def anonymous_ip(self, ip_address):
+    def anonymous_ip(self, ip_address: IPAddress) -> AnonymousIP:
         """Get the AnonymousIP object for the IP address.
 
         :param ip_address: IPv4 or IPv6 address as a string.
 
         :returns: :py:class:`geoip2.models.AnonymousIP` object
 
         """
-        return self._flat_model_for(
-            geoip2.models.AnonymousIP, "GeoIP2-Anonymous-IP", ip_address
+        return cast(
+            AnonymousIP,
+            self._flat_model_for(
+                geoip2.models.AnonymousIP, "GeoIP2-Anonymous-IP", ip_address
+            ),
         )
 
-    def asn(self, ip_address):
+    def asn(self, ip_address: IPAddress) -> ASN:
         """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)
+        return cast(
+            ASN, self._flat_model_for(geoip2.models.ASN, "GeoLite2-ASN", ip_address)
+        )
 
-    def connection_type(self, ip_address):
+    def connection_type(self, ip_address: IPAddress) -> ConnectionType:
         """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
 
         """
-        return self._flat_model_for(
-            geoip2.models.ConnectionType, "GeoIP2-Connection-Type", ip_address
+        return cast(
+            ConnectionType,
+            self._flat_model_for(
+                geoip2.models.ConnectionType, "GeoIP2-Connection-Type", ip_address
+            ),
         )
 
-    def domain(self, ip_address):
+    def domain(self, ip_address: IPAddress) -> Domain:
         """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
 
         """
-        return self._flat_model_for(geoip2.models.Domain, "GeoIP2-Domain", ip_address)
+        return cast(
+            Domain,
+            self._flat_model_for(geoip2.models.Domain, "GeoIP2-Domain", ip_address),
+        )
 
-    def enterprise(self, ip_address):
+    def enterprise(self, ip_address: IPAddress) -> Enterprise:
         """Get the Enterprise object for the IP address.
 
         :param ip_address: IPv4 or IPv6 address as a string.
 
         :returns: :py:class:`geoip2.models.Enterprise` object
 
         """
-        return self._model_for(geoip2.models.Enterprise, "Enterprise", ip_address)
+        return cast(
+            Enterprise,
+            self._model_for(geoip2.models.Enterprise, "Enterprise", ip_address),
+        )
 
-    def isp(self, ip_address):
+    def isp(self, ip_address: IPAddress) -> ISP:
         """Get the ISP object for the IP address.
 
         :param ip_address: IPv4 or IPv6 address as a string.
 
         :returns: :py:class:`geoip2.models.ISP` object
 
         """
-        return self._flat_model_for(geoip2.models.ISP, "GeoIP2-ISP", ip_address)
+        return cast(
+            ISP, self._flat_model_for(geoip2.models.ISP, "GeoIP2-ISP", ip_address)
+        )
 
-    def _get(self, database_type, ip_address):
+    def _get(self, database_type: str, ip_address: IPAddress) -> Any:
         if database_type not in self._db_type:
             caller = inspect.stack()[2][3]
             raise TypeError(
@@ -199,27 +231,39 @@ def _get(self, database_type, ip_address):
             )
         return record, prefix_len
 
-    def _model_for(self, model_class, types, ip_address):
+    def _model_for(
+        self,
+        model_class: Union[Type[Country], Type[Enterprise], Type[City]],
+        types: str,
+        ip_address: IPAddress,
+    ) -> Union[Country, Enterprise, City]:
         (record, prefix_len) = self._get(types, ip_address)
         traits = record.setdefault("traits", {})
         traits["ip_address"] = ip_address
         traits["prefix_len"] = prefix_len
         return model_class(record, locales=self._locales)
 
-    def _flat_model_for(self, model_class, types, ip_address):
+    def _flat_model_for(
+        self,
+        model_class: Union[
+            Type[Domain], Type[ISP], Type[ConnectionType], Type[ASN], Type[AnonymousIP]
+        ],
+        types: str,
+        ip_address: IPAddress,
+    ) -> Union[ConnectionType, ISP, AnonymousIP, Domain, ASN]:
         (record, prefix_len) = self._get(types, ip_address)
         record["ip_address"] = ip_address
         record["prefix_len"] = prefix_len
         return model_class(record)
 
-    def metadata(self):
+    def metadata(self) -> maxminddb.reader.Metadata:
         """The metadata for the open database.
 
         :returns: :py:class:`maxminddb.reader.Metadata` object
         """
         return self._db_reader.metadata()
 
-    def close(self):
+    def close(self) -> None:
         """Closes the GeoIP2 database."""
 
         self._db_reader.close()

geoip2/errors.py

@@ -4,6 +4,8 @@
 
 """
 
+from typing import Optional
+
 
 class GeoIP2Error(RuntimeError):
     """There was a generic error in GeoIP2.
@@ -33,7 +35,9 @@ class HTTPError(GeoIP2Error):
 
     """
 
-    def __init__(self, message, http_status=None, uri=None):
+    def __init__(
+        self, message: str, http_status: Optional[int] = None, uri: Optional[str] = None
+    ) -> None:
         super(HTTPError, self).__init__(message)
         self.http_status = http_status
         self.uri = uri

geoip2/mixins.py

@@ -1,14 +1,15 @@
 """This package contains utility mixins"""
 # pylint: disable=too-few-public-methods
 from abc import ABCMeta
+from typing import Any
 
 
 class SimpleEquality(object):
     """Naive __dict__ equality mixin"""
 
     __metaclass__ = ABCMeta
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         return isinstance(other, self.__class__) and self.__dict__ == other.__dict__
 
     def __ne__(self, other):

geoip2/models.py

@@ -13,6 +13,7 @@
 # pylint: disable=too-many-instance-attributes,too-few-public-methods
 import ipaddress
 from abc import ABCMeta
+from typing import Any, Dict, List, Optional, Union
 
 import geoip2.records
 from geoip2.mixins import SimpleEquality
@@ -66,7 +67,9 @@ class Country(SimpleEquality):
 
     """
 
-    def __init__(self, raw_response, locales=None):
+    def __init__(
+        self, raw_response: Dict[str, Any], locales: Optional[List[str]] = None
+    ) -> None:
         if locales is None:
             locales = ["en"]
         self._locales = locales
@@ -88,7 +91,7 @@ def __init__(self, raw_response, locales=None):
         self.traits = geoip2.records.Traits(**raw_response.get("traits", {}))
         self.raw = raw_response
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "{module}.{class_name}({data}, {locales})".format(
             module=self.__module__,
             class_name=self.__class__.__name__,
@@ -160,7 +163,9 @@ class City(Country):
 
     """
 
-    def __init__(self, raw_response, locales=None):
+    def __init__(
+        self, raw_response: Dict[str, Any], locales: Optional[List[str]] = None
+    ) -> None:
         super(City, self).__init__(raw_response, locales)
         self.city = geoip2.records.City(locales, **raw_response.get("city", {}))
         self.location = geoip2.records.Location(**raw_response.get("location", {}))
@@ -303,13 +308,13 @@ class SimpleModel(SimpleEquality):
 
     __metaclass__ = ABCMeta
 
-    def __init__(self, raw):
+    def __init__(self, raw: Dict[str, Union[bool, str, int]]) -> None:
         self.raw = raw
         self._network = None
         self._prefix_len = raw.get("prefix_len")
         self.ip_address = raw.get("ip_address")
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         # pylint: disable=no-member
         return "{module}.{class_name}({data})".format(
             module=self.__module__,
@@ -318,7 +323,7 @@ def __repr__(self):
         )
 
     @property
-    def network(self):
+    def network(self) -> Optional[Union[ipaddress.IPv4Network, ipaddress.IPv6Network]]:
         """The network for the record"""
         # This code is duplicated for performance reasons
         # pylint: disable=duplicate-code
@@ -391,7 +396,7 @@ class AnonymousIP(SimpleModel):
       :type: ipaddress.IPv4Network or ipaddress.IPv6Network
     """
 
-    def __init__(self, raw):
+    def __init__(self, raw: Dict[str, Union[bool, str, int]]) -> None:
         super(AnonymousIP, self).__init__(raw)
         self.is_anonymous = raw.get("is_anonymous", False)
         self.is_anonymous_vpn = raw.get("is_anonymous_vpn", False)
@@ -434,7 +439,7 @@ class ASN(SimpleModel):
     """
 
     # pylint:disable=too-many-arguments
-    def __init__(self, raw):
+    def __init__(self, raw: Dict[str, Union[str, int]]) -> None:
         super(ASN, self).__init__(raw)
         self.autonomous_system_number = raw.get("autonomous_system_number")
         self.autonomous_system_organization = raw.get("autonomous_system_organization")
@@ -473,7 +478,7 @@ class ConnectionType(SimpleModel):
       :type: ipaddress.IPv4Network or ipaddress.IPv6Network
     """
 
-    def __init__(self, raw):
+    def __init__(self, raw: Dict[str, Union[str, int]]) -> None:
         super(ConnectionType, self).__init__(raw)
         self.connection_type = raw.get("connection_type")
 
@@ -505,7 +510,7 @@ class Domain(SimpleModel):
 
     """
 
-    def __init__(self, raw):
+    def __init__(self, raw: Dict[str, Union[str, int]]) -> None:
         super(Domain, self).__init__(raw)
         self.domain = raw.get("domain")
 
@@ -556,7 +561,7 @@ class ISP(ASN):
     """
 
     # pylint:disable=too-many-arguments
-    def __init__(self, raw):
+    def __init__(self, raw: Dict[str, Union[str, int]]) -> None:
         super(ISP, self).__init__(raw)
         self.isp = raw.get("isp")
         self.organization = raw.get("organization")