sqlglot.executor.table

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