Support ipaddress objects in "get" method

extension/maxminddb.c

@@ -107,9 +107,64 @@ static int Reader_init(PyObject *self, PyObject *args, PyObject *kwds)
 static PyObject *Reader_get(PyObject *self, PyObject *args)
 {
     char *ip_address = NULL;
+    PyObject *object = NULL;
+#if PY_MAJOR_VERSION >= 3
+    PyObject *bytes = NULL;
+#endif
 
     Reader_obj *mmdb_obj = (Reader_obj *)self;
-    if (!PyArg_ParseTuple(args, "s", &ip_address)) {
+
+    if (PyArg_ParseTuple(args, "s", &ip_address)) {
+        // pass
+    } else if (PyErr_Clear(), PyArg_ParseTuple(args, "O", &object)) {
+        PyObject* module = PyImport_ImportModule("ipaddress");
+        if (module == NULL) {
+            return NULL;
+        }
+
+        PyObject* module_dict = PyModule_GetDict(module);
+        if (module_dict == NULL) {
+            return NULL;
+        }
+
+        PyObject* ipaddress_IPv4Address = PyDict_GetItemString(
+            module_dict, "IPv4Address");
+        int is_ipaddress_object = 0;
+
+        if (ipaddress_IPv4Address != NULL) {
+            is_ipaddress_object = PyObject_IsInstance(ipaddress_IPv4Address,
+                object);
+        }
+
+        PyObject* ipaddress_IPv6Address;
+        if (!is_ipaddress_object &&
+            (ipaddress_IPv6Address = PyDict_GetItemString(module_dict,
+                "IPv6Address")) != NULL) {
+            is_ipaddress_object = PyObject_IsInstance(ipaddress_IPv6Address,
+                object);
+        }
+
+        PyErr_Clear();
+        PyObject *str;
+
+        if (!is_ipaddress_object) {
+            PyErr_SetString(PyExc_TypeError, "IP address must be a string,"
+                " ipaddress.IPv4Address or ipaddress.IPv6Address object");
+        } else if ((str = PyObject_Str(object)) != NULL) {
+#if PY_MAJOR_VERSION >= 3
+            bytes = PyUnicode_AsEncodedString(str, "UTF-8", "strict");
+            ip_address = PyBytes_AS_STRING(bytes);
+#else
+            ip_address = PyString_AsString(str);
+#endif
+        }
+
+        Py_DECREF(module);
+    } else {
+        return NULL;
+    }
+
+    if (ip_address == NULL) {
         return NULL;
     }
 
@@ -131,6 +186,10 @@ static PyObject *Reader_get(PyObject *self, PyObject *args)
         PyErr_Format(PyExc_ValueError,
                      "'%s' does not appear to be an IPv4 or IPv6 address.",
                      ip_address);
+
+#if PY_MAJOR_VERSION >= 3
+        Py_XDECREF(bytes);
+#endif
         return NULL;
     }
 
@@ -143,6 +202,9 @@ static PyObject *Reader_get(PyObject *self, PyObject *args)
         }
         PyErr_Format(exception, "Error looking up %s. %s",
                      ip_address, MMDB_strerror(mmdb_error));
+#if PY_MAJOR_VERSION >= 3
+        Py_XDECREF(bytes);
+#endif
         return NULL;
     }
 
@@ -157,9 +219,16 @@ static PyObject *Reader_get(PyObject *self, PyObject *args)
                      "Error while looking up data for %s. %s",
                      ip_address, MMDB_strerror(status));
         MMDB_free_entry_data_list(entry_data_list);
+#if PY_MAJOR_VERSION >= 3
+        Py_XDECREF(bytes);
+#endif
         return NULL;
     }
 
+#if PY_MAJOR_VERSION >= 3
+    Py_XDECREF(bytes);
+#endif
+
     MMDB_entry_data_list_s *original_entry_data_list = entry_data_list;
     PyObject *py_obj = from_entry_data_list(&entry_data_list);
     MMDB_free_entry_data_list(original_entry_data_list);

maxminddb/reader.py

@@ -13,6 +13,7 @@
     # pylint: disable=invalid-name
     mmap = None
 
+import ipaddress
 import struct
 
 from maxminddb.compat import (byte_from_int, compat_ip_address, string_type,
@@ -105,11 +106,14 @@ def get(self, ip_address):
         Arguments:
         ip_address -- an IP address in the standard string notation
         """
-        if not isinstance(ip_address, string_type):
+        if isinstance(ip_address,
+                      (ipaddress.IPv4Address, ipaddress.IPv6Address)):
+            address = ip_address
+        elif not isinstance(ip_address, string_type):
             raise TypeError('argument 1 must be %s, not %s' %
                             (string_type_name, type(ip_address).__name__))
-
-        address = compat_ip_address(ip_address)
+        else:
+            address = compat_ip_address(ip_address)
 
         if address.version == 6 and self._metadata.ip_version == 4:
             raise ValueError(

tests/reader_test.py

@@ -3,6 +3,7 @@
 
 from __future__ import unicode_literals
 
+import ipaddress
 import logging
 import mock
 import os
@@ -125,9 +126,7 @@ def test_no_extension_exception(self):
     def test_ip_object_lookup(self):
         reader = open_database('tests/data/test-data/GeoIP2-City-Test.mmdb',
                                self.mode)
-        with self.assertRaisesRegex(TypeError,
-                                    "must be str(?:ing)?, not IPv6Address"):
-            reader.get(compat_ip_address('2001:220::'))
+        reader.get(compat_ip_address('2001:220::'))
         reader.close()
 
     def test_broken_database(self):
@@ -184,7 +183,7 @@ def test_no_constructor_args(self):
     def test_too_many_get_args(self):
         reader = open_database(
             'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
-        self.assertRaises(TypeError, reader.get, ('1.1.1.1', 'blah'))
+        self.assertRaises(TypeError, reader.get, '1.1.1.1', 'blah')
         reader.close()
 
     def test_no_get_args(self):
@@ -362,6 +361,11 @@ def _check_ip_v4(self, reader, file_name):
                 'found expected data record for ' + key_address + ' in ' +
                 file_name)
 
+            self.assertEqual(
+                data, reader.get(ipaddress.ip_address(key_address)),
+                'found expected data record for ' + key_address + ' in ' +
+                file_name)
+
         for ip in ['1.1.1.33', '255.254.253.123']:
             self.assertIsNone(reader.get(ip))
 
@@ -391,8 +395,14 @@ def _check_ip_v6(self, reader, file_name):
                              'found expected data record for ' + key_address +
                              ' in ' + file_name)
 
+            self.assertEqual({'ip': value_address},
+                             reader.get(ipaddress.ip_address(key_address)),
+                             'found expected data record for ' + key_address +
+                             ' in ' + file_name)
+
         for ip in ['1.1.1.33', '255.254.253.123', '89fa::']:
             self.assertIsNone(reader.get(ip))
+            self.assertIsNone(reader.get(ipaddress.ip_address(ip)))
 
 
 def has_maxminddb_extension():