Skip to content

Commit c161311

Browse files
committed
Use duck typing and make error messages consistent
Also, add a missing return when getting packed attribute fails in extension.
1 parent aaa1ef9 commit c161311

File tree

3 files changed

+28
-15
lines changed

3 files changed

+28
-15
lines changed

extension/maxminddb.c

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -185,11 +185,13 @@ static int ip_converter(PyObject *obj, struct sockaddr_storage *ip_address)
185185
Py_ssize_t len = strlen(ipstr);
186186
#endif
187187
if (!ipstr) {
188-
PyErr_SetString(PyExc_ValueError, "invalid string");
188+
PyErr_SetString(PyExc_TypeError,
189+
"argument 1 contains an invalid string");
189190
return 0;
190191
}
191192
if (strlen(ipstr) != (size_t)len) {
192-
PyErr_SetString(PyExc_ValueError, "embedded null character");
193+
PyErr_SetString(PyExc_TypeError,
194+
"argument 1 contains an embedded null character");
193195
return 0;
194196
}
195197

@@ -214,13 +216,16 @@ static int ip_converter(PyObject *obj, struct sockaddr_storage *ip_address)
214216
}
215217
PyObject *packed = PyObject_GetAttrString(obj, "packed");
216218
if (!packed) {
217-
PyErr_SetString(PyExc_ValueError, "error about object type");
219+
PyErr_SetString(PyExc_TypeError,
220+
"argument 1 must be a string or ipaddress object");
221+
return 0;
218222
}
219223
Py_ssize_t len;
220224
char *bytes;
221225
int status = PyBytes_AsStringAndSize(packed, &bytes, &len);
222226
if (status == -1) {
223-
PyErr_SetString(PyExc_ValueError, "cannot get bytes");
227+
PyErr_SetString(PyExc_TypeError,
228+
"argument 1 must be a string or ipaddress object");
224229
Py_DECREF(packed);
225230
return 0;
226231
}
@@ -241,7 +246,8 @@ static int ip_converter(PyObject *obj, struct sockaddr_storage *ip_address)
241246
return 1;
242247
}
243248
default:
244-
PyErr_SetString(PyExc_ValueError, "unexpected packed length");
249+
PyErr_SetString(PyExc_ValueError,
250+
"argument 1 returned an unexpected packed length for address");
245251
Py_DECREF(packed);
246252
return 0;
247253
}

maxminddb/reader.py

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -104,26 +104,25 @@ def get(self, ip_address):
104104
Arguments:
105105
ip_address -- an IP address in the standard string notation
106106
"""
107-
if isinstance(ip_address,
108-
(ipaddress.IPv4Address, ipaddress.IPv6Address)):
109-
address = ip_address
110-
elif isinstance(ip_address, string_type):
107+
if isinstance(ip_address, string_type):
111108
address = compat_ip_address(ip_address)
112109
else:
113-
raise TypeError('argument 1 must be %s, not %s' %
114-
(string_type_name, type(ip_address).__name__))
110+
address = ip_address
111+
112+
try:
113+
packed_address = bytearray(address.packed)
114+
except AttributeError:
115+
raise TypeError('argument 1 must be a string or ipaddress object')
115116

116117
if address.version == 6 and self._metadata.ip_version == 4:
117118
raise ValueError(
118119
'Error looking up {0}. You attempted to look up '
119120
'an IPv6 address in an IPv4-only database.'.format(ip_address))
120-
pointer = self._find_address_in_tree(address)
121+
pointer = self._find_address_in_tree(packed_address)
121122

122123
return self._resolve_data_pointer(pointer) if pointer else None
123124

124-
def _find_address_in_tree(self, ip_address):
125-
packed = bytearray(ip_address.packed)
126-
125+
def _find_address_in_tree(self, packed):
127126
bit_count = len(packed) * 8
128127
node = self._start_node(bit_count)
129128

tests/reader_test.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,14 @@ def test_no_get_args(self):
193193
reader.get()
194194
reader.close()
195195

196+
def test_incorrect_get_arg_type(self):
197+
reader = open_database('tests/data/test-data/GeoIP2-City-Test.mmdb',
198+
self.mode)
199+
with self.assertRaisesRegex(
200+
TypeError, "argument 1 must be a string or ipaddress object"):
201+
reader.get(1)
202+
reader.close()
203+
196204
def test_metadata_args(self):
197205
reader = open_database(
198206
'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)

0 commit comments

Comments
 (0)