urllib3

1
from __future__ import annotations
2

3
import datetime
4
import os
5
import time
6
import typing
7
from contextlib import contextmanager
8

9
import pytest
10

11
try:
12
    import zoneinfo
13
except ImportError:
14
    # Python < 3.9
15
    from backports import zoneinfo  # type: ignore[no-redef]
16

17

18
@contextmanager
19
def stub_timezone_ctx(tzname: str | None) -> typing.Generator[None, None, None]:
20
    """
21
    Switch to a locally-known timezone specified by `tzname`.
22
    On exit, restore the previous timezone.
23
    If `tzname` is `None`, do nothing.
24
    """
25
    if tzname is None:
26
        yield
27
        return
28

29
    # Only supported on Unix
30
    if not hasattr(time, "tzset"):
31
        pytest.skip("Timezone patching is not supported")
32

33
    # Make sure the new timezone exists
34
    try:
35
        zoneinfo.ZoneInfo(tzname)
36
    except zoneinfo.ZoneInfoNotFoundError:
37
        raise ValueError(f"Invalid timezone specified: {tzname!r}")
38

39
    # Get the current timezone
40
    old_tzname = datetime.datetime.now().astimezone().tzname()
41
    if old_tzname is None:
42
        raise OSError("Cannot determine current timezone")
43

44
    os.environ["TZ"] = tzname
45
    time.tzset()
46
    yield
47
    os.environ["TZ"] = old_tzname
48
    time.tzset()
49