cpython/Lib/test/test_zoneinfo/_support.py

101 lines
3.1 KiB
Python

import contextlib
import functools
import sys
import threading
import unittest
from test.support import import_fresh_module
OS_ENV_LOCK = threading.Lock()
TZPATH_LOCK = threading.Lock()
TZPATH_TEST_LOCK = threading.Lock()
def call_once(f):
"""Decorator that ensures a function is only ever called once."""
lock = threading.Lock()
cached = functools.lru_cache(None)(f)
@functools.wraps(f)
def inner():
with lock:
return cached()
return inner
@call_once
def get_modules():
"""Retrieve two copies of zoneinfo: pure Python and C accelerated.
Because this function manipulates the import system in a way that might
be fragile or do unexpected things if it is run many times, it uses a
`call_once` decorator to ensure that this is only ever called exactly
one time — in other words, when using this function you will only ever
get one copy of each module rather than a fresh import each time.
"""
import zoneinfo as c_module
py_module = import_fresh_module("zoneinfo", blocked=["_zoneinfo"])
return py_module, c_module
@contextlib.contextmanager
def set_zoneinfo_module(module):
"""Make sure sys.modules["zoneinfo"] refers to `module`.
This is necessary because `pickle` will refuse to serialize
an type calling itself `zoneinfo.ZoneInfo` unless `zoneinfo.ZoneInfo`
refers to the same object.
"""
NOT_PRESENT = object()
old_zoneinfo = sys.modules.get("zoneinfo", NOT_PRESENT)
sys.modules["zoneinfo"] = module
yield
if old_zoneinfo is not NOT_PRESENT:
sys.modules["zoneinfo"] = old_zoneinfo
else: # pragma: nocover
sys.modules.pop("zoneinfo")
class ZoneInfoTestBase(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.klass = cls.module.ZoneInfo
super().setUpClass()
@contextlib.contextmanager
def tzpath_context(self, tzpath, block_tzdata=True, lock=TZPATH_LOCK):
def pop_tzdata_modules():
tzdata_modules = {}
for modname in list(sys.modules):
if modname.split(".", 1)[0] != "tzdata": # pragma: nocover
continue
tzdata_modules[modname] = sys.modules.pop(modname)
return tzdata_modules
with lock:
if block_tzdata:
# In order to fully exclude tzdata from the path, we need to
# clear the sys.modules cache of all its contents — setting the
# root package to None is not enough to block direct access of
# already-imported submodules (though it will prevent new
# imports of submodules).
tzdata_modules = pop_tzdata_modules()
sys.modules["tzdata"] = None
old_path = self.module.TZPATH
try:
self.module.reset_tzpath(tzpath)
yield
finally:
if block_tzdata:
sys.modules.pop("tzdata")
for modname, module in tzdata_modules.items():
sys.modules[modname] = module
self.module.reset_tzpath(old_path)