fix: context sharing when branching

Author: ringoldsdev

laygo/pipeline.py

@@ -12,9 +12,9 @@
 
 from laygo.context import IContextManager
 from laygo.context.parallel import ParallelContextManager
+from laygo.context.types import IContextHandle
 from laygo.helpers import is_context_aware
 from laygo.transformers.transformer import Transformer
-from laygo.transformers.transformer import passthrough_chunks
 
 T = TypeVar("T")
 U = TypeVar("U")
@@ -62,7 +62,7 @@ def __init__(self, *data: Iterable[T], context_manager: IContextManager | None =
     self.processed_data: Iterator = iter(self.data_source)
 
     # Rule 1: Pipeline creates a simple context manager by default.
-    self.context_manager = context_manager or ParallelContextManager()
+    self.context_manager = context_manager if context_manager is not None else ParallelContextManager()
 
   def __del__(self) -> None:
     """Clean up the context manager when the pipeline is destroyed."""
@@ -93,16 +93,6 @@ def context(self, ctx: dict[str, Any]) -> "Pipeline[T]":
     self.context_manager.update(ctx)
     return self
 
-  def _sync_context_back(self) -> None:
-    """Synchronize the final pipeline context back to the original context reference.
-
-    This is called after processing is complete to update the original
-    context with any changes made during pipeline execution.
-    """
-    # 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.
 
@@ -170,7 +160,7 @@ def apply[U](
     match transformer:
       case Transformer():
         # Pass the pipeline's context manager to the transformer
-        self.processed_data = transformer(self.processed_data, self.context_manager)  # type: ignore
+        self.processed_data = transformer(self.processed_data, context=self.context_manager)  # type: ignore
       case _ if callable(transformer):
         if is_context_aware(transformer):
           self.processed_data = transformer(self.processed_data, self.context_manager)  # type: ignore
@@ -318,14 +308,13 @@ def branch(
     branches: dict[str, Transformer[T, Any]],
     batch_size: int = 1000,
     max_batch_buffer: int = 1,
-    use_queue_chunks: bool = True,
   ) -> tuple[dict[str, list[Any]], dict[str, Any]]:
     """Forks the pipeline into multiple branches for concurrent, parallel processing.
 
     This is a **terminal operation** that implements a fan-out pattern where
     the entire dataset is copied to each branch for independent processing.
-    Each branch processes the complete dataset concurrently using separate
-    transformers, and results are collected and returned in a dictionary.
+    Each branch gets its own Pipeline instance with isolated context management,
+    and results are collected and returned in a dictionary.
 
     Args:
         branches: A dictionary where keys are branch names (str) and values
@@ -336,13 +325,12 @@ def branch(
         max_batch_buffer: The maximum number of batches to buffer for each
                           branch queue. Controls memory usage and creates
                           backpressure. Defaults to 1.
-        use_queue_chunks: Whether to use passthrough chunking for the
-                          transformers. When True, batches are processed
-                          as chunks. Defaults to True.
 
     Returns:
-        A dictionary where keys are the branch names and values are lists
-        of all items processed by that branch's transformer.
+        A tuple containing:
+        - A dictionary where keys are the branch names and values are lists
+          of all items processed by that branch's transformer.
+        - A merged dictionary of all context values from all branches.
 
     Note:
         This operation consumes the pipeline's iterator, making subsequent
@@ -372,32 +360,41 @@ def producer() -> None:
       for q in queues:
         q.put(None)
 
-    def consumer(transformer: Transformer, queue: Queue) -> list[Any]:
-      """Consumes batches from a queue and runs them through a transformer."""
+    def consumer(
+      transformer: Transformer, queue: Queue, context_handle: IContextHandle
+    ) -> tuple[list[Any], dict[str, Any]]:
+      """Consumes batches from a queue and processes them through a dedicated pipeline."""
 
       def stream_from_queue() -> Iterator[T]:
         while (batch := queue.get()) is not None:
-          yield batch
+          yield from batch
+
+      # Create a new pipeline for this branch but share the parent's context manager
+      # This ensures all branches share the same context
+      branch_pipeline = Pipeline(stream_from_queue(), context_manager=context_handle.create_proxy())  # type: ignore
 
-      if use_queue_chunks:
-        transformer = transformer.set_chunker(passthrough_chunks)
+      # Apply the transformer to the branch pipeline and get results
+      result_list, branch_context = branch_pipeline.apply(transformer).to_list()
 
-      result_iterator = transformer(stream_from_queue(), self.context_manager)  # type: ignore
-      return list(result_iterator)
+      return result_list, branch_context
 
     with ThreadPoolExecutor(max_workers=num_branches + 1) as executor:
       executor.submit(producer)
 
       future_to_name = {
-        executor.submit(consumer, transformer, queues[i]): name for i, (name, transformer) in enumerate(branch_items)
+        executor.submit(consumer, transformer, queues[i], self.context_manager.get_handle()): name
+        for i, (name, transformer) in enumerate(branch_items)
       }
 
+      # Collect results - context is shared through the same context manager
       for future in as_completed(future_to_name):
         name = future_to_name[future]
         try:
-          final_results[name] = future.result()
-        except Exception as e:
-          print(f"Branch '{name}' raised an exception: {e}")
+          result_list, branch_context = future.result()
+          final_results[name] = result_list
+        except Exception:
           final_results[name] = []
 
-    return final_results, self.context_manager.to_dict()
+    # After all threads complete, get the final context state
+    final_context = self.context_manager.to_dict()
+    return final_results, final_context

laygo/transformers/http.py

@@ -124,7 +124,7 @@ def __call__(self, data: Iterable[In], context: IContextManager | None = None) -
     Returns:
         An iterator over the processed data.
     """
-    run_context = context or self._default_context
+    run_context = self._default_context
 
     self._finalize_config()
 

laygo/transformers/parallel.py

@@ -124,7 +124,7 @@ def from_transformer[T, U](
 
   def __call__(self, data: Iterable[In], context: IContextManager | None = None) -> Iterator[Out]:
     """Execute the transformer by distributing chunks to a process pool."""
-    run_context = context or self._default_context
+    run_context = context if context is not None else self._default_context
 
     # Get the picklable handle from the context manager.
     context_handle = run_context.get_handle()

laygo/transformers/threaded.py

@@ -119,7 +119,7 @@ def __call__(self, data: Iterable[In], context: IContextManager | None = None) -
     Returns:
         An iterator over the transformed data.
     """
-    run_context = context or self._default_context
+    run_context = context if context is not None else self._default_context
 
     # Since threads share memory, we can pass the context manager directly.
     # No handle/proxy mechanism is needed, but the locking inside

laygo/transformers/transformer.py

@@ -345,8 +345,9 @@ def __call__(self, data: Iterable[In], context: IContextManager | None = None) -
     Returns:
         An iterator over the transformed data.
     """
+
     # Use the provided context by reference, or default to a simple context.
-    run_context = context or self._default_context
+    run_context = context if context is not None else self._default_context
 
     try:
       for chunk in self._chunk_generator(data):

tests/test_pipeline.py

@@ -498,5 +498,5 @@ def context_modifier_b(chunk: list[int], ctx: IContextManager) -> list[int]:
     assert result["branch_b"] == [3, 6, 9]
 
     # Context values should reflect the actual chunk sizes processed
-    assert context["branch_a_processed"] == 3
     assert context["branch_b_processed"] == 3
+    assert context["branch_a_processed"] == 3