sqlglot.executor.table

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot.helper import dict_depth
  6from sqlglot.schema import AbstractMappingSchema
  7
  8
  9class Table:
 10    def __init__(self, columns, rows=None, column_range=None):
 11        self.columns = tuple(columns)
 12        self.column_range = column_range
 13        self.reader = RowReader(self.columns, self.column_range)
 14        self.rows = rows or []
 15        if rows:
 16            assert len(rows[0]) == len(self.columns)
 17        self.range_reader = RangeReader(self)
 18
 19    def add_columns(self, *columns: str) -> None:
 20        self.columns += columns
 21        if self.column_range:
 22            self.column_range = range(
 23                self.column_range.start, self.column_range.stop + len(columns)
 24            )
 25        self.reader = RowReader(self.columns, self.column_range)
 26
 27    def append(self, row):
 28        assert len(row) == len(self.columns)
 29        self.rows.append(row)
 30
 31    def pop(self):
 32        self.rows.pop()
 33
 34    @property
 35    def width(self):
 36        return len(self.columns)
 37
 38    def __len__(self):
 39        return len(self.rows)
 40
 41    def __iter__(self):
 42        return TableIter(self)
 43
 44    def __getitem__(self, index):
 45        self.reader.row = self.rows[index]
 46        return self.reader
 47
 48    def __repr__(self):
 49        columns = tuple(
 50            column
 51            for i, column in enumerate(self.columns)
 52            if not self.column_range or i in self.column_range
 53        )
 54        widths = {column: len(column) for column in columns}
 55        lines = [" ".join(column for column in columns)]
 56
 57        for i, row in enumerate(self):
 58            if i > 10:
 59                break
 60
 61            lines.append(
 62                " ".join(
 63                    str(row[column]).rjust(widths[column])[0 : widths[column]] for column in columns
 64                )
 65            )
 66        return "\n".join(lines)
 67
 68
 69class TableIter:
 70    def __init__(self, table):
 71        self.table = table
 72        self.index = -1
 73
 74    def __iter__(self):
 75        return self
 76
 77    def __next__(self):
 78        self.index += 1
 79        if self.index < len(self.table):
 80            return self.table[self.index]
 81        raise StopIteration
 82
 83
 84class RangeReader:
 85    def __init__(self, table):
 86        self.table = table
 87        self.range = range(0)
 88
 89    def __len__(self):
 90        return len(self.range)
 91
 92    def __getitem__(self, column):
 93        return (self.table[i][column] for i in self.range)
 94
 95
 96class RowReader:
 97    def __init__(self, columns, column_range=None):
 98        self.columns = {
 99            column: i for i, column in enumerate(columns) if not column_range or i in column_range
100        }
101        self.row = None
102
103    def __getitem__(self, column):
104        return self.row[self.columns[column]]
105
106
107class Tables(AbstractMappingSchema[Table]):
108    pass
109
110
111def ensure_tables(d: t.Optional[t.Dict]) -> Tables:
112    return Tables(_ensure_tables(d))
113
114
115def _ensure_tables(d: t.Optional[t.Dict]) -> t.Dict:
116    if not d:
117        return {}
118
119    depth = dict_depth(d)
120
121    if depth > 1:
122        return {k: _ensure_tables(v) for k, v in d.items()}
123
124    result = {}
125    for name, table in d.items():
126        if isinstance(table, Table):
127            result[name] = table
128        else:
129            columns = tuple(table[0]) if table else ()
130            rows = [tuple(row[c] for c in columns) for row in table]
131            result[name] = Table(columns=columns, rows=rows)
132
133    return result