feat: MelleaSession.register

Author: guicho271828

mellea/stdlib/session.py

@@ -3,6 +3,9 @@
 from __future__ import annotations
 
 import contextvars
+import functools
+import inspect
+from collections.abc import Callable
 from copy import copy
 from typing import Any, Literal, overload
 
@@ -236,6 +239,43 @@ def cleanup(self) -> None:
         if hasattr(self.backend, "close"):
             self.backend.close()  # type: ignore
 
+    @classmethod
+    def register(cls, fn: Callable, set_context: bool = True):
+        """Registers fn as a new method to MelleaSession.
+
+        The function fn must accept `backend` and `context` arguments.
+        """
+
+        def postprocess(self, r):
+            if set_context:
+                if isinstance(r, SamplingResult):
+                    self.ctx = r.result_ctx
+                    return r
+                else:
+                    result, context = r
+                    self.ctx = context
+                    return result
+            else:
+                return r
+
+        if inspect.iscoroutinefunction(fn):
+
+            @functools.wraps(fn)
+            async def wrapper(self, *args, **kwargs):
+                return postprocess(
+                    self,
+                    await fn(backend=self.backend, context=self.ctx, *args, **kwargs),
+                )
+        else:
+
+            @functools.wraps(fn)
+            def wrapper(self, *args, **kwargs):
+                return postprocess(
+                    self, fn(backend=self.backend, context=self.ctx, *args, **kwargs)
+                )
+
+        setattr(cls, fn.__name__, wrapper)
+
     @overload
     def act(
         self,