ai-content-maker/.venv/Lib/site-packages/wasabi/tables.py

159 lines
6.6 KiB
Python

import os
from itertools import zip_longest
from typing import Collection, Dict, Iterable, List, Optional, Sequence, Union
from typing import cast
from .compat import Literal
from .util import COLORS
from .util import color as _color
from .util import supports_ansi
ALIGN_MAP = {"l": "<", "r": ">", "c": "^"}
def table(
# fmt: off
data: Union[Collection, Dict],
header: Optional[Iterable] = None,
footer: Optional[Iterable] = None,
divider: bool = False,
widths: Union[Iterable[int], Literal["auto"]] = "auto",
max_col: int = 30,
spacing: int = 3,
aligns: Optional[Union[Iterable[Literal["r", "c", "l"]], Literal["r", "c", "l"]]] = None,
multiline: bool = False,
env_prefix: str = "WASABI",
color_values: Optional[Dict] = None,
fg_colors: Optional[Iterable] = None,
bg_colors: Optional[Iterable] = None,
# fmt: on
) -> str:
"""Format tabular data.
data (Union[Collection, Dict]): The data to render. Either a list of lists (one per
row) or a dict for two-column tables.
header (Optional[Iterable]): Optional header columns.
footer (Optional[Iterable]): Optional footer columns.
divider (bool): Show a divider line between header/footer and body.
widths (Union[Iterable[int], Literal['auto']]): Column widths in order. If "auto", widths
will be calculated automatically based on the largest value.
max_col (int): Maximum column width.
spacing (int): Spacing between columns, in spaces.
aligns (Optional[Union[Iterable[str], str]]): Optional column alignments
in order. 'l' (left, default), 'r' (right) or 'c' (center). If a string,
value is used for all columns.
multiline (bool): If a cell value is a list of a tuple, render it on
multiple lines, with one value per line.
env_prefix (str): Prefix for environment variables, e.g.
WASABI_LOG_FRIENDLY.
color_values (Optional[Dict]): Optional color values to add or overwrite, name mapped to value.
fg_colors (Optional[Iterable]): Optional foreground colors, one per column. None can be specified
for individual columns to retain the default foreground color.
bg_colors (Optional[Iterable]): Optional background colors, one per column. None can be specified
for individual columns to retain the default background color.
RETURNS (str): The formatted table.
"""
if fg_colors is not None or bg_colors is not None:
colors = dict(COLORS)
if color_values is not None:
colors.update(color_values)
if fg_colors is not None:
fg_colors = [colors.get(fg_color, fg_color) for fg_color in fg_colors]
if bg_colors is not None:
bg_colors = [colors.get(bg_color, bg_color) for bg_color in bg_colors]
if isinstance(data, dict):
data = list(data.items())
if multiline:
zipped_data = []
for i, item in enumerate(data):
vals = [v if isinstance(v, (list, tuple)) else [v] for v in item]
zipped_data.extend(list(zip_longest(*vals, fillvalue="")))
if i < len(data) - 1:
zipped_data.append(tuple(["" for i in item]))
data = zipped_data
if widths == "auto":
widths = _get_max_widths(data, header, footer, max_col)
settings = {
"widths": widths,
"spacing": spacing,
"aligns": aligns,
"env_prefix": env_prefix,
"fg_colors": fg_colors,
"bg_colors": bg_colors,
}
divider_row = row(["-" * width for width in widths], **settings) # type: ignore
rows = []
if header:
rows.append(row(header, **settings)) # type: ignore
if divider:
rows.append(divider_row)
for i, item in enumerate(data):
rows.append(row(item, **settings)) # type: ignore
if footer:
if divider:
rows.append(divider_row)
rows.append(row(footer, **settings)) # type: ignore
return "\n{}\n".format("\n".join(rows))
def row(
data: Collection,
widths: Union[Sequence[int], int, Literal["auto"]] = "auto",
spacing: int = 3,
aligns: Optional[Union[Sequence[Literal["r", "c", "l"]], str]] = None,
env_prefix: str = "WASABI",
fg_colors: Optional[Sequence] = None,
bg_colors: Optional[Sequence] = None,
) -> str:
"""Format data as a table row.
data (Collection): The individual columns to format.
widths (Union[Sequence[int], int, Literal['auto']]): Column widths, either one integer for all
columns or an iterable of values. If "auto", widths will be calculated
automatically based on the largest value.
spacing (int): Spacing between columns, in spaces.
aligns (Optional[Union[Sequence[Literal['r', 'c', 'l']], str]]): Optional column
alignments in order. 'l' (left, default), 'r' (right) or 'c' (center).
If a string, value is used for all columns.
env_prefix (str): Prefix for environment variables, e.g.
WASABI_LOG_FRIENDLY.
fg_colors (Optional[Sequence]): Optional foreground colors for the columns, in order. None can be
specified for individual columns to retain the default foreground color.
bg_colors (Optional[Sequence]): Optional background colors for the columns, in order. None can be
specified for individual columns to retain the default background color.
RETURNS (str): The formatted row.
"""
env_log_friendly = os.getenv("{}_LOG_FRIENDLY".format(env_prefix), False)
show_colors = (
supports_ansi()
and not env_log_friendly
and (fg_colors is not None or bg_colors is not None)
)
cols: List[str] = []
_aligns = (
[aligns for _ in data] if isinstance(aligns, str) else cast(List[str], aligns)
)
if not hasattr(widths, "__iter__") and widths != "auto": # single number
widths = cast(List[int], [widths for _ in range(len(data))])
for i, col in enumerate(data):
align = ALIGN_MAP.get(_aligns[i] if _aligns and i < len(_aligns) else "l")
col_width = len(col) if widths == "auto" else cast(List[int], widths)[i]
tpl = "{:%s%d}" % (align, col_width)
col = tpl.format(str(col))
if show_colors:
fg = fg_colors[i] if fg_colors is not None else None
bg = bg_colors[i] if bg_colors is not None else None
col = _color(col, fg=fg, bg=bg)
cols.append(col)
return (" " * spacing).join(cols)
def _get_max_widths(data, header, footer, max_col):
all_data = list(data)
if header:
all_data.append(header)
if footer:
all_data.append(footer)
widths = [[len(str(col)) for col in item] for item in all_data]
return [min(max(w), max_col) for w in list(zip(*widths))]