chore: transformer factory and tests

Author: ringoldsdev

laygo/__init__.py

@@ -7,14 +7,20 @@
 from laygo.pipeline import Pipeline
 from laygo.transformers.http import HTTPTransformer
 from laygo.transformers.parallel import ParallelTransformer
+from laygo.transformers.parallel import createParallelTransformer
 from laygo.transformers.threaded import ThreadedTransformer
+from laygo.transformers.threaded import createThreadedTransformer
 from laygo.transformers.transformer import Transformer
+from laygo.transformers.transformer import createTransformer
 
 __all__ = [
   "Pipeline",
   "Transformer",
+  "createTransformer",
   "ThreadedTransformer",
+  "createThreadedTransformer",
   "ParallelTransformer",
+  "createParallelTransformer",
   "HTTPTransformer",
   "PipelineContext",
   "ErrorHandler",

laygo/transformers/http.py

@@ -30,14 +30,30 @@
 U = TypeVar("U")
 
 
+def createHTTPTransformer[T](
+  _type_hint: type[T],
+  base_url: str,
+  chunk_size: int | None = None,
+  endpoint: str | None = None,
+  max_workers: int = 4,
+) -> "HTTPTransformer[T, T]":
+  """Create a new identity parallel transformer with an explicit type hint."""
+  return HTTPTransformer[T, T](
+    base_url=base_url,
+    endpoint=endpoint,
+    max_workers=max_workers,
+    chunk_size=chunk_size,
+  )
+
+
 class HTTPTransformer(Transformer[In, Out]):
   """
   A self-sufficient, chainable transformer that manages its own
   distributed execution and worker endpoint definition.
   """
 
-  def __init__(self, base_url: str, endpoint: str | None = None, max_workers: int = 8):
-    super().__init__()
+  def __init__(self, base_url: str, endpoint: str | None = None, max_workers: int = 8, chunk_size: int | None = None):
+    super().__init__(chunk_size=chunk_size)
     self.base_url = base_url.rstrip("/")
     self.endpoint = endpoint
     self.max_workers = max_workers

laygo/transformers/parallel.py

@@ -1,58 +1,70 @@
-"""Parallel transformer implementation using multiple threads."""
+"""Parallel transformer implementation using multiple processes and loky."""
 
 from collections import deque
 from collections.abc import Callable
 from collections.abc import Iterable
 from collections.abc import Iterator
+from collections.abc import MutableMapping
 from concurrent.futures import FIRST_COMPLETED
 from concurrent.futures import Future
-from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import wait
 import copy
-from functools import partial
 import itertools
-import threading
+import multiprocessing as mp
 from typing import Any
 from typing import Union
 from typing import overload
 
+from loky import ProcessPoolExecutor
+
 from laygo.errors import ErrorHandler
 from laygo.helpers import PipelineContext
-from laygo.transformers.transformer import DEFAULT_CHUNK_SIZE
 from laygo.transformers.transformer import ChunkErrorHandler
 from laygo.transformers.transformer import InternalTransformer
 from laygo.transformers.transformer import PipelineFunction
 from laygo.transformers.transformer import Transformer
 
 
-class ParallelPipelineContextType(PipelineContext):
-  """A specific context type for parallel transformers that includes a lock."""
+def _process_chunk_for_multiprocessing[In, Out](
+  transformer: InternalTransformer[In, Out],
+  shared_context: MutableMapping[str, Any],
+  chunk: list[In],
+) -> list[Out]:
+  """
+  Top-level function to process a single chunk.
+  'loky' will use cloudpickle to serialize the 'transformer' object.
+  """
+  return transformer(chunk, shared_context)  # type: ignore
+
 
-  lock: threading.Lock
+def createParallelTransformer[T](
+  _type_hint: type[T],
+  max_workers: int = 4,
+  ordered: bool = True,
+  chunk_size: int | None = None,
+) -> "ParallelTransformer[T, T]":
+  """Create a new identity parallel transformer with an explicit type hint."""
+  return ParallelTransformer[T, T](
+    max_workers=max_workers,
+    ordered=ordered,
+    chunk_size=chunk_size,
+    transformer=None,
+  )
 
 
 class ParallelTransformer[In, Out](Transformer[In, Out]):
   """
-  A transformer that executes operations concurrently using multiple threads.
+  A transformer that executes operations concurrently using multiple processes.
+  It uses 'loky' to support dynamically created transformation logic.
   """
 
   def __init__(
     self,
     max_workers: int = 4,
     ordered: bool = True,
-    chunk_size: int = DEFAULT_CHUNK_SIZE,
+    chunk_size: int | None = None,
     transformer: InternalTransformer[In, Out] | None = None,
   ):
-    """
-    Initialize the parallel transformer.
-
-    Args:
-        max_workers: Maximum number of worker threads.
-        ordered: If True, results are yielded in order. If False, results
-                 are yielded as they complete.
-        chunk_size: Size of data chunks to process.
-        transformer: The transformation logic chain.
-    """
     super().__init__(chunk_size, transformer)
     self.max_workers = max_workers
     self.ordered = ordered
@@ -65,18 +77,6 @@ def from_transformer[T, U](
     max_workers: int = 4,
     ordered: bool = True,
   ) -> "ParallelTransformer[T, U]":
-    """
-    Create a ParallelTransformer from an existing Transformer's logic.
-
-    Args:
-        transformer: The base transformer to copy the transformation logic from.
-        chunk_size: Optional chunk size override.
-        max_workers: Maximum number of worker threads.
-        ordered: If True, results are yielded in order.
-
-    Returns:
-        A new ParallelTransformer with the same transformation logic.
-    """
     return cls(
       chunk_size=chunk_size or transformer.chunk_size,
       transformer=copy.deepcopy(transformer.transformer),  # type: ignore
@@ -85,73 +85,101 @@ def from_transformer[T, U](
     )
 
   def __call__(self, data: Iterable[In], context: PipelineContext | None = None) -> Iterator[Out]:
-    """
-    Executes the transformer on data concurrently.
-
-    A new `threading.Lock` is created and added to the context for each call
-    to ensure execution runs are isolated and thread-safe.
-    """
-    # Determine the context for this run, passing it by reference as requested.
-    run_context = context or self.context
-    # Add a per-call lock for thread safety.
-    run_context["lock"] = threading.Lock()
-
-    def process_chunk(chunk: list[In], shared_context: PipelineContext) -> list[Out]:
-      """
-      Process a single chunk by passing the chunk and context explicitly
-      to the transformer chain. This is safer and avoids mutating self.
-      """
-      return self.transformer(chunk, shared_context)
-
-    # Create a partial function with the run_context "baked in".
-    process_chunk_with_context = partial(process_chunk, shared_context=run_context)
-
-    def _ordered_generator(chunks_iter: Iterator[list[In]], executor: ThreadPoolExecutor) -> Iterator[list[Out]]:
-      """Generate results in their original order."""
-      futures: deque[Future[list[Out]]] = deque()
-      for _ in range(self.max_workers + 1):
-        try:
-          chunk = next(chunks_iter)
-          futures.append(executor.submit(process_chunk_with_context, chunk))
-        except StopIteration:
-          break
-      while futures:
-        yield futures.popleft().result()
+    """Executes the transformer on data concurrently using processes."""
+    with mp.Manager() as manager:
+      initial_ctx_data = context if context is not None else self.context
+      shared_context = manager.dict(initial_ctx_data)
+
+      if "lock" not in shared_context:
+        shared_context["lock"] = manager.Lock()
+
+      try:
+        with ProcessPoolExecutor(max_workers=self.max_workers) as executor:
+          chunks_to_process = self._chunk_generator(data)
+          gen_func = self._ordered_generator if self.ordered else self._unordered_generator
+          processed_chunks_iterator = gen_func(chunks_to_process, executor, shared_context)
+
+          for result_chunk in processed_chunks_iterator:
+            yield from result_chunk
+      finally:
+        if context is not None:
+          final_context_state = dict(shared_context)
+          final_context_state.pop("lock", None)
+          # FIX 2: Do not clear the context, just update it.
+          # This allows chained transformers to merge their context results.
+          # context.clear()
+          context.update(final_context_state)
+
+  # ... The rest of the file remains the same ...
+  def _ordered_generator(
+    self,
+    chunks_iter: Iterator[list[In]],
+    executor: ProcessPoolExecutor,
+    shared_context: MutableMapping[str, Any],
+  ) -> Iterator[list[Out]]:
+    """Generate results in their original order."""
+    futures: deque[Future[list[Out]]] = deque()
+    for _ in range(self.max_workers + 1):
+      try:
+        chunk = next(chunks_iter)
+        futures.append(
+          executor.submit(
+            _process_chunk_for_multiprocessing,
+            self.transformer,
+            shared_context,
+            chunk,
+          )
+        )
+      except StopIteration:
+        break
+    while futures:
+      yield futures.popleft().result()
+      try:
+        chunk = next(chunks_iter)
+        futures.append(
+          executor.submit(
+            _process_chunk_for_multiprocessing,
+            self.transformer,
+            shared_context,
+            chunk,
+          )
+        )
+      except StopIteration:
+        continue
+
+  def _unordered_generator(
+    self,
+    chunks_iter: Iterator[list[In]],
+    executor: ProcessPoolExecutor,
+    shared_context: MutableMapping[str, Any],
+  ) -> Iterator[list[Out]]:
+    """Generate results as they complete."""
+    futures = {
+      executor.submit(
+        _process_chunk_for_multiprocessing,
+        self.transformer,
+        shared_context,
+        chunk,
+      )
+      for chunk in itertools.islice(chunks_iter, self.max_workers + 1)
+    }
+    while futures:
+      done, futures = wait(futures, return_when=FIRST_COMPLETED)
+      for future in done:
+        yield future.result()
         try:
           chunk = next(chunks_iter)
-          futures.append(executor.submit(process_chunk_with_context, chunk))
+          futures.add(
+            executor.submit(
+              _process_chunk_for_multiprocessing,
+              self.transformer,
+              shared_context,
+              chunk,
+            )
+          )
         except StopIteration:
           continue
 
-    def _unordered_generator(chunks_iter: Iterator[list[In]], executor: ThreadPoolExecutor) -> Iterator[list[Out]]:
-      """Generate results as they complete."""
-      futures = {
-        executor.submit(process_chunk_with_context, chunk)
-        for chunk in itertools.islice(chunks_iter, self.max_workers + 1)
-      }
-      while futures:
-        done, futures = wait(futures, return_when=FIRST_COMPLETED)
-        for future in done:
-          yield future.result()
-          try:
-            chunk = next(chunks_iter)
-            futures.add(executor.submit(process_chunk_with_context, chunk))
-          except StopIteration:
-            continue
-
-    def result_iterator_manager() -> Iterator[Out]:
-      """Manage the thread pool and yield flattened results."""
-      with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-        chunks_to_process = self._chunk_generator(data)
-        gen_func = _ordered_generator if self.ordered else _unordered_generator
-        processed_chunks_iterator = gen_func(chunks_to_process, executor)
-        for result_chunk in processed_chunks_iterator:
-          yield from result_chunk
-
-    return result_iterator_manager()
-
-  # --- Overridden Chaining Methods to Preserve Type ---
-
   def on_error(self, handler: ChunkErrorHandler[In, Out] | ErrorHandler) -> "ParallelTransformer[In, Out]":
     super().on_error(handler)
     return self
@@ -172,7 +200,9 @@ def flatten[T](self: "ParallelTransformer[In, tuple[T, ...]]") -> "ParallelTrans
   def flatten[T](self: "ParallelTransformer[In, set[T]]") -> "ParallelTransformer[In, T]": ...
   def flatten[T](  # type: ignore
     self: Union[
-      "ParallelTransformer[In, list[T]]", "ParallelTransformer[In, tuple[T, ...]]", "ParallelTransformer[In, set[T]]"
+      "ParallelTransformer[In, list[T]]",
+      "ParallelTransformer[In, tuple[T, ...]]",
+      "ParallelTransformer[In, set[T]]",
     ],
   ) -> "ParallelTransformer[In, T]":
     super().flatten()  # type: ignore

laygo/transformers/threaded.py

@@ -31,6 +31,21 @@ class ThreadedPipelineContextType(PipelineContext):
   lock: threading.Lock
 
 
+def createThreadedTransformer[T](
+  _type_hint: type[T],
+  max_workers: int = 4,
+  ordered: bool = True,
+  chunk_size: int = DEFAULT_CHUNK_SIZE,
+) -> "ThreadedTransformer[T, T]":
+  """Create a new identity threaded transformer with an explicit type hint."""
+  return ThreadedTransformer[T, T](
+    max_workers=max_workers,
+    ordered=ordered,
+    chunk_size=chunk_size,
+    transformer=None,
+  )
+
+
 class ThreadedTransformer[In, Out](Transformer[In, Out]):
   """
   A transformer that executes operations concurrently using multiple threads.
@@ -40,7 +55,7 @@ def __init__(
     self,
     max_workers: int = 4,
     ordered: bool = True,
-    chunk_size: int = DEFAULT_CHUNK_SIZE,
+    chunk_size: int | None = None,
     transformer: InternalTransformer[In, Out] | None = None,
   ):
     """