Commit 988fbb32 by William Tisäter

Fix Python 2encoding bug when reading CityLite from mmap or memory

parent ef94ca08
......@@ -50,9 +50,11 @@ from pygeoip.const import PY2, PY3
from pygeoip.timezone import time_zone_by_country_and_region
STANDARD = const.STANDARD
MMAP_CACHE = const.MMAP_CACHE
MEMORY_CACHE = const.MEMORY_CACHE
STANDARD = const.STANDARD
ENCODING = const.ENCODING
class GeoIPError(Exception):
......@@ -115,7 +117,7 @@ class GeoIP(GeoIPBase):
f.close()
else:
self._filehandle = codecs.open(filename, 'rb', 'iso-8859-1')
self._filehandle = codecs.open(filename, 'rb', ENCODING)
self._lock = Lock()
self._setup_segments()
......@@ -129,6 +131,7 @@ class GeoIP(GeoIPBase):
Supported databases:
* COUNTRY_EDITION
* COUNTRY_EDITION_V6
* REGION_EDITION_REV0
* REGION_EDITION_REV1
* CITY_EDITION_REV0
......@@ -148,9 +151,14 @@ class GeoIP(GeoIPBase):
for i in range(const.STRUCTURE_INFO_MAX_SIZE):
chars = chr(255) * 3
flag = 'unicode_escape'
delim = self._filehandle.read(3)
if (delim == chars) if PY3 else (delim == unicode(chars, flag)):
if PY2:
chars = chars.decode(ENCODING)
if type(delim) is str:
delim = delim.decode(ENCODING)
if delim == chars:
byte = self._filehandle.read(1)
self._databaseType = ord(byte)
......
......@@ -402,3 +402,4 @@ US_OFFSET = 1
CANADA_OFFSET = 677
WORLD_OFFSET = 1353
FIPS_RANGE = 360
ENCODING = 'iso-8859-1'
......@@ -51,6 +51,8 @@ class TestGeoIPCityFunctions(unittest.TestCase):
self.gb_region_data = {'region_name': 'N7', 'country_code': 'GB'}
self.gic = pygeoip.GeoIP(CITY_DB_PATH)
self.gic_mem = pygeoip.GeoIP(CITY_DB_PATH, pygeoip.MEMORY_CACHE)
self.gic_mmap = pygeoip.GeoIP(CITY_DB_PATH, pygeoip.MMAP_CACHE)
def testCountryCodeByName(self):
us_code = self.gic.country_code_by_name(self.us_hostname)
......@@ -94,6 +96,13 @@ class TestGeoIPCityFunctions(unittest.TestCase):
self.assertEqual(us_region, self.us_region_data)
self.assertEqual(gb_region, self.gb_region_data)
def testCacheMethods(self):
mem_record = self.gic_mem.record_by_addr(self.us_ip)
mmap_record = self.gic_mmap.record_by_addr(self.us_ip)
self.assertEqual(mem_record['city'], self.us_record_data['city'])
self.assertEqual(mmap_record['city'], self.us_record_data['city'])
def testRecordByAddr(self):
equal_keys = ('city', 'region_name', 'area_code', 'country_code3',
'postal_code', 'dma_code', 'country_code',
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment