Add IPv6 support to country_name_by_*

pygeoip/__init__.py

@@ -365,6 +365,18 @@ class GeoIP(GeoIPBase):
 
         return record
 
+    def _gethostbyname(self, hostname):
+        if self._databaseType in const.IPV6_EDITIONS:
+            try:
+                response = socket.getaddrinfo(hostname, 0, socket.AF_INET6)
+            except socket.gaierror:
+                return ''
+            family, socktype, proto, canonname, sockaddr = response[0]
+            address, port, flow, scope = sockaddr
+            return address
+        else:
+            return socket.gethostbyname(hostname)
+
     def id_by_addr(self, addr):
         """
         Get the country index.
@@ -427,24 +439,9 @@ class GeoIP(GeoIPBase):
         @return: 2-letter country code
         @rtype: str
         """
-        if self._databaseType == const.COUNTRY_EDITION_V6:
-            return self._country_code_by_name_v6(hostname)
-        else:
-            return self._country_code_by_name_v4(hostname)
-
-    def _country_code_by_name_v4(self, hostname):
-        addr = socket.gethostbyname(hostname)
+        addr = self._gethostbyname(hostname)
         return self.country_code_by_addr(addr)
 
-    def _country_code_by_name_v6(self, hostname):
-        try:
-            response = socket.getaddrinfo(hostname, 0, socket.AF_INET6)
-        except socket.gaierror:
-            return ''
-        family, socktype, proto, canonname, sockaddr = response[0]
-        address, port, flow, scope = sockaddr
-        return self.country_code_by_addr(address)
-
     def country_name_by_addr(self, addr):
         """
         Returns full country name for specified IP address.
@@ -456,7 +453,8 @@ class GeoIP(GeoIPBase):
         @rtype: str
         """
         try:
-            if self._databaseType == const.COUNTRY_EDITION:
+            COUNTRY_EDITIONS = (const.COUNTRY_EDITION, const.COUNTRY_EDITION_V6)
+            if self._databaseType in COUNTRY_EDITIONS:
                 return const.COUNTRY_NAMES[self.id_by_addr(addr)]
             elif self._databaseType in const.CITY_EDITIONS:
                 return self.record_by_addr(addr)['country_name']
@@ -476,7 +474,7 @@ class GeoIP(GeoIPBase):
         @return: country name
         @rtype: str
         """
-        addr = socket.gethostbyname(hostname)
+        addr = self._gethostbyname(hostname)
         return self.country_name_by_addr(addr)
 
     def org_by_addr(self, addr):

pygeoip/const.py

@@ -369,6 +369,7 @@ NETSPEED_EDITION = 11
 COUNTRY_EDITION_V6 = 12
 
 # Collection of databases
+IPV6_EDITIONS = (COUNTRY_EDITION_V6,)
 CITY_EDITIONS = (CITY_EDITION_REV0, CITY_EDITION_REV1)
 REGION_EDITIONS = (REGION_EDITION_REV0, REGION_EDITION_REV1)
 REGION_CITY_EDITIONS = REGION_EDITIONS + CITY_EDITIONS

tests/test_country.py

@@ -2,7 +2,7 @@
 import unittest
 
 import pygeoip
-from tests.config import COUNTRY_DB_PATH
+from tests.config import COUNTRY_DB_PATH, COUNTRY_V6_DB_PATH
 
 
 class TestGeoIPCountryFunctions(unittest.TestCase):
@@ -13,38 +13,52 @@ class TestGeoIPCountryFunctions(unittest.TestCase):
         self.us_ip = '64.233.161.99'
         self.gb_ip = '212.58.253.68'
 
+        self.ie6_hostname = 'google.com'
+        self.ie6_ip = '2a00:1450:400f:800::1002'
+
+        self.ie_code = 'IE'
         self.us_code = 'US'
         self.gb_code = 'GB'
 
+        self.ie_name = 'Ireland'
         self.us_name = 'United States'
         self.gb_name = 'United Kingdom'
 
         self.gi = pygeoip.GeoIP(COUNTRY_DB_PATH)
+        self.gi6 = pygeoip.GeoIP(COUNTRY_V6_DB_PATH)
 
     def testCountryCodeByName(self):
         us_code = self.gi.country_code_by_name(self.us_hostname)
         gb_code = self.gi.country_code_by_name(self.gb_hostname)
+        ie6_code = self.gi6.country_code_by_name(self.ie6_hostname)
 
         self.assertEqual(us_code, self.us_code)
         self.assertEqual(gb_code, self.gb_code)
+        self.assertEqual(ie6_code, self.ie_code)
 
     def testCountryCodeByAddr(self):
         us_code = self.gi.country_code_by_addr(self.us_ip)
         gb_code = self.gi.country_code_by_addr(self.gb_ip)
+        ie6_code = self.gi6.country_code_by_addr(self.ie6_ip)
         
         self.assertEqual(us_code, self.us_code)
         self.assertEqual(gb_code, self.gb_code)
+        self.assertEqual(ie6_code, self.ie_code)
 
     def testCountryNameByName(self):
         us_name = self.gi.country_name_by_name(self.us_hostname)
         gb_name = self.gi.country_name_by_name(self.gb_hostname)
+        ie6_name = self.gi6.country_name_by_name(self.ie6_hostname)
 
         self.assertEqual(us_name, self.us_name)
         self.assertEqual(gb_name, self.gb_name)
+        self.assertEqual(ie6_name, self.ie_name)
 
     def testCountryNameByAddr(self):
         us_name = self.gi.country_name_by_addr(self.us_ip)
         gb_name = self.gi.country_name_by_addr(self.gb_ip)
+        ie6_name = self.gi6.country_name_by_addr(self.ie6_ip)
 
         self.assertEqual(us_name, self.us_name)
         self.assertEqual(gb_name, self.gb_name)
+        self.assertEqual(ie6_name, self.ie_name)