Merge pull request #935 from neo4j-contrib/housekeeping/more-type-hinting

Housekeeping/more type hinting

Author: tonioo

Changelog

@@ -1,6 +1,7 @@
 Version 6.1.0 2026-01
 * Add new exists operator in filter()
 * Add allow_reload to config to allow node redefinition in hot reload environments for development purposes
+* Improve/Fix type hints
 
 Version 6.0.1 2025-12
 * Make async iterator fully async, like : async for node in MyNodeClass.nodes

neomodel/async_/match.py

@@ -2,7 +2,9 @@
 import re
 import string
 from dataclasses import dataclass
-from typing import Any, AsyncIterator, Optional, Union
+from typing import Any, AsyncIterator, Iterable, Union
+
+from typing_extensions import Self
 
 from neomodel._async_compat.util import AsyncUtil
 from neomodel.async_ import relationship_manager
@@ -1305,7 +1307,7 @@ async def process_stream(stream_iterator):
             else:
                 # Create a session for streaming
                 # Note: We need to keep the session open during iteration
-                async with adb.driver.session(
+                async with adb.driver.session(  # type: ignore
                     database=adb._database_name,
                     impersonated_user=adb.impersonated_user,
                 ) as session:
@@ -1427,7 +1429,7 @@ async def check_contains(self, obj: AsyncStructuredNode | Any) -> bool:
 
         raise ValueError("Expecting StructuredNode instance")
 
-    async def get_item(self, key: int | slice) -> Optional["AsyncBaseSet"]:
+    async def get_item(self, key: int | slice) -> Self | AsyncStructuredNode:
         if isinstance(key, slice):
             if key.stop and key.start:
                 self.limit = key.stop - key.start
@@ -1439,13 +1441,12 @@ async def get_item(self, key: int | slice) -> Optional["AsyncBaseSet"]:
 
             return self
 
-        if isinstance(key, int):
-            self.skip = key
-            self.limit = 1
+        self.skip = key
+        self.limit = 1
 
-            ast = await self.query_cls(self).build_ast()
-            _first_item = [node async for node in ast._execute()][0]
-            return _first_item
+        ast = await self.query_cls(self).build_ast()
+        _first_item = [node async for node in ast._execute()][0]
+        return _first_item
 
 
 @dataclass
@@ -1636,7 +1637,7 @@ async def _get(
         results = [node async for node in ast._execute(lazy)]
         return results
 
-    async def get(self, lazy: bool = False, **kwargs: Any) -> Any:
+    async def get(self, lazy: bool = False, **kwargs: Any) -> AsyncStructuredNode:
         """
         Retrieve one node from the set matching supplied parameters
         :param lazy: False by default, specify True to get nodes with id only without the parameters.
@@ -1650,7 +1651,7 @@ async def get(self, lazy: bool = False, **kwargs: Any) -> Any:
             raise self.source_class.DoesNotExist(repr(kwargs))
         return result[0]
 
-    async def get_or_none(self, **kwargs: Any) -> Any:
+    async def get_or_none(self, **kwargs: Any) -> AsyncStructuredNode | None:
         """
         Retrieve a node from the set matching supplied parameters or return none
 
@@ -1662,7 +1663,7 @@ async def get_or_none(self, **kwargs: Any) -> Any:
         except self.source_class.DoesNotExist:
             return None
 
-    async def first(self, **kwargs: Any) -> Any:
+    async def first(self, **kwargs: Any) -> AsyncStructuredNode:
         """
         Retrieve the first node from the set matching supplied parameters
 
@@ -1675,7 +1676,7 @@ async def first(self, **kwargs: Any) -> Any:
         else:
             raise self.source_class.DoesNotExist(repr(kwargs))
 
-    async def first_or_none(self, **kwargs: Any) -> Any:
+    async def first_or_none(self, **kwargs: Any) -> Self | None:
         """
         Retrieve the first node from the set matching supplied parameters or return none
 
@@ -1688,7 +1689,7 @@ async def first_or_none(self, **kwargs: Any) -> Any:
             pass
         return None
 
-    def filter(self, *args: Any, **kwargs: Any) -> "AsyncBaseSet":
+    def filter(self, *args: Any, **kwargs: Any) -> Self:
         """
         Apply filters to the existing nodes in the set.
 
@@ -1754,7 +1755,7 @@ def filter(self, *args: Any, **kwargs: Any) -> "AsyncBaseSet":
 
         return self
 
-    def exclude(self, *args: Any, **kwargs: Any) -> "AsyncBaseSet":
+    def exclude(self, *args: Any, **kwargs: Any) -> Self:
         """
         Exclude nodes from the NodeSet via filters.
 
@@ -1768,13 +1769,13 @@ def exclude(self, *args: Any, **kwargs: Any) -> "AsyncBaseSet":
     @deprecated(
         "This method is deprecated and set to be removed in a future release. Please use .filter(has_rel__exists=True) instead."
     )
-    def has(self, **kwargs: Any) -> "AsyncBaseSet":
+    def has(self, **kwargs: Any) -> Self:
         must_match, dont_match = process_has_args(self.source_class, kwargs)
         self.must_match.update(must_match)
         self.dont_match.update(dont_match)
         return self
 
-    def order_by(self, *props: Any) -> "AsyncBaseSet":
+    def order_by(self, *props: Any) -> Self:
         """
         Order by properties. Prepend with minus to do descending. Pass None to
         remove ordering.
@@ -1820,14 +1821,12 @@ def _register_relation_to_fetch(
             item.alias = alias
         return item
 
-    def unique_variables(self, *paths: str) -> "AsyncNodeSet":
+    def unique_variables(self, *paths: str) -> Self:
         """Generate unique variable names for the given paths."""
         self._unique_variables = list(paths)
         return self
 
-    def traverse(
-        self, *paths: tuple[str, ...], **aliased_paths: dict
-    ) -> "AsyncNodeSet":
+    def traverse(self, *paths: tuple[str, ...], **aliased_paths: dict) -> Self:
         """Specify a set of paths to traverse."""
         relations = []
         for path in paths:
@@ -1839,7 +1838,7 @@ def traverse(
         self.relations_to_fetch = relations
         return self
 
-    def annotate(self, *vars: tuple, **aliased_vars: tuple) -> "AsyncNodeSet":
+    def annotate(self, *vars: tuple, **aliased_vars: tuple) -> Self:
         """Annotate node set results with extra variables."""
 
         def register_extra_var(
@@ -1925,10 +1924,11 @@ async def resolve_subgraph(self) -> list:
 
     async def subquery(
         self,
-        nodeset: "AsyncNodeSet",
+        nodeset: Self,
         return_set: list[str],
-        initial_context: list[str] | None = None,
-    ) -> "AsyncNodeSet":
+        initial_context: list[str | NodeNameResolver | RelationNameResolver | RawCypher]
+        | None = None,
+    ) -> Self:
         """Add a subquery to this node set.
 
         A subquery is a regular cypher query but executed within the context of a CALL
@@ -1959,9 +1959,9 @@ async def subquery(
             ):
                 raise RuntimeError(f"Variable '{var}' is not returned by subquery.")
         if initial_context:
-            for var in initial_context:
-                if not isinstance(var, str) and not isinstance(
-                    var, (NodeNameResolver, RelationNameResolver, RawCypher)
+            for context_var in initial_context:
+                if not isinstance(context_var, str) and not isinstance(
+                    context_var, (NodeNameResolver, RelationNameResolver, RawCypher)
                 ):
                     raise ValueError(
                         f"Wrong variable specified in initial context, should be a string or an instance of NodeNameResolver or RelationNameResolver"
@@ -1981,7 +1981,7 @@ def intermediate_transform(
         vars: dict[str, Transformation],
         distinct: bool = False,
         ordering: list | None = None,
-    ) -> "AsyncNodeSet":
+    ) -> Self:
         if not vars:
             raise ValueError(
                 "You must provide one variable at least when calling intermediate_transform()"
@@ -2057,7 +2057,7 @@ def __init__(self, source: Any, name: str, definition: dict) -> None:
         self.name = name
         self.filters: list = []
 
-    def match(self, **kwargs: Any) -> "AsyncTraversal":
+    def match(self, **kwargs: dict[str, Any]) -> "AsyncTraversal":
         """
         Traverse relationships with properties matching the given parameters.
 

neomodel/async_/relationship_manager.py

@@ -44,7 +44,7 @@ def is_direct_subclass(obj: Any, classinfo: Any) -> bool:
     return False
 
 
-class AsyncRelationshipManager(object):
+class AsyncRelationshipManager:
     """
     Base class for all relationships managed through neomodel.
 
@@ -350,7 +350,7 @@ def _new_traversal(self) -> AsyncTraversal:
         return AsyncTraversal(self.source, self.name, self.definition)
 
     # The methods below simply proxy the match engine.
-    async def get(self, **kwargs: Any) -> AsyncNodeSet:
+    async def get(self, **kwargs: Any):
         """
         Retrieve a related node with the matching node properties.
 
@@ -359,7 +359,7 @@ async def get(self, **kwargs: Any) -> AsyncNodeSet:
         """
         return await AsyncNodeSet(self._new_traversal()).get(**kwargs)
 
-    async def get_or_none(self, **kwargs: dict) -> AsyncNodeSet:
+    async def get_or_none(self, **kwargs: dict):
         """
         Retrieve a related node with the matching node properties or return None.
 
@@ -417,7 +417,7 @@ async def single(self) -> Optional["AsyncStructuredNode"]:
         except IndexError:
             return None
 
-    def match(self, **kwargs: dict) -> AsyncNodeSet:
+    def match(self, **kwargs: Any) -> AsyncNodeSet:
         """
         Return set of nodes who's relationship properties match supplied args
 
@@ -457,7 +457,7 @@ class AsyncRelationshipDefinition:
     def __init__(
         self,
         relation_type: str,
-        cls_name: str,
+        cls_name: str | type,
         direction: int,
         manager: type[AsyncRelationshipManager] = AsyncRelationshipManager,
         model: type[AsyncStructuredRel] | None = None,
@@ -513,9 +513,11 @@ def __init__(
                 adb._NODE_CLASS_REGISTRY[label_set] = model
 
     def _validate_class(
-        self, cls_name: str, model: type[AsyncStructuredRel] | None = None
+        self,
+        cls_name: str | type[AsyncStructuredNode],
+        model: type[AsyncStructuredRel] | None = None,
     ) -> None:
-        if not isinstance(cls_name, (str, object)):
+        if not isinstance(cls_name, str) and not isinstance(cls_name, type):
             raise ValueError("Expected class name or class got " + repr(cls_name))
 
         if model and not issubclass(model, (AsyncStructuredRel,)):
@@ -642,7 +644,7 @@ class AsyncZeroOrMore(AsyncRelationshipManager):
 class AsyncRelationshipTo(AsyncRelationshipDefinition):
     def __init__(
         self,
-        cls_name: str,
+        cls_name: str | type,
         relation_type: str,
         cardinality: type[AsyncRelationshipManager] = AsyncZeroOrMore,
         model: type[AsyncStructuredRel] | None = None,
@@ -659,7 +661,7 @@ def __init__(
 class AsyncRelationshipFrom(AsyncRelationshipDefinition):
     def __init__(
         self,
-        cls_name: str,
+        cls_name: str | type,
         relation_type: str,
         cardinality: type[AsyncRelationshipManager] = AsyncZeroOrMore,
         model: type[AsyncStructuredRel] | None = None,
@@ -676,7 +678,7 @@ def __init__(
 class AsyncRelationship(AsyncRelationshipDefinition):
     def __init__(
         self,
-        cls_name: str,
+        cls_name: str | type,
         relation_type: str,
         cardinality: type[AsyncRelationshipManager] = AsyncZeroOrMore,
         model: type[AsyncStructuredRel] | None = None,

neomodel/properties.py

@@ -137,8 +137,8 @@ def __init__(
         db_property: str | None = None,
         label: str | None = None,
         help_text: str | None = None,
-        **kwargs: dict[str, Any],
-    ):
+        **kwargs: Any,
+    ) -> None:
         if default is not None and required:
             raise ValueError(
                 "The arguments `required` and `default` are mutually exclusive."
@@ -273,7 +273,7 @@ def __init__(
         choices: Any | None = None,
         max_length: int | None = None,
         **kwargs: Any,
-    ):
+    ) -> None:
         if max_length is not None:
             if choices is not None:
                 raise ValueError(