diff --git a/rocky/rocky/views/mixins.py b/rocky/rocky/views/mixins.py index 3b7871380dd..2a21b60a5f2 100644 --- a/rocky/rocky/views/mixins.py +++ b/rocky/rocky/views/mixins.py @@ -3,6 +3,7 @@ from datetime import datetime, timezone from functools import cached_property from operator import attrgetter +from typing import TypedDict import structlog from account.mixins import OrganizationView @@ -47,6 +48,12 @@ class OriginData(BaseModel): params: dict[str, str] | None = None +class Origins(TypedDict): + declarations: list[OriginData] + observations: list[OriginData] + inferences: list[OriginData] + + class OOIAttributeError(AttributeError): pass @@ -99,11 +106,11 @@ def get_origins( self, reference: Reference, organization: Organization, - ) -> tuple[list[OriginData], list[OriginData], list[OriginData]]: + ) -> Origins: declarations: list[OriginData] = [] observations: list[OriginData] = [] inferences: list[OriginData] = [] - results = declarations, observations, inferences + results: Origins = {"declarations": declarations, "observations": observations, "inferences": inferences} try: origins = self.octopoes_api_connector.list_origins(self.observed_at, result=reference) diff --git a/rocky/rocky/views/ooi_detail.py b/rocky/rocky/views/ooi_detail.py index fac902817f7..fb086b97d28 100644 --- a/rocky/rocky/views/ooi_detail.py +++ b/rocky/rocky/views/ooi_detail.py @@ -93,10 +93,10 @@ def get_context_data(self, **kwargs): context["boefjes"] = [boefje for boefje in boefjes if boefje.scan_level.value <= max_level] context["ooi"] = self.ooi - declarations, observations, inferences = self.get_origins(self.ooi.reference, self.organization) + context.update(self.get_origins(self.ooi.reference, self.organization)) inference_params = self.octopoes_api_connector.list_origin_parameters( - {inference.origin.id for inference in inferences}, + {inference.origin.id for inference in context["inferences"]}, self.observed_at, ) inference_params_per_inference = defaultdict(list) @@ -104,13 +104,9 @@ def get_context_data(self, **kwargs): inference_params_per_inference[inference_param.origin_id].append(inference_param) inference_origin_params: list[tuple] = [] - for inference in inferences: + for inference in context["inferences"]: inference_origin_params.append((inference, inference_params_per_inference[inference.origin.id])) - context["declarations"] = declarations - - context["observations"] = observations - context["inferences"] = inferences context["inference_origin_params"] = inference_origin_params context["member"] = self.organization_member