Coverage for src / sql_tool / formatters / json.py: 90%

1"""JSON formatter for QueryResult output.""" 

2 

3from __future__ import annotations 

4 

5import json 

6from typing import TYPE_CHECKING, Any 

7 

8from sql_tool.formatters.base import registry 

9 

10if TYPE_CHECKING: 

11 from collections.abc import Iterator 

12 

13 from sql_tool.core.models import QueryResult 

14 

15 

16def _serialize_value(val: Any) -> Any: 

17 if isinstance(val, (int, float, str, bool, type(None))): 

18 return val 

19 return str(val) 

20 

21 

22class JSONFormatter: 

23 def __init__(self, compact: bool = False) -> None: 

24 self.compact = compact 

25 

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

27 rows_as_dicts = [ 

28 { 

29 col.name: _serialize_value(val) 

30 for col, val in zip(result.columns, row, strict=True) 

31 } 

32 for row in result.rows 

33 ] 

34 

35 if self.compact: 

36 yield json.dumps(rows_as_dicts, default=str) 

37 else: 

38 yield json.dumps(rows_as_dicts, indent=2, default=str) 

39 

40 

41registry.register("json", JSONFormatter)