chore: implemented use of the new context manager

Author: ringoldsdev

laygo/context/__init__.py

@@ -0,0 +1,19 @@
+"""
+Laygo Context Management Package.
+
+This package provides different strategies for managing state (context)
+within a data pipeline, from simple in-memory dictionaries to
+process-safe managers for parallel execution.
+"""
+
+from .parallel import ParallelContextManager
+from .simple import SimpleContextManager
+from .types import IContextHandle
+from .types import IContextManager
+
+__all__ = [
+  "IContextManager",
+  "IContextHandle",
+  "SimpleContextManager",
+  "ParallelContextManager",
+]

laygo/errors.py

@@ -1,11 +1,11 @@
 from collections.abc import Callable
 
-from laygo.helpers import PipelineContext
+from laygo.context.types import IContextManager
 
-ChunkErrorHandler = Callable[[list, Exception, PipelineContext], None]
+ChunkErrorHandler = Callable[[list, Exception, IContextManager], None]
 
 
-def raise_error(chunk: list, error: Exception, context: PipelineContext) -> None:
+def raise_error(chunk: list, error: Exception, context: IContextManager) -> None:
   """Handler that always re-raises the error, stopping execution.
 
   This is a default error handler that provides fail-fast behavior by
@@ -47,7 +47,7 @@ def on_error(self, handler: ChunkErrorHandler) -> "ErrorHandler":
     self._handlers.insert(0, handler)
     return self
 
-  def handle(self, chunk: list, error: Exception, context: PipelineContext) -> None:
+  def handle(self, chunk: list, error: Exception, context: IContextManager) -> None:
     """Execute all handlers in the chain.
 
     Handlers are executed in reverse order of addition. Execution stops

laygo/helpers.py

@@ -3,10 +3,15 @@
 from typing import Any
 from typing import TypeGuard
 
+from laygo.context.types import IContextManager
+
 
 class PipelineContext(dict[str, Any]):
   """Generic, untyped context available to all pipeline operations.
 
+  DEPRECATED: This class is deprecated and will be removed in a future version.
+  Use IContextManager implementations (SimpleContextManager, etc.) instead.
+
   A dictionary-based context that can store arbitrary data shared across
   pipeline operations. This allows passing state and configuration between
   different stages of data processing.
@@ -16,14 +21,14 @@ class PipelineContext(dict[str, Any]):
 
 
 # Define the specific callables for clarity
-ContextAwareCallable = Callable[[Any, PipelineContext], Any]
-ContextAwareReduceCallable = Callable[[Any, Any, PipelineContext], Any]
+ContextAwareCallable = Callable[[Any, IContextManager], Any]
+ContextAwareReduceCallable = Callable[[Any, Any, IContextManager], Any]
 
 
 def is_context_aware(func: Callable[..., Any]) -> TypeGuard[ContextAwareCallable]:
   """Check if a function is context-aware by inspecting its signature.
 
-  A context-aware function accepts a PipelineContext as its second parameter,
+  A context-aware function accepts an IContextManager as its second parameter,
   allowing it to access shared state during pipeline execution.
 
   Args:
@@ -40,7 +45,7 @@ def is_context_aware_reduce(func: Callable[..., Any]) -> TypeGuard[ContextAwareR
   """Check if a reduce function is context-aware by inspecting its signature.
 
   A context-aware reduce function accepts an accumulator, current value,
-  and PipelineContext as its three parameters.
+  and IContextManager as its three parameters.
 
   Args:
       func: The reduce function to inspect for context awareness.

laygo/pipeline.py

@@ -5,13 +5,13 @@
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import as_completed
 import itertools
-import multiprocessing as mp
 from queue import Queue
 from typing import Any
 from typing import TypeVar
 from typing import overload
 
-from laygo.helpers import PipelineContext
+from laygo.context import IContextManager
+from laygo.context import SimpleContextManager
 from laygo.helpers import is_context_aware
 from laygo.transformers.transformer import Transformer
 from laygo.transformers.transformer import passthrough_chunks
@@ -44,12 +44,14 @@ class Pipeline[T]:
       pipeline effectively single-use unless the data source is re-initialized.
   """
 
-  def __init__(self, *data: Iterable[T]) -> None:
+  def __init__(self, *data: Iterable[T], context_manager: IContextManager | None = None) -> None:
     """Initialize a pipeline with one or more data sources.
 
     Args:
         *data: One or more iterable data sources. If multiple sources are
                provided, they will be chained together.
+        context_manager: An instance of a class that implements IContextManager.
+                         If None, a SimpleContextManager is used by default.
 
     Raises:
         ValueError: If no data sources are provided.
@@ -59,25 +61,16 @@ def __init__(self, *data: Iterable[T]) -> None:
     self.data_source: Iterable[T] = itertools.chain.from_iterable(data) if len(data) > 1 else data[0]
     self.processed_data: Iterator = iter(self.data_source)
 
-    # Always create a shared context with multiprocessing manager
-    self._manager = mp.Manager()
-    self.ctx = self._manager.dict()
-    # Add a shared lock to the context for safe concurrent updates
-    self.ctx["lock"] = self._manager.Lock()
-
-    # Store reference to original context for final synchronization
-    self._original_context_ref: PipelineContext | None = None
+    # Rule 1: Pipeline creates a simple context manager by default.
+    self.context_manager = context_manager or SimpleContextManager()
 
   def __del__(self) -> None:
-    """Clean up the multiprocessing manager when the pipeline is destroyed."""
-    try:
-      self._sync_context_back()
-      self._manager.shutdown()
-    except Exception:
-      pass
+    """Clean up the context manager when the pipeline is destroyed."""
+    if hasattr(self, "context_manager"):
+      self.context_manager.shutdown()
 
-  def context(self, ctx: PipelineContext) -> "Pipeline[T]":
-    """Update the pipeline context and store a reference to the original context.
+  def context(self, ctx: dict[str, Any]) -> "Pipeline[T]":
+    """Update the pipeline's context manager with values from a dictionary.
 
     The provided context will be used during pipeline execution and any
     modifications made by transformers will be synchronized back to the
@@ -96,10 +89,7 @@ def context(self, ctx: PipelineContext) -> "Pipeline[T]":
         automatically synchronized back to the original context object
         when the pipeline is destroyed or processing completes.
     """
-    # Store reference to the original context
-    self._original_context_ref = ctx
-    # Copy the context data to the pipeline's shared context
-    self.ctx.update(ctx)
+    self.context_manager.update(ctx)
     return self
 
   def _sync_context_back(self) -> None:
@@ -108,12 +98,9 @@ def _sync_context_back(self) -> None:
     This is called after processing is complete to update the original
     context with any changes made during pipeline execution.
     """
-    if self._original_context_ref is not None:
-      # Copy the final context state back to the original context reference
-      final_context_state = dict(self.ctx)
-      final_context_state.pop("lock", None)  # Remove non-serializable lock
-      self._original_context_ref.clear()
-      self._original_context_ref.update(final_context_state)
+    # This method is kept for backward compatibility but is no longer needed
+    # since we use the context manager directly
+    pass
 
   def transform[U](self, t: Callable[[Transformer[T, T]], Transformer[T, U]]) -> "Pipeline[U]":
     """Apply a transformation using a lambda function.
@@ -146,13 +133,13 @@ def apply[U](self, transformer: Transformer[T, U]) -> "Pipeline[U]": ...
   def apply[U](self, transformer: Callable[[Iterable[T]], Iterator[U]]) -> "Pipeline[U]": ...
 
   @overload
-  def apply[U](self, transformer: Callable[[Iterable[T], PipelineContext], Iterator[U]]) -> "Pipeline[U]": ...
+  def apply[U](self, transformer: Callable[[Iterable[T], IContextManager], Iterator[U]]) -> "Pipeline[U]": ...
 
   def apply[U](
     self,
     transformer: Transformer[T, U]
     | Callable[[Iterable[T]], Iterator[U]]
-    | Callable[[Iterable[T], PipelineContext], Iterator[U]],
+    | Callable[[Iterable[T], IContextManager], Iterator[U]],
   ) -> "Pipeline[U]":
     """Apply a transformer to the current data source.
 
@@ -181,10 +168,11 @@ def apply[U](
     """
     match transformer:
       case Transformer():
-        self.processed_data = transformer(self.processed_data, self.ctx)  # type: ignore
+        # Pass the pipeline's context manager to the transformer
+        self.processed_data = transformer(self.processed_data, self.context_manager)  # type: ignore
       case _ if callable(transformer):
         if is_context_aware(transformer):
-          self.processed_data = transformer(self.processed_data, self.ctx)  # type: ignore
+          self.processed_data = transformer(self.processed_data, self.context_manager)  # type: ignore
         else:
           self.processed_data = transformer(self.processed_data)  # type: ignore
       case _:
@@ -256,12 +244,12 @@ def consumer(transformer: Transformer, queue: Queue) -> list[Any]:
 
       def stream_from_queue() -> Iterator[T]:
         while (batch := queue.get()) is not None:
-          yield batch
+          yield from batch
 
       if use_queue_chunks:
         transformer = transformer.set_chunker(passthrough_chunks)
 
-      result_iterator = transformer(stream_from_queue(), self.ctx)  # type: ignore
+      result_iterator = transformer(stream_from_queue(), self.context_manager)  # type: ignore
       return list(result_iterator)
 
     with ThreadPoolExecutor(max_workers=num_branches + 1) as executor:

laygo/transformers/http.py

@@ -16,8 +16,9 @@
 
 import requests
 
+from laygo.context import IContextManager
+from laygo.context import SimpleContextManager
 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
@@ -85,6 +86,8 @@ def __init__(
     self.max_workers = max_workers
     self.session = requests.Session()
     self._worker_url: str | None = None
+    # HTTP transformers always use a simple context manager to avoid serialization issues
+    self._default_context = SimpleContextManager()
 
   def _finalize_config(self) -> None:
     """Determine the final worker URL, generating one if needed.
@@ -107,19 +110,22 @@ def _finalize_config(self) -> None:
     self.endpoint = path.lstrip("/")
     self._worker_url = f"{self.base_url}/{self.endpoint}"
 
-  def __call__(self, data: Iterable[In], context: PipelineContext | None = None) -> Iterator[Out]:
+  def __call__(self, data: Iterable[In], context: IContextManager | None = None) -> Iterator[Out]:
     """Execute distributed processing on the data (CLIENT-SIDE).
 
     This method is called by the Pipeline to start distributed processing.
     It sends data chunks to worker endpoints via HTTP.
 
     Args:
         data: The input data to process.
-        context: Optional pipeline context (currently not used in HTTP mode).
+        context: Optional pipeline context. HTTP transformers always use their
+                internal SimpleContextManager regardless of the provided context.
 
     Returns:
         An iterator over the processed data.
     """
+    run_context = context or self._default_context
+
     self._finalize_config()
 
     def process_chunk(chunk: list) -> list:
@@ -143,18 +149,24 @@ def process_chunk(chunk: list) -> list:
         print(f"Error calling worker {self._worker_url}: {e}")
         return []
 
-    with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-      chunk_iterator = self._chunk_generator(data)
-      futures = {executor.submit(process_chunk, chunk) for chunk in itertools.islice(chunk_iterator, self.max_workers)}
-      while futures:
-        done, futures = wait(futures, return_when=FIRST_COMPLETED)
-        for future in done:
-          yield from future.result()
-          try:
-            new_chunk = next(chunk_iterator)
-            futures.add(executor.submit(process_chunk, new_chunk))
-          except StopIteration:
-            continue
+    try:
+      with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+        chunk_iterator = self._chunk_generator(data)
+        futures = {
+          executor.submit(process_chunk, chunk) for chunk in itertools.islice(chunk_iterator, self.max_workers)
+        }
+        while futures:
+          done, futures = wait(futures, return_when=FIRST_COMPLETED)
+          for future in done:
+            yield from future.result()
+            try:
+              new_chunk = next(chunk_iterator)
+              futures.add(executor.submit(process_chunk, new_chunk))
+            except StopIteration:
+              continue
+    finally:
+      # Always clean up our context since we always use the default one
+      run_context.shutdown()
 
   def get_route(self):
     """Get the route configuration for registering this transformer as a worker.
@@ -167,7 +179,7 @@ def get_route(self):
     """
     self._finalize_config()
 
-    def worker_view_func(chunk: list, context: PipelineContext):
+    def worker_view_func(chunk: list, context: IContextManager):
       """The actual worker logic for this transformer.
 
       Args:
@@ -226,6 +238,6 @@ def catch[U](
     super().catch(sub_pipeline_builder, on_error)
     return self  # type: ignore
 
-  def short_circuit(self, function: Callable[[PipelineContext], bool | None]) -> "HTTPTransformer[In, Out]":
+  def short_circuit(self, function: Callable[[IContextManager], bool | None]) -> "HTTPTransformer[In, Out]":
     super().short_circuit(function)
     return self