Merge pull request #60 from pganssle/add_available_timezones

Add available timezones

docs/zoneinfo.rst

@@ -341,6 +341,29 @@ pickled in an environment with a different version of the time zone data.
 Functions
 ---------
 
+.. function:: available_timezones()
+
+    Get a set containing all the valid keys for IANA time zones available
+    anywhere on the time zone path. This is recalculated on every call to the
+    function.
+
+    This function only includes canonical zone names and does not include
+    "special" zones such as those under the ``posix/`` and ``right/``
+    directories, or the ``posixrules`` zone.
+
+    .. caution::
+
+        This function may open a large number of files, as the best way to
+        determine if a file on the time zone path is a valid time zone is to
+        read the "magic string" at the beginning.
+
+    .. note::
+
+        These values are not designed to be exposed to end-users; for user
+        facing elements, applications should use something like CLDR (the
+        Unicode Common Locale Data Repository) to get more user-friendly
+        strings. See also the cautionary note on :attr:`ZoneInfo.key`.
+
 .. function:: reset_tzpath(to=None)
 
     Sets or resets the time zone search path (:data:`TZPATH`) for the module.
@@ -355,6 +378,7 @@ Functions
     :exc:`ValueError` will be raised if something other than an absolute path
     is passed.
 
+
 Globals
 -------
 

src/zoneinfo/__init__.py

@@ -1,6 +1,7 @@
 __all__ = [
     "ZoneInfo",
     "reset_tzpath",
+    "available_timezones",
     "TZPATH",
     "ZoneInfoNotFoundError",
     "InvalidTZPathWarning",
@@ -16,6 +17,7 @@
     from ._zoneinfo import ZoneInfo
 
 reset_tzpath = _tzpath.reset_tzpath
+available_timezones = _tzpath.available_timezones
 InvalidTZPathWarning = _tzpath.InvalidTZPathWarning
 
 

src/zoneinfo/__init__.pyi

@@ -24,5 +24,6 @@ class ZoneInfo(tzinfo):
 def reset_tzpath(
     to: Optional[Sequence[Union[os.PathLike, str]]] = None
 ) -> None: ...
+def available_timezones() -> typing.Set[str]: ...
 
 TZPATH: Sequence[str]

src/zoneinfo/_tzpath.py

@@ -108,6 +108,71 @@ def _validate_tzfile_path(path, _base=_TEST_PATH):
 del _TEST_PATH
 
 
+def available_timezones():
+    """Returns a set containing all available time zones.
+
+    .. caution::
+
+        This may attempt to open a large number of files, since the best way to
+        determine if a given file on the time zone search path is to open it
+        and check for the "magic string" at the beginning.
+    """
+    from importlib import resources
+
+    valid_zones = set()
+
+    # Start with loading from the tzdata package if it exists: this has a
+    # pre-assembled list of zones that only requires opening one file.
+    try:
+        with resources.open_text("tzdata", "zones") as f:
+            for zone in f:
+                zone = zone.strip()
+                if zone:
+                    valid_zones.add(zone)
+    except (ImportError, FileNotFoundError):
+        pass
+
+    def valid_key(fpath):
+        try:
+            with open(fpath, "rb") as f:
+                return f.read(4) == b"TZif"
+        except Exception:  # pragma: nocover
+            return False
+
+    for tz_root in TZPATH:
+        if not os.path.exists(tz_root):
+            continue
+
+        for root, dirnames, files in os.walk(tz_root):
+            if root == tz_root:
+                # right/ and posix/ are special directories and shouldn't be
+                # included in the output of available zones
+                if "right" in dirnames:
+                    dirnames.remove("right")
+                if "posix" in dirnames:
+                    dirnames.remove("posix")
+
+            for file in files:
+                fpath = os.path.join(root, file)
+
+                key = os.path.relpath(fpath, start=tz_root)
+                if os.sep != "/":  # pragma: nocover
+                    key = key.replace(os.sep, "/")
+
+                if not key or key in valid_zones:
+                    continue
+
+                if valid_key(fpath):
+                    valid_zones.add(key)
+
+    if "posixrules" in valid_zones:
+        # posixrules is a special symlink-only time zone where it exists, it
+        # should not be included in the output
+        valid_zones.remove("posixrules")
+
+    return valid_zones
+
+
 class InvalidTZPathWarning(RuntimeWarning):
     """Warning raised if an invalid path is specified in PYTHONTZPATH."""
 

tests/_support.py

@@ -71,11 +71,35 @@ def setUpClass(cls):
         super().setUpClass()
 
     @contextlib.contextmanager
-    def tzpath_context(self, tzpath, lock=TZPATH_LOCK):
+    def tzpath_context(self, tzpath, block_tzdata=True, lock=TZPATH_LOCK):
+        def pop_tzdata_modules():
+            tzdata_modules = {}
+            for modname in list(sys.modules):
+                if modname.split(".", 1)[0] != "tzdata":  # pragma: nocover
+                    continue
+
+                tzdata_modules[modname] = sys.modules.pop(modname)
+
+            return tzdata_modules
+
         with lock:
+            if block_tzdata:
+                # In order to fully exclude tzdata from the path, we need to
+                # clear the sys.modules cache of all its contents — setting the
+                # root package to None is not enough to block direct access of
+                # already-imported submodules (though it will prevent new
+                # imports of submodules).
+                tzdata_modules = pop_tzdata_modules()
+                sys.modules["tzdata"] = None
+
             old_path = self.module.TZPATH
             try:
                 self.module.reset_tzpath(tzpath)
                 yield
             finally:
+                if block_tzdata:
+                    sys.modules.pop("tzdata")
+                    for modname, module in tzdata_modules.items():
+                        sys.modules[modname] = module
+
                 self.module.reset_tzpath(old_path)

tests/test_zoneinfo.py

@@ -16,6 +16,7 @@
 import tempfile
 import unittest
 from datetime import date, datetime, time, timedelta, timezone
+from functools import cached_property
 
 from . import _support as test_support
 from ._support import (
@@ -72,10 +73,18 @@ class TzPathUserMixin:
     def tzpath(self):  # pragma: nocover
         return None
 
+    @property
+    def block_tzdata(self):
+        return True
+
     def setUp(self):
         with contextlib.ExitStack() as stack:
             stack.enter_context(
-                self.tzpath_context(self.tzpath, lock=TZPATH_TEST_LOCK)
+                self.tzpath_context(
+                    self.tzpath,
+                    block_tzdata=self.block_tzdata,
+                    lock=TZPATH_TEST_LOCK,
+                )
             )
             self.addCleanup(stack.pop_all().close)
 
@@ -522,6 +531,10 @@ class TZDataTests(ZoneInfoTest):
     def tzpath(self):
         return []
 
+    @property
+    def block_tzdata(self):
+        return False
+
     def zone_from_key(self, key):
         return self.klass(key=key)
 
@@ -1628,6 +1641,32 @@ class CTzPathTest(TzPathTest):
 class TestModule(ZoneInfoTestBase):
     module = py_zoneinfo
 
+    @property
+    def zoneinfo_data(self):
+        return ZONEINFO_DATA
+
+    @cached_property
+    def _UTC_bytes(self):
+        zone_file = self.zoneinfo_data.path_from_key("UTC")
+        with open(zone_file, "rb") as f:
+            return f.read()
+
+    def touch_zone(self, key, tz_root):
+        """Creates a valid TZif file at key under the zoneinfo root tz_root.
+
+        tz_root must exist, but all folders below that will be created.
+        """
+        if not os.path.exists(tz_root):  # pragma: nocover
+            raise FileNotFoundError(f"{tz_root} does not exist.")
+
+        root_dir, *tail = key.rsplit("/", 1)
+        if tail:  # If there's no tail, then the first component isn't a dir
+            os.makedirs(os.path.join(tz_root, root_dir), exist_ok=True)
+
+        zonefile_path = os.path.join(tz_root, key)
+        with open(zonefile_path, "wb") as f:
+            f.write(self._UTC_bytes)
+
     def test_getattr_error(self):
         with self.assertRaises(AttributeError):
             self.module.NOATTRIBUTE
@@ -1648,6 +1687,79 @@ def test_dir_unique(self):
 
         self.assertCountEqual(module_dir, module_unique)
 
+    def test_available_timezones(self):
+        with self.tzpath_context([self.zoneinfo_data.tzpath]):
+            self.assertTrue(self.zoneinfo_data.keys)  # Sanity check
+
+            available_keys = self.module.available_timezones()
+            zoneinfo_keys = set(self.zoneinfo_data.keys)
+
+            # If tzdata is not present, zoneinfo_keys == available_keys,
+            # otherwise it should be a subset.
+            union = zoneinfo_keys & available_keys
+            self.assertEqual(zoneinfo_keys, union)
+
+    def test_available_timezones_weirdzone(self):
+        with tempfile.TemporaryDirectory() as td:
+            # Make a fictional zone at "Mars/Olympus_Mons"
+            self.touch_zone("Mars/Olympus_Mons", td)
+
+            with self.tzpath_context([td]):
+                available_keys = self.module.available_timezones()
+                self.assertIn("Mars/Olympus_Mons", available_keys)
+
+    def test_folder_exclusions(self):
+        expected = {
+            "America/Los_Angeles",
+            "America/Santiago",
+            "America/Indiana/Indianapolis",
+            "UTC",
+            "Europe/Paris",
+            "Europe/London",
+            "Asia/Tokyo",
+            "Australia/Sydney",
+        }
+
+        base_tree = list(expected)
+        posix_tree = [f"posix/{x}" for x in base_tree]
+        right_tree = [f"right/{x}" for x in base_tree]
+
+        cases = [
+            ("base_tree", base_tree),
+            ("base_and_posix", base_tree + posix_tree),
+            ("base_and_right", base_tree + right_tree),
+            ("all_trees", base_tree + right_tree + posix_tree),
+        ]
+
+        with tempfile.TemporaryDirectory() as td:
+            for case_name, tree in cases:
+                tz_root = os.path.join(td, case_name)
+                os.mkdir(tz_root)
+
+                for key in tree:
+                    self.touch_zone(key, tz_root)
+
+                with self.tzpath_context([tz_root]):
+                    with self.subTest(case_name):
+                        actual = self.module.available_timezones()
+                        self.assertEqual(actual, expected)
+
+    def test_exclude_posixrules(self):
+        expected = {
+            "America/New_York",
+            "Europe/London",
+        }
+
+        tree = list(expected) + ["posixrules"]
+
+        with tempfile.TemporaryDirectory() as td:
+            for key in tree:
+                self.touch_zone(key, td)
+
+            with self.tzpath_context([td]):
+                actual = self.module.available_timezones()
+                self.assertEqual(actual, expected)
+
 
 class CTestModule(TestModule):
     module = c_zoneinfo