# Copyright Kevin Deldycke <kevin@deldycke.com> and contributors.
#
# This program is Free Software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
"""Collection of table rendering utilities."""
from __future__ import annotations
import csv
import json
from enum import Enum
from functools import partial
from gettext import gettext as _
from io import StringIO
import click
import tabulate
from boltons.strutils import strip_ansi
from tabulate import DataRow
from tabulate import TableFormat as TabulateTableFormat
from . import EnumChoice, echo
from .parameters import ExtraOption
TYPE_CHECKING = False
if TYPE_CHECKING:
from collections.abc import Callable, Sequence
tabulate.MIN_PADDING = 0
"""Neutralize spurious double-spacing in table rendering."""
tabulate._table_formats.update( # type: ignore[attr-defined]
{
"aligned": TabulateTableFormat(
lineabove=None,
linebelowheader=None,
linebetweenrows=None,
linebelow=None,
headerrow=DataRow("", " ", ""),
datarow=DataRow("", " ", ""),
padding=0,
with_header_hide=None,
),
},
)
"""Custom table formats registered with tabulate.
``aligned``
A minimal format with single-space column separators and no borders or decorations.
Similar to ``plain`` but more compact (single space instead of double space between
columns). Useful for bar plugin output or other contexts requiring minimal formatting.
"""
# Patch the ``github`` format to support alignment colons in separator rows, matching
# the ``pipe`` format. Backport of https://github.com/astanin/python-tabulate/pull/410
_fmts = tabulate._table_formats # type: ignore[attr-defined]
_fmts["github"] = _fmts["pipe"]
MARKUP_FORMATS = frozenset(
{
TableFormat.ASCIIDOC,
TableFormat.CSV,
TableFormat.CSV_EXCEL,
TableFormat.CSV_EXCEL_TAB,
TableFormat.CSV_UNIX,
TableFormat.GITHUB,
TableFormat.HJSON,
TableFormat.HTML,
TableFormat.JIRA,
TableFormat.JSON,
TableFormat.JSON5,
TableFormat.JSONC,
TableFormat.LATEX,
TableFormat.LATEX_BOOKTABS,
TableFormat.LATEX_LONGTABLE,
TableFormat.LATEX_RAW,
TableFormat.MEDIAWIKI,
TableFormat.MOINMOIN,
TableFormat.ORGTBL,
TableFormat.PIPE,
TableFormat.RST,
TableFormat.TEXTILE,
TableFormat.TOML,
TableFormat.TSV,
TableFormat.UNSAFEHTML,
TableFormat.XML,
TableFormat.YAML,
TableFormat.YOUTRACK,
},
)
"""Subset of table formats that are considered as markup rendering."""
DEFAULT_FORMAT = TableFormat.ROUNDED_OUTLINE
"""Default table format, if none is specified."""
RECORD_KEY = "record"
"""Key used for each record in structured formats that require named containers
(TOML ``[[record]]``, XML ``<record>``)."""
XML_ROOT_KEY = "records"
"""Root element name for XML table output."""
SERIALIZATION_FORMATS = frozenset(
{
TableFormat.HJSON,
TableFormat.JSON,
TableFormat.JSON5,
TableFormat.JSONC,
TableFormat.TOML,
TableFormat.XML,
TableFormat.YAML,
},
)
"""Structured serialization formats whose renderers escape raw ESC bytes, making
post-render ``strip_ansi()`` ineffective."""
def _get_csv_dialect(table_format: TableFormat | None = None) -> str:
"""Extract, validate and normalize CSV dialect ID from format.
Defaults to ``excel`` rendering, like in Python's csv module.
"""
dialect = "excel"
# Extract dialect ID from table format, if any.
if table_format:
format_id = table_format.value
assert format_id.startswith("csv")
parts = format_id.split("-", 1)
assert parts[0] == "csv"
if len(parts) > 1:
dialect = parts[1]
csv.get_dialect(dialect) # Validate dialect.
return dialect
def _render_csv(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
table_format: TableFormat | None = None,
**kwargs,
) -> str:
"""Render a table in CSV format.
.. note::
StringIO is used to capture CSV output in memory. `Hard-coded to default to
UTF-8 <https://github.com/python/cpython/blob/9291095/Lib/_pyio.py#L2652>`_.
"""
defaults = {"dialect": _get_csv_dialect(table_format)}
defaults.update(kwargs)
with StringIO(newline="") as output:
writer = csv.writer(output, **defaults) # type: ignore[arg-type]
if headers:
writer.writerow(headers)
writer.writerows(table_data)
return output.getvalue()
def _rows_as_dicts(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
) -> list[dict[str, str | None]] | list[list[str | None]]:
"""Convert table data to a list of dicts keyed by headers.
Falls back to a list of lists when no headers are provided.
"""
if headers:
return [{str(k): v for k, v in zip(headers, row)} for row in table_data]
return [list(row) for row in table_data]
def _render_json(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
**kwargs,
) -> str:
"""Render a table as JSON."""
data = _rows_as_dicts(table_data, headers)
defaults: dict = {"ensure_ascii": False, "indent": 2}
defaults.update(kwargs)
return json.dumps(data, **defaults) + "\n"
def _render_yaml(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
**kwargs,
) -> str:
"""Render a table as YAML.
Requires the ``pyyaml`` package (installable via the ``[yaml]`` extra).
"""
try:
import yaml
except ImportError as exc:
msg = (
"PyYAML is required for YAML table output."
" Install it with: pip install click-extra[yaml]"
)
raise ImportError(msg) from exc
data = _rows_as_dicts(table_data, headers)
defaults: dict = {"allow_unicode": True, "default_flow_style": False}
defaults.update(kwargs)
return str(yaml.dump(data, **defaults))
def _render_toml(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
**kwargs,
) -> str:
"""Render a table as TOML using array-of-tables syntax.
``None`` values are omitted (TOML has no null type). Requires the ``tomlkit``
package (installable via the ``[toml]`` extra).
"""
try:
import tomlkit
except ImportError as exc:
msg = (
"tomlkit is required for TOML table output."
" Install it with: pip install click-extra[toml]"
)
raise ImportError(msg) from exc
aot = tomlkit.aot()
for row in table_data:
t = tomlkit.table()
if headers:
for key, value in zip(headers, row):
if value is not None and key is not None:
t.add(key, value)
else:
for i, value in enumerate(row):
if value is not None:
t.add(str(i), value)
aot.append(t)
doc = tomlkit.document()
doc.add(RECORD_KEY, aot)
return tomlkit.dumps(doc)
def _render_hjson(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
**kwargs,
) -> str:
"""Render a table as HJSON.
Requires the ``hjson`` package (installable via the ``[hjson]`` extra).
"""
try:
import hjson
except ImportError as exc:
msg = (
"hjson is required for HJSON table output."
" Install it with: pip install click-extra[hjson]"
)
raise ImportError(msg) from exc
data = _rows_as_dicts(table_data, headers)
defaults: dict = {"ensure_ascii": False}
defaults.update(kwargs)
return str(hjson.dumps(data, **defaults)) + "\n"
def _render_xml(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
**kwargs,
) -> str:
"""Render a table as XML.
``None`` values are omitted. Requires the ``xmltodict`` package (installable
via the ``[xml]`` extra).
"""
try:
import xmltodict
except ImportError as exc:
msg = (
"xmltodict is required for XML table output."
" Install it with: pip install click-extra[xml]"
)
raise ImportError(msg) from exc
def _xml_safe_name(name: str) -> str:
"""Replace characters invalid in XML element names."""
safe = "".join(c if c.isalnum() or c in "_.-" else "_" for c in name)
return safe.lstrip("0123456789.-") or "_"
if headers:
records = [
{
_xml_safe_name(k): v
for k, v in zip(headers, row)
if v is not None and k is not None
}
for row in table_data
]
else:
records = [
{str(i): v for i, v in enumerate(row) if v is not None}
for row in table_data
]
defaults: dict = {
"pretty": True,
"encoding": "unicode",
"full_document": False,
}
defaults.update(kwargs)
result: str = xmltodict.unparse({XML_ROOT_KEY: {RECORD_KEY: records}}, **defaults)
return result + "\n"
def _render_vertical(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
sep_character: str = "*",
sep_length: int = 27,
**kwargs,
) -> str:
"""Re-implements ``cli-helpers``'s vertical table layout.
.. note::
See `cli-helpers source code for reference
<https://github.com/dbcli/cli_helpers/blob/v2.7.0/cli_helpers/tabular_output/vertical_table_adapter.py>`_.
.. caution::
This layout is `hard-coded to 27 asterisks to separate rows
<https://github.com/dbcli/cli_helpers/blob/c34ae9f/cli_helpers/tabular_output/vertical_table_adapter.py#L34>`_,
as in the original implementation.
"""
if not headers:
headers = []
# Calculate header lengths and pad headers in one pass.
header_length = [0 if h is None else len(h) for h in headers]
max_length = max(header_length) if header_length else 0
padded_headers = ["" if h is None else h.ljust(max_length) for h in headers]
table_lines = []
sep = sep_character * sep_length
for index, row in enumerate(table_data):
table_lines.append(f"{sep}[ {index + 1}. row ]{sep}")
for cell_label, cell_value in zip(padded_headers, row):
# Like other formats, render None as an empty string.
cell_value = "" if cell_value is None else cell_value
table_lines.append(f"{cell_label} | {cell_value}")
return "\n".join(table_lines)
def _render_tabulate(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
table_format: TableFormat | None = None,
**kwargs,
) -> str:
"""Render a table with ``tabulate``.
Default format is ``TableFormat.ROUNDED_OUTLINE``.
"""
if not headers:
headers = ()
if not table_format:
table_format = DEFAULT_FORMAT
defaults = {
"disable_numparse": True,
"numalign": None,
# tabulate()'s format ID uses underscores instead of dashes.
"tablefmt": table_format.value.replace("-", "_"),
}
defaults.update(kwargs)
return tabulate.tabulate(table_data, headers, **defaults) # type: ignore[arg-type]
def _select_table_funcs(
table_format: TableFormat | None = None,
) -> tuple[Callable[..., str], Callable[[str], None]]:
"""Returns the rendering and print functions for the given ``table_format``.
For all formats other than CSV, we relying on Click's ``echo()`` as the print
function, to benefit from its sensitivity to global colorization settings. Thanks
to this the ``--color``/``--no-color`` option is automatically supported.
For CSV formats we returns the Python standard ``print()`` function, to preserve
line terminations and avoid extra line returns.
"""
print_func = echo
match table_format:
case (
TableFormat.CSV
| TableFormat.CSV_EXCEL
| TableFormat.CSV_EXCEL_TAB
| TableFormat.CSV_UNIX
):
print_func = partial(print, end="")
return partial(_render_csv, table_format=table_format), print_func
case TableFormat.HJSON:
print_func = partial(print, end="")
return _render_hjson, print_func
case TableFormat.JSON | TableFormat.JSON5 | TableFormat.JSONC:
print_func = partial(print, end="")
return _render_json, print_func
case TableFormat.TOML:
print_func = partial(print, end="")
return _render_toml, print_func
case TableFormat.XML:
print_func = partial(print, end="")
return _render_xml, print_func
case TableFormat.YAML:
print_func = partial(print, end="")
return _render_yaml, print_func
case TableFormat.VERTICAL:
return _render_vertical, print_func
case _:
return partial(_render_tabulate, table_format=table_format), print_func
[docs]
def render_table(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
table_format: TableFormat | None = None,
**kwargs,
) -> str:
"""Render a table and return it as a string."""
render_func, _ = _select_table_funcs(table_format)
return render_func(table_data, headers, **kwargs)
def _strip_ansi_cells(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
) -> tuple[list[list[str | None]], Sequence[str | None] | None]:
"""Strip ANSI escape codes from all string cells and headers."""
cleaned_data: list[list[str | None]] = [
[strip_ansi(v) if isinstance(v, str) else v for v in row] for row in table_data
]
cleaned_headers = (
[strip_ansi(h) if isinstance(h, str) else h for h in headers]
if headers
else headers
)
return cleaned_data, cleaned_headers
[docs]
def print_table(
table_data: Sequence[Sequence[str | None]],
headers: Sequence[str | None] | None = None,
table_format: TableFormat | None = None,
**kwargs,
) -> None:
"""Render a table and print it to the console.
For markup formats, ANSI color codes are stripped from cell values before
rendering unless ``--color`` is explicitly set.
"""
# Strip ANSI codes from cell data before rendering for markup formats.
# Pre-render stripping is necessary because some renderers (JSON, YAML) escape
# raw ESC bytes, making post-render strip_ansi() ineffective.
if table_format and table_format.is_markup:
ctx = click.get_current_context(silent=True)
# Only preserve ANSI codes when --color was explicitly passed on the
# command line. The default True from ColorOption should not prevent
# stripping.
color_explicit = False
if ctx is not None:
source = ctx.get_parameter_source("color")
color_explicit = (
ctx.color is True and source == click.core.ParameterSource.COMMANDLINE
)
if not color_explicit:
table_data, headers = _strip_ansi_cells(table_data, headers)
render_func, print_func = _select_table_funcs(table_format)
print_func(render_func(table_data, headers, **kwargs))