fix: http transformer types and added tests

Author: ringoldsdev

laygo/errors.py

@@ -1,6 +1,6 @@
 from collections.abc import Callable
 
-from laygo import PipelineContext
+from laygo.helpers import PipelineContext
 
 ChunkErrorHandler = Callable[[list, Exception, PipelineContext], None]
 

laygo/pipeline.py

@@ -6,9 +6,9 @@
 from typing import TypeVar
 from typing import overload
 
-from laygo import PipelineContext
-from laygo import Transformer
+from laygo.helpers import PipelineContext
 from laygo.helpers import is_context_aware
+from laygo.transformers.transformer import Transformer
 
 T = TypeVar("T")
 PipelineFunction = Callable[[T], Any]

laygo/transformers/http.py

@@ -1,7 +1,8 @@
 """
-The final, self-sufficient DistributedTransformer.
+The final, self-sufficient DistributedTransformer with corrected typing.
 """
 
+from collections.abc import Callable
 from collections.abc import Iterable
 from collections.abc import Iterator
 from concurrent.futures import FIRST_COMPLETED
@@ -10,14 +11,26 @@
 import hashlib
 import itertools
 import pickle
+from typing import Any
+from typing import TypeVar
+from typing import Union
+from typing import overload
 
 import requests
 
-from laygo import PipelineContext
-from laygo import Transformer
+from laygo.errors import ErrorHandler
+from laygo.helpers import PipelineContext
+from laygo.transformers.transformer import ChunkErrorHandler
+from laygo.transformers.transformer import PipelineFunction
+from laygo.transformers.transformer import Transformer
 
+In = TypeVar("In")
+Out = TypeVar("Out")
+T = TypeVar("T")
+U = TypeVar("U")
 
-class HTTPTransformer(Transformer):
+
+class HTTPTransformer(Transformer[In, Out]):
   """
   A self-sufficient, chainable transformer that manages its own
   distributed execution and worker endpoint definition.
@@ -29,32 +42,39 @@ def __init__(self, base_url: str, endpoint: str | None = None, max_workers: int
     self.endpoint = endpoint
     self.max_workers = max_workers
     self.session = requests.Session()
-    self._worker_url: str
+    self._worker_url: str | None = None
 
   def _finalize_config(self):
     """Determines the final worker URL, generating one if needed."""
-    if self._worker_url:
+    if hasattr(self, "_worker_url") and self._worker_url:
       return
 
     if self.endpoint:
       path = self.endpoint
     else:
-      # Using pickle to serialize the function chain and hashing for a unique ID
+      if not self.transformer:
+        raise ValueError("Cannot determine endpoint for an empty transformer.")
       serialized_logic = pickle.dumps(self.transformer)
       hash_id = hashlib.sha1(serialized_logic).hexdigest()[:16]
       path = f"/autogen/{hash_id}"
 
     self.endpoint = path.lstrip("/")
     self._worker_url = f"{self.base_url}/{self.endpoint}"
 
-  def __call__(self, data: Iterable, context=None) -> Iterator:
+  # --- Original HTTPTransformer Methods ---
+
+  def __call__(self, data: Iterable[In], context: PipelineContext | None = None) -> Iterator[Out]:
     """CLIENT-SIDE: Called by the Pipeline to start distributed processing."""
     self._finalize_config()
 
     def process_chunk(chunk: list) -> list:
       """Target for a thread: sends one chunk to the worker."""
       try:
-        response = self.session.post(self._worker_url, json=chunk, timeout=300)
+        response = self.session.post(
+          self._worker_url,  # type: ignore
+          json=chunk,
+          timeout=300,
+        )
         response.raise_for_status()
         return response.json()
       except requests.RequestException as e:
@@ -78,14 +98,58 @@ def get_route(self):
     """
     Function that returns the route for the worker.
     This is used to register the worker in a Flask app or similar.
-
-    Returns:
-        A tuple containing the endpoint and the worker function.
     """
     self._finalize_config()
 
     def worker_view_func(chunk: list, context: PipelineContext):
-      """The actual Flask view function for this transformer's logic."""
+      """The actual worker logic for this transformer."""
       return self.transformer(chunk, context)
 
     return (f"/{self.endpoint}", worker_view_func)
+
+  # --- Overridden Chaining Methods to Preserve Type ---
+
+  def on_error(self, handler: ChunkErrorHandler[In, Out] | ErrorHandler) -> "HTTPTransformer[In, Out]":
+    super().on_error(handler)
+    return self
+
+  def map[U](self, function: PipelineFunction[Out, U]) -> "HTTPTransformer[In, U]":
+    super().map(function)
+    return self  # type: ignore
+
+  def filter(self, predicate: PipelineFunction[Out, bool]) -> "HTTPTransformer[In, Out]":
+    super().filter(predicate)
+    return self
+
+  @overload
+  def flatten[T](self: "HTTPTransformer[In, list[T]]") -> "HTTPTransformer[In, T]": ...
+  @overload
+  def flatten[T](self: "HTTPTransformer[In, tuple[T, ...]]") -> "HTTPTransformer[In, T]": ...
+  @overload
+  def flatten[T](self: "HTTPTransformer[In, set[T]]") -> "HTTPTransformer[In, T]": ...
+  def flatten[T](
+    self: Union["HTTPTransformer[In, list[T]]", "HTTPTransformer[In, tuple[T, ...]]", "HTTPTransformer[In, set[T]]"],
+  ) -> "HTTPTransformer[In, T]":
+    super().flatten()
+    return self  # type: ignore
+
+  def tap(self, function: PipelineFunction[Out, Any]) -> "HTTPTransformer[In, Out]":
+    super().tap(function)
+    return self
+
+  def apply[T](self, t: Callable[["HTTPTransformer[In, Out]"], "Transformer[In, T]"]) -> "HTTPTransformer[In, T]":
+    # Note: The type hint for `t` is slightly adjusted to reflect it receives an HTTPTransformer
+    super().apply(t)  # type: ignore
+    return self  # type: ignore
+
+  def catch[U](
+    self,
+    sub_pipeline_builder: Callable[[Transformer[Out, Out]], Transformer[Out, U]],
+    on_error: ChunkErrorHandler[Out, U] | None = None,
+  ) -> "HTTPTransformer[In, U]":
+    super().catch(sub_pipeline_builder, on_error)
+    return self  # type: ignore
+
+  def short_circuit(self, function: Callable[[PipelineContext], bool | None]) -> "HTTPTransformer[In, Out]":
+    super().short_circuit(function)
+    return self

laygo/transformers/parallel.py

@@ -12,10 +12,10 @@
 import itertools
 import threading
 
-from laygo import PipelineContext
-from laygo import Transformer
+from laygo.helpers import PipelineContext
 from laygo.transformers.transformer import DEFAULT_CHUNK_SIZE
 from laygo.transformers.transformer import InternalTransformer
+from laygo.transformers.transformer import Transformer
 
 
 class ParallelPipelineContextType(PipelineContext):

laygo/transformers/transformer.py

@@ -9,8 +9,8 @@
 from typing import Union
 from typing import overload
 
-from laygo import ErrorHandler
-from laygo import PipelineContext
+from laygo.errors import ErrorHandler
+from laygo.helpers import PipelineContext
 from laygo.helpers import is_context_aware
 from laygo.helpers import is_context_aware_reduce
 

pyproject.toml

@@ -43,6 +43,7 @@ dev = [
     "pytest>=7.0.0",
     "ruff>=0.1.0",
     "twine>=4.0.0",
+    "requests-mock>=1.12.1",
 ]
 
 [tool.ruff]

tests/test_http_transformer.py

@@ -0,0 +1,59 @@
+# Assuming the classes from your latest example are in a file named `pipeline_lib.py`
+# This includes Pipeline, Transformer, and your HTTPTransformer.
+import requests_mock
+
+from laygo import HTTPTransformer
+from laygo import Pipeline
+from laygo import PipelineContext
+
+
+class TestHTTPTransformer:
+  """
+  Test suite for the HTTPTransformer class.
+  """
+
+  def test_distributed_transformer_with_mock(self):
+    """
+    Tests the HTTPTransformer by mocking the worker endpoint.
+    This test validates that the client-side of the transformer correctly
+    calls the endpoint and processes the response from the (mocked) worker.
+    """
+    # 1. Define the transformer's properties
+    base_url = "http://mock-worker.com"
+    endpoint = "/process/data"
+    worker_url = f"{base_url}{endpoint}"
+
+    # 2. Define the transformer and its logic using the chainable API.
+    # This single instance holds both the client and server logic.
+    http_transformer = (
+      HTTPTransformer(base_url=base_url, endpoint=endpoint).map(lambda x: x * 2).filter(lambda x: x > 10)
+    )
+
+    # Set a small chunk_size to ensure the client makes multiple requests
+    http_transformer.chunk_size = 4
+
+    # 3. Get the worker's logic from the transformer itself
+    # The `get_route` method provides the exact function the worker would run.
+    _, worker_view_func = http_transformer.get_route()
+
+    # 4. Configure the mock endpoint to use the real worker logic
+    def mock_response(request, context):
+      """The behavior of the mocked Flask endpoint."""
+      input_chunk = request.json()
+      # Call the actual view function logic obtained from get_route()
+      # We pass None for the context as it's not used in this simple case.
+      output_chunk = worker_view_func(chunk=input_chunk, context=PipelineContext())
+      return output_chunk
+
+    # Use requests_mock context manager
+    with requests_mock.Mocker() as m:
+      m.post(worker_url, json=mock_response)
+
+      # 5. Run the standard Pipeline with the configured transformer
+      initial_data = list(range(10))  # [0, 1, 2, ..., 9]
+      pipeline = Pipeline(initial_data).apply(http_transformer)
+      result = pipeline.to_list()
+
+      # 6. Assert the final result
+      expected_result = [12, 14, 16, 18]
+      assert sorted(result) == sorted(expected_result)