refactor: use TypedDict

Author: Thomas Carroll

tests/functional/e2e/testcases/parser.py

@@ -1,12 +1,17 @@
 import json
 from pathlib import Path
-from typing import cast
+from typing import TypedDict, cast
 
 # Path to the shared testcases directory in athena-protobufs
 _REPO_ROOT = Path(__file__).parent.parent.parent.parent.parent
 TESTCASES_DIR = _REPO_ROOT / "athena-protobufs" / "testcases"
 
 
+class TestCases(TypedDict):
+    classification_labels: list[str]
+    images: list[list[str | list[float]]]
+
+
 class AthenaTestCase:
     def __init__(
         self,
@@ -28,13 +33,9 @@ def load_test_cases(dirname: str = "benign_model") -> list[AthenaTestCase]:
     with Path.open(
         Path(TESTCASES_DIR / dirname / "expected_outputs.json"),
     ) as f:
-        test_cases = cast(
-            "dict[str, list[str] | list[list[str | list[float]]]]", json.load(f)
-        )
-    classification_labels = cast(
-        "list[str]", test_cases["classification_labels"]
-    )
-    images = cast("list[list[str | list[float]]]", test_cases["images"])
+        test_cases: TestCases = cast("TestCases", json.load(f))
+    classification_labels = test_cases["classification_labels"]
+    images = test_cases["images"]
     return [
         AthenaTestCase(
             str(

tests/test_classify_single.py

@@ -8,11 +8,6 @@
 import grpc.aio
 import numpy as np
 import pytest
-
-from resolver_athena_client.client.athena_client import AthenaClient
-from resolver_athena_client.client.athena_options import AthenaOptions
-from resolver_athena_client.client.exceptions import AthenaError
-from resolver_athena_client.client.models import ImageData
 from resolver_athena_client.generated.athena.models_pb2 import (
     Classification,
     ClassificationError,
@@ -25,6 +20,11 @@
     RequestEncoding,
 )
 
+from resolver_athena_client.client.athena_client import AthenaClient
+from resolver_athena_client.client.athena_options import AthenaOptions
+from resolver_athena_client.client.exceptions import AthenaError
+from resolver_athena_client.client.models import ImageData
+
 
 @pytest.fixture
 def athena_options() -> AthenaOptions: