243 行
5.5 KiB
Python
243 行
5.5 KiB
Python
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
import os
|
|
from pathlib import Path
|
|
from shutil import rmtree
|
|
import tempfile
|
|
from typing import (
|
|
IO,
|
|
Any,
|
|
Iterator,
|
|
)
|
|
import uuid
|
|
|
|
import numpy as np
|
|
|
|
from pandas import set_option
|
|
|
|
from pandas.io.common import get_handle
|
|
|
|
|
|
@contextmanager
|
|
def decompress_file(path, compression) -> Iterator[IO[bytes]]:
|
|
"""
|
|
Open a compressed file and return a file object.
|
|
|
|
Parameters
|
|
----------
|
|
path : str
|
|
The path where the file is read from.
|
|
|
|
compression : {'gzip', 'bz2', 'zip', 'xz', 'zstd', None}
|
|
Name of the decompression to use
|
|
|
|
Returns
|
|
-------
|
|
file object
|
|
"""
|
|
with get_handle(path, "rb", compression=compression, is_text=False) as handle:
|
|
yield handle.handle
|
|
|
|
|
|
@contextmanager
|
|
def set_timezone(tz: str) -> Iterator[None]:
|
|
"""
|
|
Context manager for temporarily setting a timezone.
|
|
|
|
Parameters
|
|
----------
|
|
tz : str
|
|
A string representing a valid timezone.
|
|
|
|
Examples
|
|
--------
|
|
>>> from datetime import datetime
|
|
>>> from dateutil.tz import tzlocal
|
|
>>> tzlocal().tzname(datetime(2021, 1, 1)) # doctest: +SKIP
|
|
'IST'
|
|
|
|
>>> with set_timezone('US/Eastern'):
|
|
... tzlocal().tzname(datetime(2021, 1, 1))
|
|
...
|
|
'EST'
|
|
"""
|
|
import os
|
|
import time
|
|
|
|
def setTZ(tz):
|
|
if tz is None:
|
|
try:
|
|
del os.environ["TZ"]
|
|
except KeyError:
|
|
pass
|
|
else:
|
|
os.environ["TZ"] = tz
|
|
time.tzset()
|
|
|
|
orig_tz = os.environ.get("TZ")
|
|
setTZ(tz)
|
|
try:
|
|
yield
|
|
finally:
|
|
setTZ(orig_tz)
|
|
|
|
|
|
@contextmanager
|
|
def ensure_clean(filename=None, return_filelike: bool = False, **kwargs: Any):
|
|
"""
|
|
Gets a temporary path and agrees to remove on close.
|
|
|
|
This implementation does not use tempfile.mkstemp to avoid having a file handle.
|
|
If the code using the returned path wants to delete the file itself, windows
|
|
requires that no program has a file handle to it.
|
|
|
|
Parameters
|
|
----------
|
|
filename : str (optional)
|
|
suffix of the created file.
|
|
return_filelike : bool (default False)
|
|
if True, returns a file-like which is *always* cleaned. Necessary for
|
|
savefig and other functions which want to append extensions.
|
|
**kwargs
|
|
Additional keywords are passed to open().
|
|
|
|
"""
|
|
folder = Path(tempfile.gettempdir())
|
|
|
|
if filename is None:
|
|
filename = ""
|
|
filename = str(uuid.uuid4()) + filename
|
|
path = folder / filename
|
|
|
|
path.touch()
|
|
|
|
handle_or_str: str | IO = str(path)
|
|
if return_filelike:
|
|
kwargs.setdefault("mode", "w+b")
|
|
handle_or_str = open(path, **kwargs)
|
|
|
|
try:
|
|
yield handle_or_str
|
|
finally:
|
|
if not isinstance(handle_or_str, str):
|
|
handle_or_str.close()
|
|
if path.is_file():
|
|
path.unlink()
|
|
|
|
|
|
@contextmanager
|
|
def ensure_clean_dir() -> Iterator[str]:
|
|
"""
|
|
Get a temporary directory path and agrees to remove on close.
|
|
|
|
Yields
|
|
------
|
|
Temporary directory path
|
|
"""
|
|
directory_name = tempfile.mkdtemp(suffix="")
|
|
try:
|
|
yield directory_name
|
|
finally:
|
|
try:
|
|
rmtree(directory_name)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
@contextmanager
|
|
def ensure_safe_environment_variables() -> Iterator[None]:
|
|
"""
|
|
Get a context manager to safely set environment variables
|
|
|
|
All changes will be undone on close, hence environment variables set
|
|
within this contextmanager will neither persist nor change global state.
|
|
"""
|
|
saved_environ = dict(os.environ)
|
|
try:
|
|
yield
|
|
finally:
|
|
os.environ.clear()
|
|
os.environ.update(saved_environ)
|
|
|
|
|
|
@contextmanager
|
|
def with_csv_dialect(name, **kwargs) -> Iterator[None]:
|
|
"""
|
|
Context manager to temporarily register a CSV dialect for parsing CSV.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the dialect.
|
|
kwargs : mapping
|
|
The parameters for the dialect.
|
|
|
|
Raises
|
|
------
|
|
ValueError : the name of the dialect conflicts with a builtin one.
|
|
|
|
See Also
|
|
--------
|
|
csv : Python's CSV library.
|
|
"""
|
|
import csv
|
|
|
|
_BUILTIN_DIALECTS = {"excel", "excel-tab", "unix"}
|
|
|
|
if name in _BUILTIN_DIALECTS:
|
|
raise ValueError("Cannot override builtin dialect.")
|
|
|
|
csv.register_dialect(name, **kwargs)
|
|
try:
|
|
yield
|
|
finally:
|
|
csv.unregister_dialect(name)
|
|
|
|
|
|
@contextmanager
|
|
def use_numexpr(use, min_elements=None) -> Iterator[None]:
|
|
from pandas.core.computation import expressions as expr
|
|
|
|
if min_elements is None:
|
|
min_elements = expr._MIN_ELEMENTS
|
|
|
|
olduse = expr.USE_NUMEXPR
|
|
oldmin = expr._MIN_ELEMENTS
|
|
set_option("compute.use_numexpr", use)
|
|
expr._MIN_ELEMENTS = min_elements
|
|
try:
|
|
yield
|
|
finally:
|
|
expr._MIN_ELEMENTS = oldmin
|
|
set_option("compute.use_numexpr", olduse)
|
|
|
|
|
|
class RNGContext:
|
|
"""
|
|
Context manager to set the numpy random number generator speed. Returns
|
|
to the original value upon exiting the context manager.
|
|
|
|
Parameters
|
|
----------
|
|
seed : int
|
|
Seed for numpy.random.seed
|
|
|
|
Examples
|
|
--------
|
|
with RNGContext(42):
|
|
np.random.randn()
|
|
"""
|
|
|
|
def __init__(self, seed) -> None:
|
|
self.seed = seed
|
|
|
|
def __enter__(self) -> None:
|
|
|
|
self.start_state = np.random.get_state()
|
|
np.random.seed(self.seed)
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
|
|
|
np.random.set_state(self.start_state)
|