codekingpro/portable-devtools

1import struct2 3 4def load_tzdata(key):5    from importlib import resources6 7    components = key.split("/")8    package_name = ".".join(["tzdata.zoneinfo"] + components[:-1])9    resource_name = components[-1]10 11    try:12        path = resources.files(package_name).joinpath(resource_name)13        # gh-85702: Prevent PermissionError on Windows14        if path.is_dir():15            raise IsADirectoryError16        return path.open("rb")17    except (ImportError, FileNotFoundError, UnicodeEncodeError, IsADirectoryError):18        # There are four types of exception that can be raised that all amount19        # to "we cannot find this key":20        #21        # ImportError: If package_name doesn't exist (e.g. if tzdata is not22        #   installed, or if there's an error in the folder name like23        #   Amrica/New_York)24        # FileNotFoundError: If resource_name doesn't exist in the package25        #   (e.g. Europe/Krasnoy)26        # UnicodeEncodeError: If package_name or resource_name are not UTF-8,27        #   such as keys containing a surrogate character.28        # IsADirectoryError: If package_name without a resource_name specified.29        raise ZoneInfoNotFoundError(f"No time zone found with key {key}")30 31 32def load_data(fobj):33    header = _TZifHeader.from_file(fobj)34 35    if header.version == 1:36        time_size = 437        time_type = "l"38    else:39        # Version 2+ has 64-bit integer transition times40        time_size = 841        time_type = "q"42 43        # Version 2+ also starts with a Version 1 header and data, which44        # we need to skip now45        skip_bytes = (46            header.timecnt * 5  # Transition times and types47            + header.typecnt * 6  # Local time type records48            + header.charcnt  # Time zone designations49            + header.leapcnt * 8  # Leap second records50            + header.isstdcnt  # Standard/wall indicators51            + header.isutcnt  # UT/local indicators52        )53 54        fobj.seek(skip_bytes, 1)55 56        # Now we need to read the second header, which is not the same57        # as the first58        header = _TZifHeader.from_file(fobj)59 60    typecnt = header.typecnt61    timecnt = header.timecnt62    charcnt = header.charcnt63 64    # The data portion starts with timecnt transitions and indices65    if timecnt:66        trans_list_utc = struct.unpack(67            f">{timecnt}{time_type}", fobj.read(timecnt * time_size)68        )69        trans_idx = struct.unpack(f">{timecnt}B", fobj.read(timecnt))70 71        if max(trans_idx) >= typecnt:72            raise ValueError("Invalid transition index found while reading TZif: "73                            f"{max(trans_idx)}")74    else:75        trans_list_utc = ()76        trans_idx = ()77 78    # Read the ttinfo struct, (utoff, isdst, abbrind)79    if typecnt:80        utcoff, isdst, abbrind = zip(81            *(struct.unpack(">lbb", fobj.read(6)) for i in range(typecnt))82        )83    else:84        utcoff = ()85        isdst = ()86        abbrind = ()87 88    # Now read the abbreviations. They are null-terminated strings, indexed89    # not by position in the array but by position in the unsplit90    # abbreviation string. I suppose this makes more sense in C, which uses91    # null to terminate the strings, but it's inconvenient here...92    abbr_vals = {}93    abbr_chars = fobj.read(charcnt)94 95    def get_abbr(idx):96        # Gets a string starting at idx and running until the next \x0097        #98        # We cannot pre-populate abbr_vals by splitting on \x00 because there99        # are some zones that use subsets of longer abbreviations, like so:100        #101        #  LMT\x00AHST\x00HDT\x00102        #103        # Where the idx to abbr mapping should be:104        #105        # {0: "LMT", 4: "AHST", 5: "HST", 9: "HDT"}106        if idx not in abbr_vals:107            span_end = abbr_chars.find(b"\x00", idx)108            abbr_vals[idx] = abbr_chars[idx:span_end].decode()109 110        return abbr_vals[idx]111 112    abbr = tuple(get_abbr(idx) for idx in abbrind)113 114    # The remainder of the file consists of leap seconds (currently unused) and115    # the standard/wall and ut/local indicators, which are metadata we don't need.116    # In version 2 files, we need to skip the unnecessary data to get at the TZ string:117    if header.version >= 2:118        # Each leap second record has size (time_size + 4)119        skip_bytes = header.isutcnt + header.isstdcnt + header.leapcnt * 12120        fobj.seek(skip_bytes, 1)121 122        c = fobj.read(1)  # Should be \n123        assert c == b"\n", c124 125        line = fobj.readline()126        if not line.endswith(b"\n"):127            raise ValueError("Invalid TZif file: unexpected end of file")128        tz_str = line.rstrip(b"\n")129    else:130        tz_str = None131 132    return trans_idx, trans_list_utc, utcoff, isdst, abbr, tz_str133 134 135class _TZifHeader:136    __slots__ = [137        "version",138        "isutcnt",139        "isstdcnt",140        "leapcnt",141        "timecnt",142        "typecnt",143        "charcnt",144    ]145 146    def __init__(self, *args):147        for attr, val in zip(self.__slots__, args, strict=True):148            setattr(self, attr, val)149 150    @classmethod151    def from_file(cls, stream):152        # The header starts with a 4-byte "magic" value153        if stream.read(4) != b"TZif":154            raise ValueError("Invalid TZif file: magic not found")155 156        _version = stream.read(1)157        if _version == b"\x00":158            version = 1159        else:160            version = int(_version)161        stream.read(15)162 163        args = (version,)164 165        # Slots are defined in the order that the bytes are arranged166        args = args + struct.unpack(">6l", stream.read(24))167 168        return cls(*args)169 170 171class ZoneInfoNotFoundError(KeyError):172    """Exception raised when a ZoneInfo key is not found."""173