Adds ability to set a custom zoneinfo directory.

sdispater · sdispater · commit 3ba44d094ecb · 2016-11-22T10:10:58.000-05:00
diff --git a/README.rst b/README.rst
@@ -36,6 +36,17 @@ If you just want to know the path to a specific timezone file, you may use the `
 
     tz_path('Europe/Paris')
 
+By default, ``pytzdata`` will use the bundled timezone database, however you can set
+a custom directory that holds the timezone files using the ``set_directory`` function:
+
+.. code-block:: python
+
+    import pytzdata
+
+    pytzdata.set_directory('/custom/zoneinfo')
+
+You can also set the ``PYTZDATA_TZDATADIR`` environment variable to set a custom directory.
+
 
 Release
 =======
diff --git a/pytzdata/__init__.py b/pytzdata/__init__.py
@@ -6,6 +6,14 @@
 from ._compat import FileNotFoundError
 
 
+DEFAULT_DIRECTORY = os.path.join(
+    os.path.dirname(__file__),
+    'zoneinfo'
+)
+
+_DIRECTORY = os.getenv('PYTZDATA_TZDATADIR', DEFAULT_DIRECTORY)
+
+
 def tz_file(name):
     """
     Open a timezone file from the zoneinfo subdir for reading.
@@ -50,10 +58,18 @@ def tz_path(name):
         if part == os.path.pardir or os.path.sep in part:
             raise ValueError('Bad path segment: %r' % part)
 
-    filepath = os.path.join(os.path.dirname(__file__),
-                            'zoneinfo', *name_parts)
+    filepath = os.path.join(_DIRECTORY, *name_parts)
 
     if not os.path.exists(filepath):
         raise TimezoneNotFound('Timezone {} not found at {}'.format(name, filepath))
 
     return filepath
+
+
+def set_directory(directory=None):
+    global _DIRECTORY
+
+    if directory is None:
+        directory = os.getenv('PYTZDATA_TZDATADIR', DEFAULT_DIRECTORY)
+
+    _DIRECTORY = directory
diff --git a/tests/__init__.py b/tests/__init__.py
@@ -1 +1,11 @@
 # -*- coding: utf-8 -*-
+
+from pytzdata import set_directory
+
+
+def setup_module(module):
+    set_directory()
+
+
+def teardown_module(module):
+    set_directory()
diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py
@@ -0,0 +1,2 @@
+# -*- coding: utf-8 -*-
+
diff --git a/tests/fixtures/tz/Europe/Paris b/tests/fixtures/tz/Europe/Paris
diff --git a/tests/test_set_directory.py b/tests/test_set_directory.py
@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+
+import os
+import pytest
+from pytzdata import set_directory, tz_path, TimezoneNotFound
+
+
+fixtures_path = os.path.join(os.path.dirname(__file__), 'fixtures', 'tz')
+
+
+def setup_module(module):
+    if 'PYTZDATA_TZDATADIR' in os.environ:
+        del os.environ['PYTZDATA_TZDATADIR']
+
+    set_directory()
+
+
+def teardown_module(module):
+    if 'PYTZDATA_TZDATADIR' in os.environ:
+        del os.environ['PYTZDATA_TZDATADIR']
+
+    set_directory()
+
+
+def test_set_directory():
+    set_directory(fixtures_path)
+
+    assert tz_path('Europe/Paris') == os.path.join(fixtures_path, 'Europe/Paris')
+
+    with pytest.raises(TimezoneNotFound):
+        tz_path('America/New_York')
+
+    here = os.path.realpath(os.path.dirname(__file__))
+    filepath = os.path.realpath(
+        os.path.join(here, '..', 'pytzdata', 'zoneinfo', 'America', 'New_York')
+    )
+
+    set_directory()
+
+    assert tz_path('America/New_York') == filepath
+
+
+def test_env_variable():
+    os.environ['PYTZDATA_TZDATADIR'] = fixtures_path
+    set_directory()
+
+    assert tz_path('Europe/Paris') == os.path.join(fixtures_path, 'Europe/Paris')
+
+    with pytest.raises(TimezoneNotFound):
+        tz_path('America/New_York')
+
+    del os.environ['PYTZDATA_TZDATADIR']
+
+    here = os.path.realpath(os.path.dirname(__file__))
+    filepath = os.path.realpath(
+        os.path.join(here, '..', 'pytzdata', 'zoneinfo', 'America', 'New_York')
+    )
+
+    set_directory()
+
+    assert tz_path('America/New_York') == filepath
diff --git a/tests/test_tz_file.py b/tests/test_tz_file.py
@@ -3,10 +3,24 @@
 import os
 import pytest
 
-from pytzdata import tz_file
+from pytzdata import tz_file, set_directory
 from pytzdata.exceptions import TimezoneNotFound
 
 
+def setup_module(module):
+    if 'PYTZDATA_TZDATADIR' in os.environ:
+        del os.environ['PYTZDATA_TZDATADIR']
+
+    set_directory()
+
+
+def teardown_module(module):
+    if 'PYTZDATA_TZDATADIR' in os.environ:
+        del os.environ['PYTZDATA_TZDATADIR']
+
+    set_directory()
+
+
 def test_tz_file():
     here = os.path.realpath(os.path.dirname(__file__))
     filepath = os.path.realpath(
diff --git a/tests/test_tz_path.py b/tests/test_tz_path.py
@@ -3,10 +3,24 @@
 import os
 import pytest
 
-from pytzdata import tz_path
+from pytzdata import tz_path, set_directory
 from pytzdata.exceptions import TimezoneNotFound
 
 
+def setup_module(module):
+    if 'PYTZDATA_TZDATADIR' in os.environ:
+        del os.environ['PYTZDATA_TZDATADIR']
+
+    set_directory()
+
+
+def teardown_module(module):
+    if 'PYTZDATA_TZDATADIR' in os.environ:
+        del os.environ['PYTZDATA_TZDATADIR']
+
+    set_directory()
+
+
 def test_tz_path():
     here = os.path.realpath(os.path.dirname(__file__))
     filepath = os.path.realpath(
