Start properly typing processing_results and arguments_list (#839)

services/processing/intermediate.py

@@ -1,5 +1,4 @@
 from concurrent.futures import ThreadPoolExecutor
-from dataclasses import dataclass
 
 import orjson
 import sentry_sdk
@@ -8,20 +7,9 @@
 from shared.reports.resources import Report
 
 from services.archive import ArchiveService
-from services.processing.metrics import INTERMEDIATE_REPORT_SIZE
 
-
-@dataclass
-class IntermediateReport:
-    upload_id: int
-    """
-    The `Upload` id for which this report was loaded.
-    """
-
-    report: EditableReport
-    """
-    The loaded Report.
-    """
+from .metrics import INTERMEDIATE_REPORT_SIZE
+from .types import IntermediateReport
 
 
 @sentry_sdk.trace

services/processing/merging.py

@@ -1,5 +1,3 @@
-from dataclasses import dataclass
-
 import sentry_sdk
 from shared.reports.editable import EditableReport, EditableReportFile
 from shared.reports.enums import UploadState
@@ -8,23 +6,10 @@
 from sqlalchemy.orm import Session as DbSession
 
 from database.models.reports import Upload
-from services.processing.intermediate import IntermediateReport
 from services.report import delete_uploads_by_sessionid
 from services.report.raw_upload_processor import clear_carryforward_sessions
 
-
-@dataclass
-class MergeResult:
-    session_mapping: dict[int, int]
-    """
-    This is a mapping from the input `upload_id` to the output `session_id`
-    as it exists in the merged "master Report".
-    """
-
-    deleted_sessions: set[int]
-    """
-    The Set of carryforwarded `session_id`s that have been removed from the "master Report".
-    """
+from .types import IntermediateReport, MergeResult
 
 
 @sentry_sdk.trace

services/processing/processing.py

@@ -1,6 +1,5 @@
 import logging
 from collections.abc import Callable
-from typing import NotRequired, TypedDict
 
 import sentry_sdk
 from celery.exceptions import CeleryError
@@ -11,23 +10,14 @@
 from database.models.reports import Upload
 from helpers.reports import delete_archive_setting
 from services.archive import ArchiveService
-from services.processing.intermediate import save_intermediate_report
-from services.processing.state import ProcessingState
 from services.report import ProcessingError, RawReportInfo, Report, ReportService
 from services.report.parser.types import VersionOneParsedRawReport
 
-log = logging.getLogger(__name__)
-
-
-class UploadArguments(TypedDict):
-    # TODO(swatinem): migrate this over to `upload_id`
-    upload_pk: int
+from .intermediate import save_intermediate_report
+from .state import ProcessingState
+from .types import ProcessingResult, UploadArguments
 
-
-class ProcessingResult(TypedDict):
-    arguments: UploadArguments
-    successful: bool
-    error: NotRequired[dict]
+log = logging.getLogger(__name__)
 
 
 @sentry_sdk.trace
@@ -58,7 +48,9 @@ def process_upload(
     report_service = ReportService(commit_yaml)
     archive_service = report_service.get_archive_service(commit.repository)
 
-    result = ProcessingResult(arguments=arguments, successful=False)
+    result = ProcessingResult(
+        upload_id=upload_id, arguments=arguments, successful=False
+    )
 
     try:
         report = Report()

services/processing/types.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+from typing import NotRequired, TypedDict
+
+from shared.reports.editable import EditableReport
+
+
+class UploadArguments(TypedDict):
+    # TODO(swatinem): migrate this over to `upload_id`
+    upload_pk: int
+
+
+class ProcessingResult(TypedDict):
+    upload_id: int
+    arguments: UploadArguments
+    successful: bool
+    error: NotRequired[dict]
+
+
+@dataclass
+class IntermediateReport:
+    upload_id: int
+    """
+    The `Upload` id for which this report was loaded.
+    """
+
+    report: EditableReport
+    """
+    The loaded Report.
+    """
+
+
+@dataclass
+class MergeResult:
+    session_mapping: dict[int, int]
+    """
+    This is a mapping from the input `upload_id` to the output `session_id`
+    as it exists in the merged "master Report".
+    """
+
+    deleted_sessions: set[int]
+    """
+    The Set of carryforwarded `session_id`s that have been removed from the "master Report".
+    """