Add missing docstring types

Author: edan-bainglass

aiida_restapi/common/query.py

@@ -58,11 +58,16 @@ def query_params(
     """Parse query parameters into a structured object.
 
     :param filters: AiiDA QueryBuilder filters as JSON string.
+    :type filters: str | None
     :param order_by: Comma-separated string of fields to sort by.
+    :type order_by: str | None
     :param page_size: Number of results per page.
+    :type page_size: pdt.PositiveInt
     :param page: Page number.
+    :type page: pdt.PositiveInt
     :return: Structured query parameters.
-    :raises HTTPException: If filters cannot be parsed as JSON.
+    :rtype: QueryParams
+    :raises HTTPException: If arguments cannot be parsed as JSON.
     """
     query_filters: dict[str, t.Any] = {}
     query_order_by: str | list[str] | dict[str, t.Any] | None = None

aiida_restapi/models/node.py

@@ -121,8 +121,6 @@ class NodeModelRegistry:
     """Registry for AiiDA REST API node models.
 
     This class maintains mappings of node types and their corresponding Pydantic models.
-
-    :ivar ModelUnion: A union type of all AiiDA node Pydantic models, discriminated by the `node_type` field.
     """
 
     def __init__(self) -> None:
@@ -136,22 +134,30 @@ def get_node_types(self) -> list[str]:
         """Get the list of registered node class names.
 
         :return: List of node class names.
+        :rtype: list[str]
         """
         return list(self._models.keys())
 
     def get_node_class_name(self, node_type: str) -> str:
         """Get the AiiDA node class name for a given node type.
 
         :param node_type: The AiiDA node type string.
+        :type node_type: str
         :return: The corresponding node class name.
+        :rtype: str
         """
         return node_type.rsplit('.', 2)[-2]
 
     def get_model(self, node_type: str, which: t.Literal['get', 'post'] = 'get') -> type[Node.Model]:
         """Get the Pydantic model class for a given node type.
 
         :param node_type: The AiiDA node type string.
+        :type node_type: str
+        :param which: Specify whether to get the 'get' or 'post' model.
         :return: The corresponding Pydantic model class.
+        :rtype: type[Node.Model]
+        :raises MissingEntryPointError: If the node type is not registered.
+        :raises KeyError: If the specified model type is unknown.
         """
         if (Model := self._models.get(node_type)) is None:
             raise MissingEntryPointError(f'Unknown node type: {node_type}')
@@ -163,7 +169,9 @@ def _get_node_post_model(self, node_cls: Node) -> type[Node.Model]:
         """Return a patched Model for the given node class with a literal `node_type` field.
 
         :param node_cls: The AiiDA node class.
+        :type node_cls: Node
         :return: The patched ORM Node model.
+        :rtype: type[Node.Model]
         """
         Model = node_cls.CreateModel
         node_type = node_cls.class_node_type
@@ -197,6 +205,7 @@ def _get_post_models(self) -> tuple[type[Node.Model], ...]:
         """Get a union type of all node 'post' models.
 
         :return: A union type of all node 'post' models.
+        :rtype: tuple[type[Node.Model], ...]
         """
         post_models = [model_dict['post'] for model_dict in self._models.values()]
         return tuple(post_models)

aiida_restapi/routers/server.py

@@ -217,7 +217,9 @@ def _get_route_parts(route: Route) -> tuple[str | None, set[str], str]:
     """Return the parts of a route: path, group, methods, description.
 
     :param route: A FastAPI/Starlette Route object.
+    :type route: Route
     :return: A tuple containing the group, methods, and description of the route.
+    :rtype: tuple[str | None, set[str], str]
     """
     prefix = re.escape(API_CONFIG['PREFIX'])
     match = re.match(rf'^{prefix}/([^/]+)/?.*', route.path)

aiida_restapi/routers/submit.py

@@ -24,7 +24,9 @@ def process_inputs(inputs: dict[str, t.Any]) -> dict[str, t.Any]:
     A node UUID is indicated by the key ending with the suffix ``.uuid``.
 
     :param inputs: The inputs dictionary.
+    :type inputs: dict[str, t.Any]
     :returns: The deserialized inputs dictionary.
+    :rtype: dict[str, t.Any]
     """
     uuid_suffix = '.uuid'
     results = {}
@@ -61,7 +63,9 @@ def validate_inputs(cls, inputs: dict[str, t.Any]) -> dict[str, t.Any]:
         """Process the inputs dictionary.
 
         :param inputs: The inputs to validate.
+        :type inputs: dict[str, t.Any]
         :returns: The validated inputs.
+        :rtype: dict[str, t.Any]
         """
         return process_inputs(inputs)