Coverage for src / sql_tool / formatters / table.py: 94%

1"""Rich table formatter for QueryResult output.""" 

2 

3from __future__ import annotations 

4 

5import shutil 

6from io import StringIO 

7from typing import TYPE_CHECKING 

8 

9from rich.console import Console 

10from rich.table import Table 

11 

12from sql_tool.formatters.base import registry 

13 

14if TYPE_CHECKING: 

15 from collections.abc import Iterator 

16 

17 from sql_tool.core.models import QueryResult 

18 

19_NO_RESULTS = "No results" 

20 

21 

22def _truncate(value: str, width: int) -> str: 

23 if len(value) <= width: 

24 return value 

25 return value[: width - 1] + "…" 

26 

27 

28class TableFormatter: 

29 def __init__(self, width: int = 40) -> None: 

30 self.width = width 

31 

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

33 if not result.rows: 

34 yield _NO_RESULTS 

35 return 

36 

37 table = Table(show_edge=True, pad_edge=True) 

38 for col in result.columns: 

39 table.add_column(col.name, no_wrap=True) 

40 

41 for row in result.rows: 

42 table.add_row( 

43 *(_truncate(str(v) if v is not None else "", self.width) for v in row) 

44 ) 

45 

46 buf = StringIO() 

47 term_width = shutil.get_terminal_size((120, 24)).columns 

48 console = Console(file=buf, force_terminal=True, width=term_width) 

49 console.print(table) 

50 yield buf.getvalue().rstrip("\n") 

51 

52 

53registry.register("table", TableFormatter)