Add memory mode and expose the modes. Closes #10

Author: oschwald

HISTORY.rst

@@ -3,6 +3,16 @@
 History
 -------
 
+1.1.0 (2014-12-XX)
+++++++++++++++++++
+
+* The pure Python reader now supports an optional file and memory mode in
+  addition to the existing memory-map mode. If your Python does not provide
+  the ``mmap`` module, the file mode will be used by default.
+* The preferred method for opening a database is now the ``open_database``
+  function in ``maxminddb``. This function now takes an optional file
+  ``mode``.
+
 1.0.0 (2014-09-22)
 ++++++++++++++++++
 

README.rst

@@ -40,13 +40,22 @@ provide `free GeoLite2 databases
 <http://dev.maxmind.com/geoip/geoip2/geolite2>`_. These files must be
 decompressed with ``gunzip``.
 
-After you have obtained a database and importing the module, you must create a
-``Reader`` object, providing the path to the file as the first argument to the
-constructor. After doing this, you may call the ``get`` method with an IP
-address on the object. This method will return the corresponding values for
-the IP address from the database (e.g., a dictionary for GeoIP2/GeoLite2
-databases). If the database does not contain a record for that IP address, the
-method will return ``None``.
+After you have obtained a database and importing the module, call
+``open_database`` with a path to the database as the first argument.
+Optionally, you may pass a mode as the second arguments. The modes are
+exported from ``maxminddb``. Valid modes are:
+
+* MODE_MMAP_EXT - use the C extension with memory map.
+* MODE_MMAP - read from memory map. Pure Python.
+* MODE_FILE - read database as standard file. Pure Python.
+* MODE_MEMORY - load database into memory. Pure Python.
+* MODE_AUTO - try MODE_MMAP_EXT, MODE_MMAP, MODE_FILE in that order. Default.
+
+The ``open_database`` function returns a ``Reader`` object. To look up an IP
+address, use the ``get`` method on this object. The method will return the
+corresponding values for the IP address from the database (e.g., a dictionary
+for GeoIP2/GeoLite2 databases). If the database does not contain a record for
+that IP address, the method will return ``None``.
 
 Example
 -------
@@ -55,12 +64,10 @@ Example
 
     >>> import maxminddb
     >>>
-    >>> reader = maxminddb.Reader('GeoLite2-City.mmdb')
+    >>> reader = maxminddb.open_database('GeoLite2-City.mmdb')
     >>> reader.get('1.1.1.1')
     {'country': ... }
     >>>
-    >>> # The optional 'close' method will release the resources to the
-    >>> # system immediately.
     >>> reader.close()
 
 Exceptions

maxminddb/__init__.py

@@ -1,17 +1,42 @@
 # pylint:disable=C0111
 import os
 
+import maxminddb.reader
+
 try:
-    if os.environ.get('MAXMINDDB_PURE_PYTHON') == '1':
-        raise ImportError()
-    # pylint: disable=no-name-in-module
-    from maxminddb.extension import Reader, InvalidDatabaseError
-except ImportError as import_error:
-    if os.environ.get('MAXMINDDB_PURE_PYTHON') == '0':
-        raise import_error
-    from maxminddb.decoder import InvalidDatabaseError
-    from maxminddb.reader import Reader
+    import maxminddb.extension
+except ImportError:
+    maxminddb.extension = None
+
+from maxminddb.const import (MODE_AUTO, MODE_MMAP, MODE_MMAP_EXT, MODE_FILE,
+                             MODE_MEMORY)
+from maxminddb.decoder import InvalidDatabaseError
+
+
+def open_database(database, mode=MODE_AUTO):
+    """Open a Maxmind DB database
+
+    Arguments:
+        database -- A path to a valid MaxMind DB file such as a GeoIP2
+                    database file.
+        mode -- mode to open the database with. Valid mode are:
+            * MODE_MMAP_EXT - use the C extension with memory map.
+            * MODE_MMAP - read from memory map. Pure Python.
+            * MODE_FILE - read database as standard file. Pure Python.
+            * MODE_MEMORY - load database into memory. Pure Python.
+            * MODE_AUTO - tries MODE_MMAP_EXT, MODE_MMAP, MODE_FILE in that
+                          order. Default mode.
+    """
+    if (mode == MODE_AUTO and maxminddb.extension) or mode == MODE_MMAP_EXT:
+        return maxminddb.extension.Reader(database)
+    elif mode in (MODE_AUTO, MODE_MMAP, MODE_FILE, MODE_MEMORY):
+        return maxminddb.reader.Reader(database, mode)
+    raise ValueError('Unsupported open mode: {}', mode)
+
 
+def Reader(database):  # pylint: disable=invalid-name
+    """This exists for backwards compatibility. Use open_database instead"""
+    open_database(database)
 
 __title__ = 'maxminddb'
 __version__ = '1.0.0'

maxminddb/const.py

@@ -0,0 +1,7 @@
+"""Constants used in the API"""
+
+MODE_AUTO = 0
+MODE_MMAP_EXT = 1
+MODE_MMAP = 2
+MODE_FILE = 4
+MODE_MEMORY = 8

maxminddb/extension/maxminddb.c

@@ -49,11 +49,20 @@ static PyObject *from_uint128(const MMDB_entry_data_list_s *entry_data_list);
     #  define UNUSED(x) UNUSED_ ## x
 #endif
 
-static int Reader_init(PyObject *self, PyObject *args, PyObject *UNUSED(kwds))
+static int Reader_init(PyObject *self, PyObject *args, PyObject *kwds)
 {
     char *filename;
+    int mode = 0;
 
-    if (!PyArg_ParseTuple(args, "s", &filename)) {
+    static char *kwlist[] = {"database", "mode", NULL};
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s|i", kwlist, &filename, &mode)) {
+        return -1;
+    }
+
+    if (mode != 0 && mode != 1) {
+        PyErr_Format(PyExc_ValueError, "Unsupported open mode (%i). Only "
+            "MODE_AUTO and MODE_MMAP_EXT are supported by this extension.",
+            mode);
         return -1;
     }
 
@@ -540,13 +549,21 @@ MOD_INIT(extension){
     }
     PyModule_AddObject(m, "extension", (PyObject *)&Metadata_Type);
 
-    MaxMindDB_error = PyErr_NewException("extension.InvalidDatabaseError", NULL,
-                                         NULL);
+    PyObject* error_mod = PyImport_ImportModule("maxminddb.errors");
+    if (error_mod == NULL) {
+        RETURN_MOD_INIT(NULL);
+    }
+
+    MaxMindDB_error = PyObject_GetAttrString(error_mod, "InvalidDatabaseError");
+    Py_DECREF(error_mod);
+
     if (MaxMindDB_error == NULL) {
         RETURN_MOD_INIT(NULL);
     }
 
     Py_INCREF(MaxMindDB_error);
+
+    /* We primarily add it to the module for backwards compatibility */
     PyModule_AddObject(m, "InvalidDatabaseError", MaxMindDB_error);
 
     RETURN_MOD_INIT(m);

maxminddb/file.py

@@ -21,9 +21,8 @@ def __getitem__(self, key):
         else:
             raise TypeError("Invalid argument type.")
 
-    def rfind(self, needle, beg):
-        """Reverse find needle from beg"""
-        start = max(0, beg)
+    def rfind(self, needle, start):
+        """Reverse find needle from start"""
         self._handle.seek(start)
         pos = self._handle.read(self._size - start - 1).rfind(needle)
         if pos == -1:

maxminddb/reader.py

@@ -16,6 +16,7 @@
 import struct
 
 from maxminddb.compat import byte_from_int, int_from_byte, ipaddress
+from maxminddb.const import MODE_AUTO, MODE_MMAP, MODE_FILE, MODE_MEMORY
 from maxminddb.decoder import Decoder
 from maxminddb.errors import InvalidDatabaseError
 from maxminddb.file import FileBuffer
@@ -33,24 +34,41 @@ class Reader(object):
 
     _ipv4_start = None
 
-    def __init__(self, database):
+    def __init__(self, database, mode=MODE_AUTO):
         """Reader for the MaxMind DB file format
 
         Arguments:
         database -- A path to a valid MaxMind DB file such as a GeoIP2
                     database file.
+        mode -- mode to open the database with. Valid mode are:
+            * MODE_MMAP - read from memory map.
+            * MODE_FILE - read database as standard file.
+            * MODE_MEMORY - load database into memory.
+            * MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Default.
         """
-        if mmap:
+        if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP:
             with open(database, 'rb') as db_file:
                 self._buffer = mmap.mmap(
                     db_file.fileno(), 0, access=mmap.ACCESS_READ)
-        else:
+                self._buffer_size = self._buffer.size()
+        elif mode in (MODE_AUTO, MODE_FILE):
             self._buffer = FileBuffer(database)
+            self._buffer_size = self._buffer.size()
+        elif mode == MODE_MEMORY:
+            with open(database, 'rb') as db_file:
+                self._buffer = db_file.read()
+                self._buffer_size = len(self._buffer)
+        else:
+            raise ValueError('Unsupported open mode ({0}). Only MODE_AUTO, '
+                             ' MODE_FILE, and MODE_MEMORY are support by the pure Python '
+                             'Reader'.format(mode))
 
         metadata_start = self._buffer.rfind(self._METADATA_START_MARKER,
-                                            self._buffer.size() - 128 * 1024)
+                                            max(0, self._buffer_size
+                                                - 128 * 1024))
 
         if metadata_start == -1:
+            self.close()
             raise InvalidDatabaseError('Error opening database file ({0}). '
                                        'Is this a valid MaxMind DB file?'
                                        ''.format(database))
@@ -149,7 +167,7 @@ def _resolve_data_pointer(self, pointer):
         resolved = pointer - self._metadata.node_count + \
             self._metadata.search_tree_size
 
-        if resolved > self._buffer.size():
+        if resolved > self._buffer_size:
             raise InvalidDatabaseError(
                 "The MaxMind DB file's search tree is corrupt")
 
@@ -158,7 +176,8 @@ def _resolve_data_pointer(self, pointer):
 
     def close(self):
         """Closes the MaxMind DB file and returns the resources to the system"""
-        self._buffer.close()
+        if type(self._buffer) not in (str, bytes):
+            self._buffer.close()
 
 
 class Metadata(object):