xfstestsdb: Create a Table class

This class is designed to make it easy to print out data formatted as a
table. It was inspired by the prettytable Python module but is tailored
for how I intend to use it, including support for color schemes.

I use two different character sets for drawing borders. If color is
disabled, then I use basic ascii "+" and "-" characters. However, if
colors are enabled then I use characters from the Unicode Box Drawing
block to draw a table with rounded corners.

Finally, each table cell is formatted by calling the "do_format_cell()"
function. Table subclasses can override this function to implement their
own coloring schemes.

Signed-off-by: Anna Schumaker <anna@nowheycreamery.com>
This commit is contained in:
Anna Schumaker 2023-02-02 15:47:59 -05:00
parent 6529eb61f8
commit 456cef1090
2 changed files with 272 additions and 0 deletions

167
tests/test_table.py Normal file
View File

@ -0,0 +1,167 @@
# Copyright 2023 (c) Anna Schumaker.
"""Tests our table printing code."""
import io
import unittest
import unittest.mock
import xfstestsdb.table
class TestTableBorders(unittest.TestCase):
"""Test the Table Borders character set."""
def test_ascii_borders(self):
"""Test ascii borders."""
self.assertDictEqual(xfstestsdb.table.get_borders(True),
xfstestsdb.table.ASCII_BORDERS)
self.assertDictEqual(xfstestsdb.table.ASCII_BORDERS,
{"horizontal": "-",
"vertical": "|",
"top-left": "+-",
"top-center": "-+-",
"top-right": "-+",
"mid-left": "+-",
"mid-center": "-+-",
"mid-right": "-+",
"bot-left": "+-",
"bot-center": "-+-",
"bot-right": "-+"})
def test_unicode_borders(self):
"""Test unicode borders."""
self.assertDictEqual(xfstestsdb.table.get_borders(False),
xfstestsdb.table.UNICODE_BORDERS)
self.assertDictEqual(xfstestsdb.table.UNICODE_BORDERS,
{"horizontal": "",
"vertical": "",
"top-left": "╭─",
"top-center": "─┬─",
"top-right": "─╮",
"mid-left": "├─",
"mid-center": "─┼─",
"mid-right": "─┤",
"bot-left": "╰─",
"bot-center": "─┴─",
"bot-right": "─╯"})
class TestTable(unittest.TestCase):
"""Test the Table class."""
def test_align_funcs(self):
"""Test the alignment dictionary."""
self.assertDictEqual(xfstestsdb.table.ALIGN,
{"l": str.ljust, "c": str.center, "r": str.rjust})
def test_init(self):
"""Test initializing a table."""
tbl = xfstestsdb.table.Table(["Test", "Columns"])
self.assertIsInstance(tbl.colors, xfstestsdb.colors.ColorScheme)
self.assertDictEqual(tbl.borders, xfstestsdb.table.ASCII_BORDERS)
self.assertListEqual(tbl.columns, ["Test", "Columns"])
self.assertListEqual(tbl.align, ["l", "l"])
self.assertListEqual(tbl.maxlens, [4, 7])
self.assertListEqual(tbl.rows, [])
def test_init_align(self):
"""Test initializing a table with an alignment value."""
tbl = xfstestsdb.table.Table(["Test", "Some", "Columns"], ["c", "r"])
self.assertListEqual(tbl.align, ["c", "r", "l"])
tbl = xfstestsdb.table.Table(["Test", "Align"], ["c", "r", "r", "l"])
self.assertListEqual(tbl.align, ["c", "r"])
def test_add_column(self):
"""Test adding a single column to the table."""
tbl = xfstestsdb.table.Table(["Test"])
tbl.add_column("Add")
self.assertListEqual(tbl.columns, ["Test", "Add"])
self.assertListEqual(tbl.maxlens, [4, 3])
self.assertListEqual(tbl.align, ["l", "l"])
tbl.add_row(1, 2)
tbl.add_column("Column", align="c")
self.assertListEqual(tbl.columns, ["Test", "Add", "Column"])
self.assertListEqual(tbl.maxlens, [4, 3, 6])
self.assertListEqual(tbl.align, ["l", "l", "c"])
self.assertListEqual(tbl.rows, [["1", "2", ""]])
def test_add_row(self):
"""Test adding a row to a table."""
tbl = xfstestsdb.table.Table(["Test", "Columns"])
tbl.add_row("Row", 1)
self.assertListEqual(tbl.rows, [["Row", "1"]])
self.assertListEqual(tbl.maxlens, [4, 7])
tbl.add_row("Longer", 123456789)
self.assertListEqual(tbl.rows, [["Row", "1"],
["Longer", "123456789"]])
self.assertListEqual(tbl.maxlens, [6, 9])
tbl.add_row("Too", "Many", "Columns")
self.assertListEqual(tbl.rows, [["Row", "1"],
["Longer", "123456789"],
["Too", "Many"]])
self.assertListEqual(tbl.maxlens, [6, 9])
tbl.add_row("Too Short")
self.assertListEqual(tbl.rows, [["Row", "1"],
["Longer", "123456789"],
["Too", "Many"],
["Too Short", ""]])
tbl.add_row("None", None)
self.assertListEqual(tbl.rows, [["Row", "1"],
["Longer", "123456789"],
["Too", "Many"],
["Too Short", ""],
["None", ""]])
tbl.add_row()
self.assertListEqual(tbl.rows, [["Row", "1"],
["Longer", "123456789"],
["Too", "Many"],
["Too Short", ""],
["None", ""],
["", ""]])
@unittest.mock.patch("sys.stdout", new_callable=io.StringIO)
def test_print_empty(self, mock_stdout: io.StringIO):
"""Test printing an empty table."""
print(xfstestsdb.table.Table(["Test", "Columns"]))
self.assertEqual(mock_stdout.getvalue(), "\n")
@unittest.mock.patch("sys.stdout", new_callable=io.StringIO)
def test_print_default_align_unicode(self, mock_stdout: io.StringIO):
"""Test printing a table with default alignment and unicode borders."""
tbl = xfstestsdb.table.Table(["Test", "Table", "Columns"],
color_scheme="dark")
tbl.add_row(1, 2, 3)
print()
print(tbl)
self.assertEqual(mock_stdout.getvalue(),
"""
Test Table Columns
1 2 3
""")
@unittest.mock.patch("sys.stdout", new_callable=io.StringIO)
def test_print_align_ascii(self, mock_stdout: io.StringIO):
"""Test printing with custom alignment and ascii borders."""
tbl = xfstestsdb.table.Table(["Left", "Center", "Right"],
["l", "c", "r"], color_scheme="none")
tbl.add_row("abcdefg", "hijklmnopq", "rstuvwxyz")
tbl.add_row(123, 4567, 890)
print()
print(tbl)
self.assertEqual(mock_stdout.getvalue(),
"""
+---------+------------+-----------+
| Left | Center | Right |
+---------+------------+-----------+
| abcdefg | hijklmnopq | rstuvwxyz |
| 123 | 4567 | 890 |
+---------+------------+-----------+
""")

105
xfstestsdb/table.py Normal file
View File

@ -0,0 +1,105 @@
# Copyright 2023 (c) Anna Schumaker.
"""Custom table printing code, inspired by prettytable."""
import typing
from . import colors
ALIGN = {"l": str.ljust, "c": str.center, "r": str.rjust}
ASCII_BORDERS = {"horizontal": "-", "vertical": "|",
"top-left": "+-", "top-center": "-+-", "top-right": "-+",
"mid-left": "+-", "mid-center": "-+-", "mid-right": "-+",
"bot-left": "+-", "bot-center": "-+-", "bot-right": "-+"}
UNICODE_BORDERS = {"horizontal": "", "vertical": "",
"top-left": "╭─", "top-center": "─┬─", "top-right": "─╮",
"mid-left": "├─", "mid-center": "─┼─", "mid-right": "─┤",
"bot-left": "╰─", "bot-center": "─┴─", "bot-right": "─╯"}
def get_borders(ascii: bool) -> dict:
"""Get the borders to use for this table."""
return ASCII_BORDERS if ascii else UNICODE_BORDERS
class Table:
"""Organizes data into a nice-looking table."""
def __init__(self, columns: list[str],
align: list[str] | None = None,
color_scheme: str = "none") -> None:
"""Initialize a Table object."""
self.colors = colors.get_colors(color_scheme)
self.borders = get_borders(self.colors.name == "none")
self.columns = columns
self.maxlens = [len(col) for col in columns]
self.align = self.__pad_clamp_list([] if align is None else align, "l")
self.rows = []
def __gen_falign(self) -> typing.Generator:
for (align, maxlen) in zip(self.align, self.maxlens):
match align:
case "l": yield f"<{maxlen}"
case "c": yield f"^{maxlen}"
case "r": yield f">{maxlen}"
def __pad_clamp_list(self, list: typing.Iterable, pad: str) -> list[str]:
new = ["" if i is None else str(i) for i in list][:len(self.columns)]
return new + [pad] * (len(self.columns) - len(new))
def __format_border(self, type: str) -> str:
[left, center, right] = [self.borders[f"{type}-{dir}"]
for dir in ("left", "center", "right")]
cols = [self.borders["horizontal"] * len for len in self.maxlens]
return self.colors.format(f"{left}{center.join(cols)}{right}",
color="table-border", bgcolor="table-bg")
def __format_line(self, cells: list[str]) -> str:
border = self.colors.format(self.borders["vertical"],
color="table-border", bgcolor="table-bg")
return f"{border}{border.join(cells)}{border}"
def __format_headers(self) -> str:
formatted = [Table.do_format_cell(self, 0, f" {col:{align}} ",
"fg-header", bold=True)
for col, align in zip(self.columns, self.__gen_falign())]
return self.__format_line(formatted)
def __format_row(self, rownum: int, row: list[str]) -> str:
formatted = [self.do_format_cell(rownum, f" {cell:{align}} ",
"fg-odd" if rownum % 2 else "fg-even")
for cell, align in zip(row, self.__gen_falign())]
return self.__format_line(formatted)
def __repr__(self) -> str:
"""Generate the string representation for the Table."""
if len(self.rows) == 0:
return ""
lines = [self.__format_border("top"),
self.__format_headers(),
self.__format_border("mid")]
for i, row in enumerate(self.rows):
lines.append(self.__format_row(i, row))
lines.append(self.__format_border("bot"))
return "\n".join(lines)
def add_column(self, name: str, align: str = "l") -> None:
"""Add a column to the table."""
self.columns.append(name)
self.maxlens.append(len(name))
self.align.append(align)
for row in self.rows:
row.append("")
def add_row(self, *args) -> None:
"""Add a row to the Table."""
row = self.__pad_clamp_list(args, "")
self.maxlens = [max(a, len(b)) for (a, b) in zip(self.maxlens, row)]
self.rows.append(row)
def do_format_cell(self, rownum: int, text: str, color: str,
bold: bool = False,
dark: bool | str = "fg-dark") -> str:
"""Format the text for this cell and return the resulting string."""
return self.colors.format(text, color=color, bgcolor="table-bg",
bold=bold, dark=dark)