Rollback changes to TransformFunction

Author: gwaramadze

quixstreams/core/stream/functions/transform.py

@@ -1,13 +1,7 @@
 from typing import Any, Literal, Union, cast, overload
 
 from .base import StreamFunction
-from .types import (
-    TransformCallback,
-    TransformExpandedCallback,
-    TransformWallClockCallback,
-    TransformWallClockExpandedCallback,
-    VoidExecutor,
-)
+from .types import TransformCallback, TransformExpandedCallback, VoidExecutor
 
 __all__ = ("TransformFunction",)
 
@@ -29,81 +23,28 @@ class TransformFunction(StreamFunction):
 
     @overload
     def __init__(
-        self,
-        func: TransformCallback,
-        expand: Literal[False] = False,
-        wall_clock: Literal[False] = False,
-    ) -> None: ...
-
-    @overload
-    def __init__(
-        self,
-        func: TransformExpandedCallback,
-        expand: Literal[True],
-        wall_clock: Literal[False] = False,
-    ) -> None: ...
-
-    @overload
-    def __init__(
-        self,
-        func: TransformWallClockCallback,
-        expand: Literal[False] = False,
-        wall_clock: Literal[True] = True,
+        self, func: TransformCallback, expand: Literal[False] = False
     ) -> None: ...
 
     @overload
     def __init__(
-        self,
-        func: TransformWallClockExpandedCallback,
-        expand: Literal[True],
-        wall_clock: Literal[True],
+        self, func: TransformExpandedCallback, expand: Literal[True]
     ) -> None: ...
 
     def __init__(
         self,
-        func: Union[
-            TransformCallback,
-            TransformExpandedCallback,
-            TransformWallClockCallback,
-            TransformWallClockExpandedCallback,
-        ],
+        func: Union[TransformCallback, TransformExpandedCallback],
         expand: bool = False,
-        wall_clock: bool = False,
     ):
         super().__init__(func)
 
-        self.func: Union[
-            TransformCallback,
-            TransformExpandedCallback,
-            TransformWallClockCallback,
-            TransformWallClockExpandedCallback,
-        ]
+        self.func: Union[TransformCallback, TransformExpandedCallback]
         self.expand = expand
-        self.wall_clock = wall_clock
 
     def get_executor(self, *child_executors: VoidExecutor) -> VoidExecutor:
         child_executor = self._resolve_branching(*child_executors)
 
-        if self.expand and self.wall_clock:
-            wall_clock_expanded_func = cast(
-                TransformWallClockExpandedCallback, self.func
-            )
-
-            def wrapper(
-                value: Any,
-                key: Any,
-                timestamp: int,
-                headers: Any,
-            ):
-                for (
-                    new_value,
-                    new_key,
-                    new_timestamp,
-                    new_headers,
-                ) in wall_clock_expanded_func(timestamp):
-                    child_executor(new_value, new_key, new_timestamp, new_headers)
-
-        elif self.expand:
+        if self.expand:
             expanded_func = cast(TransformExpandedCallback, self.func)
 
             def wrapper(
@@ -116,22 +57,8 @@ def wrapper(
                 for new_value, new_key, new_timestamp, new_headers in result:
                     child_executor(new_value, new_key, new_timestamp, new_headers)
 
-        elif self.wall_clock:
-            wall_clock_func = cast(TransformWallClockCallback, self.func)
-
-            def wrapper(
-                value: Any,
-                key: Any,
-                timestamp: int,
-                headers: Any,
-            ):
-                new_value, new_key, new_timestamp, new_headers = wall_clock_func(
-                    timestamp
-                )
-                child_executor(new_value, new_key, new_timestamp, new_headers)
-
         else:
-            regular_func = cast(TransformCallback, self.func)
+            func = cast(TransformCallback, self.func)
 
             def wrapper(
                 value: Any,
@@ -140,7 +67,7 @@ def wrapper(
                 headers: Any,
             ):
                 # Execute a function on a single value and return its result
-                new_value, new_key, new_timestamp, new_headers = regular_func(
+                new_value, new_key, new_timestamp, new_headers = func(
                     value, key, timestamp, headers
                 )
                 child_executor(new_value, new_key, new_timestamp, new_headers)

quixstreams/core/stream/functions/types.py

@@ -14,8 +14,6 @@
     "FilterWithMetadataCallback",
     "TransformCallback",
     "TransformExpandedCallback",
-    "TransformWallClockCallback",
-    "TransformWallClockExpandedCallback",
 )
 
 
@@ -37,10 +35,6 @@ def __bool__(self) -> bool: ...
 TransformExpandedCallback = Callable[
     [Any, Any, int, Any], Iterable[Tuple[Any, Any, int, Any]]
 ]
-TransformWallClockCallback = Callable[[int], Tuple[Any, Any, int, Any]]
-TransformWallClockExpandedCallback = Callable[
-    [int], Iterable[Tuple[Any, Any, int, Any]]
-]
 
 StreamCallback = Union[
     ApplyCallback,

quixstreams/core/stream/stream.py

@@ -30,8 +30,6 @@
     TransformCallback,
     TransformExpandedCallback,
     TransformFunction,
-    TransformWallClockCallback,
-    TransformWallClockExpandedCallback,
     UpdateCallback,
     UpdateFunction,
     UpdateWithMetadataCallback,
@@ -256,7 +254,6 @@ def add_transform(
         func: TransformCallback,
         *,
         expand: Literal[False] = False,
-        wall_clock: Literal[False] = False,
     ):
         pass
 
@@ -266,41 +263,14 @@ def add_transform(
         func: TransformExpandedCallback,
         *,
         expand: Literal[True],
-        wall_clock: Literal[False] = False,
     ):
         pass
 
-    @overload
-    def add_transform(
-        self,
-        func: TransformWallClockCallback,
-        *,
-        expand: Literal[False] = False,
-        wall_clock: Literal[True],
-    ):
-        pass
-
-    @overload
     def add_transform(
         self,
-        func: TransformWallClockExpandedCallback,
-        *,
-        expand: Literal[True],
-        wall_clock: Literal[True],
-    ):
-        pass
-
-    def add_transform(
-        self,
-        func: Union[
-            TransformCallback,
-            TransformExpandedCallback,
-            TransformWallClockCallback,
-            TransformWallClockExpandedCallback,
-        ],
+        func: Union[TransformCallback, TransformExpandedCallback],
         *,
         expand: bool = False,
-        wall_clock: bool = False,
     ) -> "Stream":
         """
         Add a "transform" function to the Stream, that will mutate the input value.
@@ -316,11 +286,9 @@ def add_transform(
         :param expand: if True, expand the returned iterable into individual items
             downstream. If returned value is not iterable, `TypeError` will be raised.
             Default - `False`.
-        :param wall_clock: if True, the callback is expected to accept timestamp only.
-            Default - `False`.
         :return: a new Stream derived from the current one
         """
-        return self._add(TransformFunction(func, expand=expand, wall_clock=wall_clock))  # type: ignore[call-overload]
+        return self._add(TransformFunction(func, expand=expand))  # type: ignore[call-overload]
 
     def merge(self, other: "Stream") -> "Stream":
         """