diff --git a/.pylintrc b/.pylintrc index 53a1201..dd54a45 100644 --- a/.pylintrc +++ b/.pylintrc @@ -4,3 +4,4 @@ disable=C0330 [BASIC] no-docstring-rgx=_.* +extension-pkg-allow-list=maxminddb.extension diff --git a/maxminddb/__init__.py b/maxminddb/__init__.py index 582d209..5574cf0 100644 --- a/maxminddb/__init__.py +++ b/maxminddb/__init__.py @@ -1,29 +1,42 @@ # pylint:disable=C0111 import os -from typing import AnyStr, IO, Union +from typing import IO, AnyStr, Union, cast -import maxminddb.reader - -try: - import maxminddb.extension -except ImportError: - maxminddb.extension = None # type: ignore - -from maxminddb.const import ( +from .const import ( MODE_AUTO, - MODE_MMAP, - MODE_MMAP_EXT, + MODE_FD, MODE_FILE, MODE_MEMORY, - MODE_FD, + MODE_MMAP, + MODE_MMAP_EXT, ) -from maxminddb.decoder import InvalidDatabaseError -from maxminddb.reader import Reader as PyReader +from .decoder import InvalidDatabaseError +from .reader import Reader + +try: + # pylint: disable=import-self + from . import extension as _extension +except ImportError: + _extension = None # type: ignore[assignment] + + +__all__ = [ + "InvalidDatabaseError", + "MODE_AUTO", + "MODE_FD", + "MODE_FILE", + "MODE_MEMORY", + "MODE_MMAP", + "MODE_MMAP_EXT", + "Reader", + "open_database", +] def open_database( - database: Union[AnyStr, int, os.PathLike, IO], mode: int = MODE_AUTO -) -> Union[PyReader, "maxminddb.extension.Reader"]: + database: Union[AnyStr, int, os.PathLike, IO], + mode: int = MODE_AUTO, +) -> Reader: """Open a MaxMind DB database Arguments: @@ -39,21 +52,32 @@ def open_database( * MODE_AUTO - tries MODE_MMAP_EXT, MODE_MMAP, MODE_FILE in that order. Default mode. """ - has_extension = maxminddb.extension and hasattr(maxminddb.extension, "Reader") - if (mode == MODE_AUTO and has_extension) or mode == MODE_MMAP_EXT: - if not has_extension: - raise ValueError( - "MODE_MMAP_EXT requires the maxminddb.extension module to be available" - ) - return maxminddb.extension.Reader(database) - if mode in (MODE_AUTO, MODE_MMAP, MODE_FILE, MODE_MEMORY, MODE_FD): - return PyReader(database, mode) - raise ValueError(f"Unsupported open mode: {mode}") - - -def Reader(database): # pylint: disable=invalid-name - """This exists for backwards compatibility. Use open_database instead""" - return open_database(database) + if mode not in ( + MODE_AUTO, + MODE_FD, + MODE_FILE, + MODE_MEMORY, + MODE_MMAP, + MODE_MMAP_EXT, + ): + raise ValueError(f"Unsupported open mode: {mode}") + + has_extension = _extension and hasattr(_extension, "Reader") + use_extension = has_extension if mode == MODE_AUTO else mode == MODE_MMAP_EXT + + if not use_extension: + return Reader(database, mode) + + if not has_extension: + raise ValueError( + "MODE_MMAP_EXT requires the maxminddb.extension module to be available" + ) + + # The C type exposes the same API as the Python Reader, so for type + # checking purposes, pretend it is one. (Ideally this would be a subclass + # of, or share a common parent class with, the Python Reader + # implementation.) + return cast(Reader, _extension.Reader(database, mode)) __title__ = "maxminddb" diff --git a/tests/reader_test.py b/tests/reader_test.py index 30acdbc..9b9fe5a 100644 --- a/tests/reader_test.py +++ b/tests/reader_test.py @@ -250,8 +250,8 @@ def test_opening_path(self): self.assertEqual(reader.metadata().database_type, "MaxMind DB Decoder Test") def test_no_extension_exception(self): - real_extension = maxminddb.extension - maxminddb.extension = None + real_extension = maxminddb._extension + maxminddb._extension = None with self.assertRaisesRegex( ValueError, "MODE_MMAP_EXT requires the maxminddb.extension module to be available", @@ -259,7 +259,7 @@ def test_no_extension_exception(self): open_database( "tests/data/test-data/MaxMind-DB-test-decoder.mmdb", MODE_MMAP_EXT ) - maxminddb.extension = real_extension + maxminddb._extension = real_extension def test_broken_database(self): reader = open_database(