Minimize use of to_model, especially when fetching many

Author: edan-bainglass

aiida_restapi/routers/computers.py

@@ -22,7 +22,7 @@
 
 @read_router.get(
     '/schema',
-    response_model=dict,
+    response_model=dict[str, t.Any],
     responses={
         422: {'model': errors.RequestValidationError},
     },
@@ -32,7 +32,7 @@ async def get_computers_schema(
         t.Literal['get', 'post'],
         Query(description='Type of schema to retrieve: "get" or "post"'),
     ] = 'get',
-) -> dict:
+) -> dict[str, t.Any]:
     """Get JSON schema for AiiDA computers."""
     return service.get_schema(which=which)
 
@@ -61,7 +61,7 @@ async def get_computers(
         query.QueryParams,
         Depends(query.query_params),
     ],
-) -> PaginatedResults[orm.Computer.Model]:
+) -> PaginatedResults[dict[str, t.Any]]:
     """Get AiiDA computers with optional filtering, sorting, and/or pagination."""
     return service.get_many(query_params)
 
@@ -78,7 +78,7 @@ async def get_computers(
     },
 )
 @with_dbenv()
-async def get_computer(pk: str) -> orm.Computer.Model:
+async def get_computer(pk: str) -> dict[str, t.Any]:
     """Get AiiDA computer by pk."""
     return service.get_one(pk)
 
@@ -112,6 +112,6 @@ async def get_computer_metadata(pk: str) -> dict[str, t.Any]:
 async def create_computer(
     computer_model: orm.Computer.CreateModel,
     current_user: t.Annotated[UserInDB, Depends(get_current_active_user)],
-) -> orm.Computer.Model:
+) -> dict[str, t.Any]:
     """Create new AiiDA computer."""
     return service.add_one(computer_model)

aiida_restapi/routers/groups.py

@@ -22,7 +22,7 @@
 
 @read_router.get(
     '/schema',
-    response_model=dict,
+    response_model=dict[str, t.Any],
     responses={
         422: {'model': errors.RequestValidationError},
     },
@@ -32,7 +32,7 @@ async def get_groups_schema(
         t.Literal['get', 'post'],
         Query(description='Type of schema to retrieve: "get" or "post"'),
     ] = 'get',
-) -> dict:
+) -> dict[str, t.Any]:
     """Get JSON schema for AiiDA groups."""
     return service.get_schema(which=which)
 
@@ -61,7 +61,7 @@ async def get_groups(
         query.QueryParams,
         Depends(query.query_params),
     ],
-) -> PaginatedResults[orm.Group.Model]:
+) -> PaginatedResults[dict[str, t.Any]]:
     """Get AiiDA groups with optional filtering, sorting, and/or pagination."""
     return service.get_many(query_params)
 
@@ -78,7 +78,7 @@ async def get_groups(
     },
 )
 @with_dbenv()
-async def get_group(uuid: str) -> orm.Group.Model:
+async def get_group(uuid: str) -> dict[str, t.Any]:
     """Get AiiDA group by uuid."""
     return service.get_one(uuid)
 
@@ -95,9 +95,9 @@ async def get_group(uuid: str) -> orm.Group.Model:
     },
 )
 @with_dbenv()
-async def get_group_user(uuid: str) -> orm.User.Model:
+async def get_group_user(uuid: str) -> dict[str, t.Any]:
     """Get the user associated with a group."""
-    return t.cast(orm.User.Model, service.get_related_one(uuid, orm.User))
+    return service.get_related_one(uuid, orm.User)
 
 
 @read_router.get(
@@ -118,7 +118,7 @@ async def get_group_nodes(
         query.QueryParams,
         Depends(query.query_params),
     ],
-) -> PaginatedResults[orm.Node.Model]:
+) -> PaginatedResults[dict[str, t.Any]]:
     """Get the nodes of a group."""
     return service.get_related_many(uuid, orm.Node, query_params)
 
@@ -152,6 +152,6 @@ async def get_group_extras(uuid: str) -> dict[str, t.Any]:
 async def create_group(
     group_model: orm.Group.CreateModel,
     current_user: t.Annotated[UserInDB, Depends(get_current_active_user)],
-) -> orm.Group.Model:
+) -> dict[str, t.Any]:
     """Create new AiiDA group."""
     return service.add_one(group_model)

aiida_restapi/routers/nodes.py

@@ -39,7 +39,7 @@
 
 @read_router.get(
     '/schema',
-    response_model=dict,
+    response_model=dict[str, t.Any],
     responses={
         422: {'model': errors.InvalidNodeTypeError},
     },
@@ -56,7 +56,7 @@ async def get_nodes_schema(
         t.Literal['get', 'post'],
         Query(description='Type of schema to retrieve'),
     ] = 'get',
-) -> dict:
+) -> dict[str, t.Any]:
     """Get JSON schema for the base AiiDA node 'get' model."""
     if not node_type:
         return orm.Node.Model.model_json_schema()
@@ -120,7 +120,7 @@ async def get_nodes(
         query.QueryParams,
         Depends(query.query_params),
     ],
-) -> PaginatedResults[orm.Node.Model]:
+) -> PaginatedResults[dict[str, t.Any]]:
     """Get AiiDA nodes with optional filtering, sorting, and/or pagination."""
     return service.get_many(query_params)
 
@@ -158,7 +158,7 @@ async def get_node_types() -> list:
     },
 )
 @with_dbenv()
-async def get_node(uuid: str) -> orm.Node.Model:
+async def get_node(uuid: str) -> dict[str, t.Any]:
     """Get AiiDA node by uuid."""
     return service.get_one(uuid)
 
@@ -175,9 +175,9 @@ async def get_node(uuid: str) -> orm.Node.Model:
     },
 )
 @with_dbenv()
-async def get_node_user(uuid: str) -> orm.User.Model:
+async def get_node_user(uuid: str) -> dict[str, t.Any]:
     """Get the user associated with a node."""
-    return t.cast(orm.User.Model, service.get_related_one(uuid, orm.User))
+    return service.get_related_one(uuid, orm.User)
 
 
 @read_router.get(
@@ -192,9 +192,9 @@ async def get_node_user(uuid: str) -> orm.User.Model:
     },
 )
 @with_dbenv()
-async def get_node_computer(uuid: str) -> orm.Computer.Model:
+async def get_node_computer(uuid: str) -> dict[str, t.Any]:
     """Get the computer associated with a node."""
-    return t.cast(orm.Computer.Model, service.get_related_one(uuid, orm.Computer))
+    return service.get_related_one(uuid, orm.Computer)
 
 
 @read_router.get(
@@ -215,7 +215,7 @@ async def get_node_groups(
         query.QueryParams,
         Depends(query.query_params),
     ],
-) -> PaginatedResults[orm.Group.Model]:
+) -> PaginatedResults[dict[str, t.Any]]:
     """Get the groups of a node."""
     return service.get_related_many(uuid, orm.Group, query_params)
 

aiida_restapi/routers/submit.py

@@ -84,7 +84,7 @@ def validate_inputs(cls, inputs: dict[str, t.Any]) -> dict[str, t.Any]:
 async def submit_process(
     process: ProcessSubmitModel,
     current_user: t.Annotated[UserInDB, Depends(get_current_active_user)],
-) -> orm.Node.Model:
+) -> dict[str, t.Any]:
     """Submit new AiiDA process."""
     try:
         entry_point_process = load_entry_point_from_string(process.entry_point)
@@ -94,4 +94,4 @@ async def submit_process(
         process_node = engine.submit(entry_point_process, **process.inputs)
     except Exception as exception:
         raise exceptions.InputValidationError(str(exception)) from exception
-    return t.cast(orm.Node.Model, process_node.to_model())
+    return process_node.serialize(minimal=True)

aiida_restapi/routers/users.py

@@ -22,7 +22,7 @@
 
 @read_router.get(
     '/schema',
-    response_model=dict,
+    response_model=dict[str, t.Any],
     responses={
         422: {'model': errors.RequestValidationError},
     },
@@ -32,7 +32,7 @@ async def get_users_schema(
         t.Literal['get', 'post'],
         Query(description='Type of schema to retrieve: "get" or "post"'),
     ] = 'get',
-) -> dict:
+) -> dict[str, t.Any]:
     """Get JSON schema for AiiDA users."""
     return service.get_schema(which=which)
 
@@ -61,7 +61,7 @@ async def get_users(
         query.QueryParams,
         Depends(query.query_params),
     ],
-) -> PaginatedResults[orm.User.Model]:
+) -> PaginatedResults[dict[str, t.Any]]:
     """Get AiiDA users with optional filtering, sorting, and/or pagination."""
     return service.get_many(query_params)
 
@@ -76,7 +76,7 @@ async def get_users(
     },
 )
 @with_dbenv()
-async def get_user(pk: int) -> orm.User.Model:
+async def get_user(pk: int) -> dict[str, t.Any]:
     """Get AiiDA user by pk."""
     return service.get_one(pk)
 
@@ -95,6 +95,6 @@ async def get_user(pk: int) -> orm.User.Model:
 async def create_user(
     user_model: orm.User.CreateModel,
     current_user: t.Annotated[UserInDB, Depends(get_current_active_user)],
-) -> orm.User.Model:
+) -> dict[str, t.Any]:
     """Create new AiiDA user."""
     return service.add_one(user_model)