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

1from __future__ import annotations 

2 

3import gzip 

4import io 

5import tarfile 

6from typing import ( 

7 TYPE_CHECKING, 

8 Any, 

9) 

10import zipfile 

11 

12from pandas.compat._optional import import_optional_dependency 

13 

14import pandas as pd 

15 

16if TYPE_CHECKING: 

17 from collections.abc import Callable 

18 from pathlib import Path 

19 

20 from pandas import ( 

21 DataFrame, 

22 Series, 

23 ) 

24 

25# ------------------------------------------------------------------ 

26# File-IO 

27 

28 

29def round_trip_pickle(obj: Any, tmp_path: Path) -> DataFrame | Series: 

30 """ 

31 Pickle an object and then read it again. 

32 

33 Parameters 

34 ---------- 

35 obj : any object 

36 The object to pickle and then re-read. 

37 tmp_path : Path 

38 The path where the pickled object is written and then read. 

39 

40 Returns 

41 ------- 

42 pandas object 

43 The original object that was pickled and then re-read. 

44 """ 

45 pd.to_pickle(obj, tmp_path) 

46 return pd.read_pickle(tmp_path) 

47 

48 

49def round_trip_pathlib( 

50 writer: Callable[[Path], Any], reader: Callable[[Path], Any], tmp_path: Path 

51) -> DataFrame | Series: 

52 """ 

53 Write an object to file specified by a pathlib.Path and read it back 

54 

55 Parameters 

56 ---------- 

57 writer : callable bound to pandas object 

58 IO writing function (e.g. DataFrame.to_csv ) 

59 reader : callable 

60 IO reading function (e.g. pd.read_csv ) 

61 tmp_path : Path 

62 The path where the object is written and then read. 

63 

64 Returns 

65 ------- 

66 pandas object 

67 The original object that was serialized and then re-read. 

68 """ 

69 writer(tmp_path) 

70 obj = reader(tmp_path) 

71 return obj 

72 

73 

74def write_to_compressed( 

75 compression: str, path: str, data: bytes, dest: str = "test" 

76) -> None: 

77 """ 

78 Write data to a compressed file. 

79 

80 Parameters 

81 ---------- 

82 compression : {'gzip', 'bz2', 'zip', 'xz', 'zstd'} 

83 The compression type to use. 

84 path : str 

85 The file path to write the data. 

86 data : str 

87 The data to write. 

88 dest : str, default "test" 

89 The destination file (for ZIP only) 

90 

91 Raises 

92 ------ 

93 ValueError : An invalid compression value was passed in. 

94 """ 

95 args: tuple[Any, ...] = (data,) 

96 mode = "wb" 

97 method = "write" 

98 compress_method: Callable[..., Any] 

99 

100 if compression == "zip": 

101 compress_method = zipfile.ZipFile 

102 mode = "w" 

103 args = (dest, data) 

104 method = "writestr" 

105 elif compression == "tar": 

106 compress_method = tarfile.TarFile 

107 mode = "w" 

108 file = tarfile.TarInfo(name=dest) 

109 bytes = io.BytesIO(data) 

110 file.size = len(data) 

111 args = (file, bytes) 

112 method = "addfile" 

113 elif compression == "gzip": 

114 compress_method = gzip.GzipFile 

115 elif compression == "bz2": 

116 import bz2 

117 

118 compress_method = bz2.BZ2File 

119 elif compression == "zstd": 

120 compress_method = import_optional_dependency("zstandard").open 

121 elif compression == "xz": 

122 import lzma 

123 

124 compress_method = lzma.LZMAFile 

125 else: 

126 raise ValueError(f"Unrecognized compression type: {compression}") 

127 

128 # error: No overload variant of "ZipFile" matches argument types "str", "str" 

129 # error: No overload variant of "BZ2File" matches argument types "str", "str" 

130 # error: Argument "mode" to "TarFile" has incompatible type "str"; 

131 # expected "Literal['r', 'a', 'w', 'x'] 

132 with compress_method(path, mode=mode) as f: # type: ignore[call-overload, arg-type] 

133 getattr(f, method)(*args)