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:
parent
6529eb61f8
commit
456cef1090
|
@ -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 |
|
||||
+---------+------------+-----------+
|
||||
""")
|
|
@ -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)
|
Loading…
Reference in New Issue