1"""
2Script to automatically generate a JSON file containing time zone information.
3
4This is done to allow "pinning" a small subset of the tzdata in the tests,
5since we are testing properties of a file that may be subject to change. For
6example, the behavior in the far future of any given zone is likely to change,
7but "does this give the right answer for this file in 2040" is still an
8important property to test.
9
10This must be run from a computer with zoneinfo data installed.
11"""
12from __future__ import annotations
13
14import base64
15import functools
16import json
17import lzma
18import pathlib
19import textwrap
20import typing
21
22import zoneinfo
23
24KEYS = [
25    "Africa/Abidjan",
26    "Africa/Casablanca",
27    "America/Los_Angeles",
28    "America/Santiago",
29    "Asia/Tokyo",
30    "Australia/Sydney",
31    "Europe/Dublin",
32    "Europe/Lisbon",
33    "Europe/London",
34    "Pacific/Kiritimati",
35    "UTC",
36]
37
38TEST_DATA_LOC = pathlib.Path(__file__).parent
39
40
41@functools.lru_cache(maxsize=None)
42def get_zoneinfo_path() -> pathlib.Path:
43    """Get the first zoneinfo directory on TZPATH containing the "UTC" zone."""
44    key = "UTC"
45    for path in map(pathlib.Path, zoneinfo.TZPATH):
46        if (path / key).exists():
47            return path
48    else:
49        raise OSError("Cannot find time zone data.")
50
51
52def get_zoneinfo_metadata() -> typing.Dict[str, str]:
53    path = get_zoneinfo_path()
54
55    tzdata_zi = path / "tzdata.zi"
56    if not tzdata_zi.exists():
57        # tzdata.zi is necessary to get the version information
58        raise OSError("Time zone data does not include tzdata.zi.")
59
60    with open(tzdata_zi, "r") as f:
61        version_line = next(f)
62
63    _, version = version_line.strip().rsplit(" ", 1)
64
65    if (
66        not version[0:4].isdigit()
67        or len(version) < 5
68        or not version[4:].isalpha()
69    ):
70        raise ValueError(
71            "Version string should be YYYYx, "
72            + "where YYYY is the year and x is a letter; "
73            + f"found: {version}"
74        )
75
76    return {"version": version}
77
78
79def get_zoneinfo(key: str) -> bytes:
80    path = get_zoneinfo_path()
81
82    with open(path / key, "rb") as f:
83        return f.read()
84
85
86def encode_compressed(data: bytes) -> typing.List[str]:
87    compressed_zone = lzma.compress(data)
88    raw = base64.b85encode(compressed_zone)
89
90    raw_data_str = raw.decode("utf-8")
91
92    data_str = textwrap.wrap(raw_data_str, width=70)
93    return data_str
94
95
96def load_compressed_keys() -> typing.Dict[str, typing.List[str]]:
97    output = {key: encode_compressed(get_zoneinfo(key)) for key in KEYS}
98
99    return output
100
101
102def update_test_data(fname: str = "zoneinfo_data.json") -> None:
103    TEST_DATA_LOC.mkdir(exist_ok=True, parents=True)
104
105    # Annotation required: https://github.com/python/mypy/issues/8772
106    json_kwargs: typing.Dict[str, typing.Any] = dict(
107        indent=2, sort_keys=True,
108    )
109
110    compressed_keys = load_compressed_keys()
111    metadata = get_zoneinfo_metadata()
112    output = {
113        "metadata": metadata,
114        "data": compressed_keys,
115    }
116
117    with open(TEST_DATA_LOC / fname, "w") as f:
118        json.dump(output, f, **json_kwargs)
119
120
121if __name__ == "__main__":
122    update_test_data()
123