Merge pull request #29 from hongquan/master

Add context manager support

Author: oschwald

geoip2/database.py

@@ -83,6 +83,12 @@ def __init__(self, filename, locales=None, mode=MODE_AUTO):
         self._db_reader = maxminddb.open_database(filename, mode)
         self._locales = locales
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
     def country(self, ip_address):
         """Get the Country object for the IP address
 

tests/database_test.py

@@ -149,6 +149,12 @@ def test_isp(self):
 
         reader.close()
 
+    def test_context_manager(self):
+        with geoip2.database.Reader(
+            'tests/data/test-data/GeoIP2-Country-Test.mmdb') as reader:
+            record = reader.country('81.2.69.160')
+            self.assertEqual(record.traits.ip_address, '81.2.69.160')
+
 
 @unittest.skipUnless(maxminddb.extension, 'No C extension module found. Skipping tests')
 class TestExtensionReader(BaseTestReader, unittest.TestCase):