Coverage for polypandas/io.py: 77%

1"""Data export and import utilities for polypandas."""

2

3 from pathlib import Path

4 from typing import Any, Dict, List

5

6 from polypandas.exceptions import PolypandasError

7 from polypandas.protocols import is_pandas_available

8

9

10 class DataIOError(PolypandasError):

11 """Raised when data I/O operations fail."""

12

13 pass

14

15

16 def save_as_parquet(df: Any, path: str, **kwargs: Any) -> None:

17 """Save DataFrame as Parquet file.

18

19 Args:

20 df: pandas DataFrame to save.

21 path: Output path for Parquet file.

22 **kwargs: Additional arguments for df.to_parquet() (e.g. index=False).

23

24 Raises:

25 DataIOError: If save operation fails.

26 """

27 if not is_pandas_available():

28 raise DataIOError("pandas is required for Parquet operations")

29

30 import pandas as pd

31

32 if not isinstance(df, pd.DataFrame):

33 raise DataIOError("Expected a pandas DataFrame")

34

35 try:

36 Path(path).parent.mkdir(parents=True, exist_ok=True)

37 df.to_parquet(path, **kwargs)

38 except Exception as e:

39 raise DataIOError(f"Failed to save Parquet file: {e}") from e

40

41

42 def save_as_json(df: Any, path: str, **kwargs: Any) -> None:

43 """Save DataFrame as JSON file."""

44 if not is_pandas_available():

45 raise DataIOError("pandas is required for JSON operations")

46

47 try:

48 Path(path).parent.mkdir(parents=True, exist_ok=True)

49 df.to_json(path, **kwargs)

50 except Exception as e:

51 raise DataIOError(f"Failed to save JSON file: {e}") from e

52

53

54 def save_as_csv(

55 df: Any,

56 path: str,

57 header: bool = True,

58 **kwargs: Any,

59) -> None:

60 """Save DataFrame as CSV file."""

61 if not is_pandas_available():

62 raise DataIOError("pandas is required for CSV operations")

63

64 try:

65 Path(path).parent.mkdir(parents=True, exist_ok=True)

66 df.to_csv(path, header=header, **kwargs)

67 except Exception as e:

68 raise DataIOError(f"Failed to save CSV file: {e}") from e

69

70

71 def load_parquet(path: str, **kwargs: Any) -> Any:

72 """Load DataFrame from Parquet file."""

73 if not is_pandas_available():

74 raise DataIOError("pandas is required for Parquet operations")

75

76 import pandas as pd

77

78 try:

79 return pd.read_parquet(path, **kwargs)

80 except Exception as e:

81 raise DataIOError(f"Failed to load Parquet file: {e}") from e

82

83

84 def load_json(path: str, **kwargs: Any) -> Any:

85 """Load DataFrame from JSON file."""

86 if not is_pandas_available():

87 raise DataIOError("pandas is required for JSON operations")

88

89 import pandas as pd

90

91 try:

92 return pd.read_json(path, **kwargs)

93 except Exception as e:

94 raise DataIOError(f"Failed to load JSON file: {e}") from e

95

96

97 def load_csv(

98 path: str,

99 header: int = 0,

100 **kwargs: Any,

101) -> Any:

102 """Load DataFrame from CSV file."""

103 if not is_pandas_available():

104 raise DataIOError("pandas is required for CSV operations")

105

106 import pandas as pd

107

108 try:

109 return pd.read_csv(path, header=header, **kwargs)

110 except Exception as e:

111 raise DataIOError(f"Failed to load CSV file: {e}") from e

112

113

114 def load_and_validate(

115 path: str,

116 expected_schema: Any = None,

117 validate_schema: bool = True,

118) -> Any:

119 """Load data file and optionally validate against expected dtype dict.

120

121 Args:

122 path: Path to data file (.parquet, .json, or .csv).

123 expected_schema: Optional dict of column name -> expected dtype (string).

124 validate_schema: Whether to validate that columns and dtypes match.

125

126 Returns:

127 Loaded pandas DataFrame.

128 """

129 if not is_pandas_available():

130 raise DataIOError("pandas is required for data loading")

131

132 file_path = Path(path)

133 suffix = file_path.suffix.lower()

134

135 if suffix == ".parquet":

136 df = load_parquet(path)

137 elif suffix == ".json":

138 df = load_json(path)

139 elif suffix in (".csv", ".txt"):

140 df = load_csv(path)

141 else:

142 raise DataIOError(f"Unsupported file format: {suffix}. Supported: .parquet, .json, .csv")

143

144 if validate_schema and expected_schema is not None:

145 from polypandas.testing import assert_column_exists

146

147 for col in expected_schema:

148 assert_column_exists(df, col)

149 # Optionally check dtypes

150 for col, dtype in expected_schema.items():

151 if col in df.columns and str(df[col].dtype) != str(dtype):

152 raise DataIOError(f"Column '{col}' has dtype {df[col].dtype}, expected {dtype}")

153

154 return df

155

156

157 def save_dicts_as_json(data: List[Dict[str, Any]], path: str) -> None:

158 """Save list of dictionaries as JSON lines file. Works without pandas."""

159 import json

160

161 try:

162 file_path = Path(path)

163 file_path.parent.mkdir(parents=True, exist_ok=True)

164 with open(file_path, "w") as f:

165 for record in data:

166 json.dump(record, f, default=str)

167 f.write("\n")

168 except Exception as e:

169 raise DataIOError(f"Failed to save JSON file: {e}") from e

170

171

172 def load_dicts_from_json(path: str) -> List[Dict[str, Any]]:

173 """Load list of dictionaries from JSON lines file. Works without pandas."""

174 import json

175

176 try:

177 file_path = Path(path)

178 if not file_path.exists():

179 raise DataIOError(f"File not found: {path}")

180 data = []

181 with open(file_path) as f:

182 for line in f:

183 if line.strip():

184 data.append(json.loads(line))

185 return data

186 except DataIOError:

187 raise

188 except Exception as e:

189 raise DataIOError(f"Failed to load JSON file: {e}") from e