106 lines
4.5 KiB
Python
106 lines
4.5 KiB
Python
# 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)
|