Coverage for src/crump/tabular_file.py: 74%

1"""Abstract base classes for tabular file formats (CSV, Parquet, etc.)."""

2

3 from __future__ import annotations

4

5 from abc import ABC, abstractmethod

6 from collections.abc import Iterator

7 from pathlib import Path

8 from typing import Any

9

10 from .file_types import InputFileType, OutputFileType

11

12

13 class TabularFileReader(ABC):

14 """Abstract base class for reading tabular file formats.

15

16 Provides a common interface for reading different tabular file formats

17 like CSV and Parquet. Designed to work as a context manager and iterator,

18 similar to csv.DictReader.

19

20 Example usage:

21 with TabularFileReader(file_path) as reader:

22 print(f"Columns: {reader.fieldnames}")

23 for row in reader:

24 print(row) # row is a dict

25 """

26

27 def __init__(self, file_path: str | Path):

28 """Initialize the reader with a file path.

29

30 Args:

31 file_path: Path to the file to read

32 """

33 self.file_path = Path(file_path)

34

35 @abstractmethod

36 def __enter__(self) -> TabularFileReader:

37 """Enter context manager and prepare for reading.

38

39 Returns:

40 Self for use in with statement

41 """

42 pass

43

44 @abstractmethod

45 def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:

46 """Exit context manager and cleanup resources.

47

48 Args:

49 exc_type: Exception type if an error occurred

50 exc_val: Exception value if an error occurred

51 exc_tb: Exception traceback if an error occurred

52 """

53 pass

54

55 @property

56 @abstractmethod

57 def fieldnames(self) -> list[str]:

58 """Get column names from the file.

59

60 Returns:

61 List of column names

62 """

63 pass

64

65 @abstractmethod

66 def __iter__(self) -> Iterator[dict[str, Any]]:

67 """Iterate through rows as dictionaries.

68

69 Yields:

70 Dictionary mapping column names to values for each row

71 """

72 pass

73

74

75 class TabularFileWriter(ABC):

76 """Abstract base class for writing tabular file formats.

77

78 Provides a common interface for writing different tabular file formats

79 like CSV and Parquet. Designed to work as a context manager, similar to

80 csv.writer.

81

82 Example usage:

83 with TabularFileWriter(file_path, append=False) as writer:

84 writer.writerow(['col1', 'col2']) # header

85 writer.writerow(['val1', 'val2']) # data row

86 """

87

88 def __init__(self, file_path: str | Path, append: bool = False):

89 """Initialize the writer with a file path.

90

91 Args:

92 file_path: Path to the file to write

93 append: If True, append to existing file. If False, overwrite.

94 """

95 self.file_path = Path(file_path)

96 self.append = append

97

98 @abstractmethod

99 def __enter__(self) -> TabularFileWriter:

100 """Enter context manager and prepare for writing.

101

102 Returns:

103 Self for use in with statement

104 """

105 pass

106

107 @abstractmethod

108 def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:

109 """Exit context manager and cleanup resources.

110

111 Args:

112 exc_type: Exception type if an error occurred

113 exc_val: Exception value if an error occurred

114 exc_tb: Exception traceback if an error occurred

115 """

116 pass

117

118 @abstractmethod

119 def writerow(self, row: list[Any]) -> None:

120 """Write a single row to the file.

121

122 Args:

123 row: List of values to write (can be header or data row)

124 """

125 pass

126

127

128 def create_reader(

129 file_path: str | Path, file_format: InputFileType | str | None = None

130) -> TabularFileReader:

131 """Factory function to create appropriate reader based on file format.

132

133 Args:

134 file_path: Path to the file to read

135 file_format: File format (InputFileType enum, 'csv', or 'parquet').

136 If None, auto-detect from extension. Defaults to CSV for unknown extensions.

137

138 Returns:

139 TabularFileReader instance for the file format

140

141 Raises:

142 ValueError: If file format is not supported

143 """

144 from .csv_file import CsvFileReader

145 from .parquet_file import ParquetFileReader

146

147 path = Path(file_path)

148

149 # Convert string to enum if needed, or auto-detect from extension

150 if file_format is None:

151 # Auto-detect with CSV fallback for unknown extensions (like .cdf)

152 try:

153 format_enum = InputFileType.from_path(str(path))

154 except ValueError:

155 format_enum = InputFileType.CSV

156 elif isinstance(file_format, str):

157 # Convert string to enum

158 format_str = file_format.lower()

159 if format_str == "csv":

160 format_enum = InputFileType.CSV

161 elif format_str == "parquet":

162 format_enum = InputFileType.PARQUET

163 else:

164 raise ValueError(f"Unsupported file format: {file_format}")

165 else:

166 format_enum = file_format

167

168 # Create appropriate reader (CDF not supported for reading tabular data)

169 if format_enum == InputFileType.CSV:

170 return CsvFileReader(path)

171 elif format_enum == InputFileType.PARQUET:

172 return ParquetFileReader(path)

173 else:

174 raise ValueError(f"Cannot read {format_enum.value} files as tabular data")

175

176

177 def create_writer(

178 file_path: str | Path, file_format: OutputFileType | str | None = None, append: bool = False

179) -> TabularFileWriter:

180 """Factory function to create appropriate writer based on file format.

181

182 Args:

183 file_path: Path to the file to write

184 file_format: File format (OutputFileType enum, 'csv', or 'parquet').

185 If None, auto-detect from extension.

186 append: If True, append to existing file. If False, overwrite.

187

188 Returns:

189 TabularFileWriter instance for the file format

190

191 Raises:

192 ValueError: If file format is not supported or cannot be detected

193 """

194 from .csv_file import CsvFileWriter

195 from .parquet_file import ParquetFileWriter

196

197 path = Path(file_path)

198

199 # Convert string to enum if needed, or auto-detect from extension

200 if file_format is None:

201 format_enum = OutputFileType.from_path(str(path))

202 elif isinstance(file_format, str):

203 # Convert string to enum

204 format_str = file_format.lower()

205 if format_str == "csv":

206 format_enum = OutputFileType.CSV

207 elif format_str == "parquet":

208 format_enum = OutputFileType.PARQUET

209 else:

210 raise ValueError(f"Unsupported file format: {file_format}")

211 else:

212 format_enum = file_format

213

214 # Create appropriate writer

215 if format_enum == OutputFileType.CSV:

216 return CsvFileWriter(path, append=append)

217 elif format_enum == OutputFileType.PARQUET:

218 return ParquetFileWriter(path, append=append)

219 else:

220 raise ValueError(f"Unsupported file format: {format_enum.value}")