Support for with-statement (context manager) for the reader object

maxminddb/extension/maxminddb.c

@@ -12,6 +12,7 @@ static PyObject *MaxMindDB_error;
 typedef struct {
     PyObject_HEAD               /* no semicolon */
     MMDB_s *mmdb;
+    PyObject *closed;
 } Reader_obj;
 
 typedef struct {
@@ -99,6 +100,7 @@ static int Reader_init(PyObject *self, PyObject *args, PyObject *kwds)
     }
 
     mmdb_obj->mmdb = mmdb;
+    mmdb_obj->closed = Py_False;
     return 0;
 }
 
@@ -209,6 +211,28 @@ static PyObject *Reader_close(PyObject *self, PyObject *UNUSED(args))
         mmdb_obj->mmdb = NULL;
     }
 
+    mmdb_obj->closed = Py_True;
+
+    Py_RETURN_NONE;
+}
+
+static PyObject *Reader__enter__(PyObject *self, PyObject *UNUSED(args))
+{
+    Reader_obj *mmdb_obj = (Reader_obj *)self;
+
+    if(mmdb_obj->closed == Py_True) {
+        PyErr_SetString(PyExc_ValueError,
+                        "Attempt to reopen a closed MaxMind DB.");
+        return NULL;
+    }
+
+    Py_INCREF(self);
+    return (PyObject *)self;
+}
+
+static PyObject *Reader__exit__(PyObject *self, PyObject *UNUSED(args))
+{
+    Reader_close(self, NULL);
     Py_RETURN_NONE;
 }
 
@@ -455,16 +479,24 @@ static PyMethodDef Reader_methods[] = {
     { "metadata", Reader_metadata, METH_NOARGS,
       "Returns metadata object for database" },
     { "close",    Reader_close,    METH_NOARGS, "Closes database"},
+    { "__exit__", Reader__exit__, METH_VARARGS, "Called when exiting a with-context. Calls close"},
+    { "__enter__", Reader__enter__, METH_NOARGS, "Called when entering a with-context."},
     { NULL,       NULL,            0,           NULL        }
 };
 
+static PyMemberDef Reader_members[] = {
+    { "closed", T_OBJECT, offsetof(Reader_obj, closed), READONLY, NULL },
+    { NULL, 0, 0, 0, NULL }
+};
+
 static PyTypeObject Reader_Type = {
     PyVarObject_HEAD_INIT(NULL, 0)
     .tp_basicsize = sizeof(Reader_obj),
     .tp_dealloc = Reader_dealloc,
     .tp_doc = "Reader object",
     .tp_flags = Py_TPFLAGS_DEFAULT,
     .tp_methods = Reader_methods,
+    .tp_members = Reader_members,
     .tp_name = "Reader",
     .tp_init = Reader_init,
 };

maxminddb/reader.py

@@ -45,7 +45,6 @@ def __init__(self, database, mode=MODE_AUTO):
             * MODE_MEMORY - load database into memory.
             * MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Default.
         """
-        # pylint: disable=redefined-variable-type
         if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP:
             with open(database, 'rb') as db_file:
                 self._buffer = mmap.mmap(
@@ -81,6 +80,7 @@ def __init__(self, database, mode=MODE_AUTO):
 
         self._decoder = Decoder(self._buffer, self._metadata.search_tree_size +
                                 self._DATA_SECTION_SEPARATOR_SIZE)
+        self.closed = False
 
     def metadata(self):
         """Return the metadata associated with the MaxMind DB file"""
@@ -181,6 +181,15 @@ def close(self):
         # pylint: disable=unidiomatic-typecheck
         if type(self._buffer) not in (str, bytes):
             self._buffer.close()
+        self.closed = True
+
+    def __exit__(self, *args):
+        self.close()
+
+    def __enter__(self):
+        if self.closed:
+            raise ValueError('Attempt to reopen a closed MaxMind DB')
+        return self
 
 
 class Metadata(object):

tests/reader_test.py

@@ -201,6 +201,33 @@ def test_closed_get(self):
                                'Attempt to read from a closed MaxMind DB.'
                                '|closed', reader.get, ('1.1.1.1'))
 
+    def test_with_statement(self):
+        filename = 'tests/data/test-data/MaxMind-DB-test-ipv4-24.mmdb'
+        with open_database(filename, self.mode) as reader:
+            self._check_ip_v4(reader, filename)
+        self.assertEqual(reader.closed, True)
+
+    def test_with_statement_close(self):
+        filename = 'tests/data/test-data/MaxMind-DB-test-ipv4-24.mmdb'
+        reader = open_database(filename, self.mode)
+        reader.close()
+
+        def use_with(reader):
+            with reader:
+                pass
+
+        self.assertRaisesRegex(ValueError, 'Attempt to reopen a closed MaxMind DB',
+                               use_with, reader)
+
+
+    def test_closed(self):
+        reader = open_database(
+            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
+        self.assertEqual(reader.closed, False)
+        reader.close()
+        self.assertEqual(reader.closed, True)
+
+
     # XXX - Figure out whether we want to have the same behavior on both the
     #       extension and the pure Python reader. If we do, the pure Python
     #       reader will need to throw an exception or the extension will need