Fix typing issues

Author: dedeswim

src/prompt_siren/build_images.py

@@ -21,7 +21,7 @@
 import tempfile
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Literal
+from typing import Literal, TypeVar
 
 import click
 from pydantic import BaseModel
@@ -440,7 +440,10 @@ async def build_all_specs(self, specs: list[ImageBuildSpec]) -> list[BuildError]
         return errors
 
 
-def _maybe_override_cache_dir(config: BaseModel, cache_dir: str | None) -> BaseModel:
+T = TypeVar("T", bound=BaseModel)
+
+
+def _maybe_override_cache_dir(config: T, cache_dir: str | None) -> T:
     """Override cache_dir on configs that define it."""
     if cache_dir is None:
         return config

tests/test_build_images.py

@@ -6,11 +6,11 @@
 from unittest.mock import AsyncMock
 
 import pytest
-from pydantic import BaseModel
-from prompt_siren.build_images import ImageBuilder, _maybe_override_cache_dir
+from prompt_siren.build_images import _maybe_override_cache_dir, ImageBuilder
 from prompt_siren.datasets.swebench_dataset.config import SwebenchDatasetConfig
 from prompt_siren.sandbox_managers.docker.plugins.errors import DockerClientError
 from prompt_siren.sandbox_managers.image_spec import BuildImageSpec
+from pydantic import BaseModel
 
 
 class MockDockerClient: