Merge pull request #31 from microsoft/language_fixes

Language fixes

Author: niclasko

flowquery-py/src/graph/database.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, Optional, Union
+from typing import Any, AsyncIterator, Dict, List, Optional, Union
 
 from ..parsing.ast_node import ASTNode
 from .node import Node
@@ -48,35 +48,57 @@ def add_relationship(self, relationship: 'Relationship', statement: ASTNode) ->
         physical = PhysicalRelationship()
         physical.type = relationship.type
         physical.statement = statement
+        if relationship.source is not None:
+            physical.source = relationship.source
+        if relationship.target is not None:
+            physical.target = relationship.target
         Database._relationships[relationship.type] = physical
 
     def get_relationship(self, relationship: 'Relationship') -> Optional['PhysicalRelationship']:
         """Gets a relationship from the database."""
         return Database._relationships.get(relationship.type) if relationship.type else None
 
-    async def schema(self) -> list[dict[str, Any]]:
+    def get_relationships(self, relationship: 'Relationship') -> list['PhysicalRelationship']:
+        """Gets multiple physical relationships for ORed types."""
+        result = []
+        for rel_type in relationship.types:
+            physical = Database._relationships.get(rel_type)
+            if physical:
+                result.append(physical)
+        return result
+
+    async def schema(self) -> List[Dict[str, Any]]:
         """Returns the graph schema with node/relationship labels and sample data."""
-        result: list[dict[str, Any]] = []
+        return [item async for item in self._schema()]
 
+    async def _schema(self) -> AsyncIterator[Dict[str, Any]]:
+        """Async generator for graph schema with node/relationship labels and sample data."""
         for label, physical_node in Database._nodes.items():
             records = await physical_node.data()
-            entry: dict[str, Any] = {"kind": "node", "label": label}
+            entry: Dict[str, Any] = {"kind": "Node", "label": label}
             if records:
                 sample = {k: v for k, v in records[0].items() if k != "id"}
-                if sample:
+                properties = list(sample.keys())
+                if properties:
+                    entry["properties"] = properties
                     entry["sample"] = sample
-            result.append(entry)
+            yield entry
 
         for rel_type, physical_rel in Database._relationships.items():
             records = await physical_rel.data()
-            entry_rel: dict[str, Any] = {"kind": "relationship", "type": rel_type}
+            entry_rel: Dict[str, Any] = {
+                "kind": "Relationship",
+                "type": rel_type,
+                "from_label": physical_rel.source.label if physical_rel.source else None,
+                "to_label": physical_rel.target.label if physical_rel.target else None,
+            }
             if records:
                 sample = {k: v for k, v in records[0].items() if k not in ("left_id", "right_id")}
-                if sample:
+                properties = list(sample.keys())
+                if properties:
+                    entry_rel["properties"] = properties
                     entry_rel["sample"] = sample
-            result.append(entry_rel)
-
-        return result
+            yield entry_rel
 
     async def get_data(self, element: Union['Node', 'Relationship']) -> Union['NodeData', 'RelationshipData']:
         """Gets data for a node or relationship."""
@@ -87,6 +109,17 @@ async def get_data(self, element: Union['Node', 'Relationship']) -> Union['NodeD
             data = await node.data()
             return NodeData(data)
         elif isinstance(element, Relationship):
+            if len(element.types) > 1:
+                physicals = self.get_relationships(element)
+                if not physicals:
+                    raise ValueError(f"No physical relationships found for types {', '.join(element.types)}")
+                all_records = []
+                for i, physical in enumerate(physicals):
+                    records = await physical.data()
+                    type_name = element.types[i]
+                    for record in records:
+                        all_records.append({**record, "_type": type_name})
+                return RelationshipData(all_records)
             relationship = self.get_relationship(element)
             if relationship is None:
                 raise ValueError(f"Physical relationship not found for type {element.type}")

flowquery-py/src/graph/relationship.py

@@ -19,7 +19,7 @@ class Relationship(ASTNode):
     def __init__(self) -> None:
         super().__init__()
         self._identifier: Optional[str] = None
-        self._type: Optional[str] = None
+        self._types: List[str] = []
         self._hops: Hops = Hops()
         self._source: Optional['Node'] = None
         self._target: Optional['Node'] = None
@@ -39,11 +39,19 @@ def identifier(self, value: str) -> None:
 
     @property
     def type(self) -> Optional[str]:
-        return self._type
+        return self._types[0] if self._types else None
 
     @type.setter
     def type(self, value: str) -> None:
-        self._type = value
+        self._types = [value]
+
+    @property
+    def types(self) -> List[str]:
+        return self._types
+
+    @types.setter
+    def types(self, value: List[str]) -> None:
+        self._types = value
 
     @property
     def hops(self) -> Hops:

flowquery-py/src/graph/relationship_data.py

@@ -25,11 +25,12 @@ def find(self, id: str, hop: int = 0, direction: str = "right") -> bool:
         return self._find(id, hop, key)
 
     def properties(self) -> Optional[Dict[str, Any]]:
-        """Get properties of current relationship, excluding left_id and right_id."""
+        """Get properties of current relationship, excluding left_id, right_id, and _type."""
         current = self.current()
         if current:
             props = dict(current)
             props.pop("left_id", None)
             props.pop("right_id", None)
+            props.pop("_type", None)
             return props
         return None

flowquery-py/src/graph/relationship_match_collector.py

@@ -28,9 +28,15 @@ def push(self, relationship: 'Relationship', traversal_id: str = "") -> Relation
         """Push a new match onto the collector."""
         start_node_value = relationship.source.value() if relationship.source else None
         rel_data = relationship.get_data()
+        current_record = rel_data.current() if rel_data else None
+        default_type = relationship.type or ""
+        if current_record and isinstance(current_record, dict):
+            actual_type = current_record.get('_type', default_type)
+        else:
+            actual_type = default_type
         rel_props: Dict[str, Any] = (rel_data.properties() or {}) if rel_data else {}
         match: RelationshipMatchRecord = {
-            "type": relationship.type or "",
+            "type": actual_type,
             "startNode": start_node_value or {},
             "endNode": None,
             "properties": rel_props,

flowquery-py/src/graph/relationship_reference.py

@@ -10,8 +10,8 @@ class RelationshipReference(Relationship):
     def __init__(self, relationship: Relationship, referred: ASTNode) -> None:
         super().__init__()
         self._referred = referred
-        if relationship.type:
-            self.type = relationship.type
+        if relationship.types:
+            self.types = relationship.types
 
     @property
     def referred(self) -> ASTNode:

flowquery-py/src/parsing/functions/__init__.py

@@ -39,6 +39,7 @@
 from .to_json import ToJson
 from .to_lower import ToLower
 from .to_string import ToString
+from .trim import Trim
 from .type_ import Type
 from .value_holder import ValueHolder
 
@@ -78,6 +79,7 @@
     "ToJson",
     "ToLower",
     "ToString",
+    "Trim",
     "Type",
     "Functions",
     "Schema",

flowquery-py/src/parsing/functions/predicate_sum.py

@@ -1,6 +1,6 @@
 """PredicateSum function."""
 
-from typing import Any, Optional
+from typing import Any
 
 from .function_metadata import FunctionDef
 from .predicate_function import PredicateFunction
@@ -41,12 +41,9 @@ def value(self) -> Any:
         if array is None or not isinstance(array, list):
             raise ValueError("Invalid array for sum function")
 
-        _sum: Optional[Any] = None
+        _sum: int = 0
         for item in array:
             self._value_holder.holder = item
             if self.where is None or self.where.value():
-                if _sum is None:
-                    _sum = self._return.value()
-                else:
-                    _sum += self._return.value()
+                _sum += self._return.value()
         return _sum

flowquery-py/src/parsing/functions/schema.py

@@ -9,23 +9,27 @@
 @FunctionDef({
     "description": (
         "Returns the graph schema listing all nodes and relationships "
-        "with a sample of their data."
+        "with their properties and a sample of their data."
     ),
     "category": "async",
     "parameters": [],
     "output": {
-        "description": "Schema entry with kind, label/type, and optional sample data",
+        "description": "Schema entry with label/type, properties, and optional sample data",
         "type": "object",
     },
     "examples": [
-        "CALL schema() YIELD kind, label, type, sample RETURN kind, label, type, sample",
+        "CALL schema() YIELD label, type, from_label, to_label, properties, sample "
+        "RETURN label, type, from_label, to_label, properties, sample",
     ],
 })
 class Schema(AsyncFunction):
     """Returns the graph schema of the database.
 
-    Lists all nodes and relationships with their labels/types and a sample
-    of their data (excluding id from nodes, left_id and right_id from relationships).
+    Lists all nodes and relationships with their labels/types, properties,
+    and a sample of their data (excluding id from nodes, left_id and right_id from relationships).
+
+    Nodes: {label, properties, sample}
+    Relationships: {type, from_label, to_label, properties, sample}
     """
 
     async def generate(self) -> AsyncGenerator[Any, None]:

flowquery-py/src/parsing/functions/trim.py

@@ -0,0 +1,35 @@
+"""Trim function."""
+
+from typing import Any
+
+from .function import Function
+from .function_metadata import FunctionDef
+
+
+@FunctionDef({
+    "description": "Removes leading and trailing whitespace from a string",
+    "category": "scalar",
+    "parameters": [
+        {"name": "text", "description": "String to trim", "type": "string"}
+    ],
+    "output": {"description": "Trimmed string", "type": "string", "example": "hello"},
+    "examples": [
+        "WITH '  hello  ' AS s RETURN trim(s)",
+        "WITH '\\tfoo\\n' AS s RETURN trim(s)"
+    ]
+})
+class Trim(Function):
+    """Trim function.
+
+    Removes leading and trailing whitespace from a string.
+    """
+
+    def __init__(self) -> None:
+        super().__init__("trim")
+        self._expected_parameter_count = 1
+
+    def value(self) -> Any:
+        val = self.get_children()[0].value()
+        if not isinstance(val, str):
+            raise ValueError("Invalid argument for trim function: expected a string")
+        return val.strip()

flowquery-py/src/parsing/operations/group_by.py

@@ -122,6 +122,8 @@ def generate_results(
                 self.mappers[mapper_index].overridden = child.value
                 yield from self.generate_results(mapper_index + 1, child)
         else:
+            if node.elements is None:
+                node.elements = [reducer.element() for reducer in self.reducers]
             if node.elements:
                 for i, element in enumerate(node.elements):
                     self.reducers[i].overridden = element.value