chore: implemented context managers

Author: ringoldsdev

laygo/context/parallel.py

@@ -0,0 +1,133 @@
+"""
+A context manager for parallel and distributed processing using
+multiprocessing.Manager to share state across processes.
+"""
+
+from collections.abc import Iterator
+import multiprocessing as mp
+from multiprocessing.managers import BaseManager
+from multiprocessing.managers import DictProxy
+from multiprocessing.synchronize import Lock
+from typing import Any
+
+from laygo.context.types import IContextHandle
+from laygo.context.types import IContextManager
+
+
+class _ParallelStateManager(BaseManager):
+  """A custom manager to expose a shared dictionary and lock."""
+
+  pass
+
+
+class ParallelContextHandle(IContextHandle):
+  """
+  A lightweight, picklable "blueprint" for recreating a connection to the
+  shared context in a different process.
+  """
+
+  def __init__(self, address: tuple[str, int], manager_class: type["ParallelContextManager"]):
+    self.address = address
+    self.manager_class = manager_class
+
+  def create_proxy(self) -> "IContextManager":
+    """
+    Creates a new instance of the ParallelContextManager in "proxy" mode
+    by initializing it with this handle.
+    """
+    return self.manager_class(handle=self)
+
+
+class ParallelContextManager(IContextManager):
+  """
+  A context manager that uses a background multiprocessing.Manager to enable
+  state sharing across different processes.
+
+  This single class operates in two modes:
+  1. Server Mode (when created normally): It starts and manages the background
+     server process that holds the shared state.
+  2. Proxy Mode (when created with a handle): It acts as a client, connecting
+     to an existing server process to manipulate the shared state.
+  """
+
+  def __init__(self, initial_context: dict[str, Any] | None = None, handle: ParallelContextHandle | None = None):
+    """
+    Initializes the manager. If a handle is provided, it initializes in
+    proxy mode; otherwise, it starts a new server.
+    """
+    if handle:
+      # --- PROXY MODE INITIALIZATION ---
+      # This instance is a client connecting to an existing server.
+      self._is_proxy = True
+      self._manager_server = None  # Proxies do not own the server process.
+
+      manager = _ParallelStateManager(address=handle.address)
+      manager.connect()
+      self._manager = manager
+
+    else:
+      # --- SERVER MODE INITIALIZATION ---
+      # This is the main instance that owns the server process.
+      self._is_proxy = False
+      manager = mp.Manager()  # type: ignore
+      _ParallelStateManager.register("get_dict", callable=lambda: manager.dict(initial_context or {}))
+      _ParallelStateManager.register("get_lock", callable=lambda: manager.Lock())
+
+      self._manager_server = _ParallelStateManager(address=("", 0))
+      self._manager_server.start()
+      self._manager = self._manager_server
+
+    # Common setup for both modes
+    self._shared_dict: DictProxy = self._manager.get_dict()  # type: ignore
+    self._lock: Lock = self._manager.get_lock()  # type: ignore
+
+  def get_handle(self) -> ParallelContextHandle:
+    """
+    Returns a picklable handle for reconstruction in a worker.
+    Only the main server instance can generate handles.
+    """
+    if self._is_proxy or not self._manager_server:
+      raise TypeError("Cannot get a handle from a proxy context instance.")
+
+    return ParallelContextHandle(
+      address=self._manager_server.address,  # type: ignore
+      manager_class=self.__class__,  # Pass its own class for reconstruction
+    )
+
+  def shutdown(self) -> None:
+    """
+    Shuts down the background manager process.
+    This is a no-op for proxy instances, as only the main instance
+    should control the server's lifecycle.
+    """
+    if not self._is_proxy and self._manager_server:
+      self._manager_server.shutdown()
+
+  def __enter__(self) -> "ParallelContextManager":
+    """Acquires the lock for use in a 'with' statement."""
+    self._lock.acquire()
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    """Releases the lock."""
+    self._lock.release()
+
+  def __getitem__(self, key: str) -> Any:
+    with self._lock:
+      return self._shared_dict[key]
+
+  def __setitem__(self, key: str, value: Any) -> None:
+    with self._lock:
+      self._shared_dict[key] = value
+
+  def __delitem__(self, key: str) -> None:
+    with self._lock:
+      del self._shared_dict[key]
+
+  def __iter__(self) -> Iterator[str]:
+    with self._lock:
+      return iter(list(self._shared_dict.keys()))
+
+  def __len__(self) -> int:
+    with self._lock:
+      return len(self._shared_dict)

laygo/context/simple.py

@@ -0,0 +1,89 @@
+"""
+A simple, dictionary-based context manager for single-process pipelines.
+"""
+
+from collections.abc import Iterator
+from typing import Any
+
+from laygo.context.types import IContextHandle
+from laygo.context.types import IContextManager
+
+
+class SimpleContextHandle(IContextHandle):
+  """
+  A handle for the SimpleContextManager that provides a reference back to the
+  original manager instance.
+
+  In a single-process environment, the "proxy" is the manager itself, ensuring
+  all transformers in a chain share the exact same context dictionary.
+  """
+
+  def __init__(self, manager_instance: "IContextManager"):
+    self._manager_instance = manager_instance
+
+  def create_proxy(self) -> "IContextManager":
+    """
+    Returns the original SimpleContextManager instance.
+
+    This ensures that in a non-distributed pipeline, all chained transformers
+    operate on the same shared dictionary.
+    """
+    return self._manager_instance
+
+
+class SimpleContextManager(IContextManager):
+  """
+  A basic context manager that uses a standard Python dictionary for state.
+
+  This manager is suitable for single-threaded, single-process pipelines where
+  no state needs to be shared across process boundaries. It is the default
+  context manager for a Laygo pipeline.
+  """
+
+  def __init__(self, initial_context: dict[str, Any] | None = None) -> None:
+    """
+    Initializes the context manager with an optional dictionary.
+
+    Args:
+        initial_context: An optional dictionary to populate the context with.
+    """
+    self._context = dict(initial_context or {})
+
+  def get_handle(self) -> IContextHandle:
+    """
+    Returns a handle that holds a reference back to this same instance.
+    """
+    return SimpleContextHandle(self)
+
+  def __enter__(self) -> "SimpleContextManager":
+    """
+    Provides 'with' statement compatibility. No lock is needed for this
+    simple, single-threaded context manager.
+    """
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    """
+    Provides 'with' statement compatibility. No lock is needed for this
+    simple, single-threaded context manager.
+    """
+    pass
+
+  def __getitem__(self, key: str) -> Any:
+    return self._context[key]
+
+  def __setitem__(self, key: str, value: Any) -> None:
+    self._context[key] = value
+
+  def __delitem__(self, key: str) -> None:
+    del self._context[key]
+
+  def __iter__(self) -> Iterator[str]:
+    return iter(self._context)
+
+  def __len__(self) -> int:
+    return len(self._context)
+
+  def shutdown(self) -> None:
+    """No-op for the simple context manager."""
+    pass

laygo/context/types.py

@@ -0,0 +1,68 @@
+"""
+Defines the abstract base classes for context management in Laygo.
+
+This module provides the core interfaces (IContextHandle and IContextManager)
+that all context managers must implement, ensuring a consistent API for
+state management across different execution environments (simple, threaded, parallel).
+"""
+
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import MutableMapping
+from typing import Any
+
+
+class IContextHandle(ABC):
+  """
+  An abstract base class for a picklable handle to a context manager.
+
+  A handle contains the necessary information for a worker process to
+  reconstruct a connection (a proxy) to the shared context.
+  """
+
+  @abstractmethod
+  def create_proxy(self) -> "IContextManager":
+    """
+    Creates the appropriate context proxy instance from the handle's data.
+
+    This method is called within a worker process to establish its own
+    connection to the shared state.
+
+    Returns:
+        An instance of an IContextManager proxy.
+    """
+    raise NotImplementedError
+
+
+class IContextManager(MutableMapping[str, Any], ABC):
+  """
+  Abstract base class for managing shared state (context) in a pipeline.
+
+  This class defines the contract for all context managers, ensuring they
+  provide a dictionary-like interface for state manipulation by inheriting
+  from `collections.abc.MutableMapping`. It also includes methods for
+  distribution (get_handle) and resource management (shutdown).
+  """
+
+  @abstractmethod
+  def get_handle(self) -> IContextHandle:
+    """
+    Returns a picklable handle for connecting from a worker process.
+
+    This handle is serialized and sent to distributed workers, which then
+    use it to create a proxy to the shared context.
+
+    Returns:
+        A picklable IContextHandle instance.
+    """
+    raise NotImplementedError
+
+  @abstractmethod
+  def shutdown(self) -> None:
+    """
+    Performs final synchronization and cleans up any resources.
+
+    This method is responsible for releasing connections, shutting down
+    background processes, or any other cleanup required by the manager.
+    """
+    raise NotImplementedError