diff --git a/src/crate/client/cursor.py b/src/crate/client/cursor.py index 26a9faad0..7b2686869 100644 --- a/src/crate/client/cursor.py +++ b/src/crate/client/cursor.py @@ -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): """