Improved test, adjusted custom metrics

Author: anuunchin

dlt/extract/incremental/__init__.py

@@ -1,6 +1,16 @@
 import os
 from datetime import datetime  # noqa: I251
-from typing import Generic, ClassVar, Any, Optional, Type, Dict, Union, Literal, Tuple
+from typing import (
+    Generic,
+    ClassVar,
+    Any,
+    Optional,
+    Type,
+    Dict,
+    Union,
+    Literal,
+    Tuple,
+)
 
 import inspect
 from functools import wraps
@@ -20,6 +30,7 @@
     is_optional_type,
     is_subclass,
     TColumnNames,
+    TypedDict,
 )
 from dlt.common.configuration import configspec, ConfigurationValueError
 from dlt.common.configuration.specs import BaseConfiguration
@@ -63,8 +74,17 @@
     pandas = None
 
 
+class IncrementalCustomMetrics(TypedDict, total=False):
+    unfiltered_items_count: int
+    unfiltered_batches_count: int
+    initial_unique_hashes_count: int
+    final_unique_hashes_count: int
+
+
 @configspec
-class Incremental(ItemTransform[TDataItem], BaseConfiguration, Generic[TCursorValue]):
+class Incremental(
+    ItemTransform[TDataItem, IncrementalCustomMetrics], BaseConfiguration, Generic[TCursorValue]
+):
     """Adds incremental extraction for a resource by storing a cursor value in persistent state.
 
     The cursor could for example be a timestamp for when the record was created and you can use this to load only
@@ -191,8 +211,12 @@ def __init__(
         """Bound pipe"""
         self.range_start = range_start
         self.range_end = range_end
-        # Initialize custom metrics
-        BaseItemTransform.__init__(self)
+        self._custom_metrics: IncrementalCustomMetrics = {
+            "unfiltered_items_count": 0,
+            "unfiltered_batches_count": 0,
+            "initial_unique_hashes_count": 0,
+            "final_unique_hashes_count": 0,
+        }
 
     @property
     def primary_key(self) -> Optional[TTableHintTemplate[TColumnNames]]:
@@ -570,13 +594,8 @@ def __call__(self, rows: TDataItems, meta: Any = None) -> Optional[TDataItems]:
             return rows
 
         # collect metrics
-        self.custom_metrics["unfiltered_items_count"] = self.custom_metrics.get(
-            "unfiltered_items_count", 0
-        ) + count_rows_in_items(rows)
-        self.custom_metrics["unfiltered_batches_count"] = (
-            self.custom_metrics.get("unfiltered_batches_count", 0) + 1
-        )
-        self.custom_metrics["unique_hashes_count"] = len(self.get_state().get("unique_hashes", []))
+        self.custom_metrics["unfiltered_items_count"] += count_rows_in_items(rows)
+        self.custom_metrics["unfiltered_batches_count"] += 1
 
         transformer = self._get_transform(rows)
         if isinstance(rows, list):
@@ -599,6 +618,10 @@ def __call__(self, rows: TDataItems, meta: Any = None) -> Optional[TDataItems]:
         # writing back state
         self._cached_state["last_value"] = transformer.last_value
 
+        initial_hash_list = self._cached_state.get("unique_hashes")
+        initial_hash_count = len(initial_hash_list) if initial_hash_list else 0
+        self.custom_metrics["initial_unique_hashes_count"] = initial_hash_count
+
         if transformer.boundary_deduplication:
             # compute hashes for new last rows
             # NOTE: object transform uses last_rows to pass rows to dedup, arrow computes
@@ -607,11 +630,11 @@ def __call__(self, rows: TDataItems, meta: Any = None) -> Optional[TDataItems]:
                 transformer.compute_unique_value(row, self.primary_key)
                 for row in transformer.last_rows
             )
-            initial_hash_count = len(self._cached_state.get("unique_hashes", []))
             # add directly computed hashes
             unique_hashes.update(transformer.unique_hashes)
             self._cached_state["unique_hashes"] = list(unique_hashes)
             final_hash_count = len(self._cached_state["unique_hashes"])
+            self.custom_metrics["final_unique_hashes_count"] = final_hash_count
 
             self._check_duplicate_cursor_threshold(initial_hash_count, final_hash_count)
         return rows
@@ -636,7 +659,7 @@ def _check_duplicate_cursor_threshold(
 TIncrementalConfig = Union[Incremental[Any], IncrementalArgs]
 
 
-class IncrementalResourceWrapper(ItemTransform[TDataItem]):
+class IncrementalResourceWrapper(ItemTransform[TDataItem, IncrementalCustomMetrics]):
     placement_affinity: ClassVar[float] = 1  # stick to end
 
     _incremental: Optional[Incremental[Any]] = None
@@ -798,8 +821,8 @@ def allow_external_schedulers(self, value: bool) -> None:
             self._incremental.allow_external_schedulers = value
 
     @property
-    def custom_metrics(self) -> Dict[str, Any]:
-        """Returns custom metrics of the Incremental object itself"""
+    def custom_metrics(self) -> IncrementalCustomMetrics:
+        """Returns custom metrics of the Incremental object itself if exists"""
         if self._incremental:
             return self._incremental.custom_metrics
         return {}

dlt/extract/items_transform.py

@@ -11,6 +11,8 @@
     Optional,
     Union,
     Dict,
+    TypeVar,
+    cast,
 )
 
 from dlt.common.data_writers.writers import count_rows_in_items
@@ -32,16 +34,19 @@
 ItemTransformFunc = Union[ItemTransformFunctionWithMeta[TAny], ItemTransformFunctionNoMeta[TAny]]
 
 
-class BaseItemTransform:
+TCustomMetrics = TypeVar("TCustomMetrics", covariant=True)
+
+
+class BaseItemTransform(Generic[TCustomMetrics]):
     def __init__(self) -> None:
-        self._custom_metrics: Dict[str, Any] = {}
+        self._custom_metrics: TCustomMetrics = cast(TCustomMetrics, {})
 
     @property
-    def custom_metrics(self) -> Dict[str, Any]:
+    def custom_metrics(self) -> TCustomMetrics:
         return self._custom_metrics
 
 
-class ItemTransform(BaseItemTransform, ABC, Generic[TAny]):
+class ItemTransform(BaseItemTransform[TCustomMetrics], ABC, Generic[TAny, TCustomMetrics]):
     _f_meta: ItemTransformFunctionWithMeta[TAny] = None
     _f: ItemTransformFunctionNoMeta[TAny] = None
 
@@ -58,7 +63,9 @@ def __init__(self, transform_f: ItemTransformFunc[TAny]) -> None:
         else:  # TODO: do better check
             self._f_meta = transform_f  # type: ignore
 
-    def bind(self: "ItemTransform[TAny]", pipe: SupportsPipe) -> "ItemTransform[TAny]":
+    def bind(
+        self: "ItemTransform[TAny, TCustomMetrics]", pipe: SupportsPipe
+    ) -> "ItemTransform[TAny, TCustomMetrics]":
         return self
 
     @abstractmethod
@@ -67,7 +74,7 @@ def __call__(self, item: TDataItems, meta: Any = None) -> Optional[TDataItems]:
         pass
 
 
-class FilterItem(ItemTransform[bool]):
+class FilterItem(ItemTransform[bool, Dict[str, Any]]):
     # mypy needs those to type correctly
     _f_meta: ItemTransformFunctionWithMeta[bool]
     _f: ItemTransformFunctionNoMeta[bool]
@@ -93,7 +100,7 @@ def __call__(self, item: TDataItems, meta: Any = None) -> Optional[TDataItems]:
                 return item if self._f(item) else None
 
 
-class MapItem(ItemTransform[TDataItem]):
+class MapItem(ItemTransform[TDataItem, Dict[str, Any]]):
     # mypy needs those to type correctly
     _f_meta: ItemTransformFunctionWithMeta[TDataItem]
     _f: ItemTransformFunctionNoMeta[TDataItem]
@@ -115,7 +122,7 @@ def __call__(self, item: TDataItems, meta: Any = None) -> Optional[TDataItems]:
                 return self._f(item)
 
 
-class YieldMapItem(ItemTransform[Iterator[TDataItem]]):
+class YieldMapItem(ItemTransform[Iterator[TDataItem], Dict[str, Any]]):
     # mypy needs those to type correctly
     _f_meta: ItemTransformFunctionWithMeta[TDataItem]
     _f: ItemTransformFunctionNoMeta[TDataItem]
@@ -138,7 +145,7 @@ def __call__(self, item: TDataItems, meta: Any = None) -> Optional[TDataItems]:
                 yield from self._f(item)
 
 
-class ValidateItem(ItemTransform[TDataItem]):
+class ValidateItem(ItemTransform[TDataItem, Dict[str, Any]]):
     """Base class for validators of data items.
 
     Subclass should implement the `__call__` method to either return the data item(s) or raise `extract.exceptions.ValidationError`.
@@ -149,12 +156,12 @@ class ValidateItem(ItemTransform[TDataItem]):
 
     table_name: str
 
-    def bind(self, pipe: SupportsPipe) -> ItemTransform[TDataItem]:
+    def bind(self, pipe: SupportsPipe) -> ItemTransform[TDataItem, Dict[str, Any]]:
         self.table_name = pipe.name
         return self
 
 
-class LimitItem(ItemTransform[TDataItem]):
+class LimitItem(ItemTransform[TDataItem, Dict[str, Any]]):
     placement_affinity: ClassVar[float] = 1.1  # stick to end right behind incremental
 
     def __init__(

dlt/extract/pipe.py

@@ -12,6 +12,7 @@
     Iterator,
     List,
     Tuple,
+    Dict,
 )
 
 from dlt.common.reflection.inspect import isasyncgenfunction, isgeneratorfunction
@@ -42,7 +43,7 @@
 )
 
 
-class ForkPipe(ItemTransform[ResolvablePipeItem]):
+class ForkPipe(ItemTransform[ResolvablePipeItem, Dict[str, Any]]):
     placement_affinity: ClassVar[float] = 2
 
     def __init__(self, pipe: "Pipe", step: int = -1, copy_on_fork: bool = False) -> None: