mypy

Author: pedroslopez

airbyte_cdk/manifest_runner/command_processor/processor.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Iterable, List, Mapping, Tuple
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple
 
 from airbyte_protocol_dataclasses.models import (
     AirbyteCatalog,
@@ -13,7 +13,7 @@
 from airbyte_cdk.connector_builder.models import StreamRead
 from airbyte_cdk.connector_builder.test_reader import TestReader
 from airbyte_cdk.entrypoint import AirbyteEntrypoint
-from airbyte_cdk.models.airbyte_protocol import (
+from airbyte_cdk.models import (
     AirbyteStateMessage,
     ConfiguredAirbyteCatalog,
 )
@@ -78,7 +78,7 @@ def check_connection(
     def discover(
         self,
         config: Mapping[str, Any],
-    ) -> AirbyteCatalog | None:
+    ) -> Optional[AirbyteCatalog]:
         """
         Discover the catalog from the source.
         """
@@ -91,11 +91,11 @@ def discover(
     def _get_messages_by_type(
         self,
         messages: Iterable[AirbyteMessage],
-    ) -> Mapping[str, Iterable[AirbyteMessage]]:
+    ) -> Dict[str, List[AirbyteMessage]]:
         """
         Group messages by type.
         """
-        grouped = {}
+        grouped: Dict[str, List[AirbyteMessage]] = {}
         for message in messages:
             msg_type = message.type
             if msg_type not in grouped:
@@ -106,7 +106,7 @@ def _get_messages_by_type(
     def _get_connection_status(
         self,
         messages_by_type: Mapping[str, List[AirbyteMessage]],
-    ) -> AirbyteConnectionStatus | None:
+    ) -> Optional[AirbyteConnectionStatus]:
         """
         Get the connection status from the messages.
         """
@@ -116,7 +116,7 @@ def _get_connection_status(
     def _get_catalog(
         self,
         messages_by_type: Mapping[str, List[AirbyteMessage]],
-    ) -> AirbyteCatalog:
+    ) -> Optional[AirbyteCatalog]:
         """
         Get the catalog from the messages.
         """

airbyte_cdk/manifest_runner/command_processor/utils.py

@@ -1,6 +1,6 @@
 from typing import Any, Mapping
 
-from airbyte_cdk.models.airbyte_protocol import (
+from airbyte_cdk.models import (
     AirbyteStream,
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteStream,

airbyte_cdk/manifest_runner/routers/health.py

@@ -1,3 +1,5 @@
+from typing import Dict
+
 from fastapi import APIRouter
 
 router = APIRouter(
@@ -7,5 +9,5 @@
 
 
 @router.get("/")
-def health():
+def health() -> Dict[str, str]:
     return {"status": "ok"}

airbyte_cdk/manifest_runner/routers/manifest.py

@@ -1,6 +1,6 @@
 import hashlib
 from dataclasses import asdict
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Mapping
 
 import jsonschema
 from fastapi import APIRouter, Depends, HTTPException
@@ -12,13 +12,15 @@
     DiscoverResponse,
 )
 from airbyte_cdk.models import AirbyteStateMessageSerializer
+from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
 from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
     INJECTED_COMPONENTS_PY,
     INJECTED_COMPONENTS_PY_CHECKSUMS,
 )
 
 from ..api_models import (
     FullResolveRequest,
+    Manifest,
     ManifestResponse,
     ResolveRequest,
     StreamRead,
@@ -29,7 +31,9 @@
 from ..command_processor.utils import build_catalog, build_source
 
 
-def safe_build_source(manifest_dict, config_dict):
+def safe_build_source(
+    manifest_dict: Mapping[str, Any], config_dict: Mapping[str, Any]
+) -> ManifestDeclarativeSource:
     """Wrapper around build_source that converts ValidationError to HTTPException."""
     try:
         return build_source(manifest_dict, config_dict)
@@ -97,7 +101,7 @@ def discover(request: DiscoverRequest) -> DiscoverResponse:
 def resolve(request: ResolveRequest) -> ManifestResponse:
     """Resolve a manifest to its final configuration."""
     source = safe_build_source(request.manifest.model_dump(), {})
-    return ManifestResponse(manifest=source.resolved_manifest)
+    return ManifestResponse(manifest=Manifest(**source.resolved_manifest))
 
 
 @router.post("/full_resolve", operation_id="fullResolve")
@@ -125,4 +129,4 @@ def full_resolve(request: FullResolveRequest) -> ManifestResponse:
         streams.extend(generated_streams_list)
 
     manifest["streams"] = streams
-    return ManifestResponse(manifest=manifest)
+    return ManifestResponse(manifest=Manifest(**manifest))