From 641729d499a567e5109f4f0aaed419e15aab52de Mon Sep 17 00:00:00 2001 From: Nick Williams Date: Tue, 7 Apr 2015 16:00:16 -0500 Subject: [PATCH] Implemented Issue #13 - Add reader mode information to reader metadata object. --- maxminddb/__init__.py | 2 +- maxminddb/extension/maxminddb.c | 27 +++++++++++++++++++++++++-- maxminddb/reader.py | 8 ++++++++ 3 files changed, 34 insertions(+), 3 deletions(-) diff --git a/maxminddb/__init__.py b/maxminddb/__init__.py index 50d8f54..b58ef8f 100644 --- a/maxminddb/__init__.py +++ b/maxminddb/__init__.py @@ -29,7 +29,7 @@ def open_database(database, mode=MODE_AUTO): """ if (mode == MODE_AUTO and maxminddb.extension and hasattr(maxminddb.extension, 'Reader')) or mode == MODE_MMAP_EXT: - return maxminddb.extension.Reader(database) + return maxminddb.extension.Reader(database, mode) elif mode in (MODE_AUTO, MODE_MMAP, MODE_FILE, MODE_MEMORY): return maxminddb.reader.Reader(database, mode) raise ValueError('Unsupported open mode: {0}'.format(mode)) diff --git a/maxminddb/extension/maxminddb.c b/maxminddb/extension/maxminddb.c index 9e4d45e..c69f8ae 100644 --- a/maxminddb/extension/maxminddb.c +++ b/maxminddb/extension/maxminddb.c @@ -12,6 +12,7 @@ static PyObject *MaxMindDB_error; typedef struct { PyObject_HEAD /* no semicolon */ MMDB_s *mmdb; + PyObject *mode; } Reader_obj; typedef struct { @@ -25,6 +26,8 @@ typedef struct { PyObject *languages; PyObject *node_count; PyObject *record_size; + PyObject *mode_auto; + PyObject *mode; } Metadata_obj; static PyObject *from_entry_data_list(MMDB_entry_data_list_s **entry_data_list); @@ -99,6 +102,7 @@ static int Reader_init(PyObject *self, PyObject *args, PyObject *kwds) } mmdb_obj->mmdb = mmdb; + mmdb_obj->mode = PyInt_FromLong((long)mode); return 0; } @@ -192,10 +196,14 @@ static PyObject *Reader_metadata(PyObject *self, PyObject *UNUSED(args)) return NULL; } + Py_INCREF(mmdb_obj->mode); + PyDict_SetItemString(metadata_dict, "mode", mmdb_obj->mode); + PyObject *metadata = PyObject_Call((PyObject *)&Metadata_Type, args, metadata_dict); Py_DECREF(metadata_dict); + Py_DECREF(mmdb_obj->mode); return metadata; } @@ -225,6 +233,7 @@ static void Reader_dealloc(PyObject *self) static int Metadata_init(PyObject *self, PyObject *args, PyObject *kwds) { + int mode = 0; PyObject *binary_format_major_version, *binary_format_minor_version, @@ -246,10 +255,11 @@ static int Metadata_init(PyObject *self, PyObject *args, PyObject *kwds) "languages", "node_count", "record_size", + "mode", NULL }; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOOOOOOOO", kwlist, + if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOOOOOOOOi", kwlist, &binary_format_major_version, &binary_format_minor_version, &build_epoch, @@ -258,7 +268,8 @@ static int Metadata_init(PyObject *self, PyObject *args, PyObject *kwds) &ip_version, &languages, &node_count, - &record_size)) { + &record_size, + &mode)) { return -1; } @@ -273,6 +284,12 @@ static int Metadata_init(PyObject *self, PyObject *args, PyObject *kwds) obj->languages = languages; obj->node_count = node_count; obj->record_size = record_size; + if (mode == 0) { + obj->mode_auto = PyBool_FromLong(1); + } else { + obj->mode_auto = PyBool_FromLong(0); + } + obj->mode = PyString_FromString("MODE_MMAP_EXT"); Py_INCREF(obj->binary_format_major_version); Py_INCREF(obj->binary_format_minor_version); @@ -299,6 +316,8 @@ static void Metadata_dealloc(PyObject *self) Py_DECREF(obj->languages); Py_DECREF(obj->node_count); Py_DECREF(obj->record_size); + Py_DECREF(obj->mode_auto); + Py_DECREF(obj->mode); PyObject_Del(self); } @@ -493,6 +512,10 @@ static PyMemberDef Metadata_members[] = { READONLY, NULL }, { "record_size", T_OBJECT, offsetof(Metadata_obj, record_size), READONLY, NULL }, + { "mode_auto", T_OBJECT, offsetof(Metadata_obj, mode_auto), + READONLY, NULL }, + { "mode", T_OBJECT, offsetof(Metadata_obj, mode), + READONLY, NULL }, { NULL, 0, 0, 0, NULL } }; /* *INDENT-ON* */ diff --git a/maxminddb/reader.py b/maxminddb/reader.py index 6c0b24f..e39dab1 100644 --- a/maxminddb/reader.py +++ b/maxminddb/reader.py @@ -46,18 +46,22 @@ def __init__(self, database, mode=MODE_AUTO): * MODE_MEMORY - load database into memory. * MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Default. """ + metadata_mode_auto = mode == MODE_AUTO 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) self._buffer_size = self._buffer.size() + metadata_mode = 'MODE_MMAP' elif mode in (MODE_AUTO, MODE_FILE): self._buffer = FileBuffer(database) self._buffer_size = self._buffer.size() + metadata_mode = 'MODE_FILE' elif mode == MODE_MEMORY: with open(database, 'rb') as db_file: self._buffer = db_file.read() self._buffer_size = len(self._buffer) + metadata_mode = 'MODE_MEMORY' else: raise ValueError('Unsupported open mode ({0}). Only MODE_AUTO, ' ' MODE_FILE, and MODE_MEMORY are support by the pure Python ' @@ -76,6 +80,8 @@ def __init__(self, database, mode=MODE_AUTO): metadata_start += len(self._METADATA_START_MARKER) metadata_decoder = Decoder(self._buffer, metadata_start) (metadata, _) = metadata_decoder.decode(metadata_start) + metadata['mode_auto'] = metadata_mode_auto + metadata['mode'] = metadata_mode self._metadata = Metadata(**metadata) # pylint: disable=star-args self._decoder = Decoder(self._buffer, self._metadata.search_tree_size @@ -200,6 +206,8 @@ def __init__(self, **kwargs): 'binary_format_minor_version'] self.build_epoch = kwargs['build_epoch'] self.description = kwargs['description'] + self.mode_auto = kwargs.get('mode_auto', None) + self.mode = kwargs.get('mode', None) @property def node_byte_size(self):