Add type for build log message

Author: buurro

src/fastapi_cloud_cli/commands/deploy.py

@@ -343,7 +343,7 @@ def _wait_for_deployment(
             for log in client.stream_build_logs(deployment.id):
                 time_elapsed = time.monotonic() - started_at
 
-                if log.type == BuildLogType.message and log.message:
+                if log.type == BuildLogType.message:
                     progress.log(Text.from_ansi(log.message.rstrip()))
 
                 if log.type == BuildLogType.complete:

src/fastapi_cloud_cli/utils/api.py

@@ -3,10 +3,11 @@
 from contextlib import contextmanager
 from datetime import timedelta
 from enum import Enum
-from typing import ContextManager, Generator, Optional
+from typing import ContextManager, Generator, Literal, Optional
 
 import httpx
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, Field, TypeAdapter, ValidationError
+from typing_extensions import Annotated
 
 from fastapi_cloud_cli import __version__
 from fastapi_cloud_cli.config import Settings
@@ -18,7 +19,8 @@
 BUILD_LOG_TIMEOUT = timedelta(minutes=5)
 
 
-class BuildLogError(Exception): ...
+class BuildLogError(Exception):
+    pass
 
 
 class BuildLogType(str, Enum):
@@ -30,11 +32,28 @@ class BuildLogType(str, Enum):
 
 
 class BuildLogLine(BaseModel):
-    type: BuildLogType
+    type: Literal[
+        BuildLogType.complete,
+        BuildLogType.failed,
+        BuildLogType.timeout,
+        BuildLogType.heartbeat,
+    ]
     message: Optional[str] = None
     id: Optional[str] = None
 
 
+class BuildLogLineMessage(BaseModel):
+    type: Literal[BuildLogType.message] = BuildLogType.message
+    message: str
+    id: Optional[str] = None
+
+
+# check pydantic v1
+BuildLogAdapter = TypeAdapter[BuildLogLineMessage | BuildLogLine](
+    Annotated[BuildLogLineMessage | BuildLogLine, Field(discriminator="type")]
+)
+
+
 @contextmanager
 def attempt(attempt_number: int) -> Generator[None, None, None]:
     def _backoff() -> None:
@@ -108,7 +127,7 @@ def __init__(self) -> None:
 
     def stream_build_logs(
         self, deployment_id: str
-    ) -> Generator[BuildLogLine, None, None]:
+    ) -> Generator[BuildLogLineMessage | BuildLogLine, None, None]:
         last_id = None
 
         for attempt in attempts(BUILD_LOG_MAX_RETRIES, BUILD_LOG_TIMEOUT):
@@ -158,9 +177,11 @@ def stream_build_logs(
         # Exhausted retries without getting any response
         raise BuildLogError(f"Failed after {BUILD_LOG_MAX_RETRIES} attempts")
 
-    def _parse_log_line(self, line: str) -> Optional[BuildLogLine]:
+    def _parse_log_line(
+        self, line: str
+    ) -> Optional[BuildLogLineMessage | BuildLogLine]:
         try:
-            return BuildLogLine.model_validate_json(line)
+            return BuildLogAdapter.validate_json(line)
         except ValidationError as e:
             logger.debug("Skipping malformed log: %s (error: %s)", line[:100], e)
             return None