Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/_testing/contexts.py: 29%

1from __future__ import annotations 

2 

3from contextlib import ( 

4 AbstractContextManager, 

5 contextmanager, 

6) 

7import os 

8import sys 

9from typing import ( 

10 IO, 

11 TYPE_CHECKING, 

12 Any, 

13) 

14 

15from pandas.compat import CHAINED_WARNING_DISABLED 

16from pandas.errors import ChainedAssignmentError 

17 

18from pandas.io.common import get_handle 

19 

20if TYPE_CHECKING: 

21 from collections.abc import Generator 

22 import warnings 

23 

24 from pandas._typing import ( 

25 BaseBuffer, 

26 CompressionOptions, 

27 FilePath, 

28 ) 

29 

30 

31@contextmanager 

32def decompress_file( 

33 path: FilePath | BaseBuffer, compression: CompressionOptions 

34) -> Generator[IO[bytes]]: 

35 """ 

36 Open a compressed file and return a file object. 

37 

38 Parameters 

39 ---------- 

40 path : str 

41 The path where the file is read from. 

42 

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

44 Name of the decompression to use 

45 

46 Returns 

47 ------- 

48 file object 

49 """ 

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

51 yield handle.handle 

52 

53 

54@contextmanager 

55def set_timezone(tz: str) -> Generator[None]: 

56 """ 

57 Context manager for temporarily setting a timezone. 

58 

59 Parameters 

60 ---------- 

61 tz : str 

62 A string representing a valid timezone. 

63 

64 Examples 

65 -------- 

66 >>> from datetime import datetime 

67 >>> from dateutil.tz import tzlocal 

68 >>> tzlocal().tzname(datetime(2021, 1, 1)) # doctest: +SKIP 

69 'IST' 

70 

71 >>> with set_timezone("US/Eastern"): 

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

73 'EST' 

74 """ 

75 import time 

76 

77 def setTZ(tz: str | None) -> None: 

78 if hasattr(time, "tzset"): 

79 if tz is None: 

80 try: 

81 del os.environ["TZ"] 

82 except KeyError: 

83 pass 

84 else: 

85 os.environ["TZ"] = tz 

86 # Next line allows typing checks to pass on Windows 

87 if sys.platform != "win32": 

88 time.tzset() 

89 

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

91 setTZ(tz) 

92 try: 

93 yield 

94 finally: 

95 setTZ(orig_tz) 

96 

97 

98@contextmanager 

99def with_csv_dialect(name: str, **kwargs: Any) -> Generator[None]: 

100 """ 

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

102 

103 Parameters 

104 ---------- 

105 name : str 

106 The name of the dialect. 

107 kwargs : mapping 

108 The parameters for the dialect. 

109 

110 Raises 

111 ------ 

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

113 

114 See Also 

115 -------- 

116 csv : Python's CSV library. 

117 """ 

118 import csv 

119 

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

121 

122 if name in _BUILTIN_DIALECTS: 

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

124 

125 csv.register_dialect(name, **kwargs) 

126 try: 

127 yield 

128 finally: 

129 csv.unregister_dialect(name) 

130 

131 

132def raises_chained_assignment_error( 

133 extra_warnings: tuple[type[Warning], ...] = (), 

134 extra_match: tuple[str | None, ...] = (), 

135) -> AbstractContextManager[list[warnings.WarningMessage] | None]: 

136 from pandas._testing import assert_produces_warning 

137 

138 if CHAINED_WARNING_DISABLED: 

139 if not extra_warnings: 

140 from contextlib import nullcontext 

141 

142 return nullcontext() 

143 else: 

144 return assert_produces_warning( 

145 extra_warnings, 

146 match=extra_match, 

147 ) 

148 else: 

149 warning = ChainedAssignmentError 

150 match = ( 

151 "A value is being set on a copy of a DataFrame or Series " 

152 "through chained assignment" 

153 ) 

154 if extra_warnings: 

155 warning = (warning, *extra_warnings) # type: ignore[assignment] 

156 return assert_produces_warning( 

157 warning, 

158 match=(match, *extra_match), 

159 )