_common.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164
  1. import struct
  2. def load_tzdata(key):
  3. from importlib import resources
  4. components = key.split("/")
  5. package_name = ".".join(["tzdata.zoneinfo"] + components[:-1])
  6. resource_name = components[-1]
  7. try:
  8. return resources.files(package_name).joinpath(resource_name).open("rb")
  9. except (ImportError, FileNotFoundError, UnicodeEncodeError):
  10. # There are three types of exception that can be raised that all amount
  11. # to "we cannot find this key":
  12. #
  13. # ImportError: If package_name doesn't exist (e.g. if tzdata is not
  14. # installed, or if there's an error in the folder name like
  15. # Amrica/New_York)
  16. # FileNotFoundError: If resource_name doesn't exist in the package
  17. # (e.g. Europe/Krasnoy)
  18. # UnicodeEncodeError: If package_name or resource_name are not UTF-8,
  19. # such as keys containing a surrogate character.
  20. raise ZoneInfoNotFoundError(f"No time zone found with key {key}")
  21. def load_data(fobj):
  22. header = _TZifHeader.from_file(fobj)
  23. if header.version == 1:
  24. time_size = 4
  25. time_type = "l"
  26. else:
  27. # Version 2+ has 64-bit integer transition times
  28. time_size = 8
  29. time_type = "q"
  30. # Version 2+ also starts with a Version 1 header and data, which
  31. # we need to skip now
  32. skip_bytes = (
  33. header.timecnt * 5 # Transition times and types
  34. + header.typecnt * 6 # Local time type records
  35. + header.charcnt # Time zone designations
  36. + header.leapcnt * 8 # Leap second records
  37. + header.isstdcnt # Standard/wall indicators
  38. + header.isutcnt # UT/local indicators
  39. )
  40. fobj.seek(skip_bytes, 1)
  41. # Now we need to read the second header, which is not the same
  42. # as the first
  43. header = _TZifHeader.from_file(fobj)
  44. typecnt = header.typecnt
  45. timecnt = header.timecnt
  46. charcnt = header.charcnt
  47. # The data portion starts with timecnt transitions and indices
  48. if timecnt:
  49. trans_list_utc = struct.unpack(
  50. f">{timecnt}{time_type}", fobj.read(timecnt * time_size)
  51. )
  52. trans_idx = struct.unpack(f">{timecnt}B", fobj.read(timecnt))
  53. else:
  54. trans_list_utc = ()
  55. trans_idx = ()
  56. # Read the ttinfo struct, (utoff, isdst, abbrind)
  57. if typecnt:
  58. utcoff, isdst, abbrind = zip(
  59. *(struct.unpack(">lbb", fobj.read(6)) for i in range(typecnt))
  60. )
  61. else:
  62. utcoff = ()
  63. isdst = ()
  64. abbrind = ()
  65. # Now read the abbreviations. They are null-terminated strings, indexed
  66. # not by position in the array but by position in the unsplit
  67. # abbreviation string. I suppose this makes more sense in C, which uses
  68. # null to terminate the strings, but it's inconvenient here...
  69. abbr_vals = {}
  70. abbr_chars = fobj.read(charcnt)
  71. def get_abbr(idx):
  72. # Gets a string starting at idx and running until the next \x00
  73. #
  74. # We cannot pre-populate abbr_vals by splitting on \x00 because there
  75. # are some zones that use subsets of longer abbreviations, like so:
  76. #
  77. # LMT\x00AHST\x00HDT\x00
  78. #
  79. # Where the idx to abbr mapping should be:
  80. #
  81. # {0: "LMT", 4: "AHST", 5: "HST", 9: "HDT"}
  82. if idx not in abbr_vals:
  83. span_end = abbr_chars.find(b"\x00", idx)
  84. abbr_vals[idx] = abbr_chars[idx:span_end].decode()
  85. return abbr_vals[idx]
  86. abbr = tuple(get_abbr(idx) for idx in abbrind)
  87. # The remainder of the file consists of leap seconds (currently unused) and
  88. # the standard/wall and ut/local indicators, which are metadata we don't need.
  89. # In version 2 files, we need to skip the unnecessary data to get at the TZ string:
  90. if header.version >= 2:
  91. # Each leap second record has size (time_size + 4)
  92. skip_bytes = header.isutcnt + header.isstdcnt + header.leapcnt * 12
  93. fobj.seek(skip_bytes, 1)
  94. c = fobj.read(1) # Should be \n
  95. assert c == b"\n", c
  96. tz_bytes = b""
  97. while (c := fobj.read(1)) != b"\n":
  98. tz_bytes += c
  99. tz_str = tz_bytes
  100. else:
  101. tz_str = None
  102. return trans_idx, trans_list_utc, utcoff, isdst, abbr, tz_str
  103. class _TZifHeader:
  104. __slots__ = [
  105. "version",
  106. "isutcnt",
  107. "isstdcnt",
  108. "leapcnt",
  109. "timecnt",
  110. "typecnt",
  111. "charcnt",
  112. ]
  113. def __init__(self, *args):
  114. for attr, val in zip(self.__slots__, args, strict=True):
  115. setattr(self, attr, val)
  116. @classmethod
  117. def from_file(cls, stream):
  118. # The header starts with a 4-byte "magic" value
  119. if stream.read(4) != b"TZif":
  120. raise ValueError("Invalid TZif file: magic not found")
  121. _version = stream.read(1)
  122. if _version == b"\x00":
  123. version = 1
  124. else:
  125. version = int(_version)
  126. stream.read(15)
  127. args = (version,)
  128. # Slots are defined in the order that the bytes are arranged
  129. args = args + struct.unpack(">6l", stream.read(24))
  130. return cls(*args)
  131. class ZoneInfoNotFoundError(KeyError):
  132. """Exception raised when a ZoneInfo key is not found."""