refactor: address code review feedback on public image info, API key resolution, and serialization

app.py

@@ -10,10 +10,10 @@ from fastapi.staticfiles import StaticFiles
 from cleaner import (
     DEFAULT_MODEL,
     CleanerError,
-    _get_image_info,
     build_prompt,
     clean_garment,
-    encode_image_bytes,
+    get_image_info,
+    resolve_api_key,
 )
 
 DEFAULT_PROMPT = build_prompt("beige")
@@ -53,23 +53,25 @@ async def clean_endpoint(
     api_key: str | None = Form(default=None),
     restore_res: bool = Form(default=True),
 ):
-    effective_api_key = (api_key.strip() if api_key and api_key.strip() else None) or os.environ.get(
-        "OPENROUTER_API_KEY"
-    )
+    effective_api_key = resolve_api_key(api_key)
     if not effective_api_key:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail="API key is required. Provide it in the form or set OPENROUTER_API_KEY.",
         )
 
-    file_bytes = await file.read()
+    try:
+        file_bytes = await file.read()
+    finally:
+        await file.close()
+
     if not file_bytes:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail="Uploaded file is empty.",
         )
 
-    orig_dims, _ = _get_image_info(file_bytes)
+    orig_dims, _ = get_image_info(file_bytes)
     if orig_dims is None:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -89,21 +91,7 @@ async def clean_endpoint(
             detail=str(exc),
         )
 
-    return {
-        "image": encode_image_bytes(clean_result.image_bytes, mime=clean_result.media_type),
-        "cost": clean_result.cost,
-        "width": clean_result.width,
-        "height": clean_result.height,
-        "original_dimensions": list(clean_result.original_dimensions)
-        if clean_result.original_dimensions
-        else None,
-        "model_dimensions": list(clean_result.model_dimensions)
-        if clean_result.model_dimensions
-        else None,
-        "was_rescaled": clean_result.was_rescaled,
-        "media_type": clean_result.media_type,
-    }
-
+    return clean_result.to_payload()
 if STATIC_DIR.is_dir():
     app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
 

cleaner.py

@@ -26,7 +26,14 @@ EXTENSIONS = {
 }
 
 
-def _get_image_info(data_or_path: bytes | str) -> tuple[tuple[int, int] | None, str | None]:
+def resolve_api_key(api_key: str | None = None) -> str | None:
+    """Resolve OpenRouter API key from parameter or OPENROUTER_API_KEY environment variable."""
+    if api_key and api_key.strip():
+        return api_key.strip()
+    return os.environ.get("OPENROUTER_API_KEY")
+
+
+def get_image_info(data_or_path: bytes | str) -> tuple[tuple[int, int] | None, str | None]:
     """Extract (width, height) dimensions and mime type from bytes, file path, or data URI."""
     try:
         from PIL import Image
@@ -34,7 +41,6 @@ def _get_image_info(data_or_path: bytes | str) -> tuple[tuple[int, int] | None,
         if isinstance(data_or_path, bytes):
             raw_bytes = data_or_path
         elif isinstance(data_or_path, str) and data_or_path.startswith("data:"):
-            # Parse data URI: data:<mime>;base64,<payload>
             comma_idx = data_or_path.find(",")
             if comma_idx != -1:
                 raw_bytes = base64.b64decode(data_or_path[comma_idx + 1:])
@@ -50,6 +56,9 @@ def _get_image_info(data_or_path: bytes | str) -> tuple[tuple[int, int] | None,
         pass
     return None, None
 
+
+_get_image_info = get_image_info
+
 def build_prompt(bg_choice: str) -> str:
     bg_description = BACKGROUND_PRESETS.get(bg_choice.lower(), bg_choice)
     return (
@@ -132,6 +141,22 @@ class CleanResult:
             return self.model_dimensions != self.original_dimensions
         return False
 
+    def to_payload(self) -> dict[str, Any]:
+        """Serialize result to a dictionary suitable for API responses."""
+        return {
+            "image": encode_image_bytes(self.image_bytes, mime=self.media_type),
+            "cost": self.cost,
+            "width": self.width,
+            "height": self.height,
+            "original_dimensions": list(self.original_dimensions)
+            if self.original_dimensions
+            else None,
+            "model_dimensions": list(self.model_dimensions)
+            if self.model_dimensions
+            else None,
+            "was_rescaled": self.was_rescaled,
+            "media_type": self.media_type,
+        }
 
 def call_openrouter_images(
     image_url_or_data_uri: str,

tests/test_cleaner.py

@@ -285,3 +285,35 @@ def test_clean_garment_file(tmp_path, monkeypatch):
     assert res.width == 250
     assert res.height == 350
     assert res.cost == 0.03
+
+def test_resolve_api_key(monkeypatch):
+    import cleaner
+    monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
+    assert cleaner.resolve_api_key(None) is None
+    assert cleaner.resolve_api_key("   ") is None
+    assert cleaner.resolve_api_key("custom-key") == "custom-key"
+    monkeypatch.setenv("OPENROUTER_API_KEY", "env-key")
+    assert cleaner.resolve_api_key(None) == "env-key"
+    assert cleaner.resolve_api_key("custom-key") == "custom-key"
+
+
+def test_clean_result_to_payload():
+    raw = b"img-data"
+    res = CleanResult(
+        image_bytes=raw,
+        media_type="image/jpeg",
+        cost=0.015,
+        width=100,
+        height=200,
+        original_dimensions=(100, 200),
+        model_dimensions=(50, 100),
+    )
+    payload = res.to_payload()
+    assert payload["cost"] == 0.015
+    assert payload["width"] == 100
+    assert payload["height"] == 200
+    assert payload["original_dimensions"] == [100, 200]
+    assert payload["model_dimensions"] == [50, 100]
+    assert payload["was_rescaled"] is True
+    assert payload["media_type"] == "image/jpeg"
+    assert payload["image"].startswith("data:image/jpeg;base64,")