Class Table

tests/utils/table.py:6–77

class Table:
    def __init__(self, rows: List[List[str]], header: Optional[List[str]]):
        self.rows = rows
        self.header = header

    def __getitem__(self, item):
        if isinstance(item, slice):
            return Table(self.rows[item], header=self.header)
        return self.rows[item]

    def __iter__(self):
        yield from self.rows

    def as_horizontal(self) -> "Table":
        if self.header is not None:
            return self
        assert self.rows
        return Table(self.rows[1:], self.rows[0])

    def get_row_value(self, key) -> Optional[str]:
        """
        Assumes vertical table (each value has a separate row).
        """
        row = [r for r in self.rows if r and r[0] == key]
        if not row:
            return None
        assert len(row) == 1
        return row[0][1]

    def check_row_value(self, key: str, value: str):
        row = self.get_row_value(key)
        if row is None:
            raise Exception(f"Key `{key!r}` not found in\n{self}")
        assert row == value

    def get_column_value(self, key: str) -> Optional[List[str]]:
        """
        Assumes horizontal table (each value has a separate column).
        """
        if self.header is None:
            raise Exception(f"This table is not horizontal!\n{self}")
        if key not in self.header:
            return None

        index = self.header.index(key)
        return [row[index] for row in self.rows]

    def check_column_value(self, key: str, index: int, value: str):
        column = self.get_column_value(key)
        if not column:
            raise Exception(f"Value for key `{key!r}` not found in\n{self}")
        row = column[index]
        assert row == value

    def check_columns_value(self, keys: List[str], index: int, values: List[str]):
        assert len(keys) == len(values)
        for key, val in zip(keys, values):
            self.check_column_value(key, index, val)