fix: clarify state updates type

qua-platform · Sep 20, 2024 · 83fb120 · 83fb120
1 parent 04de4df
commit 83fb120
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/qualibrate_runner/api/routes/others.py b/qualibrate_runner/api/routes/others.py
@@ -1,4 +1,4 @@
-from typing import Annotated, Optional
+from typing import Annotated, Mapping, Optional, cast
 
 from fastapi import APIRouter, Depends, HTTPException, status
 
@@ -8,7 +8,11 @@
 from qualibrate_runner.config import (
     State,
 )
-from qualibrate_runner.core.models.last_run import LastRun, RunStatus
+from qualibrate_runner.core.models.last_run import (
+    LastRun,
+    RunStatus,
+    StateUpdate,
+)
 
 others_router = APIRouter()
 
@@ -54,5 +58,6 @@ def state_updated(
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
             detail="Unknown state update key.",
         )
+    state_updates = cast(Mapping[str, StateUpdate], state_updates)
     state_updates[key].updated = True
     return state.last_run
diff --git a/qualibrate_runner/core/models/last_run.py b/qualibrate_runner/core/models/last_run.py
@@ -41,7 +41,9 @@ class LastRun(BaseModel):
     runnable_type: RunnableType
     passed_parameters: Mapping[str, Any] = Field(default_factory=dict)
     run_result: Optional[Union[NodeRunSummary, GraphRunSummary]] = None
-    state_updates: Mapping[str, StateUpdate] = Field(default_factory=dict)
+    state_updates: Union[
+        Mapping[str, StateUpdate], Mapping[str, Mapping[str, StateUpdate]]
+    ] = Field(default_factory=dict)
     error: Optional[RunError] = None
 
     @computed_field