Fix access to result_set and QueryPolyResult construction

Author: tobias-weber

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = ipython-polypheny
-version = 0.1.2
+version = 0.1.3
 description = Access Polypheny via IPython magics
 long_description = file: README.md
 long_description_content_type = text/markdown

src/poly/poly_magic.py

@@ -145,7 +145,7 @@ def expand_variables(self, template):
 
 def separate_args(line, cell, split_str=":"):
     """
-    Finds the first occurence of an element of termination_strings in line. The line is split after this element and
+    Finds the first occurrence of an element of termination_strings in line. The line is split after this element and
     the two parts are returned.
     """
     if cell is None:

src/poly/poly_result.py

@@ -6,7 +6,7 @@
 def build_result(raw_result: Union[str, dict, List[dict]]):
     result_set = get_result_dict(raw_result)
 
-    if 'error' in result_set:
+    if result_set.get('error') is not None:
         return ErrorPolyResult(result_set)
     if 'header' not in result_set:
         return InfoPolyResult(result_set)
@@ -35,33 +35,39 @@ def get_type_from_string(type_string):
         return float
     if cleaned in ['ARRAY', 'JSON', 'DOCUMENT', 'DOCUMENT NOT NULL']:
         return json.loads
-    if cleaned.startswith('PATH') or cleaned.startswith('NODE'):
+    if cleaned.startswith('PATH') or cleaned.startswith('NODE') or cleaned.startswith('EDGE'):
         return json.loads
     return None
 
 
-def cast_data(raw_data, header):
-    data = [row[:] for row in raw_data]  # Create shallow copy
-    data_types = [get_type_from_string(col['dataType']) for col in header]
-
-    for col_idx, col_type in enumerate(data_types):
-        if col_type is None:
-            continue
-        for row_idx in range(len(data)):
-            try:
-                data[row_idx][col_idx] = col_type(data[row_idx][col_idx])
-            except (TypeError, json.JSONDecodeError):
-                pass
+def cast_data(raw_data, header, data_model):
+    if data_model in ['RELATIONAL', 'GRAPH']:
+        data = [row[:] for row in raw_data]  # Create shallow copy
+        data_types = [get_type_from_string(col['dataType']) for col in header]
+
+        for col_idx, col_type in enumerate(data_types):
+            if col_type is None:
+                continue
+            for row_idx in range(len(data)):
+                try:
+                    data[row_idx][col_idx] = col_type(data[row_idx][col_idx])
+                except (TypeError, json.JSONDecodeError):
+                    pass
+    elif data_model == 'DOCUMENT':
+        # transform documents to rows with 1 column
+        data = [[json.loads(doc[0] if isinstance(doc, list) else doc)] for doc in raw_data]
+    else:
+        raise ValueError(f'Unsupported dataModel: {data_model}')
     return data
 
 
 class QueryPolyResult(list):  # Data is stored in a nested list. For simplicity, we do not use Numpy, but Python lists
     def __init__(self, result_set):
         self.result_set = result_set
-        self.type = result_set['namespaceType']
+        self.type = result_set['dataModel']
         self._header = result_set['header']
         self.keys = [col['name'] for col in self._header]
-        self._data = cast_data(result_set['data'], self._header)
+        self._data = cast_data(result_set['data'], self._header, self.type)
         self._pretty = PrettyTable(self.keys)
         self._pretty.add_rows(self._data)
         self._pretty.set_style(PLAIN_COLUMNS)
@@ -87,7 +93,7 @@ def __init__(self, result_set):
 
     def __repr__(self):
         rows = ['Successfully executed:',
-                'Query:'.ljust(30) + str(self.result_set['generatedQuery'])
+                'Query:'.ljust(30) + str(self.result_set['query'])
                 ]
         return "\n".join(rows)
 
@@ -98,7 +104,7 @@ def __init__(self, result_set):
 
     def __repr__(self):
         rows = ['ERROR:',
-                'Query:'.ljust(30) + str(self.result_set['generatedQuery']),
+                'Query:'.ljust(30) + str(self.result_set['query']),
                 'Message:'.ljust(30) + str(self.result_set['error'])
                 ]
         return "\n".join(rows)