Coverage for src / crump / csv_file.py: 98%

1"""CSV file reader and writer implementations.""" 

2 

3from __future__ import annotations 

4 

5import csv 

6from collections.abc import Iterator 

7from pathlib import Path 

8from typing import Any 

9 

10from .tabular_file import TabularFileReader, TabularFileWriter 

11 

12 

13class CsvFileReader(TabularFileReader): 

14 """CSV file reader implementation. 

15 

16 Wraps Python's csv.DictReader to provide a consistent interface 

17 for reading tabular data files. 

18 """ 

19 

20 def __init__(self, file_path: str | Path, encoding: str = "utf-8", errors: str = "strict"): 

21 """Initialize CSV file reader. 

22 

23 Args: 

24 file_path: Path to the CSV file 

25 encoding: Character encoding (default: utf-8) 

26 errors: How to handle encoding errors (default: strict, can be 'replace') 

27 """ 

28 super().__init__(file_path) 

29 self.encoding = encoding 

30 self.errors = errors 

31 self._file_handle: Any = None 

32 self._reader: csv.DictReader[str] | None = None 

33 

34 def __enter__(self) -> CsvFileReader: 

35 """Open the CSV file and create a DictReader. 

36 

37 Returns: 

38 Self for use in with statement 

39 """ 

40 self._file_handle = open(self.file_path, encoding=self.encoding, errors=self.errors) 

41 self._reader = csv.DictReader(self._file_handle) 

42 return self 

43 

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

45 """Close the CSV file. 

46 

47 Args: 

48 exc_type: Exception type if an error occurred 

49 exc_val: Exception value if an error occurred 

50 exc_tb: Exception traceback if an error occurred 

51 """ 

52 if self._file_handle: 

53 self._file_handle.close() 

54 self._file_handle = None 

55 self._reader = None 

56 

57 @property 

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

59 """Get column names from the CSV file. 

60 

61 Returns: 

62 List of column names 

63 

64 Raises: 

65 RuntimeError: If called outside of context manager 

66 """ 

67 if self._reader is None: 

68 raise RuntimeError("Reader must be used within a context manager (with statement)") 

69 return list(self._reader.fieldnames or []) 

70 

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

72 """Iterate through CSV rows as dictionaries. 

73 

74 Yields: 

75 Dictionary mapping column names to values for each row 

76 

77 Raises: 

78 RuntimeError: If called outside of context manager 

79 """ 

80 if self._reader is None: 

81 raise RuntimeError("Reader must be used within a context manager (with statement)") 

82 yield from self._reader 

83 

84 

85class CsvFileWriter(TabularFileWriter): 

86 """CSV file writer implementation. 

87 

88 Wraps Python's csv.writer to provide a consistent interface 

89 for writing tabular data files. 

90 """ 

91 

92 def __init__(self, file_path: str | Path, append: bool = False, encoding: str = "utf-8"): 

93 """Initialize CSV file writer. 

94 

95 Args: 

96 file_path: Path to the CSV file 

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

98 encoding: Character encoding (default: utf-8) 

99 """ 

100 super().__init__(file_path, append) 

101 self.encoding = encoding 

102 self._file_handle: Any = None 

103 self._writer: Any = None 

104 

105 def __enter__(self) -> CsvFileWriter: 

106 """Open the CSV file and create a writer. 

107 

108 Returns: 

109 Self for use in with statement 

110 """ 

111 mode = "a" if self.append and self.file_path.exists() else "w" 

112 self._file_handle = open(self.file_path, mode, newline="", encoding=self.encoding) 

113 self._writer = csv.writer(self._file_handle) 

114 return self 

115 

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

117 """Close the CSV file. 

118 

119 Args: 

120 exc_type: Exception type if an error occurred 

121 exc_val: Exception value if an error occurred 

122 exc_tb: Exception traceback if an error occurred 

123 """ 

124 if self._file_handle: 

125 self._file_handle.close() 

126 self._file_handle = None 

127 self._writer = None 

128 

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

130 """Write a single row to the CSV file. 

131 

132 Args: 

133 row: List of values to write 

134 

135 Raises: 

136 RuntimeError: If called outside of context manager 

137 """ 

138 if self._writer is None: 

139 raise RuntimeError("Writer must be used within a context manager (with statement)") 

140 self._writer.writerow(row)