Coverage for src / sql_tool / formatters / csv.py: 91%

1"""CSV formatter for QueryResult output (RFC 4180 compliant).""" 

2 

3from __future__ import annotations 

4 

5import csv 

6from io import StringIO 

7from typing import TYPE_CHECKING 

8 

9from sql_tool.formatters.base import registry 

10 

11if TYPE_CHECKING: 

12 from collections.abc import Iterator 

13 

14 from sql_tool.core.models import QueryResult 

15 

16 

17def _write_row(values: list[str]) -> str: 

18 buf = StringIO() 

19 writer = csv.writer(buf) 

20 writer.writerow(values) 

21 return buf.getvalue().rstrip("\r\n") 

22 

23 

24class CSVFormatter: 

25 def __init__(self, no_header: bool = False) -> None: 

26 self.no_header = no_header 

27 

28 def format(self, result: QueryResult) -> Iterator[str]: 

29 if not self.no_header: 

30 yield _write_row([col.name for col in result.columns]) 

31 

32 for row in result.rows: 

33 yield _write_row([str(v) if v is not None else "" for v in row]) 

34 

35 

36registry.register("csv", CSVFormatter)