From 8608557baec40b36d79ab47379fe735b69e0f2fa Mon Sep 17 00:00:00 2001 From: Javier Sancho Date: Fri, 11 May 2018 12:15:38 +0200 Subject: [PATCH] Return columns and rows ready to be shown on datasette --- datasette_pytables/__init__.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/datasette_pytables/__init__.py b/datasette_pytables/__init__.py index 93be0f9..d44c9f9 100644 --- a/datasette_pytables/__init__.py +++ b/datasette_pytables/__init__.py @@ -69,8 +69,9 @@ class Connection: else: table_rows = table.iterrows() + # Prepare rows if len(fields) == 1 and fields[0] == 'count(*)': - rows.append(Row({'count(*)': table.nrows})) + rows.append(Row({fields[0]: table.nrows})) else: for table_row in table_rows: row = Row() @@ -84,10 +85,16 @@ class Connection: row[field] = table_row[field] rows.append(row) - description = ((col,) for col in table.colnames) + # Prepare query description + for field in fields: + if field == '*': + for col in table.colnames: + description.append((col,)) + else: + description.append((field,)) if truncate: - return rows, truncated, description + return rows, truncated, tuple(description) else: return rows @@ -95,3 +102,6 @@ class Row(OrderedDict): def __getitem__(self, label): if type(label) is int: return super(OrderedDict, self).__getitem__(list(self.keys())[label]) + + def __iter__(self): + return self.values().__iter__() -- 2.39.2