From e527ec8abe0849e784ce100f53c2736986b670ae Mon Sep 17 00:00:00 2001 From: Paul Ganssle Date: Sun, 17 May 2020 21:55:11 -0400 Subject: bpo-40536: Add zoneinfo.available_timezones (GH-20158) This was not specified in the PEP, but it will likely be a frequently requested feature if it's not included. This includes only the "canonical" zones, not a simple listing of every valid value of `key` that can be passed to `Zoneinfo`, because it seems likely that that's what people will want. --- Doc/library/zoneinfo.rst | 23 +++++ Lib/test/test_zoneinfo/_support.py | 26 ++++- Lib/test/test_zoneinfo/test_zoneinfo.py | 114 ++++++++++++++++++++- Lib/zoneinfo/__init__.py | 2 + Lib/zoneinfo/_tzpath.py | 65 ++++++++++++ .../2020-05-17-14-00-12.bpo-40536.FCpoRA.rst | 2 + 6 files changed, 230 insertions(+), 2 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2020-05-17-14-00-12.bpo-40536.FCpoRA.rst diff --git a/Doc/library/zoneinfo.rst b/Doc/library/zoneinfo.rst index d8e2796..1b6f2e7 100644 --- a/Doc/library/zoneinfo.rst +++ b/Doc/library/zoneinfo.rst @@ -337,6 +337,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. diff --git a/Lib/test/test_zoneinfo/_support.py b/Lib/test/test_zoneinfo/_support.py index 6bd8d8d..0fe162c 100644 --- a/Lib/test/test_zoneinfo/_support.py +++ b/Lib/test/test_zoneinfo/_support.py @@ -66,11 +66,35 @@ class ZoneInfoTestBase(unittest.TestCase): 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) diff --git a/Lib/test/test_zoneinfo/test_zoneinfo.py b/Lib/test/test_zoneinfo/test_zoneinfo.py index 05db03a..fe2c380 100644 --- a/Lib/test/test_zoneinfo/test_zoneinfo.py +++ b/Lib/test/test_zoneinfo/test_zoneinfo.py @@ -16,6 +16,7 @@ import struct 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): + 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 @@ class TestModule(ZoneInfoTestBase): 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 diff --git a/Lib/zoneinfo/__init__.py b/Lib/zoneinfo/__init__.py index 81a2d5e..f5510ee 100644 --- a/Lib/zoneinfo/__init__.py +++ b/Lib/zoneinfo/__init__.py @@ -1,6 +1,7 @@ __all__ = [ "ZoneInfo", "reset_tzpath", + "available_timezones", "TZPATH", "ZoneInfoNotFoundError", "InvalidTZPathWarning", @@ -15,6 +16,7 @@ except ImportError: # pragma: nocover from ._zoneinfo import ZoneInfo reset_tzpath = _tzpath.reset_tzpath +available_timezones = _tzpath.available_timezones InvalidTZPathWarning = _tzpath.InvalidTZPathWarning diff --git a/Lib/zoneinfo/_tzpath.py b/Lib/zoneinfo/_tzpath.py index 8cff0b1..c4c671d 100644 --- a/Lib/zoneinfo/_tzpath.py +++ b/Lib/zoneinfo/_tzpath.py @@ -102,6 +102,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.""" diff --git a/Misc/NEWS.d/next/Library/2020-05-17-14-00-12.bpo-40536.FCpoRA.rst b/Misc/NEWS.d/next/Library/2020-05-17-14-00-12.bpo-40536.FCpoRA.rst new file mode 100644 index 0000000..ba7773b --- /dev/null +++ b/Misc/NEWS.d/next/Library/2020-05-17-14-00-12.bpo-40536.FCpoRA.rst @@ -0,0 +1,2 @@ +Added the :func:`~zoneinfo.available_timezones` function to the +:mod:`zoneinfo` module. Patch by Paul Ganssle. -- cgit v0.12