codekingpro/portable-devtools

1import bisect2import calendar3import collections4import functools5import re6import weakref7from datetime import datetime, timedelta, tzinfo8 9from . import _common, _tzpath10 11EPOCH = datetime(1970, 1, 1)12EPOCHORDINAL = datetime(1970, 1, 1).toordinal()13 14# It is relatively expensive to construct new timedelta objects, and in most15# cases we're looking at the same deltas, like integer numbers of hours, etc.16# To improve speed and memory use, we'll keep a dictionary with references17# to the ones we've already used so far.18#19# Loading every time zone in the 2020a version of the time zone database20# requires 447 timedeltas, which requires approximately the amount of space21# that ZoneInfo("America/New_York") with 236 transitions takes up, so we will22# set the cache size to 512 so that in the common case we always get cache23# hits, but specifically crafted ZoneInfo objects don't leak arbitrary amounts24# of memory.25@functools.lru_cache(maxsize=512)26def _load_timedelta(seconds):27    return timedelta(seconds=seconds)28 29 30class ZoneInfo(tzinfo):31    _strong_cache_size = 832    _strong_cache = collections.OrderedDict()33    _weak_cache = weakref.WeakValueDictionary()34    __module__ = "zoneinfo"35 36    def __init_subclass__(cls):37        cls._strong_cache = collections.OrderedDict()38        cls._weak_cache = weakref.WeakValueDictionary()39 40    def __new__(cls, key):41        instance = cls._weak_cache.get(key, None)42        if instance is None:43            instance = cls._weak_cache.setdefault(key, cls._new_instance(key))44            instance._from_cache = True45 46        # Update the "strong" cache47        cls._strong_cache[key] = cls._strong_cache.pop(key, instance)48 49        if len(cls._strong_cache) > cls._strong_cache_size:50            try:51                cls._strong_cache.popitem(last=False)52            except KeyError:53                # another thread may have already emptied the cache54                pass55 56        return instance57 58    @classmethod59    def no_cache(cls, key):60        obj = cls._new_instance(key)61        obj._from_cache = False62 63        return obj64 65    @classmethod66    def _new_instance(cls, key):67        obj = super().__new__(cls)68        obj._key = key69        obj._file_path = obj._find_tzfile(key)70 71        if obj._file_path is not None:72            file_obj = open(obj._file_path, "rb")73        else:74            file_obj = _common.load_tzdata(key)75 76        with file_obj as f:77            obj._load_file(f)78 79        return obj80 81    @classmethod82    def from_file(cls, file_obj, /, key=None):83        obj = super().__new__(cls)84        obj._key = key85        obj._file_path = None86        obj._load_file(file_obj)87        obj._file_repr = repr(file_obj)88 89        # Disable pickling for objects created from files90        obj.__reduce__ = obj._file_reduce91 92        return obj93 94    @classmethod95    def clear_cache(cls, *, only_keys=None):96        if only_keys is not None:97            for key in only_keys:98                cls._weak_cache.pop(key, None)99                cls._strong_cache.pop(key, None)100 101        else:102            cls._weak_cache.clear()103            cls._strong_cache.clear()104 105    @property106    def key(self):107        return self._key108 109    def utcoffset(self, dt):110        return self._find_trans(dt).utcoff111 112    def dst(self, dt):113        return self._find_trans(dt).dstoff114 115    def tzname(self, dt):116        return self._find_trans(dt).tzname117 118    def fromutc(self, dt):119        """Convert from datetime in UTC to datetime in local time"""120 121        if not isinstance(dt, datetime):122            raise TypeError("fromutc() requires a datetime argument")123        if dt.tzinfo is not self:124            raise ValueError("dt.tzinfo is not self")125 126        timestamp = self._get_local_timestamp(dt)127        num_trans = len(self._trans_utc)128 129        if num_trans >= 1 and timestamp < self._trans_utc[0]:130            tti = self._tti_before131            fold = 0132        elif (133            num_trans == 0 or timestamp > self._trans_utc[-1]134        ) and not isinstance(self._tz_after, _ttinfo):135            tti, fold = self._tz_after.get_trans_info_fromutc(136                timestamp, dt.year137            )138        elif num_trans == 0:139            tti = self._tz_after140            fold = 0141        else:142            idx = bisect.bisect_right(self._trans_utc, timestamp)143 144            if num_trans > 1 and timestamp >= self._trans_utc[1]:145                tti_prev, tti = self._ttinfos[idx - 2 : idx]146            elif timestamp > self._trans_utc[-1]:147                tti_prev = self._ttinfos[-1]148                tti = self._tz_after149            else:150                tti_prev = self._tti_before151                tti = self._ttinfos[0]152 153            # Detect fold154            shift = tti_prev.utcoff - tti.utcoff155            fold = shift.total_seconds() > timestamp - self._trans_utc[idx - 1]156        dt += tti.utcoff157        if fold:158            return dt.replace(fold=1)159        else:160            return dt161 162    def _find_trans(self, dt):163        if dt is None:164            if self._fixed_offset:165                return self._tz_after166            else:167                return _NO_TTINFO168 169        ts = self._get_local_timestamp(dt)170 171        lt = self._trans_local[dt.fold]172 173        num_trans = len(lt)174 175        if num_trans and ts < lt[0]:176            return self._tti_before177        elif not num_trans or ts > lt[-1]:178            if isinstance(self._tz_after, _TZStr):179                return self._tz_after.get_trans_info(ts, dt.year, dt.fold)180            else:181                return self._tz_after182        else:183            # idx is the transition that occurs after this timestamp, so we184            # subtract off 1 to get the current ttinfo185            idx = bisect.bisect_right(lt, ts) - 1186            assert idx >= 0187            return self._ttinfos[idx]188 189    def _get_local_timestamp(self, dt):190        return (191            (dt.toordinal() - EPOCHORDINAL) * 86400192            + dt.hour * 3600193            + dt.minute * 60194            + dt.second195        )196 197    def __str__(self):198        if self._key is not None:199            return f"{self._key}"200        else:201            return repr(self)202 203    def __repr__(self):204        if self._key is not None:205            return f"{self.__class__.__name__}(key={self._key!r})"206        else:207            return f"{self.__class__.__name__}.from_file({self._file_repr})"208 209    def __reduce__(self):210        return (self.__class__._unpickle, (self._key, self._from_cache))211 212    def _file_reduce(self):213        import pickle214 215        raise pickle.PicklingError(216            "Cannot pickle a ZoneInfo file created from a file stream."217        )218 219    @classmethod220    def _unpickle(cls, key, from_cache, /):221        if from_cache:222            return cls(key)223        else:224            return cls.no_cache(key)225 226    def _find_tzfile(self, key):227        return _tzpath.find_tzfile(key)228 229    def _load_file(self, fobj):230        # Retrieve all the data as it exists in the zoneinfo file231        trans_idx, trans_utc, utcoff, isdst, abbr, tz_str = _common.load_data(232            fobj233        )234 235        # Infer the DST offsets (needed for .dst()) from the data236        dstoff = self._utcoff_to_dstoff(trans_idx, utcoff, isdst)237 238        # Convert all the transition times (UTC) into "seconds since 1970-01-01 local time"239        trans_local = self._ts_to_local(trans_idx, trans_utc, utcoff)240 241        # Construct `_ttinfo` objects for each transition in the file242        _ttinfo_list = [243            _ttinfo(244                _load_timedelta(utcoffset), _load_timedelta(dstoffset), tzname245            )246            for utcoffset, dstoffset, tzname in zip(utcoff, dstoff, abbr)247        ]248 249        self._trans_utc = trans_utc250        self._trans_local = trans_local251        self._ttinfos = [_ttinfo_list[idx] for idx in trans_idx]252 253        # Find the first non-DST transition254        for i in range(len(isdst)):255            if not isdst[i]:256                self._tti_before = _ttinfo_list[i]257                break258        else:259            if self._ttinfos:260                self._tti_before = self._ttinfos[0]261            else:262                self._tti_before = None263 264        # Set the "fallback" time zone265        if tz_str is not None and tz_str != b"":266            self._tz_after = _parse_tz_str(tz_str.decode())267        else:268            if not self._ttinfos and not _ttinfo_list:269                raise ValueError("No time zone information found.")270 271            if self._ttinfos:272                self._tz_after = self._ttinfos[-1]273            else:274                self._tz_after = _ttinfo_list[-1]275 276        # Determine if this is a "fixed offset" zone, meaning that the output277        # of the utcoffset, dst and tzname functions does not depend on the278        # specific datetime passed.279        #280        # We make three simplifying assumptions here:281        #282        # 1. If _tz_after is not a _ttinfo, it has transitions that might283        #    actually occur (it is possible to construct TZ strings that284        #    specify STD and DST but no transitions ever occur, such as285        #    AAA0BBB,0/0,J365/25).286        # 2. If _ttinfo_list contains more than one _ttinfo object, the objects287        #    represent different offsets.288        # 3. _ttinfo_list contains no unused _ttinfos (in which case an289        #    otherwise fixed-offset zone with extra _ttinfos defined may290        #    appear to *not* be a fixed offset zone).291        #292        # Violations to these assumptions would be fairly exotic, and exotic293        # zones should almost certainly not be used with datetime.time (the294        # only thing that would be affected by this).295        if len(_ttinfo_list) > 1 or not isinstance(self._tz_after, _ttinfo):296            self._fixed_offset = False297        elif not _ttinfo_list:298            self._fixed_offset = True299        else:300            self._fixed_offset = _ttinfo_list[0] == self._tz_after301 302    @staticmethod303    def _utcoff_to_dstoff(trans_idx, utcoffsets, isdsts):304        # Now we must transform our ttis and abbrs into `_ttinfo` objects,305        # but there is an issue: .dst() must return a timedelta with the306        # difference between utcoffset() and the "standard" offset, but307        # the "base offset" and "DST offset" are not encoded in the file;308        # we can infer what they are from the isdst flag, but it is not309        # sufficient to just look at the last standard offset, because310        # occasionally countries will shift both DST offset and base offset.311 312        typecnt = len(isdsts)313        dstoffs = [0] * typecnt  # Provisionally assign all to 0.314        dst_cnt = sum(isdsts)315        dst_found = 0316 317        for i in range(1, len(trans_idx)):318            if dst_cnt == dst_found:319                break320 321            idx = trans_idx[i]322 323            dst = isdsts[idx]324 325            # We're only going to look at daylight saving time326            if not dst:327                continue328 329            # Skip any offsets that have already been assigned330            if dstoffs[idx] != 0:331                continue332 333            dstoff = 0334            utcoff = utcoffsets[idx]335 336            comp_idx = trans_idx[i - 1]337 338            if not isdsts[comp_idx]:339                dstoff = utcoff - utcoffsets[comp_idx]340 341            if not dstoff and idx < (typecnt - 1) and i + 1 < len(trans_idx):342                comp_idx = trans_idx[i + 1]343 344                # If the following transition is also DST and we couldn't345                # find the DST offset by this point, we're going to have to346                # skip it and hope this transition gets assigned later347                if isdsts[comp_idx]:348                    continue349 350                dstoff = utcoff - utcoffsets[comp_idx]351 352            if dstoff:353                dst_found += 1354                dstoffs[idx] = dstoff355        else:356            # If we didn't find a valid value for a given index, we'll end up357            # with dstoff = 0 for something where `isdst=1`. This is obviously358            # wrong - one hour will be a much better guess than 0359            for idx in range(typecnt):360                if not dstoffs[idx] and isdsts[idx]:361                    dstoffs[idx] = 3600362 363        return dstoffs364 365    @staticmethod366    def _ts_to_local(trans_idx, trans_list_utc, utcoffsets):367        """Generate number of seconds since 1970 *in the local time*.368 369        This is necessary to easily find the transition times in local time"""370        if not trans_list_utc:371            return [[], []]372 373        # Start with the timestamps and modify in-place374        trans_list_wall = [list(trans_list_utc), list(trans_list_utc)]375 376        if len(utcoffsets) > 1:377            offset_0 = utcoffsets[0]378            offset_1 = utcoffsets[trans_idx[0]]379            if offset_1 > offset_0:380                offset_1, offset_0 = offset_0, offset_1381        else:382            offset_0 = offset_1 = utcoffsets[0]383 384        trans_list_wall[0][0] += offset_0385        trans_list_wall[1][0] += offset_1386 387        for i in range(1, len(trans_idx)):388            offset_0 = utcoffsets[trans_idx[i - 1]]389            offset_1 = utcoffsets[trans_idx[i]]390 391            if offset_1 > offset_0:392                offset_1, offset_0 = offset_0, offset_1393 394            trans_list_wall[0][i] += offset_0395            trans_list_wall[1][i] += offset_1396 397        return trans_list_wall398 399 400class _ttinfo:401    __slots__ = ["utcoff", "dstoff", "tzname"]402 403    def __init__(self, utcoff, dstoff, tzname):404        self.utcoff = utcoff405        self.dstoff = dstoff406        self.tzname = tzname407 408    def __eq__(self, other):409        return (410            self.utcoff == other.utcoff411            and self.dstoff == other.dstoff412            and self.tzname == other.tzname413        )414 415    def __repr__(self):  # pragma: nocover416        return (417            f"{self.__class__.__name__}"418            + f"({self.utcoff}, {self.dstoff}, {self.tzname})"419        )420 421 422_NO_TTINFO = _ttinfo(None, None, None)423 424 425class _TZStr:426    __slots__ = (427        "std",428        "dst",429        "start",430        "end",431        "get_trans_info",432        "get_trans_info_fromutc",433        "dst_diff",434    )435 436    def __init__(437        self, std_abbr, std_offset, dst_abbr, dst_offset, start=None, end=None438    ):439        self.dst_diff = dst_offset - std_offset440        std_offset = _load_timedelta(std_offset)441        self.std = _ttinfo(442            utcoff=std_offset, dstoff=_load_timedelta(0), tzname=std_abbr443        )444 445        self.start = start446        self.end = end447 448        dst_offset = _load_timedelta(dst_offset)449        delta = _load_timedelta(self.dst_diff)450        self.dst = _ttinfo(utcoff=dst_offset, dstoff=delta, tzname=dst_abbr)451 452        # These are assertions because the constructor should only be called453        # by functions that would fail before passing start or end454        assert start is not None, "No transition start specified"455        assert end is not None, "No transition end specified"456 457        self.get_trans_info = self._get_trans_info458        self.get_trans_info_fromutc = self._get_trans_info_fromutc459 460    def transitions(self, year):461        start = self.start.year_to_epoch(year)462        end = self.end.year_to_epoch(year)463        return start, end464 465    def _get_trans_info(self, ts, year, fold):466        """Get the information about the current transition - tti"""467        start, end = self.transitions(year)468 469        # With fold = 0, the period (denominated in local time) with the470        # smaller offset starts at the end of the gap and ends at the end of471        # the fold; with fold = 1, it runs from the start of the gap to the472        # beginning of the fold.473        #474        # So in order to determine the DST boundaries we need to know both475        # the fold and whether DST is positive or negative (rare), and it476        # turns out that this boils down to fold XOR is_positive.477        if fold == (self.dst_diff >= 0):478            end -= self.dst_diff479        else:480            start += self.dst_diff481 482        if start < end:483            isdst = start <= ts < end484        else:485            isdst = not (end <= ts < start)486 487        return self.dst if isdst else self.std488 489    def _get_trans_info_fromutc(self, ts, year):490        start, end = self.transitions(year)491        start -= self.std.utcoff.total_seconds()492        end -= self.dst.utcoff.total_seconds()493 494        if start < end:495            isdst = start <= ts < end496        else:497            isdst = not (end <= ts < start)498 499        # For positive DST, the ambiguous period is one dst_diff after the end500        # of DST; for negative DST, the ambiguous period is one dst_diff before501        # the start of DST.502        if self.dst_diff > 0:503            ambig_start = end504            ambig_end = end + self.dst_diff505        else:506            ambig_start = start507            ambig_end = start - self.dst_diff508 509        fold = ambig_start <= ts < ambig_end510 511        return (self.dst if isdst else self.std, fold)512 513 514def _post_epoch_days_before_year(year):515    """Get the number of days between 1970-01-01 and YEAR-01-01"""516    y = year - 1517    return y * 365 + y // 4 - y // 100 + y // 400 - EPOCHORDINAL518 519 520class _DayOffset:521    __slots__ = ["d", "julian", "hour", "minute", "second"]522 523    def __init__(self, d, julian, hour=2, minute=0, second=0):524        min_day = 0 + julian  # convert bool to int525        if not min_day <= d <= 365:526            raise ValueError(f"d must be in [{min_day}, 365], not: {d}")527 528        self.d = d529        self.julian = julian530        self.hour = hour531        self.minute = minute532        self.second = second533 534    def year_to_epoch(self, year):535        days_before_year = _post_epoch_days_before_year(year)536 537        d = self.d538        if self.julian and d >= 59 and calendar.isleap(year):539            d += 1540 541        epoch = (days_before_year + d) * 86400542        epoch += self.hour * 3600 + self.minute * 60 + self.second543 544        return epoch545 546 547class _CalendarOffset:548    __slots__ = ["m", "w", "d", "hour", "minute", "second"]549 550    _DAYS_BEFORE_MONTH = (551        -1,552        0,553        31,554        59,555        90,556        120,557        151,558        181,559        212,560        243,561        273,562        304,563        334,564    )565 566    def __init__(self, m, w, d, hour=2, minute=0, second=0):567        if not 1 <= m <= 12:568            raise ValueError("m must be in [1, 12]")569 570        if not 1 <= w <= 5:571            raise ValueError("w must be in [1, 5]")572 573        if not 0 <= d <= 6:574            raise ValueError("d must be in [0, 6]")575 576        self.m = m577        self.w = w578        self.d = d579        self.hour = hour580        self.minute = minute581        self.second = second582 583    @classmethod584    def _ymd2ord(cls, year, month, day):585        return (586            _post_epoch_days_before_year(year)587            + cls._DAYS_BEFORE_MONTH[month]588            + (month > 2 and calendar.isleap(year))589            + day590        )591 592    # TODO: These are not actually epoch dates as they are expressed in local time593    def year_to_epoch(self, year):594        """Calculates the datetime of the occurrence from the year"""595        # We know year and month, we need to convert w, d into day of month596        #597        # Week 1 is the first week in which day `d` (where 0 = Sunday) appears.598        # Week 5 represents the last occurrence of day `d`, so we need to know599        # the range of the month.600        first_day, days_in_month = calendar.monthrange(year, self.m)601 602        # This equation seems magical, so I'll break it down:603        # 1. calendar says 0 = Monday, POSIX says 0 = Sunday604        #    so we need first_day + 1 to get 1 = Monday -> 7 = Sunday,605        #    which is still equivalent because this math is mod 7606        # 2. Get first day - desired day mod 7: -1 % 7 = 6, so we don't need607        #    to do anything to adjust negative numbers.608        # 3. Add 1 because month days are a 1-based index.609        month_day = (self.d - (first_day + 1)) % 7 + 1610 611        # Now use a 0-based index version of `w` to calculate the w-th612        # occurrence of `d`613        month_day += (self.w - 1) * 7614 615        # month_day will only be > days_in_month if w was 5, and `w` means616        # "last occurrence of `d`", so now we just check if we over-shot the617        # end of the month and if so knock off 1 week.618        if month_day > days_in_month:619            month_day -= 7620 621        ordinal = self._ymd2ord(year, self.m, month_day)622        epoch = ordinal * 86400623        epoch += self.hour * 3600 + self.minute * 60 + self.second624        return epoch625 626 627def _parse_tz_str(tz_str):628    # The tz string has the format:629    #630    # std[offset[dst[offset],start[/time],end[/time]]]631    #632    # std and dst must be 3 or more characters long and must not contain633    # a leading colon, embedded digits, commas, nor a plus or minus signs;634    # The spaces between "std" and "offset" are only for display and are635    # not actually present in the string.636    #637    # The format of the offset is ``[+|-]hh[:mm[:ss]]``638 639    offset_str, *start_end_str = tz_str.split(",", 1)640 641    parser_re = re.compile(642        r"""643        (?P<std>[^<0-9:.+-]+|<[a-zA-Z0-9+-]+>)644        (?:645            (?P<stdoff>[+-]?\d{1,3}(?::\d{2}(?::\d{2})?)?)646            (?:647                (?P<dst>[^0-9:.+-]+|<[a-zA-Z0-9+-]+>)648                (?P<dstoff>[+-]?\d{1,3}(?::\d{2}(?::\d{2})?)?)?649            )? # dst650        )? # stdoff651        """,652        re.ASCII|re.VERBOSE653    )654 655    m = parser_re.fullmatch(offset_str)656 657    if m is None:658        raise ValueError(f"{tz_str} is not a valid TZ string")659 660    std_abbr = m.group("std")661    dst_abbr = m.group("dst")662    dst_offset = None663 664    std_abbr = std_abbr.strip("<>")665 666    if dst_abbr:667        dst_abbr = dst_abbr.strip("<>")668 669    if std_offset := m.group("stdoff"):670        try:671            std_offset = _parse_tz_delta(std_offset)672        except ValueError as e:673            raise ValueError(f"Invalid STD offset in {tz_str}") from e674    else:675        std_offset = 0676 677    if dst_abbr is not None:678        if dst_offset := m.group("dstoff"):679            try:680                dst_offset = _parse_tz_delta(dst_offset)681            except ValueError as e:682                raise ValueError(f"Invalid DST offset in {tz_str}") from e683        else:684            dst_offset = std_offset + 3600685 686        if not start_end_str:687            raise ValueError(f"Missing transition rules: {tz_str}")688 689        start_end_strs = start_end_str[0].split(",", 1)690        try:691            start, end = (_parse_dst_start_end(x) for x in start_end_strs)692        except ValueError as e:693            raise ValueError(f"Invalid TZ string: {tz_str}") from e694 695        return _TZStr(std_abbr, std_offset, dst_abbr, dst_offset, start, end)696    elif start_end_str:697        raise ValueError(f"Transition rule present without DST: {tz_str}")698    else:699        # This is a static ttinfo, don't return _TZStr700        return _ttinfo(701            _load_timedelta(std_offset), _load_timedelta(0), std_abbr702        )703 704 705def _parse_dst_start_end(dststr):706    date, *time = dststr.split("/", 1)707    type = date[:1]708    if type == "M":709        n_is_julian = False710        m = re.fullmatch(r"M(\d{1,2})\.(\d).(\d)", date, re.ASCII)711        if m is None:712            raise ValueError(f"Invalid dst start/end date: {dststr}")713        date_offset = tuple(map(int, m.groups()))714        offset = _CalendarOffset(*date_offset)715    else:716        if type == "J":717            n_is_julian = True718            date = date[1:]719        else:720            n_is_julian = False721 722        doy = int(date)723        offset = _DayOffset(doy, n_is_julian)724 725    if time:726        offset.hour, offset.minute, offset.second = _parse_transition_time(time[0])727 728    return offset729 730 731def _parse_transition_time(time_str):732    match = re.fullmatch(733        r"(?P<sign>[+-])?(?P<h>\d{1,3})(:(?P<m>\d{2})(:(?P<s>\d{2}))?)?",734        time_str,735        re.ASCII736    )737    if match is None:738        raise ValueError(f"Invalid time: {time_str}")739 740    h, m, s = (int(v or 0) for v in match.group("h", "m", "s"))741 742    if h > 167:743        raise ValueError(744            f"Hour must be in [0, 167]: {time_str}"745        )746 747    if match.group("sign") == "-":748        h, m, s = -h, -m, -s749 750    return h, m, s751 752 753def _parse_tz_delta(tz_delta):754    match = re.fullmatch(755        r"(?P<sign>[+-])?(?P<h>\d{1,3})(:(?P<m>\d{2})(:(?P<s>\d{2}))?)?",756        tz_delta,757        re.ASCII758    )759    # Anything passed to this function should already have hit an equivalent760    # regular expression to find the section to parse.761    assert match is not None, tz_delta762 763    h, m, s = (int(v or 0) for v in match.group("h", "m", "s"))764 765    total = h * 3600 + m * 60 + s766 767    if h > 24:768        raise ValueError(769            f"Offset hours must be in [0, 24]: {tz_delta}"770        )771 772    # Yes, +5 maps to an offset of -5h773    if match.group("sign") != "-":774        total = -total775 776    return total777