Coverage for /var/srv/projects/api.amasfac.comuna18.com/tmp/venv/lib/python3.9/site-packages/pandas/_testing/contexts.py: 27%
97 statements
« prev ^ index » next coverage.py v6.4.4, created at 2023-07-17 14:22 -0600
« prev ^ index » next coverage.py v6.4.4, created at 2023-07-17 14:22 -0600
1from __future__ import annotations
3from contextlib import contextmanager
4import os
5from pathlib import Path
6from shutil import rmtree
7import tempfile
8from typing import (
9 IO,
10 Any,
11 Iterator,
12)
13import uuid
15import numpy as np
17from pandas import set_option
19from pandas.io.common import get_handle
22@contextmanager
23def decompress_file(path, compression) -> Iterator[IO[bytes]]:
24 """
25 Open a compressed file and return a file object.
27 Parameters
28 ----------
29 path : str
30 The path where the file is read from.
32 compression : {'gzip', 'bz2', 'zip', 'xz', 'zstd', None}
33 Name of the decompression to use
35 Returns
36 -------
37 file object
38 """
39 with get_handle(path, "rb", compression=compression, is_text=False) as handle:
40 yield handle.handle
43@contextmanager
44def set_timezone(tz: str) -> Iterator[None]:
45 """
46 Context manager for temporarily setting a timezone.
48 Parameters
49 ----------
50 tz : str
51 A string representing a valid timezone.
53 Examples
54 --------
55 >>> from datetime import datetime
56 >>> from dateutil.tz import tzlocal
57 >>> tzlocal().tzname(datetime(2021, 1, 1)) # doctest: +SKIP
58 'IST'
60 >>> with set_timezone('US/Eastern'):
61 ... tzlocal().tzname(datetime(2021, 1, 1))
62 ...
63 'EST'
64 """
65 import os
66 import time
68 def setTZ(tz):
69 if tz is None:
70 try:
71 del os.environ["TZ"]
72 except KeyError:
73 pass
74 else:
75 os.environ["TZ"] = tz
76 time.tzset()
78 orig_tz = os.environ.get("TZ")
79 setTZ(tz)
80 try:
81 yield
82 finally:
83 setTZ(orig_tz)
86@contextmanager
87def ensure_clean(filename=None, return_filelike: bool = False, **kwargs: Any):
88 """
89 Gets a temporary path and agrees to remove on close.
91 This implementation does not use tempfile.mkstemp to avoid having a file handle.
92 If the code using the returned path wants to delete the file itself, windows
93 requires that no program has a file handle to it.
95 Parameters
96 ----------
97 filename : str (optional)
98 suffix of the created file.
99 return_filelike : bool (default False)
100 if True, returns a file-like which is *always* cleaned. Necessary for
101 savefig and other functions which want to append extensions.
102 **kwargs
103 Additional keywords are passed to open().
105 """
106 folder = Path(tempfile.gettempdir())
108 if filename is None:
109 filename = ""
110 filename = str(uuid.uuid4()) + filename
111 path = folder / filename
113 path.touch()
115 handle_or_str: str | IO = str(path)
116 if return_filelike:
117 kwargs.setdefault("mode", "w+b")
118 handle_or_str = open(path, **kwargs)
120 try:
121 yield handle_or_str
122 finally:
123 if not isinstance(handle_or_str, str):
124 handle_or_str.close()
125 if path.is_file():
126 path.unlink()
129@contextmanager
130def ensure_clean_dir() -> Iterator[str]:
131 """
132 Get a temporary directory path and agrees to remove on close.
134 Yields
135 ------
136 Temporary directory path
137 """
138 directory_name = tempfile.mkdtemp(suffix="")
139 try:
140 yield directory_name
141 finally:
142 try:
143 rmtree(directory_name)
144 except OSError:
145 pass
148@contextmanager
149def ensure_safe_environment_variables() -> Iterator[None]:
150 """
151 Get a context manager to safely set environment variables
153 All changes will be undone on close, hence environment variables set
154 within this contextmanager will neither persist nor change global state.
155 """
156 saved_environ = dict(os.environ)
157 try:
158 yield
159 finally:
160 os.environ.clear()
161 os.environ.update(saved_environ)
164@contextmanager
165def with_csv_dialect(name, **kwargs) -> Iterator[None]:
166 """
167 Context manager to temporarily register a CSV dialect for parsing CSV.
169 Parameters
170 ----------
171 name : str
172 The name of the dialect.
173 kwargs : mapping
174 The parameters for the dialect.
176 Raises
177 ------
178 ValueError : the name of the dialect conflicts with a builtin one.
180 See Also
181 --------
182 csv : Python's CSV library.
183 """
184 import csv
186 _BUILTIN_DIALECTS = {"excel", "excel-tab", "unix"}
188 if name in _BUILTIN_DIALECTS:
189 raise ValueError("Cannot override builtin dialect.")
191 csv.register_dialect(name, **kwargs)
192 try:
193 yield
194 finally:
195 csv.unregister_dialect(name)
198@contextmanager
199def use_numexpr(use, min_elements=None) -> Iterator[None]:
200 from pandas.core.computation import expressions as expr
202 if min_elements is None:
203 min_elements = expr._MIN_ELEMENTS
205 olduse = expr.USE_NUMEXPR
206 oldmin = expr._MIN_ELEMENTS
207 set_option("compute.use_numexpr", use)
208 expr._MIN_ELEMENTS = min_elements
209 try:
210 yield
211 finally:
212 expr._MIN_ELEMENTS = oldmin
213 set_option("compute.use_numexpr", olduse)
216class RNGContext:
217 """
218 Context manager to set the numpy random number generator speed. Returns
219 to the original value upon exiting the context manager.
221 Parameters
222 ----------
223 seed : int
224 Seed for numpy.random.seed
226 Examples
227 --------
228 with RNGContext(42):
229 np.random.randn()
230 """
232 def __init__(self, seed) -> None:
233 self.seed = seed
235 def __enter__(self) -> None:
237 self.start_state = np.random.get_state()
238 np.random.seed(self.seed)
240 def __exit__(self, exc_type, exc_value, traceback) -> None:
242 np.random.set_state(self.start_state)