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.Any = None,
 14        rows: t.Any = None,
 15        column_range: range | None = None,
 16    ) -> None:
 17        self.columns: t.Any = tuple(columns) if columns is not None else ()
 18        self.column_range = column_range
 19        self.reader: t.Any = RowReader(self.columns, self.column_range)
 20        self.rows: t.Any = 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.Any) -> 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) -> 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: t.Any = None) -> None:
 95        self.table: t.Any = 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=None, 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            if columns is not None
110            else {}
111        )
112        self.row = None
113
114    def __getitem__(self, column):
115        return self.row[self.columns[column]]
116
117
118class Tables(AbstractMappingSchema):
119    pass
120
121
122def ensure_tables(d: dict | None, dialect: DialectType = None) -> Tables:
123    return Tables(_ensure_tables(d, dialect=dialect))
124
125
126def _ensure_tables(d: dict | None, dialect: DialectType = None) -> dict:
127    if not d:
128        return {}
129
130    depth = dict_depth(d)
131    if depth > 1:
132        return {
133            normalize_name(k, dialect=dialect, is_table=True).name: _ensure_tables(
134                v, dialect=dialect
135            )
136            for k, v in d.items()
137        }
138
139    result = {}
140    for table_name, table in d.items():
141        table_name = normalize_name(table_name, dialect=dialect).name
142
143        if isinstance(table, Table):
144            result[table_name] = table
145        else:
146            table = [
147                {
148                    normalize_name(column_name, dialect=dialect).name: value
149                    for column_name, value in row.items()
150                }
151                for row in table
152            ]
153            column_names = tuple(column_name for column_name in table[0]) if table else ()
154            rows = [tuple(row[name] for name in column_names) for row in table]
155            result[table_name] = Table(columns=column_names, rows=rows)
156
157    return result