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 path : str, path object or file-like object, default None 

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(writer, reader, tmp_path: Path): 

50 """ 

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

52 

53 Parameters 

54 ---------- 

55 writer : callable bound to pandas object 

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

57 reader : callable 

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

59 path : str, default None 

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

61 

62 Returns 

63 ------- 

64 pandas object 

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

66 """ 

67 writer(tmp_path) 

68 obj = reader(tmp_path) 

69 return obj 

70 

71 

72def write_to_compressed(compression, path: str, data, dest: str = "test") -> None: 

73 """ 

74 Write data to a compressed file. 

75 

76 Parameters 

77 ---------- 

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

79 The compression type to use. 

80 path : str 

81 The file path to write the data. 

82 data : str 

83 The data to write. 

84 dest : str, default "test" 

85 The destination file (for ZIP only) 

86 

87 Raises 

88 ------ 

89 ValueError : An invalid compression value was passed in. 

90 """ 

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

92 mode = "wb" 

93 method = "write" 

94 compress_method: Callable 

95 

96 if compression == "zip": 

97 compress_method = zipfile.ZipFile 

98 mode = "w" 

99 args = (dest, data) 

100 method = "writestr" 

101 elif compression == "tar": 

102 compress_method = tarfile.TarFile 

103 mode = "w" 

104 file = tarfile.TarInfo(name=dest) 

105 bytes = io.BytesIO(data) 

106 file.size = len(data) 

107 args = (file, bytes) 

108 method = "addfile" 

109 elif compression == "gzip": 

110 compress_method = gzip.GzipFile 

111 elif compression == "bz2": 

112 import bz2 

113 

114 compress_method = bz2.BZ2File 

115 elif compression == "zstd": 

116 compress_method = import_optional_dependency("zstandard").open 

117 elif compression == "xz": 

118 import lzma 

119 

120 compress_method = lzma.LZMAFile 

121 else: 

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

123 

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

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

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

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

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

129 getattr(f, method)(*args)