Skip to content

Commit

Permalink
Polish datetime conversion implementation
Browse files Browse the repository at this point in the history
- Naming things
- Slight structural changes
- Use iterator instead of generator for column type flagging
- Improve inline documentation
  • Loading branch information
amotl committed Jul 26, 2022
1 parent 6a5c61c commit 6522e43
Showing 1 changed file with 25 additions and 24 deletions.
49 changes: 25 additions & 24 deletions src/crate/client/cursor.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,44 +51,45 @@ def execute(self, sql, parameters=None, bulk_parameters=None):
self._result = self.connection.client.sql(sql, parameters,
bulk_parameters)

if "rows" in self._result:
transformed_result = False
if "col_types" in self._result:
transformed_result = True
self.rows = self.result_set_transformed()
if "rows" not in self._result:
return

if "col_types" in self._result:
self.rows = iter(self._transform_result_types())

if not transformed_result:
self.rows = iter(self._result["rows"])
else:
self.rows = iter(self._result["rows"])

def result_set_transformed(self):
def _transform_result_types(self):
"""
Generator that iterates over each row from the result set
Generate row items with column values converted to their corresponding
native Python types, based on information from `col_types`.
Currently, only converting to native `datetime` objects is implemented.
"""
rows_to_convert = [True if col_type == 11 or col_type == 15 else False for col_type in
self._result["col_types"]]
datetime_column_types = [11, 15]
datetime_columns_mask = [
True if col_type in datetime_column_types else False
for col_type in self._result["col_types"]
]
for row in self._result["rows"]:
gen_flags = (flag for flag in rows_to_convert)
yield [t_row for t_row in self._transform_date_columns(row, gen_flags)]
yield list(self._transform_datetime_columns(row, iter(datetime_columns_mask)))

@staticmethod
def _transform_date_columns(row, gen_flags):
def _transform_datetime_columns(row, column_flags):
"""
Generates iterates over each value from a row and converts timestamps to pandas TIMESTAMP
Convert all designated columns to native Python `datetime` objects.
"""
for value in row:
try:
flag = next(gen_flags)
flag = next(column_flags)
except StopIteration:
break

if not flag or value is None:
yield value
else:
if value < 0:
yield None
else:
value = datetime.fromtimestamp(value / 1000)
yield value
if flag and value is not None:
value = datetime.fromtimestamp(value / 1e3)

yield value

def executemany(self, sql, seq_of_parameters):
"""
Expand Down

0 comments on commit 6522e43

Please sign in to comment.