From 0860e419e49c294e5fa8593e735b8665d54d8b05 Mon Sep 17 00:00:00 2001 From: Saugat Pachhai Date: Mon, 26 Apr 2021 17:39:23 +0545 Subject: [PATCH] ui: tabular data (#5852) * introduce tabular data structure * ignore ellipsis --- .coveragerc | 1 + dvc/compare.py | 169 ++++++++++++++++++++++++++++++ dvc/ui/table.py | 4 +- tests/unit/test_tabular_data.py | 178 ++++++++++++++++++++++++++++++++ 4 files changed, 350 insertions(+), 2 deletions(-) create mode 100644 dvc/compare.py create mode 100644 tests/unit/test_tabular_data.py diff --git a/.coveragerc b/.coveragerc index 17a064779..cca27a875 100644 --- a/.coveragerc +++ b/.coveragerc @@ -13,3 +13,4 @@ exclude_lines = raise NotImplementedError # pragma: no cover # pragma: no cover + \.\.\.$ diff --git a/dvc/compare.py b/dvc/compare.py new file mode 100644 index 000000000..fe2c71238 --- /dev/null +++ b/dvc/compare.py @@ -0,0 +1,169 @@ +from itertools import chain, repeat, zip_longest +from operator import itemgetter +from typing import ( + TYPE_CHECKING, + Any, + Dict, + ItemsView, + Iterable, + Iterator, + List, + MutableSequence, + Sequence, + Tuple, + Union, + overload, +) + +from funcy import reraise + +if TYPE_CHECKING: + from dvc.ui.table import CellT + + +class Column(List["CellT"]): + pass + + +def with_value(value, default): + return default if value is None else value + + +class TabularData(MutableSequence[Sequence["CellT"]]): + def __init__(self, columns: Sequence[str], fill_value: str = ""): + self._columns: Dict[str, Column] = {name: Column() for name in columns} + self._keys: List[str] = list(columns) + self._fill_value = fill_value + + @property + def columns(self) -> List[Column]: + return list(map(self.column, self.keys())) + + def column(self, name: str) -> Column: + return self._columns[name] + + def items(self) -> ItemsView[str, Column]: + projection = {k: self.column(k) for k in self.keys()} + return projection.items() + + def keys(self) -> List[str]: + return self._keys + + def _iter_col_row( + self, row: Sequence["CellT"] + ) -> Iterator[Tuple["CellT", Column]]: + for val, col in zip_longest(row, self.columns): + if col is None: + break + yield with_value(val, self._fill_value), col + + def append(self, value: Sequence["CellT"]) -> None: + for val, col in self._iter_col_row(value): + col.append(val) + + def extend(self, values: Iterable[Sequence["CellT"]]) -> None: + for row in values: + self.append(row) + + def insert(self, index: int, value: Sequence["CellT"]) -> None: + for val, col in self._iter_col_row(value): + col.insert(index, val) + + def __iter__(self) -> Iterator[List["CellT"]]: + return map(list, zip(*self.columns)) + + def __getattr__(self, item: str) -> Column: + with reraise(KeyError, AttributeError): + return self.column(item) + + def __getitem__(self, item: Union[int, slice]): + func = itemgetter(item) + it = map(func, self.columns) + if isinstance(item, slice): + it = map(list, zip(*it)) + return list(it) + + @overload + def __setitem__(self, item: int, value: Sequence["CellT"]) -> None: + ... + + @overload + def __setitem__( + self, item: slice, value: Iterable[Sequence["CellT"]] + ) -> None: + ... + + def __setitem__(self, item, value) -> None: + it = value + if isinstance(item, slice): + n = len(self.columns) + normalized_rows = ( + chain(val, repeat(self._fill_value, n - len(val))) + for val in value + ) + # we need to transpose those rows into columnar format + # as we work in terms of column-based arrays + it = zip(*normalized_rows) + + for i, col in self._iter_col_row(it): + col[item] = i + + def __delitem__(self, item: Union[int, slice]) -> None: + for col in self.columns: + del col[item] + + def __len__(self) -> int: + return len(self.columns[0]) + + @property + def shape(self) -> Tuple[int, int]: + return len(self.columns), len(self) + + def drop(self, *col_names: str) -> None: + for col_name in col_names: + self._keys.remove(col_name) + self._columns.pop(col_name) + + def rename(self, from_col_name: str, to_col_name: str) -> None: + self._columns[to_col_name] = self._columns.pop(from_col_name) + self._keys[self._keys.index(from_col_name)] = to_col_name + + def project(self, *col_names: str) -> None: + self.drop(*(set(self._keys) - set(col_names))) + self._keys = list(col_names) + + def to_csv(self) -> str: + import csv + from io import StringIO + + buff = StringIO() + writer = csv.writer(buff) + writer.writerow(self.keys()) + + for row in self: + writer.writerow(row) + return buff.getvalue() + + def add_column(self, name: str) -> None: + self._columns[name] = Column([self._fill_value] * len(self)) + self._keys.append(name) + + def row_from_dict(self, d: Dict[str, "CellT"]) -> None: + keys = self.keys() + for key in d: + if key not in keys: + self.add_column(key) + + row: List["CellT"] = [ + with_value(d.get(key), self._fill_value) for key in self.keys() + ] + self.append(row) + + def render(self, **kwargs: Any): + from dvc.ui import ui + + ui.table(self, headers=self.keys(), **kwargs) + + def as_dict(self): + keys = self.keys() + return [dict(zip(keys, row)) for row in self] diff --git a/dvc/ui/table.py b/dvc/ui/table.py index aaf6d10dd..54dbd7509 100644 --- a/dvc/ui/table.py +++ b/dvc/ui/table.py @@ -1,6 +1,6 @@ from contextlib import contextmanager from itertools import zip_longest -from typing import TYPE_CHECKING, Iterator, List, Sequence, Union +from typing import TYPE_CHECKING, Iterator, Sequence, Union from dvc.types import DictStrAny @@ -17,7 +17,7 @@ SHOW_MAX_WIDTH = 1024 CellT = Union[str, "Text"] # Text is mostly compatible with str Row = Sequence[CellT] -TableData = List[Row] +TableData = Sequence[Row] Headers = Sequence[str] Styles = DictStrAny diff --git a/tests/unit/test_tabular_data.py b/tests/unit/test_tabular_data.py new file mode 100644 index 000000000..e70645a36 --- /dev/null +++ b/tests/unit/test_tabular_data.py @@ -0,0 +1,178 @@ +import pytest + +from dvc.compare import TabularData + + +def test_table_empty(capsys): + td = TabularData(["Col1", "Col2", "Col3"]) + assert dict(td.items()) == {"Col1": [], "Col2": [], "Col3": []} + assert td.columns == [[], [], []] + assert td.keys() == ["Col1", "Col2", "Col3"] + assert list(td) == [] + assert td.Col1 == [] + assert td.Col2 == [] + assert td.Col3 == [] + + assert td[1:] == [] + with pytest.raises(IndexError): + _ = td[1] + + assert len(td) == 0 + assert td.shape == (3, 0) + assert td.to_csv() == """Col1,Col2,Col3\r\n""" + + td.render() + assert capsys.readouterr() == ("", "") + + td.render(rich_table=True) + assert capsys.readouterr() == ("", "") + + td.render(markdown=True) + assert capsys.readouterr() == ( + "| Col1 | Col2 | Col3 |\n|--------|--------|--------|\n\n", + "", + ) + + td.rename("Col1", "Col11") + assert td.keys() == ["Col11", "Col2", "Col3"] + + td.project("Col3", "Col11") + assert td.keys() == ["Col3", "Col11"] + + +def test_list_operations(): + td = TabularData(["col1", "col2", "col3"]) + td.append(["1", "2", "3"]) + + assert list(td) == [["1", "2", "3"]] + td.extend((["11", "12", "13"], ["21", "22", "23"])) + assert list(td) == [ + ["1", "2", "3"], + ["11", "12", "13"], + ["21", "22", "23"], + ] + td.insert(1, ["01", "02", "03"]) + assert list(td) == [ + ["1", "2", "3"], + ["01", "02", "03"], + ["11", "12", "13"], + ["21", "22", "23"], + ] + assert td.shape == (3, 4) + assert len(td) == 4 + assert td[1] == ["01", "02", "03"] + assert td[1:] == [ + ["01", "02", "03"], + ["11", "12", "13"], + ["21", "22", "23"], + ] + assert td[::-1] == [ + ["21", "22", "23"], + ["11", "12", "13"], + ["01", "02", "03"], + ["1", "2", "3"], + ] + del td[1] + assert list(td) == [ + ["1", "2", "3"], + ["11", "12", "13"], + ["21", "22", "23"], + ] + assert td.shape == (3, 3) + td[1:3] = [["51", "52", "53"], ["61", "62", "63"]] + assert list(td) == [ + ["1", "2", "3"], + ["51", "52", "53"], + ["61", "62", "63"], + ] + td[1] = ["41", "42", "43"] + assert td[1] == ["41", "42", "43"] + + del td[1:3] + assert td.shape == (3, 1) + + assert td.to_csv() == "col1,col2,col3\r\n1,2,3\r\n" + + +def test_dict_like_interfaces(): + td = TabularData(["col-1", "col-2"]) + + td.extend([["foo", "bar"], ["foobar", "foobar"]]) + assert td.keys() == ["col-1", "col-2"] + assert dict(td.items()) == { + "col-1": ["foo", "foobar"], + "col-2": ["bar", "foobar"], + } + assert td.as_dict() == [ + {"col-1": "foo", "col-2": "bar"}, + {"col-1": "foobar", "col-2": "foobar"}, + ] + + +def test_fill_value(): + td = TabularData(["col-1", "col-2", "col-3"], fill_value="?") + td.append(["foo"]) + assert list(td) == [["foo", "?", "?"]] + + td.extend( + [ + ["bar"], + ["foobar", "foobar2"], + ["f", "fo", "foo", "foob", "fooba", "foobar"], + ] + ) + assert list(td) == [ + ["foo", "?", "?"], + ["bar", "?", "?"], + ["foobar", "foobar2", "?"], + ["f", "fo", "foo"], + ] + + td.insert(1, ["lorem"]) + assert td[1] == ["lorem", "?", "?"] + + td[0] = ["lorem", "ipsum"] + assert td[0] == ["lorem", "ipsum", "?"] + + td[1:2] = [["f", "fo"]] + assert td[1:2] == [["f", "fo", "?"]] + + td.add_column("col-4") + assert td.keys() == ["col-1", "col-2", "col-3", "col-4"] + assert td[0][3] == "?" + + +def test_drop(): + td = TabularData(["col1", "col2", "col3"]) + td.append(["foo", "bar", "baz"]) + assert list(td) == [["foo", "bar", "baz"]] + td.drop("col2") + assert td.keys() == ["col1", "col3"] + assert list(td) == [["foo", "baz"]] + + +def test_row_from_dict(): + td = TabularData(["col1", "col2"]) + td.row_from_dict({"col3": "value3", "col4": "value4"}) + assert td.keys() == ["col1", "col2", "col3", "col4"] + assert dict(td.items()) == { + "col1": [""], + "col2": [""], + "col3": ["value3"], + "col4": ["value4"], + } + td.row_from_dict({"col3": "value3", "col5": "value5", "col6": "value6"}) + assert td.keys() == ["col1", "col2", "col3", "col4", "col5", "col6"] + assert dict(td.items()) == { + "col1": ["", ""], + "col2": ["", ""], + "col3": ["value3", "value3"], + "col4": ["value4", ""], + "col5": ["", "value5"], + "col6": ["", "value6"], + } + assert td.shape == (6, 2) + assert list(td) == [ + ["", "", "value3", "value4", "", ""], + ["", "", "value3", "", "value5", "value6"], + ]