Coverage for /var/srv/projects/api.amasfac.comuna18.com/tmp/venv/lib/python3.9/site-packages/pandas/_testing/contexts.py: 27%

1from __future__ import annotations 

2 

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 

14 

15import numpy as np 

16 

17from pandas import set_option 

18 

19from pandas.io.common import get_handle 

20 

21 

22@contextmanager 

23def decompress_file(path, compression) -> Iterator[IO[bytes]]: 

24 """ 

25 Open a compressed file and return a file object. 

26 

27 Parameters 

28 ---------- 

29 path : str 

30 The path where the file is read from. 

31 

32 compression : {'gzip', 'bz2', 'zip', 'xz', 'zstd', None} 

33 Name of the decompression to use 

34 

35 Returns 

36 ------- 

37 file object 

38 """ 

39 with get_handle(path, "rb", compression=compression, is_text=False) as handle: 

40 yield handle.handle 

41 

42 

43@contextmanager 

44def set_timezone(tz: str) -> Iterator[None]: 

45 """ 

46 Context manager for temporarily setting a timezone. 

47 

48 Parameters 

49 ---------- 

50 tz : str 

51 A string representing a valid timezone. 

52 

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' 

59 

60 >>> with set_timezone('US/Eastern'): 

61 ... tzlocal().tzname(datetime(2021, 1, 1)) 

62 ... 

63 'EST' 

64 """ 

65 import os 

66 import time 

67 

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() 

77 

78 orig_tz = os.environ.get("TZ") 

79 setTZ(tz) 

80 try: 

81 yield 

82 finally: 

83 setTZ(orig_tz) 

84 

85 

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. 

90 

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. 

94 

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(). 

104 

105 """ 

106 folder = Path(tempfile.gettempdir()) 

107 

108 if filename is None: 

109 filename = "" 

110 filename = str(uuid.uuid4()) + filename 

111 path = folder / filename 

112 

113 path.touch() 

114 

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) 

119 

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() 

127 

128 

129@contextmanager 

130def ensure_clean_dir() -> Iterator[str]: 

131 """ 

132 Get a temporary directory path and agrees to remove on close. 

133 

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 

146 

147 

148@contextmanager 

149def ensure_safe_environment_variables() -> Iterator[None]: 

150 """ 

151 Get a context manager to safely set environment variables 

152 

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) 

162 

163 

164@contextmanager 

165def with_csv_dialect(name, **kwargs) -> Iterator[None]: 

166 """ 

167 Context manager to temporarily register a CSV dialect for parsing CSV. 

168 

169 Parameters 

170 ---------- 

171 name : str 

172 The name of the dialect. 

173 kwargs : mapping 

174 The parameters for the dialect. 

175 

176 Raises 

177 ------ 

178 ValueError : the name of the dialect conflicts with a builtin one. 

179 

180 See Also 

181 -------- 

182 csv : Python's CSV library. 

183 """ 

184 import csv 

185 

186 _BUILTIN_DIALECTS = {"excel", "excel-tab", "unix"} 

187 

188 if name in _BUILTIN_DIALECTS: 

189 raise ValueError("Cannot override builtin dialect.") 

190 

191 csv.register_dialect(name, **kwargs) 

192 try: 

193 yield 

194 finally: 

195 csv.unregister_dialect(name) 

196 

197 

198@contextmanager 

199def use_numexpr(use, min_elements=None) -> Iterator[None]: 

200 from pandas.core.computation import expressions as expr 

201 

202 if min_elements is None: 

203 min_elements = expr._MIN_ELEMENTS 

204 

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) 

214 

215 

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. 

220 

221 Parameters 

222 ---------- 

223 seed : int 

224 Seed for numpy.random.seed 

225 

226 Examples 

227 -------- 

228 with RNGContext(42): 

229 np.random.randn() 

230 """ 

231 

232 def __init__(self, seed) -> None: 

233 self.seed = seed 

234 

235 def __enter__(self) -> None: 

236 

237 self.start_state = np.random.get_state() 

238 np.random.seed(self.seed) 

239 

240 def __exit__(self, exc_type, exc_value, traceback) -> None: 

241 

242 np.random.set_state(self.start_state)