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 contextmanager 

4import os 

5import sys 

6from typing import ( 

7 IO, 

8 TYPE_CHECKING, 

9) 

10 

11from pandas.compat import CHAINED_WARNING_DISABLED 

12from pandas.errors import ChainedAssignmentError 

13 

14from pandas.io.common import get_handle 

15 

16if TYPE_CHECKING: 

17 from collections.abc import Generator 

18 

19 from pandas._typing import ( 

20 BaseBuffer, 

21 CompressionOptions, 

22 FilePath, 

23 ) 

24 

25 

26@contextmanager 

27def decompress_file( 

28 path: FilePath | BaseBuffer, compression: CompressionOptions 

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

30 """ 

31 Open a compressed file and return a file object. 

32 

33 Parameters 

34 ---------- 

35 path : str 

36 The path where the file is read from. 

37 

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

39 Name of the decompression to use 

40 

41 Returns 

42 ------- 

43 file object 

44 """ 

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

46 yield handle.handle 

47 

48 

49@contextmanager 

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

51 """ 

52 Context manager for temporarily setting a timezone. 

53 

54 Parameters 

55 ---------- 

56 tz : str 

57 A string representing a valid timezone. 

58 

59 Examples 

60 -------- 

61 >>> from datetime import datetime 

62 >>> from dateutil.tz import tzlocal 

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

64 'IST' 

65 

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

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

68 'EST' 

69 """ 

70 import time 

71 

72 def setTZ(tz) -> None: 

73 if hasattr(time, "tzset"): 

74 if tz is None: 

75 try: 

76 del os.environ["TZ"] 

77 except KeyError: 

78 pass 

79 else: 

80 os.environ["TZ"] = tz 

81 # Next line allows typing checks to pass on Windows 

82 if sys.platform != "win32": 

83 time.tzset() 

84 

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

86 setTZ(tz) 

87 try: 

88 yield 

89 finally: 

90 setTZ(orig_tz) 

91 

92 

93@contextmanager 

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

95 """ 

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

97 

98 Parameters 

99 ---------- 

100 name : str 

101 The name of the dialect. 

102 kwargs : mapping 

103 The parameters for the dialect. 

104 

105 Raises 

106 ------ 

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

108 

109 See Also 

110 -------- 

111 csv : Python's CSV library. 

112 """ 

113 import csv 

114 

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

116 

117 if name in _BUILTIN_DIALECTS: 

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

119 

120 csv.register_dialect(name, **kwargs) 

121 try: 

122 yield 

123 finally: 

124 csv.unregister_dialect(name) 

125 

126 

127def raises_chained_assignment_error(extra_warnings=(), extra_match=()): 

128 from pandas._testing import assert_produces_warning 

129 

130 if CHAINED_WARNING_DISABLED: 

131 if not extra_warnings: 

132 from contextlib import nullcontext 

133 

134 return nullcontext() 

135 else: 

136 return assert_produces_warning( 

137 extra_warnings, 

138 match=extra_match, 

139 ) 

140 else: 

141 warning = ChainedAssignmentError 

142 match = ( 

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

144 "through chained assignment" 

145 ) 

146 if extra_warnings: 

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

148 return assert_produces_warning( 

149 warning, 

150 match=(match, *extra_match), 

151 )