Merge pull request #924 from neo4j-contrib/task/fix-typing-stubs-2

Task/fix typing stubs 2

Author: mariusconjeaud

.github/workflows/type-checking.yml

@@ -1219,6 +1219,18 @@ async def _execute(self, lazy: bool = False, dict_output: bool = False) -> Any:
                 yield result
 
 
+@dataclass
+class Path:
+    """Path traversal definition."""
+
+    value: str
+    optional: bool = False
+    include_nodes_in_return: bool = True
+    include_rels_in_return: bool = True
+    relation_filtering: bool = False
+    alias: str | None = None
+
+
 class AsyncBaseSet:
     """
     Base class for all node sets.
@@ -1229,6 +1241,10 @@ class AsyncBaseSet:
     query_cls = AsyncQueryBuilder
     source_class: type[AsyncStructuredNode]
 
+    # Attributes defined in subclasses (AsyncNodeSet)
+    _unique_variables: list[str]
+    relations_to_fetch: list[Path]
+
     async def all(self, lazy: bool = False) -> list:
         """
         Return all nodes belonging to the set
@@ -1300,18 +1316,6 @@ async def get_item(self, key: int | slice) -> Optional["AsyncBaseSet"]:
             return _first_item
 
 
-@dataclass
-class Path:
-    """Path traversal definition."""
-
-    value: str
-    optional: bool = False
-    include_nodes_in_return: bool = True
-    include_rels_in_return: bool = True
-    relation_filtering: bool = False
-    alias: str | None = None
-
-
 @dataclass
 class RelationNameResolver:
     """Helper to refer to a relation variable name.

neomodel/async_/match.py

@@ -617,6 +617,7 @@ def deflate_relationship_properties(
     :return: Dictionary mapping property names to parameter placeholders (e.g. {'since': '$since'})
     """
     rel_model = relationship.definition.get("model")
+    assert rel_model is not None, "Relationship model is required"
     tmp = rel_model(**rel_props)
 
     rel_prop = {}

neomodel/async_/relationship_manager.py

@@ -16,11 +16,13 @@
 """
 
 
+from typing import Any, Literal
 from warnings import warn
 
 try:
     # noinspection PyPackageRequirements
     from numpy import array as nparray
+    from numpy import ndarray
 except ImportError:
     warn(
         "The neomodel.integration.numpy module expects numpy to be installed "
@@ -29,7 +31,11 @@
     raise
 
 
-def to_ndarray(query_results: tuple, dtype=None, order="K"):
+def to_ndarray(
+    query_results: tuple[list[list[Any]], list[str]],
+    dtype: Any | None = None,
+    order: Literal["K", "A", "C", "F"] = "K",
+) -> ndarray:
     """Convert the results of a db.cypher_query call into a numpy array.
     Optionally, specify a datatype and/or an order for the columns.
     """

neomodel/integration/numpy.py

@@ -20,7 +20,7 @@
 
 """
 
-
+from typing import Any
 from warnings import warn
 
 try:
@@ -34,7 +34,11 @@
     raise
 
 
-def to_dataframe(query_results: tuple, index=None, dtype=None):
+def to_dataframe(
+    query_results: tuple[list[list[Any]], list[str]],
+    index: Any | None = None,
+    dtype: Any | None = None,
+) -> DataFrame:
     """Convert the results of a db.cypher_query call and associated metadata
     into a pandas DataFrame.
     Optionally, specify an index and/or a datatype for the columns.
@@ -43,10 +47,18 @@ def to_dataframe(query_results: tuple, index=None, dtype=None):
     return DataFrame(results, columns=meta, index=index, dtype=dtype)
 
 
-def to_series(query_results: tuple, field=0, index=None, dtype=None):
+def to_series(
+    query_results: tuple[list[list[Any]], list[str]],
+    field: int = 0,
+    index: Any | None = None,
+    dtype: Any | None = None,
+) -> Series:
     """Convert the results of a db.cypher_query call
     into a pandas Series for the given field.
     Optionally, specify an index and/or a datatype for the columns.
     """
     results, _ = query_results
-    return Series([record[field] for record in results], index=index, dtype=dtype)
+    if dtype is None:
+        return Series([record[field] for record in results], index=index)
+    else:
+        return Series([record[field] for record in results], index=index, dtype=dtype)

neomodel/integration/pandas.py

@@ -1215,6 +1215,18 @@ def _execute(self, lazy: bool = False, dict_output: bool = False) -> Any:
                 yield result
 
 
+@dataclass
+class Path:
+    """Path traversal definition."""
+
+    value: str
+    optional: bool = False
+    include_nodes_in_return: bool = True
+    include_rels_in_return: bool = True
+    relation_filtering: bool = False
+    alias: str | None = None
+
+
 class BaseSet:
     """
     Base class for all node sets.
@@ -1225,6 +1237,10 @@ class BaseSet:
     query_cls = QueryBuilder
     source_class: type[StructuredNode]
 
+    # Attributes defined in subclasses (AsyncNodeSet)
+    _unique_variables: list[str]
+    relations_to_fetch: list[Path]
+
     def all(self, lazy: bool = False) -> list:
         """
         Return all nodes belonging to the set
@@ -1296,18 +1312,6 @@ def __getitem__(self, key: int | slice) -> Optional["BaseSet"]:
             return _first_item
 
 
-@dataclass
-class Path:
-    """Path traversal definition."""
-
-    value: str
-    optional: bool = False
-    include_nodes_in_return: bool = True
-    include_rels_in_return: bool = True
-    relation_filtering: bool = False
-    alias: str | None = None
-
-
 @dataclass
 class RelationNameResolver:
     """Helper to refer to a relation variable name.

neomodel/sync_/match.py

@@ -597,6 +597,7 @@ def deflate_relationship_properties(
     :return: Dictionary mapping property names to parameter placeholders (e.g. {'since': '$since'})
     """
     rel_model = relationship.definition.get("model")
+    assert rel_model is not None, "Relationship model is required"
     tmp = rel_model(**rel_props)
 
     rel_prop = {}