Query result wrapper with convenient access methods
| 35 | |
| 36 | |
| 37 | class QueryResult: |
| 38 | """Query result wrapper with convenient access methods""" |
| 39 | |
| 40 | def __init__(self, data: Dict[str, Any]): |
| 41 | self._data = data |
| 42 | self.variables = data.get("variables", []) |
| 43 | # Flatten rows from nested structure |
| 44 | raw_rows = data.get("rows", []) |
| 45 | self.rows = [self._flatten_row(row) for row in raw_rows] |
| 46 | self.row_count = len(self.rows) |
| 47 | |
| 48 | def _flatten_row(self, row: Dict[str, Any]) -> Dict[str, Any]: |
| 49 | """Flatten a row from nested value structure to simple dict""" |
| 50 | if "values" not in row: |
| 51 | return row |
| 52 | |
| 53 | flattened = {} |
| 54 | for key, value_wrapper in row["values"].items(): |
| 55 | flattened[key] = self._extract_value(value_wrapper) |
| 56 | return flattened |
| 57 | |
| 58 | def _extract_value(self, value_wrapper: Any) -> Any: |
| 59 | """Extract value from Rust enum wrapper like {'String': 'foo'} or {'Number': 42}""" |
| 60 | if not isinstance(value_wrapper, dict): |
| 61 | return value_wrapper |
| 62 | |
| 63 | # Handle Rust enum variants |
| 64 | if "String" in value_wrapper: |
| 65 | return value_wrapper["String"] |
| 66 | elif "Number" in value_wrapper: |
| 67 | num = value_wrapper["Number"] |
| 68 | # Convert to int if it's a whole number |
| 69 | return int(num) if isinstance(num, float) and num.is_integer() else num |
| 70 | elif "Boolean" in value_wrapper: |
| 71 | return value_wrapper["Boolean"] |
| 72 | elif "Null" in value_wrapper: |
| 73 | return None |
| 74 | elif "List" in value_wrapper: |
| 75 | return [self._extract_value(v) for v in value_wrapper["List"]] |
| 76 | elif "Map" in value_wrapper: |
| 77 | return {k: self._extract_value(v) for k, v in value_wrapper["Map"].items()} |
| 78 | elif "Node" in value_wrapper: |
| 79 | return value_wrapper # Return node reference as-is |
| 80 | elif "Edge" in value_wrapper: |
| 81 | return value_wrapper # Return edge reference as-is |
| 82 | elif "Path" in value_wrapper: |
| 83 | return value_wrapper # Return path as-is |
| 84 | else: |
| 85 | return value_wrapper |
| 86 | |
| 87 | def __repr__(self): |
| 88 | return f"QueryResult(rows={self.row_count}, variables={self.variables})" |
| 89 | |
| 90 | def to_dict(self) -> Dict[str, Any]: |
| 91 | """Get raw dictionary representation""" |
| 92 | return self._data |
| 93 | |
| 94 | def first(self) -> Optional[Dict[str, Any]]: |