ui: tabular data (#5852)

* introduce tabular data structure

* ignore ellipsis

.coveragerc

@@ -13,3 +13,4 @@ exclude_lines =
     raise NotImplementedError
     # pragma: no cover
     #  pragma: no cover
+    \.\.\.$

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]

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
 

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"],
+    ]