Add types to internal modules

Author: oschwald

maxminddb/decoder.py

@@ -7,13 +7,23 @@
 """
 import struct
 
+from mmap import mmap
+from typing import cast, Dict, List, Tuple, Union
+
 from maxminddb.errors import InvalidDatabaseError
+from maxminddb.file import FileBuffer
+from maxminddb.types import Record
 
 
 class Decoder:  # pylint: disable=too-few-public-methods
     """Decoder for the data section of the MaxMind DB"""
 
-    def __init__(self, database_buffer, pointer_base=0, pointer_test=False):
+    def __init__(
+        self,
+        database_buffer: Union[FileBuffer, mmap, bytes],
+        pointer_base: int = 0,
+        pointer_test: bool = False,
+    ) -> None:
         """Created a Decoder for a MaxMind DB
 
         Arguments:
@@ -25,35 +35,37 @@ def __init__(self, database_buffer, pointer_base=0, pointer_test=False):
         self._buffer = database_buffer
         self._pointer_base = pointer_base
 
-    def _decode_array(self, size, offset):
+    def _decode_array(self, size: int, offset: int) -> Tuple[List[Record], int]:
         array = []
         for _ in range(size):
             (value, offset) = self.decode(offset)
             array.append(value)
         return array, offset
 
-    def _decode_boolean(self, size, offset):  # pylint: disable=no-self-use
+    def _decode_boolean(  # pylint: disable=no-self-use
+        self, size: int, offset: int
+    ) -> Tuple[bool, int]:
         return size != 0, offset
 
-    def _decode_bytes(self, size, offset):
+    def _decode_bytes(self, size: int, offset: int) -> Tuple[bytes, int]:
         new_offset = offset + size
         return self._buffer[offset:new_offset], new_offset
 
-    def _decode_double(self, size, offset):
+    def _decode_double(self, size: int, offset: int) -> Tuple[float, int]:
         self._verify_size(size, 8)
         new_offset = offset + size
         packed_bytes = self._buffer[offset:new_offset]
         (value,) = struct.unpack(b"!d", packed_bytes)
         return value, new_offset
 
-    def _decode_float(self, size, offset):
+    def _decode_float(self, size: int, offset: int) -> Tuple[float, int]:
         self._verify_size(size, 4)
         new_offset = offset + size
         packed_bytes = self._buffer[offset:new_offset]
         (value,) = struct.unpack(b"!f", packed_bytes)
         return value, new_offset
 
-    def _decode_int32(self, size, offset):
+    def _decode_int32(self, size: int, offset: int) -> Tuple[int, int]:
         if size == 0:
             return 0, offset
         new_offset = offset + size
@@ -64,15 +76,15 @@ def _decode_int32(self, size, offset):
         (value,) = struct.unpack(b"!i", packed_bytes)
         return value, new_offset
 
-    def _decode_map(self, size, offset):
-        container = {}
+    def _decode_map(self, size: int, offset: int) -> Tuple[Dict[str, Record], int]:
+        container: Dict[str, Record] = {}
         for _ in range(size):
             (key, offset) = self.decode(offset)
             (value, offset) = self.decode(offset)
-            container[key] = value
+            container[cast(str, key)] = value
         return container, offset
 
-    def _decode_pointer(self, size, offset):
+    def _decode_pointer(self, size: int, offset: int) -> Tuple[Record, int]:
         pointer_size = (size >> 3) + 1
 
         buf = self._buffer[offset : offset + pointer_size]
@@ -95,12 +107,12 @@ def _decode_pointer(self, size, offset):
         (value, _) = self.decode(pointer)
         return value, new_offset
 
-    def _decode_uint(self, size, offset):
+    def _decode_uint(self, size: int, offset: int) -> Tuple[int, int]:
         new_offset = offset + size
         uint_bytes = self._buffer[offset:new_offset]
         return int.from_bytes(uint_bytes, "big"), new_offset
 
-    def _decode_utf8_string(self, size, offset):
+    def _decode_utf8_string(self, size: int, offset: int) -> Tuple[str, int]:
         new_offset = offset + size
         return self._buffer[offset:new_offset].decode("utf-8"), new_offset
 
@@ -120,7 +132,7 @@ def _decode_utf8_string(self, size, offset):
         15: _decode_float,
     }
 
-    def decode(self, offset):
+    def decode(self, offset: int) -> Tuple[Record, int]:
         """Decode a section of the data section starting at offset
 
         Arguments:
@@ -143,7 +155,7 @@ def decode(self, offset):
         (size, new_offset) = self._size_from_ctrl_byte(ctrl_byte, new_offset, type_num)
         return decoder(self, size, new_offset)
 
-    def _read_extended(self, offset):
+    def _read_extended(self, offset: int) -> Tuple[int, int]:
         next_byte = self._buffer[offset]
         type_num = next_byte + 7
         if type_num < 7:
@@ -155,14 +167,16 @@ def _read_extended(self, offset):
         return type_num, offset + 1
 
     @staticmethod
-    def _verify_size(expected, actual):
+    def _verify_size(expected: int, actual: int) -> None:
         if expected != actual:
             raise InvalidDatabaseError(
                 "The MaxMind DB file's data section contains bad data "
                 "(unknown data type or corrupt data)"
             )
 
-    def _size_from_ctrl_byte(self, ctrl_byte, offset, type_num):
+    def _size_from_ctrl_byte(
+        self, ctrl_byte: int, offset: int, type_num: int
+    ) -> Tuple[int, int]:
         size = ctrl_byte & 0x1F
         if type_num == 1 or size < 29:
             return size, offset

maxminddb/file.py

@@ -1,6 +1,7 @@
 """For internal use only. It provides a slice-like file reader."""
 
 import os
+from typing import Union
 
 try:
     # pylint: disable=no-name-in-module
@@ -12,44 +13,44 @@
 class FileBuffer:
     """A slice-able file reader"""
 
-    def __init__(self, database):
+    def __init__(self, database: str) -> None:
         self._handle = open(database, "rb")
         self._size = os.fstat(self._handle.fileno()).st_size
         if not hasattr(os, "pread"):
             self._lock = Lock()
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: Union[slice, int]):
         if isinstance(key, slice):
             return self._read(key.stop - key.start, key.start)
         if isinstance(key, int):
             return self._read(1, key)[0]
         raise TypeError("Invalid argument type.")
 
-    def rfind(self, needle, start):
+    def rfind(self, needle: bytes, start: int) -> int:
         """Reverse find needle from start"""
         pos = self._read(self._size - start - 1, start).rfind(needle)
         if pos == -1:
             return pos
         return start + pos
 
-    def size(self):
+    def size(self) -> int:
         """Size of file"""
         return self._size
 
-    def close(self):
+    def close(self) -> None:
         """Close file"""
         self._handle.close()
 
     if hasattr(os, "pread"):
 
-        def _read(self, buffersize, offset):
+        def _read(self, buffersize: int, offset: int) -> bytes:
             """read that uses pread"""
             # pylint: disable=no-member
             return os.pread(self._handle.fileno(), buffersize, offset)
 
     else:
 
-        def _read(self, buffersize, offset):
+        def _read(self, buffersize: int, offset: int) -> bytes:
             """read with a lock
 
             This lock is necessary as after a fork, the different processes

maxminddb/reader.py

@@ -17,7 +17,7 @@
 import struct
 from ipaddress import IPv4Address, IPv6Address
 from os import PathLike
-from typing import Any, AnyStr, Dict, List, IO, Optional, Tuple, Union
+from typing import Any, AnyStr, IO, Optional, Tuple, Union
 
 from maxminddb.const import MODE_AUTO, MODE_MMAP, MODE_FILE, MODE_MEMORY, MODE_FD
 from maxminddb.decoder import Decoder
@@ -61,7 +61,7 @@ def __init__(
                 self._buffer_size = self._buffer.size()
             filename = database
         elif mode in (MODE_AUTO, MODE_FILE):
-            self._buffer = FileBuffer(database)
+            self._buffer = FileBuffer(database)  # type: ignore
             self._buffer_size = self._buffer.size()
             filename = database
         elif mode == MODE_MEMORY:
@@ -95,6 +95,12 @@ def __init__(
         metadata_start += len(self._METADATA_START_MARKER)
         metadata_decoder = Decoder(self._buffer, metadata_start)
         (metadata, _) = metadata_decoder.decode(metadata_start)
+
+        if not isinstance(metadata, dict):
+            raise InvalidDatabaseError(
+                "Error reading metadata in database file ({0}).".format(filename)
+            )
+
         self._metadata = Metadata(**metadata)  # pylint: disable=bad-option-value
 
         self._decoder = Decoder(
@@ -208,23 +214,7 @@ def _read_node(self, node_number: int, index: int) -> int:
             raise InvalidDatabaseError("Unknown record size: {0}".format(record_size))
         return struct.unpack(b"!I", node_bytes)[0]
 
-    def _resolve_data_pointer(
-        self, pointer: int
-    ) -> Union[
-        str,
-        Dict[str, str],
-        Dict[
-            str,
-            Union[
-                List[int],
-                bytes,
-                float,
-                int,
-                Dict[str, Dict[str, Union[List[int], str]]],
-                str,
-            ],
-        ],
-    ]:
+    def _resolve_data_pointer(self, pointer: int) -> Record:
         resolved = pointer - self._metadata.node_count + self._metadata.search_tree_size
 
         if resolved >= self._buffer_size: