Add type annotations

airbytehq · May 23, 2024 · d3b2591 · d3b2591
1 parent 84b40a0
commit d3b2591
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/airbyte-cdk/python/airbyte_cdk/sources/declarative/incremental/per_partition_cursor.py b/airbyte-cdk/python/airbyte_cdk/sources/declarative/incremental/per_partition_cursor.py
@@ -6,7 +6,7 @@
 from typing import Any, Callable, Iterable, Mapping, MutableMapping, Optional, Union
 
 from airbyte_cdk.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
-from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer import StreamSlicer
+from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
 from airbyte_cdk.sources.types import Record, StreamSlice, StreamState
 
 
@@ -62,7 +62,7 @@ class PerPartitionCursor(DeclarativeCursor):
     _KEY = 0
     _VALUE = 1
 
-    def __init__(self, cursor_factory: CursorFactory, partition_router: StreamSlicer):
+    def __init__(self, cursor_factory: CursorFactory, partition_router: PartitionRouter):
         self._cursor_factory = cursor_factory
         self._partition_router = partition_router
         self._cursor_per_partition: MutableMapping[str, DeclarativeCursor] = {}

diff --git a/airbyte-cdk/python/airbyte_cdk/sources/declarative/partition_routers/partition_router.py b/airbyte-cdk/python/airbyte_cdk/sources/declarative/partition_routers/partition_router.py
@@ -3,8 +3,10 @@
 #
 
 from dataclasses import dataclass
+from typing import Optional
 
 from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer import StreamSlicer
+from airbyte_cdk.sources.types import StreamState
 
 
 @dataclass
@@ -17,7 +19,7 @@ class PartitionRouter(StreamSlicer):
         get_parent_state(): Get the state of the parent streams.
     """
 
-    def set_parent_state(self, stream_state):
+    def set_parent_state(self, stream_state: StreamState) -> None:
         """
         Set the state of the parent streams.
 
@@ -26,7 +28,7 @@ def set_parent_state(self, stream_state):
         """
         pass
 
-    def get_parent_state(self):
+    def get_parent_state(self) -> Optional[StreamState]:
         """
         Get the state of the parent streams.
 

diff --git a/...dk/python/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py b/...dk/python/airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py
@@ -3,7 +3,7 @@
 #
 
 from dataclasses import InitVar, dataclass
-from typing import TYPE_CHECKING, Any, Iterable, List, Mapping, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Mapping, Optional, Union
 
 import dpath.util
 from airbyte_cdk.models import AirbyteMessage, SyncMode, Type
@@ -59,7 +59,7 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         if not self.parent_stream_configs:
             raise ValueError("SubstreamPartitionRouter needs at least 1 parent stream")
         self._parameters = parameters
-        self._parent_state = {}
+        self._parent_state: Dict[str, Any] = {}
 
     def get_request_params(
         self,