refactor: NodeContext -> RayContext

Author: blepabyte

maxray/__init__.py

@@ -1,4 +1,5 @@
-from .transforms import recompile_fn_with_transform, NodeContext
+from .nodes import NodeContext, RayContext
+from .transforms import recompile_fn_with_transform
 from .function_store import FunctionStore, set_property_on_functionlike
 
 import inspect
@@ -106,7 +107,14 @@ class TransformSettings:
     Whether to always populate `NodeContext.local_scope` with the evaluation of `locals()` at every node.
     """
 
-    def update(self, forbid_modules, restrict_modules, pass_local_scopes):
+    preserve_values: bool = True
+    """
+    If False, will attempt to unpack/destructure `x` values to support pattern matching features.
+    """
+
+    def update(
+        self, forbid_modules, restrict_modules, pass_local_scopes, preserve_values
+    ):
         def or_empty_set(s):
             return frozenset([] if s is None else s)
 
@@ -119,6 +127,7 @@ def or_empty_set(s):
             forbid_modules=self.forbid_modules.union(forbid_modules),
             restrict_modules=restrict_modules,
             pass_local_scopes=self.pass_local_scopes or pass_local_scopes,
+            preserve_values=self.preserve_values and preserve_values,
         )
 
 
@@ -372,6 +381,11 @@ def _maxray_walker_handler(x, ctx: NodeContext):
     # 2. run the active hooks
     global_write_active_token = _GLOBAL_WRITER_ACTIVE_FLAG.set(True)
     try:
+        ray = RayContext(
+            x, ctx, unpack_assignments=not _MAXRAY_TRANSFORM_SETTINGS.preserve_values
+        )
+        x = ray.value()
+
         for walk_hook in _MAXRAY_REGISTERED_HOOKS:
             # Our recompiled fn sets and unsets a contextvar whenever it is active
             if not walk_hook.active_call_state.get():
@@ -380,9 +394,9 @@ def _maxray_walker_handler(x, ctx: NodeContext):
             # Set the writer active flag
             write_active_token = walk_hook.writer_active_call_state.set(True)
             if walk_hook.mutable:
-                x = walk_hook.impl_fn(x, ctx)
+                x = walk_hook.impl_fn(x, ray)
             else:
-                walk_hook.impl_fn(x, ctx)
+                walk_hook.impl_fn(x, ray)
             walk_hook.writer_active_call_state.reset(write_active_token)
     finally:
         _GLOBAL_WRITER_ACTIVE_FLAG.reset(global_write_active_token)
@@ -394,14 +408,15 @@ def _maxray_walker_handler(x, ctx: NodeContext):
 
 
 def maxray(
-    writer: Callable[[Any, NodeContext], Any],
+    writer: Callable[[Any, RayContext], Any],
     *,
     mutable=True,
     forbid_modules=frozenset(),
     restrict_modules=None,
-    pass_scope=False,
     initial_scope={},
     assume_transformed=False,
+    pass_scope=False,
+    preserve_values=True,
 ) -> Callable[[T], T]:
     """
     A transform that recursively hooks into all further calls made within the function, so that `writer` will (in theory) observe every single expression evaluated by the Python interpreter occurring as part of the decorated function call.
@@ -419,6 +434,7 @@ def maxray(
         forbid_modules=forbid_modules,
         restrict_modules=restrict_modules,
         pass_local_scopes=pass_scope,
+        preserve_values=preserve_values,
     )
 
     # TODO: allow configuring injection of variables into exec scope

maxray/__main__.py

@@ -4,25 +4,36 @@
 
 import click
 
+from typing import Optional
+
 
 @click.group()
 def cli():
     pass
 
 
 @cli.command()
-@click.argument("file", type=str)
+@click.option("--over", type=str)
+@click.option("--new", type=str)
 @click.option("-f", "--force", is_flag=True)
 @click.option("--runner", is_flag=True)
-def template(file: str, force: bool, runner: bool):
-    path = Path(file).resolve(True)
-    assert path.suffix == ".py"
+def template(over: Optional[str], new: Optional[str], force: bool, runner: bool):
+    save_path_args = [f for f in [over, new] if f is not None]
+
+    if not save_path_args or len(save_path_args) > 1:
+        raise ValueError("Must specify exactly one of --new or --over")
+
+    save_path = Path(save_path_args[0]).resolve(False)
+    if over is not None:
+        save_path = save_path.with_name(f"over_{save_path.name}")
+
+    assert save_path.suffix == ".py"
 
     spec = find_spec("maxray.inators.template")
     assert spec is not None
     assert spec.origin is not None
 
-    template_path = path.with_name(f"over_{path.name}")
+    template_path = save_path
     if not force:
         assert not template_path.exists(), f"{template_path} exists!"
 

maxray/capture/logs.py

@@ -24,8 +24,7 @@ class CaptureLogs:
         "loc_col_start": pa.int32(),
         "loc_col_end": pa.int32(),
         # Function calls
-        "context_call_id": pa.string(),
-        "target_call_id": pa.string(),
+        "fn_compile_id": pa.string(),
         "fn": pa.string(),
         "fn_call_count": pa.int32(),
         # Source info/code
@@ -58,16 +57,13 @@ class CaptureLogs:
     @staticmethod
     def extractor(x, ctx: NodeContext):
         if isinstance(instance := CaptureLogs.instance.get(None), CaptureLogs):
-            if ctx.caller_id is not None:
-                instance.fn_sources[ctx.caller_id] = ctx.fn_context
             if ctx.source != "self":
                 instance.builder("loc_line_start").append(ctx.location[0])
                 instance.builder("loc_line_end").append(ctx.location[1])
                 instance.builder("loc_col_start").append(ctx.location[2])
                 instance.builder("loc_col_end").append(ctx.location[3])
 
-                instance.builder("context_call_id").append(ctx.fn_context.compile_id)
-                instance.builder("target_call_id").append(ctx.caller_id)
+                instance.builder("fn_compile_id").append(ctx.fn_context.compile_id)
                 instance.builder("source_file").append(ctx.fn_context.source_file)
                 instance.builder("source_module").append(ctx.fn_context.module)
                 instance.builder("source").append(ctx.source)
@@ -106,9 +102,6 @@ def schema():
         )
 
     def __init__(self, stream_to_arrow_file=None, flush_every_records: int = 10_000):
-        # Maps function UUIDs (_MAXRAY_TRANSFORM_ID) to FnContext instances
-        self.fn_sources = {}
-
         self.builders = {}
 
         if stream_to_arrow_file is not None:

maxray/inators/core.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
-from maxray.transforms import NodeContext
+from maxray.nodes import NodeContext, RayContext
 from .display import Display
 
 import ipdb
 import attrs
 
 import json
-from result import Result, Ok, Err
 from contextvars import ContextVar
 from dataclasses import dataclass
 from pathlib import Path
@@ -32,7 +31,11 @@ def __getitem__(self, key):
         return self._existing_keys[key][1]
 
     def __setitem__(self, key, value):
-        v, _old_value = self._existing_keys[key]
+        if key in self._existing_keys:
+            v, _old_value = self._existing_keys[key]
+        else:
+            v = 1
+
         self._existing_keys[key] = (v, value)
 
     def define_once(self, key, factory, /, v: int = 0):
@@ -243,70 +246,21 @@ def __getattr__(self, rewrite_cls_name: str) -> RewriteContext:
         return self.by_class[rewrite_cls_name]
 
 
-def unpack_assign_context(x, ctx):
-    match ctx.props:
-        case {"assigned": {"targets": targets}}:
-            if len(targets) > 1:
-                if inspect.isgenerator(x) or isinstance(x, (map, filter)):
-                    # Greedily consume iterators before assignment
-                    unpacked_x = tuple(iter(x))
-                else:
-                    # Otherwise for chained equality like a = b, c = it, code may rely on `a` being of the original type
-                    unpacked_x = x
-
-                # TODO: doesn't work for starred assignments: x, *y, z = iterable
-                assigned = {target: val for target, val in zip(targets, unpacked_x)}
-                return unpacked_x, assigned
-
-            elif len(targets) == 1:
-                return x, {targets[0]: x}
-            else:
-                return x, {}
-        case _:
-            return x, {}
-
-
 class LoggingEncoder(json.JSONEncoder):
     def default(self, o):
         if attrs.has(type(o)):
             return attrs.asdict(o)
         return super().default(o)
 
 
-class Ray:
+class Ray(RayContext):
     """
     Captures the state of a point (syntax node) in the source code of the original program.
 
     One instance is created for each point in the program, that is then passed to multiple handlers.
     """
 
-    INSTANCE = ContextVar("RY")
-
-    def __init__(self, x, ctx: NodeContext, *, unpack_assignments: bool):
-        if unpack_assignments:
-            self._x, self._assigned = unpack_assign_context(x, ctx)
-        self.ctx = ctx
-
-    # TODO: staticmethod to get
-
-    @contextmanager
-    def _bind(self, x, ctx):
-        # TODO: assert no previous active instance
-        reset = Ray.INSTANCE.set(self)
-        try:
-            yield Ray.INSTANCE.get()
-        finally:
-            Ray.INSTANCE.reset(reset)
-
-    @staticmethod
-    def try_get() -> Result[Ray, None]:
-        try:
-            return Ok(Ray.INSTANCE.get())
-        except LookupError:
-            return Err(None)
-
-    @staticmethod
-    def log(msg, *, level="INFO"):
+    def log(self, msg, *, level="INFO"):
         """
         Logs to Rerun with the current context if active.
         """
@@ -319,52 +273,9 @@ def log(msg, *, level="INFO"):
             case _ if attrs.has(type(msg)):
                 msg = str(msg)
 
-        match Ray.try_get():
-            case Ok(ray):
-                location = Path(ray.ctx.fn_context.source_file).name
-                line = ray.ctx.location[0]
-                rr.log(f"log/{location}:{line}", rr.TextLog(msg, level=level))
-            case Err():
-                rr.log("log/somewhere", rr.TextLog(msg, level=level))
-
-    def value(self):
-        return self._x
-
-    def locals(self):
-        match self.ctx.local_scope:
-            case dict():
-                return self.ctx.local_scope
-            case None:
-                return {}
-            case _:
-                raise TypeError(
-                    f"Unexpected type {type(self.ctx.local_scope)} for local scope"
-                )
-
-    def assigned(self):
-        return self._assigned
-
-    def iterated(self):
-        match self.ctx.props:
-            case {"iterated": {"target": target}}:
-                return [target]
-            case _:
-                return []
-
-    def returned(self): ...
-
-    def entered(self):
-        """
-        Returns:
-            {
-                source: Source code of the LHS being entered
-                as_var: Variable binding after "as", if present
-        """
-        match self.ctx.props:
-            case {"entered": entered}:
-                return entered
-            case _:
-                return {}
+        location = Path(self.ctx.fn_context.source_file).name
+        line = self.ctx.location[0]
+        rr.log(f"log/{location}:{line}", rr.TextLog(msg, level=level))
 
     def contextmanager(self, fn: Callable[[Ray], Iterator[Any]]):
         return contextmanager(partial(fn, self))