diff options
Diffstat (limited to 'sqlglot/executor/table.py')
| -rw-r--r-- | sqlglot/executor/table.py | 27 |
1 files changed, 19 insertions, 8 deletions
diff --git a/sqlglot/executor/table.py b/sqlglot/executor/table.py index 80674cb..6796740 100644 --- a/sqlglot/executor/table.py +++ b/sqlglot/executor/table.py @@ -1,10 +1,12 @@ class Table: - def __init__(self, *columns, rows=None): - self.columns = tuple(columns if isinstance(columns[0], str) else columns[0]) + def __init__(self, columns, rows=None, column_range=None): + self.columns = tuple(columns) + self.column_range = column_range + self.reader = RowReader(self.columns, self.column_range) + self.rows = rows or [] if rows: assert len(rows[0]) == len(self.columns) - self.reader = RowReader(self.columns) self.range_reader = RangeReader(self) def append(self, row): @@ -29,15 +31,22 @@ class Table: return self.reader def __repr__(self): - widths = {column: len(column) for column in self.columns} - lines = [" ".join(column for column in self.columns)] + columns = tuple( + column + for i, column in enumerate(self.columns) + if not self.column_range or i in self.column_range + ) + widths = {column: len(column) for column in columns} + lines = [" ".join(column for column in columns)] for i, row in enumerate(self): if i > 10: break lines.append( - " ".join(str(row[column]).rjust(widths[column])[0 : widths[column]] for column in self.columns) + " ".join( + str(row[column]).rjust(widths[column])[0 : widths[column]] for column in columns + ) ) return "\n".join(lines) @@ -70,8 +79,10 @@ class RangeReader: class RowReader: - def __init__(self, columns): - self.columns = {column: i for i, column in enumerate(columns)} + def __init__(self, columns, column_range=None): + self.columns = { + column: i for i, column in enumerate(columns) if not column_range or i in column_range + } self.row = None def __getitem__(self, column): |