feat: add type for projections

Author: matmoncon

pyneo4j_ogm/core/node.py

@@ -36,7 +36,12 @@
 )
 from pyneo4j_ogm.fields.settings import NodeModelSettings, RelationshipModelSettings
 from pyneo4j_ogm.logger import logger
-from pyneo4j_ogm.queries.types import MultiHopFilters, NodeFilters, QueryOptions
+from pyneo4j_ogm.queries.types import (
+    MultiHopFilters,
+    NodeFilters,
+    Projection,
+    QueryOptions,
+)
 
 if TYPE_CHECKING:
     from pyneo4j_ogm.fields.relationship_property import RelationshipProperty
@@ -271,7 +276,7 @@ async def refresh(self) -> None:
     async def find_connected_nodes(
         self,
         filters: MultiHopFilters,
-        projections: Optional[Dict[str, str]] = None,
+        projections: Optional[Projection] = None,
         options: Optional[QueryOptions] = None,
         auto_fetch_nodes: bool = False,
         auto_fetch_models: Optional[List[Union[str, Type["NodeModel"]]]] = None,
@@ -281,7 +286,7 @@ async def find_connected_nodes(
 
         Args:
             filters (MultiHopFilters): The filters to apply to the query.
-            projections (Dict[str, str], optional): The properties to project from the node. The keys define
+            projections (Projection, optional): The properties to project from the node. The keys define
                 the new keys in the projection and the value defines the model property to be projected. A invalid
                 or empty projection will result in the whole model instance being returned. Defaults to `None`.
             options (QueryOptions, optional): The options to apply to the query. Defaults to `None`.
@@ -414,7 +419,7 @@ async def find_connected_nodes(
     async def find_one(
         cls: Type[T],
         filters: NodeFilters,
-        projections: Optional[Dict[str, str]] = None,
+        projections: Optional[Projection] = None,
         auto_fetch_nodes: bool = False,
         auto_fetch_models: Optional[List[Union[str, Type["NodeModel"]]]] = None,
         raise_on_empty: bool = False,
@@ -425,7 +430,7 @@ async def find_one(
 
         Args:
             filters (NodeFilters): The filters to apply to the query.
-            projections (Dict[str, str], optional): The properties to project from the node. The keys define
+            projections (Projection, optional): The properties to project from the node. The keys define
                 the new keys in the projection and the value defines the model property to be projected. A invalid
                 or empty projection will result in the whole model instance being returned. Defaults to `None`.
             auto_fetch_nodes (bool, optional): Whether to automatically fetch connected nodes. Takes priority over the
@@ -542,7 +547,7 @@ async def find_one(
     async def find_many(
         cls: Type[T],
         filters: Optional[NodeFilters] = None,
-        projections: Optional[Dict[str, str]] = None,
+        projections: Optional[Projection] = None,
         options: Optional[QueryOptions] = None,
         auto_fetch_nodes: bool = False,
         auto_fetch_models: Optional[List[Union[str, Type["NodeModel"]]]] = None,
@@ -553,7 +558,7 @@ async def find_many(
 
         Args:
             filters (NodeFilters, optional): The filters to apply to the query. Defaults to `None`.
-            projections (Dict[str, str], optional): The properties to project from the node. The keys define
+            projections (Projection, optional): The properties to project from the node. The keys define
                 the new keys in the projection and the value defines the model property to be projected. A invalid
                 or empty projection will result in the whole model instance being returned. Defaults to `None`.
             options (QueryOptions, optional): The options to apply to the query. Defaults to `None`.

pyneo4j_ogm/core/relationship.py

@@ -23,6 +23,7 @@
 from pyneo4j_ogm.fields.settings import RelationshipModelSettings
 from pyneo4j_ogm.logger import logger
 from pyneo4j_ogm.queries.types import (
+    Projection,
     QueryOptions,
     RelationshipFilters,
     RelationshipMatchDirection,
@@ -267,7 +268,7 @@ async def end_node(self) -> Type[NodeModel]:
     async def find_one(
         cls: Type[T],
         filters: RelationshipFilters,
-        projections: Optional[Dict[str, str]] = None,
+        projections: Optional[Projection] = None,
         raise_on_empty: bool = False,
     ) -> Optional[Union[T, Dict[str, Any]]]:
         """
@@ -276,7 +277,7 @@ async def find_one(
 
         Args:
             filters (RelationshipFilters): Expressions applied to the query.
-            projections (Dict[str, str], optional): The properties to project from the relationship. The keys define
+            projections (Projection, optional): The properties to project from the relationship. The keys define
                 the new keys in the projection and the value defines the model property to be projected. A invalid
                 or empty projection will result in the whole model instance being returned. Defaults to `None`.
             raise_on_empty (bool, optional): Whether to raise a `NoResultFound` if no match is found.
@@ -343,7 +344,7 @@ async def find_one(
     async def find_many(
         cls: Type[T],
         filters: Optional[RelationshipFilters] = None,
-        projections: Optional[Dict[str, str]] = None,
+        projections: Optional[Projection] = None,
         options: Optional[QueryOptions] = None,
     ) -> List[Union[T, Dict[str, Any]]]:
         """
@@ -353,7 +354,7 @@ async def find_many(
         Args:
             filters (RelationshipFilters | None, optional): Expressions applied to the query. Defaults to
                 `None`.
-            projections (Dict[str, str], optional): The properties to project from the relationship. The keys define
+            projections (Projection, optional): The properties to project from the relationship. The keys define
                 the new keys in the projection and the value defines the model property to be projected. A invalid
                 or empty projection will result in the whole model instance being returned. Defaults to `None`.
             options (QueryOptions | None, optional): Options for modifying the query result. Defaults to `None`.

pyneo4j_ogm/fields/relationship_property.py

@@ -41,6 +41,7 @@
 )
 from pyneo4j_ogm.queries.query_builder import QueryBuilder
 from pyneo4j_ogm.queries.types import (
+    Projection,
     QueryOptions,
     RelationshipFilters,
     RelationshipPropertyFilters,
@@ -249,7 +250,7 @@ async def relationships(
         self,
         node: T,
         filters: Optional[RelationshipFilters] = None,
-        projections: Optional[Dict[str, str]] = None,
+        projections: Optional[Projection] = None,
         options: Optional[QueryOptions] = None,
     ) -> List[U]:
         """
@@ -260,7 +261,7 @@ async def relationships(
             node (T): The node to which to get the relationship.
             filters (RelationshipFilters | None, optional): Expressions applied to the query. Defaults to
                 `None`.
-            projections (Dict[str, str], optional): The properties to project from the node. The keys define
+            projections (Projection, optional): The properties to project from the node. The keys define
                 the new keys in the projection and the value defines the model property to be projected. A invalid
                 or empty projection will result in the whole model instance being returned. Defaults to `None`.
             options (QueryOptions, optional): The options to apply to the query. Defaults to `None`.
@@ -681,7 +682,7 @@ async def replace(self, old_node: T, new_node: T) -> List[U]:
     async def find_connected_nodes(
         self,
         filters: Optional[RelationshipPropertyFilters] = None,
-        projections: Optional[Dict[str, str]] = None,
+        projections: Optional[Projection] = None,
         options: Optional[QueryOptions] = None,
         auto_fetch_nodes: bool = False,
         auto_fetch_models: Optional[List[Union[str, Type["NodeModel"]]]] = None,
@@ -691,7 +692,7 @@ async def find_connected_nodes(
 
         Args:
             filters (RelationshipPropertyFilters | None, optional): Expressions applied to the query. Defaults to `None`.
-            projections (Dict[str, str], optional): The properties to project from the node. A invalid or empty
+            projections (Projection, optional): The properties to project from the node. A invalid or empty
                 projection will result in the whole model instances being returned. Defaults to `None`.
             options (QueryOptions | None, optional): Options for modifying the query result. Defaults to `None`.
             auto_fetch_nodes (bool, optional): Whether to automatically fetch connected nodes. Takes priority over the

pyneo4j_ogm/queries/query_builder.py

@@ -11,6 +11,7 @@
 from pyneo4j_ogm.queries.types import (
     MultiHopFilters,
     NodeFilters,
+    Projection,
     QueryOptions,
     RelationshipFilters,
     RelationshipMatchDirection,
@@ -367,15 +368,13 @@ def relationship_match(
             case _:
                 raise InvalidRelationshipDirection(direction)
 
-    def build_projections(
-        self, projections: Dict[str, Union[str, Literal["$elementId"], Literal["$id"]]], ref: str = "n"
-    ) -> None:
+    def build_projections(self, projections: Projection, ref: str = "n") -> None:
         """
         Builds a projection which only returns the node properties defined in the projection.
 
         Args:
             ref (str): The reference to the node. Defaults to `'n'`.
-            projections (Dict[str, str]): The projections to build.
+            projections (Projection): The projections to build.
         """
         if not isinstance(projections, dict):
             return

pyneo4j_ogm/queries/types.py

@@ -196,3 +196,7 @@ class QueryOptions(TypedDict, total=False):
     skip: Optional[int]
     sort: Optional[Union[List[str], str]]
     order: Optional[QueryOptionsOrder]
+
+
+# Interface for a projection
+Projection = Dict[str, Union[str, Literal["$elementId"], Literal["$id"]]]