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

pygeoip/__init__.py

@@ -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)
 

pygeoip/const.py

@@ -402,3 +402,4 @@ US_OFFSET = 1
 CANADA_OFFSET = 677
 WORLD_OFFSET = 1353
 FIPS_RANGE = 360
+ENCODING = 'iso-8859-1'

tests/test_city.py

@@ -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',