Fix mypy

airbytehq · Feb 12, 2025 · 1e71e63 · 1e71e63
1 parent e46a88a
commit 1e71e63
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py
@@ -1647,7 +1647,7 @@ def create_declarative_stream(
 
     def _build_stream_slicer_from_partition_router(
         self,
-        model: Union[AsyncRetrieverModel, CustomRetrieverModel, SimpleRetrieverModel],
+        model: Union[AsyncRetrieverModel, CustomRetrieverModel, SimpleRetrieverModel, StateDelegatingRetrieverModel],
         config: Config,
     ) -> Optional[PartitionRouter]:
         if (
@@ -1672,7 +1672,7 @@ def _build_stream_slicer_from_partition_router(
 
     def _build_resumable_cursor_from_paginator(
         self,
-        model: Union[AsyncRetrieverModel, CustomRetrieverModel, SimpleRetrieverModel],
+        model: Union[AsyncRetrieverModel, CustomRetrieverModel, SimpleRetrieverModel, StateDelegatingRetrieverModel],
         stream_slicer: Optional[StreamSlicer],
     ) -> Optional[StreamSlicer]:
         if hasattr(model, "paginator") and model.paginator and not stream_slicer:

diff --git a/airbyte_cdk/sources/declarative/retrievers/state_delegating_retriever.py b/airbyte_cdk/sources/declarative/retrievers/state_delegating_retriever.py
@@ -5,7 +5,7 @@
 from dataclasses import dataclass
 from typing import (
     Any,
-    MutableMapping,
+    Mapping,
 )
 
 from typing_extensions import deprecated
@@ -55,11 +55,11 @@ def retriever(self) -> Retriever:
         )
 
     @property
-    def state(self) -> MutableMapping[str, Any]:
+    def state(self) -> Mapping[str, Any]:
         return self.cursor.get_stream_state() if self.cursor else {}
 
     @state.setter
-    def state(self, value: MutableMapping[str, Any]) -> None:
+    def state(self, value: Mapping[str, Any]) -> None:
         """State setter, accept state serialized by state getter."""
         if self.cursor:
             self.cursor.set_initial_state(value)