Coverage for gws-app/gws/plugin/csv_helper/__init__.py: 84%

1"""Common csv writer helper.""" 

2 

3from typing import BinaryIO 

4 

5import decimal 

6import datetime 

7 

8import gws 

9import gws.lib.intl 

10 

11gws.ext.new.helper('csv') 

12 

13 

14class FormatConfig(gws.Config): 

15 """CSV format settings""" 

16 

17 delimiter: str = ',' 

18 """Field delimiter.""" 

19 encoding: str = 'utf8' 

20 """Text encoding.""" 

21 formulaHack: bool = True 

22 """Prepend numeric strings with an equals sign.""" 

23 quote: str = '"' 

24 """Quote character.""" 

25 quoteAll: bool = False 

26 """Quote all fields.""" 

27 rowDelimiter: str = 'LF' 

28 """Row delimiter.""" 

29 

30 

31class Config(gws.Config): 

32 """CSV helper.""" 

33 

34 format: FormatConfig 

35 """CSV format settings.""" 

36 

37 

38class Format(gws.Data): 

39 delimiter: str 

40 encoding: str 

41 formulaHack: bool 

42 quote: str 

43 quoteAll: bool 

44 rowDelimiter: str 

45 

46 

47class Object(gws.Node): 

48 format: Format 

49 

50 def configure(self) -> None: 

51 """Configure the CSV helper with format settings from config. 

52 

53 Sets up the format attribute with values from configuration or defaults. 

54 """ 

55 self.format = Format( 

56 delimiter=self.cfg('format.delimiter', default=','), 

57 encoding=self.cfg('format.encoding', default='utf8'), 

58 formulaHack=self.cfg('format.formulaHack', default=True), 

59 quote=self.cfg('format.quote', default='"'), 

60 quoteAll=self.cfg('format.quoteAll', default=False), 

61 rowDelimiter=self.cfg('format.rowDelimiter', default='LF').replace('CR', '\r').replace('LF', '\n'), 

62 ) 

63 

64 def writer(self, locale: gws.Locale, stream_to: BinaryIO = None) -> '_Writer': 

65 """Creates a new CSV Writer. 

66 

67 Args: 

68 locale: Locale to use for formatting values. 

69 stream_to: Optional binary stream to write to. If None, data is stored in memory. 

70 

71 Returns: 

72 A new _Writer instance configured with this helper's format settings. 

73 """ 

74 

75 return _Writer(self, locale, stream_to) 

76 

77 

78class _Writer: 

79 def __init__(self, helper: 'Object', locale: gws.Locale, stream_to: BinaryIO = None) -> None: 

80 """Initialize a CSV writer. 

81 

82 Args: 

83 helper: The CSV helper object containing format settings. 

84 locale: Locale to use for formatting values. 

85 stream_to: Optional binary stream to write to. If None, data is stored in memory. 

86 """ 

87 self.helper: Object = helper 

88 self.format = self.helper.format 

89 self.stream_to = stream_to 

90 self.eol = self.format.rowDelimiter.encode(self.format.encoding) 

91 

92 self.headers = [] 

93 self.str_rows = [] 

94 self.str_headers = '' 

95 

96 f = gws.lib.intl.formatters(locale) 

97 self.dateFormatter = f[0] 

98 self.timeFormatter = f[1] 

99 self.numberFormatter = f[2] 

100 

101 def write_headers(self, headers: list[str]) -> '_Writer': 

102 """Writes headers to the CSV output. 

103 

104 Args: 

105 headers: List of header column names. 

106 

107 Returns: 

108 Self for method chaining. 

109 """ 

110 

111 self.headers = headers 

112 self.str_headers = self.format.delimiter.join(self._quote(s) for s in headers) 

113 if self.stream_to: 

114 self.stream_to.write(self.str_headers.encode(self.format.encoding) + self.eol) 

115 return self 

116 

117 def write_row(self, row: list) -> '_Writer': 

118 """Writes a row of data to the CSV output. 

119 

120 Args: 

121 row: List of values to write as a single row. 

122 

123 Returns: 

124 Self for method chaining. 

125 """ 

126 

127 s = self.format.delimiter.join(self._format(v) for v in row) 

128 if self.stream_to: 

129 self.stream_to.write(s.encode(self.format.encoding) + self.eol) 

130 else: 

131 self.str_rows.append(s) 

132 return self 

133 

134 def write_dict(self, d: dict) -> '_Writer': 

135 """Writes a dict of data to the CSV output. 

136 

137 Args: 

138 d: Dictionary where keys are column names and values are the data. 

139 

140 Returns: 

141 Self for method chaining. 

142 """ 

143 

144 if not self.headers: 

145 self.write_headers(list(d.keys())) 

146 return self.write_row([d.get(h, '') for h in self.headers]) 

147 

148 def to_str(self) -> str: 

149 """Converts the headers and rows to a CSV string. 

150 

151 Returns: 

152 A string containing the complete CSV data. 

153 """ 

154 

155 rows = [] 

156 if self.headers: 

157 rows.append(self.str_headers) 

158 rows.extend(self.str_rows) 

159 return self.format.rowDelimiter.join(rows) 

160 

161 def to_bytes(self, encoding: str = None) -> bytes: 

162 """Converts the CSV data to a byte string. 

163 

164 Args: 

165 encoding: Optional encoding to use. If None, uses the format's encoding. 

166 

167 Returns: 

168 Byte string representation of the CSV data. 

169 """ 

170 

171 return self.to_str().encode(encoding or self.format.encoding, errors='replace') 

172 

173 def _format(self, val) -> str: 

174 """Format a value for CSV output according to its type. 

175 

176 Args: 

177 val: The value to format. 

178 

179 Returns: 

180 Formatted string representation of the value. 

181 """ 

182 if val is None: 

183 return self._quote('') 

184 

185 if isinstance(val, (float, decimal.Decimal)): 

186 s = self.numberFormatter.decimal(val) 

187 return self._quote(s) if self.format.quoteAll else s 

188 

189 if isinstance(val, int): 

190 s = str(val) 

191 return self._quote(s) if self.format.quoteAll else s 

192 

193 if isinstance(val, (datetime.datetime, datetime.date)): 

194 s = self.dateFormatter.short(val) 

195 return self._quote(s) 

196 

197 if isinstance(val, datetime.time): 

198 s = self.timeFormatter.short(val) 

199 return self._quote(s) 

200 

201 val = gws.u.to_str(val) 

202 

203 if val and val.isdigit() and self.format.formulaHack: 

204 val = '=' + self._quote(val) 

205 

206 return self._quote(val) 

207 

208 def _quote(self, val) -> str: 

209 """Quote a value according to CSV quoting rules. 

210 

211 Doubles any quote characters in the value and wraps the result in quotes. 

212 

213 Args: 

214 val: The value to quote. 

215 

216 Returns: 

217 Quoted string. 

218 """ 

219 q = self.format.quote 

220 s = gws.u.to_str(val).replace(q, q + q) 

221 return q + s + q