bpo-40503: Add tests and implementation for ZoneInfo

.github/workflows/coverage.yml

@@ -48,6 +48,7 @@ jobs:
         ./python -m venv .venv
         source ./.venv/bin/activate
         python -m pip install -U coverage
+        python -m pip install -r Misc/requirements-test.txt
         python -m test.pythoninfo
     - name: 'Tests with coverage'
       run: >

.travis.yml

@@ -87,6 +87,7 @@ matrix:
         # Need a venv that can parse covered code.
         - ./python -m venv venv
         - ./venv/bin/python -m pip install -U coverage
+        - ./venv/bin/python -m pip install -r Misc/requirements-test.txt
         - ./venv/bin/python -m test.pythoninfo
       script:
         # Skip tests that re-run the entire test suite.

Lib/sysconfig.py

@@ -546,6 +546,7 @@ def get_config_vars(*args):
 
         if os.name == 'nt':
             _init_non_posix(_CONFIG_VARS)
+            _CONFIG_VARS['TZPATH'] = ''
         if os.name == 'posix':
             _init_posix(_CONFIG_VARS)
         # For backward compatibility, see issue19555

Lib/test/test_zoneinfo/__init__.py

@@ -0,0 +1 @@
+from .test_zoneinfo import *

Lib/test/test_zoneinfo/__main__.py

@@ -0,0 +1,3 @@
+import unittest
+
+unittest.main('test.test_zoneinfo')

Lib/test/test_zoneinfo/_support.py

@@ -0,0 +1,76 @@
+import contextlib
+import functools
+import sys
+import threading
+import unittest
+from test.support import import_fresh_module
+
+OS_ENV_LOCK = threading.Lock()
+TZPATH_LOCK = threading.Lock()
+TZPATH_TEST_LOCK = threading.Lock()
+
+
+def call_once(f):
+    """Decorator that ensures a function is only ever called once."""
+    lock = threading.Lock()
+    cached = functools.lru_cache(None)(f)
+
+    @functools.wraps(f)
+    def inner():
+        with lock:
+            return cached()
+
+    return inner
+
+
+@call_once
+def get_modules():
+    """Retrieve two copies of zoneinfo: pure Python and C accelerated.
+
+    Because this function manipulates the import system in a way that might
+    be fragile or do unexpected things if it is run many times, it uses a
+    `call_once` decorator to ensure that this is only ever called exactly
+    one time — in other words, when using this function you will only ever
+    get one copy of each module rather than a fresh import each time.
+    """
+    import zoneinfo as c_module
+
+    py_module = import_fresh_module("zoneinfo", blocked=["_zoneinfo"])
+
+    return py_module, c_module
+
+
+@contextlib.contextmanager
+def set_zoneinfo_module(module):
+    """Make sure sys.modules["zoneinfo"] refers to `module`.
+
+    This is necessary because `pickle` will refuse to serialize
+    an type calling itself `zoneinfo.ZoneInfo` unless `zoneinfo.ZoneInfo`
+    refers to the same object.
+    """
+
+    NOT_PRESENT = object()
+    old_zoneinfo = sys.modules.get("zoneinfo", NOT_PRESENT)
+    sys.modules["zoneinfo"] = module
+    yield
+    if old_zoneinfo is not NOT_PRESENT:
+        sys.modules["zoneinfo"] = old_zoneinfo
+    else:  # pragma: nocover
+        sys.modules.pop("zoneinfo")
+
+
+class ZoneInfoTestBase(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.klass = cls.module.ZoneInfo
+        super().setUpClass()
+
+    @contextlib.contextmanager
+    def tzpath_context(self, tzpath, lock=TZPATH_LOCK):
+        with lock:
+            old_path = self.module.TZPATH
+            try:
+                self.module.reset_tzpath(tzpath)
+                yield
+            finally:
+                self.module.reset_tzpath(old_path)

Lib/test/test_zoneinfo/data/update_test_data.py

@@ -0,0 +1,122 @@
+"""
+Script to automatically generate a JSON file containing time zone information.
+
+This is done to allow "pinning" a small subset of the tzdata in the tests,
+since we are testing properties of a file that may be subject to change. For
+example, the behavior in the far future of any given zone is likely to change,
+but "does this give the right answer for this file in 2040" is still an
+important property to test.
+
+This must be run from a computer with zoneinfo data installed.
+"""
+from __future__ import annotations
+
+import base64
+import functools
+import json
+import lzma
+import pathlib
+import textwrap
+import typing
+
+import zoneinfo
+
+KEYS = [
+    "Africa/Abidjan",
+    "Africa/Casablanca",
+    "America/Los_Angeles",
+    "America/Santiago",
+    "Asia/Tokyo",
+    "Australia/Sydney",
+    "Europe/Dublin",
+    "Europe/Lisbon",
+    "Europe/London",
+    "Pacific/Kiritimati",
+    "UTC",
+]
+
+TEST_DATA_LOC = pathlib.Path(__file__).parent
+
+
+@functools.lru_cache(maxsize=None)
+def get_zoneinfo_path() -> pathlib.Path:
+    """Get the first zoneinfo directory on TZPATH containing the "UTC" zone."""
+    key = "UTC"
+    for path in map(pathlib.Path, zoneinfo.TZPATH):
+        if (path / key).exists():
+            return path
+    else:
+        raise OSError("Cannot find time zone data.")
+
+
+def get_zoneinfo_metadata() -> typing.Dict[str, str]:
+    path = get_zoneinfo_path()
+
+    tzdata_zi = path / "tzdata.zi"
+    if not tzdata_zi.exists():
+        # tzdata.zi is necessary to get the version information
+        raise OSError("Time zone data does not include tzdata.zi.")
+
+    with open(tzdata_zi, "r") as f:
+        version_line = next(f)
+
+    _, version = version_line.strip().rsplit(" ", 1)
+
+    if (
+        not version[0:4].isdigit()
+        or len(version) < 5
+        or not version[4:].isalpha()
+    ):
+        raise ValueError(
+            "Version string should be YYYYx, "
+            + "where YYYY is the year and x is a letter; "
+            + f"found: {version}"
+        )
+
+    return {"version": version}
+
+
+def get_zoneinfo(key: str) -> bytes:
+    path = get_zoneinfo_path()
+
+    with open(path / key, "rb") as f:
+        return f.read()
+
+
+def encode_compressed(data: bytes) -> typing.List[str]:
+    compressed_zone = lzma.compress(data)
+    raw = base64.b85encode(compressed_zone)
+
+    raw_data_str = raw.decode("utf-8")
+
+    data_str = textwrap.wrap(raw_data_str, width=70)
+    return data_str
+
+
+def load_compressed_keys() -> typing.Dict[str, typing.List[str]]:
+    output = {key: encode_compressed(get_zoneinfo(key)) for key in KEYS}
+
+    return output
+
+
+def update_test_data(fname: str = "zoneinfo_data.json") -> None:
+    TEST_DATA_LOC.mkdir(exist_ok=True, parents=True)
+
+    # Annotation required: https://github.com/python/mypy/issues/8772
+    json_kwargs: typing.Dict[str, typing.Any] = dict(
+        indent=2, sort_keys=True,
+    )
+
+    compressed_keys = load_compressed_keys()
+    metadata = get_zoneinfo_metadata()
+    output = {
+        "metadata": metadata,
+        "data": compressed_keys,
+    }
+
+    with open(TEST_DATA_LOC / fname, "w") as f:
+        json.dump(output, f, **json_kwargs)
+
+
+if __name__ == "__main__":
+    update_test_data()