cocoindex 0.1.43__cp311-cp311-macosx_11_0_arm64.whl → 0.1.45__cp311-cp311-macosx_11_0_arm64.whl

cocoindex/flow.py

@@ -10,14 +10,24 @@ import inspect
 import datetime
 import functools
 
-from typing import Any, Callable, Sequence, TypeVar, Generic, get_args, get_origin, Type, NamedTuple
+from typing import (
+    Any,
+    Callable,
+    Sequence,
+    TypeVar,
+    Generic,
+    get_args,
+    get_origin,
+    NamedTuple,
+    cast,
+)
 from threading import Lock
 from enum import Enum
 from dataclasses import dataclass
 from rich.text import Text
 from rich.tree import Tree
 
-from . import _engine
+from . import _engine  # type: ignore
 from . import index
 from . import op
 from . import setting
@@ -25,11 +35,12 @@ from .convert import dump_engine_object, encode_engine_value, make_engine_value_
 from .typing import encode_enriched_type
 from .runtime import execution_context
 
+
 class _NameBuilder:
     _existing_names: set[str]
     _next_name_index: dict[str, int]
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._existing_names = set()
         self._next_name_index = {}
 
@@ -51,40 +62,53 @@ class _NameBuilder:
                 return name
 
 
-_WORD_BOUNDARY_RE = re.compile('(?<!^)(?=[A-Z])')
+_WORD_BOUNDARY_RE = re.compile("(?<!^)(?=[A-Z])")
+
+
 def _to_snake_case(name: str) -> str:
-    return _WORD_BOUNDARY_RE.sub('_', name).lower()
+    return _WORD_BOUNDARY_RE.sub("_", name).lower()
+
 
 def _create_data_slice(
-        flow_builder_state: _FlowBuilderState,
-        creator: Callable[[_engine.DataScopeRef | None, str | None], _engine.DataSlice],
-        name: str | None = None) -> DataSlice:
+    flow_builder_state: _FlowBuilderState,
+    creator: Callable[[_engine.DataScopeRef | None, str | None], _engine.DataSlice],
+    name: str | None = None,
+) -> DataSlice[T]:
     if name is None:
-        return DataSlice(_DataSliceState(
-            flow_builder_state,
-            lambda target:
-                creator(target[0], target[1]) if target is not None else creator(None, None)))
+        return DataSlice(
+            _DataSliceState(
+                flow_builder_state,
+                lambda target: creator(target[0], target[1])
+                if target is not None
+                else creator(None, None),
+            )
+        )
     else:
         return DataSlice(_DataSliceState(flow_builder_state, creator(None, name)))
 
 
 def _spec_kind(spec: Any) -> str:
-    return spec.__class__.__name__
+    return cast(str, spec.__class__.__name__)
+
+
+T = TypeVar("T")
 
-T = TypeVar('T')
 
 class _DataSliceState:
     flow_builder_state: _FlowBuilderState
 
     _lazy_lock: Lock | None = None  # None means it's not lazy.
     _data_slice: _engine.DataSlice | None = None
-    _data_slice_creator: Callable[[tuple[_engine.DataScopeRef, str] | None],
-                                  _engine.DataSlice] | None = None
+    _data_slice_creator: (
+        Callable[[tuple[_engine.DataScopeRef, str] | None], _engine.DataSlice] | None
+    ) = None
 
     def __init__(
-            self, flow_builder_state: _FlowBuilderState,
-            data_slice: _engine.DataSlice | Callable[[tuple[_engine.DataScopeRef, str] | None],
-                                                     _engine.DataSlice]):
+        self,
+        flow_builder_state: _FlowBuilderState,
+        data_slice: _engine.DataSlice
+        | Callable[[tuple[_engine.DataScopeRef, str] | None], _engine.DataSlice],
+    ):
         self.flow_builder_state = flow_builder_state
 
         if isinstance(data_slice, _engine.DataSlice):
@@ -124,6 +148,7 @@ class _DataSliceState:
         # TODO: We'll support this by an identity transformer or "aliasing" in the future.
         raise ValueError("DataSlice is already attached to a field")
 
+
 class DataSlice(Generic[T]):
     """A data slice represents a slice of data in a flow. It's readonly."""
 
@@ -132,13 +157,13 @@ class DataSlice(Generic[T]):
     def __init__(self, state: _DataSliceState):
         self._state = state
 
-    def __str__(self):
+    def __str__(self) -> str:
         return str(self._state.engine_data_slice)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return repr(self._state.engine_data_slice)
 
-    def __getitem__(self, field_name: str) -> DataSlice:
+    def __getitem__(self, field_name: str) -> DataSlice[T]:
         field_slice = self._state.engine_data_slice.field(field_name)
         if field_slice is None:
             raise KeyError(field_name)
@@ -158,7 +183,9 @@ class DataSlice(Generic[T]):
         with self.row() as scope:
             f(scope)
 
-    def transform(self, fn_spec: op.FunctionSpec, *args, **kwargs) -> DataSlice:
+    def transform(
+        self, fn_spec: op.FunctionSpec, *args: Any, **kwargs: Any
+    ) -> DataSlice[T]:
         """
         Apply a function to the data slice.
         """
@@ -167,63 +194,77 @@ class DataSlice(Generic[T]):
 
         transform_args: list[tuple[Any, str | None]]
         transform_args = [(self._state.engine_data_slice, None)]
-        transform_args += [(self._state.flow_builder_state.get_data_slice(v), None) for v in args]
-        transform_args += [(self._state.flow_builder_state.get_data_slice(v), k)
-                           for (k, v) in kwargs.items()]
+        transform_args += [
+            (self._state.flow_builder_state.get_data_slice(v), None) for v in args
+        ]
+        transform_args += [
+            (self._state.flow_builder_state.get_data_slice(v), k)
+            for (k, v) in kwargs.items()
+        ]
 
         flow_builder_state = self._state.flow_builder_state
         return _create_data_slice(
             flow_builder_state,
-            lambda target_scope, name:
-                flow_builder_state.engine_flow_builder.transform(
-                    _spec_kind(fn_spec),
-                    dump_engine_object(fn_spec),
-                    transform_args,
-                    target_scope,
-                    flow_builder_state.field_name_builder.build_name(
-                        name, prefix=_to_snake_case(_spec_kind(fn_spec))+'_'),
-                ))
-
-    def call(self, func: Callable[[DataSlice], T], *args, **kwargs) -> T:
+            lambda target_scope, name: flow_builder_state.engine_flow_builder.transform(
+                _spec_kind(fn_spec),
+                dump_engine_object(fn_spec),
+                transform_args,
+                target_scope,
+                flow_builder_state.field_name_builder.build_name(
+                    name, prefix=_to_snake_case(_spec_kind(fn_spec)) + "_"
+                ),
+            ),
+        )
+
+    def call(self, func: Callable[[DataSlice[T]], T], *args: Any, **kwargs: Any) -> T:
         """
         Call a function with the data slice.
         """
         return func(self, *args, **kwargs)
 
-def _data_slice_state(data_slice: DataSlice) -> _DataSliceState:
+
+def _data_slice_state(data_slice: DataSlice[T]) -> _DataSliceState:
     return data_slice._state  # pylint: disable=protected-access
 
+
 class DataScope:
     """
     A data scope in a flow.
     It has multple fields and collectors, and allow users to add new fields and collectors.
     """
+
     _flow_builder_state: _FlowBuilderState
     _engine_data_scope: _engine.DataScopeRef
 
-    def __init__(self, flow_builder_state: _FlowBuilderState, data_scope: _engine.DataScopeRef):
+    def __init__(
+        self, flow_builder_state: _FlowBuilderState, data_scope: _engine.DataScopeRef
+    ):
         self._flow_builder_state = flow_builder_state
         self._engine_data_scope = data_scope
 
-    def __str__(self):
+    def __str__(self) -> str:
         return str(self._engine_data_scope)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return repr(self._engine_data_scope)
 
-    def __getitem__(self, field_name: str) -> DataSlice:
-        return DataSlice(_DataSliceState(
-            self._flow_builder_state,
-            self._flow_builder_state.engine_flow_builder.scope_field(
-                self._engine_data_scope, field_name)))
+    def __getitem__(self, field_name: str) -> DataSlice[T]:
+        return DataSlice(
+            _DataSliceState(
+                self._flow_builder_state,
+                self._flow_builder_state.engine_flow_builder.scope_field(
+                    self._engine_data_scope, field_name
+                ),
+            )
+        )
 
-    def __setitem__(self, field_name: str, value: DataSlice):
+    def __setitem__(self, field_name: str, value: DataSlice[T]) -> None:
         value._state.attach_to_scope(self._engine_data_scope, field_name)
 
-    def __enter__(self):
+    def __enter__(self) -> DataScope:
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         del self._engine_data_scope
 
     def add_collector(self, name: str | None = None) -> DataCollector:
@@ -233,27 +274,36 @@ class DataScope:
         return DataCollector(
             self._flow_builder_state,
             self._engine_data_scope.add_collector(
-                self._flow_builder_state.field_name_builder.build_name(name, prefix="_collector_")
-            )
+                self._flow_builder_state.field_name_builder.build_name(
+                    name, prefix="_collector_"
+                )
+            ),
         )
 
+
 class GeneratedField(Enum):
     """
     A generated field is automatically set by the engine.
     """
+
     UUID = "Uuid"
 
+
 class DataCollector:
     """A data collector is used to collect data into a collector."""
+
     _flow_builder_state: _FlowBuilderState
     _engine_data_collector: _engine.DataCollector
 
-    def __init__(self, flow_builder_state: _FlowBuilderState,
-                 data_collector: _engine.DataCollector):
+    def __init__(
+        self,
+        flow_builder_state: _FlowBuilderState,
+        data_collector: _engine.DataCollector,
+    ):
         self._flow_builder_state = flow_builder_state
         self._engine_data_collector = data_collector
 
-    def collect(self, **kwargs):
+    def collect(self, **kwargs: Any) -> None:
         """
         Collect data into the collector.
         """
@@ -268,45 +318,62 @@ class DataCollector:
                 else:
                     raise ValueError(f"Unexpected generated field: {v}")
             else:
-                regular_kwargs.append(
-                    (k, self._flow_builder_state.get_data_slice(v)))
+                regular_kwargs.append((k, self._flow_builder_state.get_data_slice(v)))
 
         self._flow_builder_state.engine_flow_builder.collect(
-            self._engine_data_collector, regular_kwargs, auto_uuid_field)
+            self._engine_data_collector, regular_kwargs, auto_uuid_field
+        )
 
-    def export(self, name: str, target_spec: op.StorageSpec, /, *,
-              primary_key_fields: Sequence[str],
-              vector_indexes: Sequence[index.VectorIndexDef] = (),
-              vector_index: Sequence[tuple[str, index.VectorSimilarityMetric]] = (),
-              setup_by_user: bool = False):
+    def export(
+        self,
+        name: str,
+        target_spec: op.StorageSpec,
+        /,
+        *,
+        primary_key_fields: Sequence[str],
+        vector_indexes: Sequence[index.VectorIndexDef] = (),
+        vector_index: Sequence[tuple[str, index.VectorSimilarityMetric]] = (),
+        setup_by_user: bool = False,
+    ) -> None:
         """
         Export the collected data to the specified target.
 
         `vector_index` is for backward compatibility only. Please use `vector_indexes` instead.
         """
         if not isinstance(target_spec, op.StorageSpec):
-            raise ValueError("export() can only be called on a CocoIndex target storage")
+            raise ValueError(
+                "export() can only be called on a CocoIndex target storage"
+            )
 
         # For backward compatibility only.
         if len(vector_indexes) == 0 and len(vector_index) > 0:
-            vector_indexes = [index.VectorIndexDef(field_name=field_name, metric=metric)
-                             for field_name, metric in vector_index]
+            vector_indexes = [
+                index.VectorIndexDef(field_name=field_name, metric=metric)
+                for field_name, metric in vector_index
+            ]
 
         index_options = index.IndexOptions(
             primary_key_fields=primary_key_fields,
             vector_indexes=vector_indexes,
         )
         self._flow_builder_state.engine_flow_builder.export(
-            name, _spec_kind(target_spec), dump_engine_object(target_spec),
-            dump_engine_object(index_options), self._engine_data_collector, setup_by_user)
+            name,
+            _spec_kind(target_spec),
+            dump_engine_object(target_spec),
+            dump_engine_object(index_options),
+            self._engine_data_collector,
+            setup_by_user,
+        )
 
 
 _flow_name_builder = _NameBuilder()
 
+
 class _FlowBuilderState:
     """
     A flow builder is used to build a flow.
     """
+
     engine_flow_builder: _engine.FlowBuilder
     field_name_builder: _NameBuilder
 
@@ -322,32 +389,40 @@ class _FlowBuilderState:
             return v._state.engine_data_slice
         return self.engine_flow_builder.constant(encode_enriched_type(type(v)), v)
 
+
 @dataclass
 class _SourceRefreshOptions:
     """
     Options for refreshing a source.
     """
+
     refresh_interval: datetime.timedelta | None = None
 
+
 class FlowBuilder:
     """
     A flow builder is used to build a flow.
     """
+
     _state: _FlowBuilderState
 
     def __init__(self, state: _FlowBuilderState):
         self._state = state
 
-    def __str__(self):
+    def __str__(self) -> str:
         return str(self._state.engine_flow_builder)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return repr(self._state.engine_flow_builder)
 
-    def add_source(self, spec: op.SourceSpec, /, *,
-            name: str | None = None,
-            refresh_interval: datetime.timedelta | None = None,
-        ) -> DataSlice:
+    def add_source(
+        self,
+        spec: op.SourceSpec,
+        /,
+        *,
+        name: str | None = None,
+        refresh_interval: datetime.timedelta | None = None,
+    ) -> DataSlice[T]:
         """
         Import a source to the flow.
         """
@@ -360,30 +435,37 @@ class FlowBuilder:
                 dump_engine_object(spec),
                 target_scope,
                 self._state.field_name_builder.build_name(
-                    name, prefix=_to_snake_case(_spec_kind(spec))+'_'),
-                dump_engine_object(_SourceRefreshOptions(refresh_interval=refresh_interval)),
+                    name, prefix=_to_snake_case(_spec_kind(spec)) + "_"
+                ),
+                dump_engine_object(
+                    _SourceRefreshOptions(refresh_interval=refresh_interval)
+                ),
             ),
-            name
+            name,
         )
 
-    def declare(self, spec: op.DeclarationSpec):
+    def declare(self, spec: op.DeclarationSpec) -> None:
         """
         Add a declaration to the flow.
         """
         self._state.engine_flow_builder.declare(dump_engine_object(spec))
 
+
 @dataclass
 class FlowLiveUpdaterOptions:
     """
     Options for live updating a flow.
     """
+
     live_mode: bool = True
     print_stats: bool = False
 
+
 class FlowLiveUpdater:
     """
     A live updater for a flow.
     """
+
     _flow: Flow
     _options: FlowLiveUpdaterOptions
     _engine_live_updater: _engine.FlowLiveUpdater | None = None
@@ -396,7 +478,7 @@ class FlowLiveUpdater:
         self.start()
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         self.abort()
         self.wait()
 
@@ -404,7 +486,7 @@ class FlowLiveUpdater:
         await self.start_async()
         return self
 
-    async def __aexit__(self, exc_type, exc_value, traceback):
+    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         self.abort()
         await self.wait_async()
 
@@ -419,7 +501,8 @@ class FlowLiveUpdater:
         Start the live updater.
         """
         self._engine_live_updater = await _engine.FlowLiveUpdater.create(
-            await self._flow.internal_flow_async(), dump_engine_object(self._options))
+            await self._flow.internal_flow_async(), dump_engine_object(self._options)
+        )
 
     def wait(self) -> None:
         """
@@ -456,22 +539,28 @@ class EvaluateAndDumpOptions:
     """
     Options for evaluating and dumping a flow.
     """
+
     output_dir: str
     use_cache: bool = True
 
+
 class Flow:
     """
     A flow describes an indexing pipeline.
     """
+
     _name: str
     _full_name: str
     _lazy_engine_flow: Callable[[], _engine.Flow]
 
-    def __init__(self, name: str, full_name: str, engine_flow_creator: Callable[[], _engine.Flow]):
+    def __init__(
+        self, name: str, full_name: str, engine_flow_creator: Callable[[], _engine.Flow]
+    ):
         self._name = name
         self._full_name = full_name
         engine_flow = None
         lock = Lock()
+
         def _lazy_engine_flow() -> _engine.Flow:
             nonlocal engine_flow, lock
             if engine_flow is None:
@@ -479,6 +568,7 @@ class Flow:
                     if engine_flow is None:
                         engine_flow = engine_flow_creator()
             return engine_flow
+
         self._lazy_engine_flow = _lazy_engine_flow
 
     def _render_spec(self, verbose: bool = False) -> Tree:
@@ -488,7 +578,7 @@ class Flow:
         spec = self._get_spec(verbose=verbose)
         tree = Tree(f"Flow: {self.full_name}", style="cyan")
 
-        def build_tree(label: str, lines: list):
+        def build_tree(label: str, lines: list[Any]) -> Tree:
             node = Tree(label=label if lines else label + " None", style="cyan")
             for line in lines:
                 child_node = node.add(Text(line.content, style="yellow"))
@@ -501,15 +591,17 @@ class Flow:
         return tree
 
     def _get_spec(self, verbose: bool = False) -> _engine.RenderedSpec:
-        return self._lazy_engine_flow().get_spec(output_mode="verbose" if verbose else "concise")
-    
+        return self._lazy_engine_flow().get_spec(
+            output_mode="verbose" if verbose else "concise"
+        )
+
     def _get_schema(self) -> list[tuple[str, str, str]]:
-        return self._lazy_engine_flow().get_schema()
+        return cast(list[tuple[str, str, str]], self._lazy_engine_flow().get_schema())
 
-    def __str__(self):
+    def __str__(self) -> str:
         return str(self._get_spec())
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return repr(self._lazy_engine_flow())
 
     @property
@@ -538,11 +630,15 @@ class Flow:
         Update the index defined by the flow.
         Once the function returns, the index is fresh up to the moment when the function is called.
         """
-        updater = await FlowLiveUpdater.create_async(self, FlowLiveUpdaterOptions(live_mode=False))
-        await updater.wait_async()
+        async with FlowLiveUpdater(
+            self, FlowLiveUpdaterOptions(live_mode=False)
+        ) as updater:
+            await updater.wait_async()
         return updater.update_stats()
 
-    def evaluate_and_dump(self, options: EvaluateAndDumpOptions):
+    def evaluate_and_dump(
+        self, options: EvaluateAndDumpOptions
+    ) -> _engine.IndexUpdateInfo:
         """
         Evaluate the flow and dump flow outputs to files.
         """
@@ -560,19 +656,26 @@ class Flow:
         """
         return await asyncio.to_thread(self.internal_flow)
 
-def _create_lazy_flow(name: str | None, fl_def: Callable[[FlowBuilder, DataScope], None]) -> Flow:
+
+def _create_lazy_flow(
+    name: str | None, fl_def: Callable[[FlowBuilder, DataScope], None]
+) -> Flow:
     """
     Create a flow without really building it yet.
     The flow will be built the first time when it's really needed.
     """
     flow_name = _flow_name_builder.build_name(name, prefix="_flow_")
     flow_full_name = get_full_flow_name(flow_name)
+
     def _create_engine_flow() -> _engine.Flow:
         flow_builder_state = _FlowBuilderState(flow_full_name)
         root_scope = DataScope(
-            flow_builder_state, flow_builder_state.engine_flow_builder.root_scope())
+            flow_builder_state, flow_builder_state.engine_flow_builder.root_scope()
+        )
         fl_def(FlowBuilder(flow_builder_state), root_scope)
-        return flow_builder_state.engine_flow_builder.build_flow(execution_context.event_loop)
+        return flow_builder_state.engine_flow_builder.build_flow(
+            execution_context.event_loop
+        )
 
     return Flow(flow_name, flow_full_name, _create_engine_flow)
 
@@ -580,28 +683,36 @@ def _create_lazy_flow(name: str | None, fl_def: Callable[[FlowBuilder, DataScope
 _flows_lock = Lock()
 _flows: dict[str, Flow] = {}
 
+
 def get_full_flow_name(name: str) -> str:
     """
     Get the full name of a flow.
     """
     return f"{setting.get_app_namespace(trailing_delimiter='.')}{name}"
 
+
 def add_flow_def(name: str, fl_def: Callable[[FlowBuilder, DataScope], None]) -> Flow:
     """Add a flow definition to the cocoindex library."""
-    if not all(c.isalnum() or c == '_' for c in name):
-        raise ValueError(f"Flow name '{name}' contains invalid characters. Only alphanumeric characters and underscores are allowed.")
+    if not all(c.isalnum() or c == "_" for c in name):
+        raise ValueError(
+            f"Flow name '{name}' contains invalid characters. Only alphanumeric characters and underscores are allowed."
+        )
     with _flows_lock:
         if name in _flows:
             raise KeyError(f"Flow with name {name} already exists")
         fl = _flows[name] = _create_lazy_flow(name, fl_def)
     return fl
 
-def flow_def(name = None) -> Callable[[Callable[[FlowBuilder, DataScope], None]], Flow]:
+
+def flow_def(
+    name: str | None = None,
+) -> Callable[[Callable[[FlowBuilder, DataScope], None]], Flow]:
     """
     A decorator to wrap the flow definition.
     """
     return lambda fl_def: add_flow_def(name or fl_def.__name__, fl_def)
 
+
 def flow_names() -> list[str]:
     """
     Get the names of all flows.
@@ -609,6 +720,7 @@ def flow_names() -> list[str]:
     with _flows_lock:
         return list(_flows.keys())
 
+
 def flows() -> dict[str, Flow]:
     """
     Get all flows.
@@ -616,6 +728,7 @@ def flows() -> dict[str, Flow]:
     with _flows_lock:
         return dict(_flows)
 
+
 def flow_by_name(name: str) -> Flow:
     """
     Get a flow by name.
@@ -623,12 +736,14 @@ def flow_by_name(name: str) -> Flow:
     with _flows_lock:
         return _flows[name]
 
+
 def ensure_all_flows_built() -> None:
     """
     Ensure all flows are built.
     """
     execution_context.run(ensure_all_flows_built_async())
 
+
 async def ensure_all_flows_built_async() -> None:
     """
     Ensure all flows are built.
@@ -636,43 +751,63 @@ async def ensure_all_flows_built_async() -> None:
     for fl in flows().values():
         await fl.internal_flow_async()
 
-def update_all_flows(options: FlowLiveUpdaterOptions) -> dict[str, _engine.IndexUpdateInfo]:
+
+def update_all_flows(
+    options: FlowLiveUpdaterOptions,
+) -> dict[str, _engine.IndexUpdateInfo]:
     """
     Update all flows.
     """
-    return execution_context.run(update_all_flows_async(options))
+    return cast(
+        dict[str, _engine.IndexUpdateInfo],
+        execution_context.run(update_all_flows_async(options)),
+    )
 
-async def update_all_flows_async(options: FlowLiveUpdaterOptions) -> dict[str, _engine.IndexUpdateInfo]:
+
+async def update_all_flows_async(
+    options: FlowLiveUpdaterOptions,
+) -> dict[str, _engine.IndexUpdateInfo]:
     """
     Update all flows.
     """
     await ensure_all_flows_built_async()
+
     async def _update_flow(name: str, fl: Flow) -> tuple[str, _engine.IndexUpdateInfo]:
         async with FlowLiveUpdater(fl, options) as updater:
             await updater.wait_async()
             return (name, updater.update_stats())
+
     fls = flows()
-    all_stats = await asyncio.gather(*(_update_flow(name, fl) for (name, fl) in fls.items()))
+    all_stats = await asyncio.gather(
+        *(_update_flow(name, fl) for (name, fl) in fls.items())
+    )
     return dict(all_stats)
 
-def _get_data_slice_annotation_type(data_slice_type: Type[DataSlice[T]]) -> Type[T] | None:
+
+def _get_data_slice_annotation_type(
+    data_slice_type: type[DataSlice[T] | inspect._empty],
+) -> type[T] | None:
     type_args = get_args(data_slice_type)
     if data_slice_type is inspect.Parameter.empty or data_slice_type is DataSlice:
         return None
     if get_origin(data_slice_type) != DataSlice or len(type_args) != 1:
         raise ValueError(f"Expect a DataSlice[T] type, but got {data_slice_type}")
-    return type_args[0]
+    return cast(type[T] | None, type_args[0])
+
 
 _transform_flow_name_builder = _NameBuilder()
 
+
 class TransformFlowInfo(NamedTuple):
     engine_flow: _engine.TransientFlow
     result_decoder: Callable[[Any], T]
 
+
 class TransformFlow(Generic[T]):
     """
     A transient transformation flow that transforms in-memory data.
     """
+
     _flow_fn: Callable[..., DataSlice[T]]
     _flow_name: str
     _flow_arg_types: list[Any]
@@ -682,21 +817,27 @@ class TransformFlow(Generic[T]):
     _lazy_flow_info: TransformFlowInfo | None = None
 
     def __init__(
-            self, flow_fn: Callable[..., DataSlice[T]],
-            flow_arg_types: Sequence[Any], /, name: str | None = None):
+        self,
+        flow_fn: Callable[..., DataSlice[T]],
+        flow_arg_types: Sequence[Any],
+        /,
+        name: str | None = None,
+    ):
         self._flow_fn = flow_fn
-        self._flow_name = _transform_flow_name_builder.build_name(name, prefix="_transform_flow_")
+        self._flow_name = _transform_flow_name_builder.build_name(
+            name, prefix="_transform_flow_"
+        )
         self._flow_arg_types = list(flow_arg_types)
         self._lazy_lock = asyncio.Lock()
 
-    def __call__(self, *args, **kwargs) -> DataSlice[T]:
+    def __call__(self, *args: Any, **kwargs: Any) -> DataSlice[T]:
         return self._flow_fn(*args, **kwargs)
 
     @property
     def _flow_info(self) -> TransformFlowInfo:
         if self._lazy_flow_info is not None:
             return self._lazy_flow_info
-        return execution_context.run(self._flow_info_async())
+        return cast(TransformFlowInfo, execution_context.run(self._flow_info_async()))
 
     async def _flow_info_async(self) -> TransformFlowInfo:
         if self._lazy_flow_info is not None:
@@ -712,35 +853,57 @@ class TransformFlow(Generic[T]):
         if len(sig.parameters) != len(self._flow_arg_types):
             raise ValueError(
                 f"Number of parameters in the flow function ({len(sig.parameters)}) "
-                f"does not match the number of argument types ({len(self._flow_arg_types)})")
+                f"does not match the number of argument types ({len(self._flow_arg_types)})"
+            )
 
-        kwargs: dict[str, DataSlice] = {}
-        for (param_name, param), param_type in zip(sig.parameters.items(), self._flow_arg_types):
-            if param.kind not in (inspect.Parameter.POSITIONAL_OR_KEYWORD,
-                                  inspect.Parameter.KEYWORD_ONLY):
-                raise ValueError(f"Parameter `{param_name}` is not a parameter can be passed by name")
+        kwargs: dict[str, DataSlice[T]] = {}
+        for (param_name, param), param_type in zip(
+            sig.parameters.items(), self._flow_arg_types
+        ):
+            if param.kind not in (
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                inspect.Parameter.KEYWORD_ONLY,
+            ):
+                raise ValueError(
+                    f"Parameter `{param_name}` is not a parameter can be passed by name"
+                )
             encoded_type = encode_enriched_type(param_type)
             if encoded_type is None:
                 raise ValueError(f"Parameter `{param_name}` has no type annotation")
-            engine_ds = flow_builder_state.engine_flow_builder.add_direct_input(param_name, encoded_type)
-            kwargs[param_name] = DataSlice(_DataSliceState(flow_builder_state, engine_ds))
+            engine_ds = flow_builder_state.engine_flow_builder.add_direct_input(
+                param_name, encoded_type
+            )
+            kwargs[param_name] = DataSlice(
+                _DataSliceState(flow_builder_state, engine_ds)
+            )
 
         output = self._flow_fn(**kwargs)
         flow_builder_state.engine_flow_builder.set_direct_output(
-            _data_slice_state(output).engine_data_slice)
-        engine_flow = await flow_builder_state.engine_flow_builder.build_transient_flow_async(execution_context.event_loop)
+            _data_slice_state(output).engine_data_slice
+        )
+        engine_flow = (
+            await flow_builder_state.engine_flow_builder.build_transient_flow_async(
+                execution_context.event_loop
+            )
+        )
         self._param_names = list(sig.parameters.keys())
 
-        engine_return_type = _data_slice_state(output).engine_data_slice.data_type().schema()
-        python_return_type = _get_data_slice_annotation_type(sig.return_annotation)
-        result_decoder = make_engine_value_decoder([], engine_return_type['type'], python_return_type)
+        engine_return_type = (
+            _data_slice_state(output).engine_data_slice.data_type().schema()
+        )
+        python_return_type: type[T] | None = _get_data_slice_annotation_type(
+            sig.return_annotation
+        )
+        result_decoder = make_engine_value_decoder(
+            [], engine_return_type["type"], python_return_type
+        )
 
         return TransformFlowInfo(engine_flow, result_decoder)
 
-    def __str__(self):
+    def __str__(self) -> str:
         return str(self._flow_info.engine_flow)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return repr(self._flow_info.engine_flow)
 
     def internal_flow(self) -> _engine.TransientFlow:
@@ -749,13 +912,13 @@ class TransformFlow(Generic[T]):
         """
         return self._flow_info.engine_flow
 
-    def eval(self, *args, **kwargs) -> T:
+    def eval(self, *args: Any, **kwargs: Any) -> T:
         """
         Evaluate the transform flow.
         """
-        return execution_context.run(self.eval_async(*args, **kwargs))
+        return cast(T, execution_context.run(self.eval_async(*args, **kwargs)))
 
-    async def eval_async(self, *args, **kwargs) -> T:
+    async def eval_async(self, *args: Any, **kwargs: Any) -> T:
         """
         Evaluate the transform flow.
         """
@@ -776,18 +939,26 @@ def transform_flow() -> Callable[[Callable[..., DataSlice[T]]], TransformFlow[T]
     """
     A decorator to wrap the transform function.
     """
-    def _transform_flow_wrapper(fn: Callable[..., DataSlice[T]]):
+
+    def _transform_flow_wrapper(fn: Callable[..., DataSlice[T]]) -> TransformFlow[T]:
         sig = inspect.signature(fn)
         arg_types = []
-        for (param_name, param) in sig.parameters.items():
-            if param.kind not in (inspect.Parameter.POSITIONAL_OR_KEYWORD,
-                                  inspect.Parameter.KEYWORD_ONLY):
-                raise ValueError(f"Parameter `{param_name}` is not a parameter can be passed by name")
-            value_type_annotation = _get_data_slice_annotation_type(param.annotation)
+        for param_name, param in sig.parameters.items():
+            if param.kind not in (
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                inspect.Parameter.KEYWORD_ONLY,
+            ):
+                raise ValueError(
+                    f"Parameter `{param_name}` is not a parameter can be passed by name"
+                )
+            value_type_annotation: type[T] | None = _get_data_slice_annotation_type(
+                param.annotation
+            )
             if value_type_annotation is None:
                 raise ValueError(
                     f"Parameter `{param_name}` for {fn} has no value type annotation. "
-                    "Please use `cocoindex.DataSlice[T]` where T is the type of the value.")
+                    "Please use `cocoindex.DataSlice[T]` where T is the type of the value."
+                )
             arg_types.append(value_type_annotation)
 
         _transform_flow = TransformFlow(fn, arg_types)