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

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

2

3 from __future__ import annotations

4

5 import csv

6 from collections.abc import Iterator

7 from pathlib import Path

8 from typing import Any

9

10 from .tabular_file import TabularFileReader, TabularFileWriter

11

12

13 class 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

85 class 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)