feat: convenience command to capture logs from external scripts

Author: blepabyte

maxray/__init__.py

@@ -260,6 +260,7 @@ def maxray(
     *,
     mutable=True,
     pass_scope=False,
+    initial_scope={},
 ):
     """
     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.
@@ -284,7 +285,7 @@ def maxray(
     #     del frame
 
     # TODO: allow configuring injection of variables into exec scope
-    caller_locals = {}
+    caller_locals = initial_scope
 
     def recursive_transform(fn):
         _MAXRAY_REGISTERED_HOOKS.append(
@@ -304,7 +305,7 @@ def recursive_transform(fn):
             match recompile_fn_with_transform(
                 fn,
                 _maxray_walker_handler,
-                initial_scope=caller_locals,
+                override_scope=caller_locals,
                 pass_scope=pass_scope,
                 is_maxray_root=True,
             ):

maxray/capture/logs.py

@@ -9,13 +9,27 @@
 from contextvars import ContextVar
 from functools import wraps
 import time
+from pathlib import Path
 
 from loguru import logger
 
 
 class CaptureLogs:
     instance = ContextVar("CaptureLogs")
 
+    LSP_NO_SHOW_TYPES = {
+        "type",
+        "function",
+        "builtin_function_or_method",
+        "staticmethod",
+        "method",
+        "module",
+        "NoneType",
+        "int",  # mostly uninteresting and makes array indexing expressions cluttered
+        "int32",
+        "int64",
+    }
+
     @staticmethod
     def extractor(x, ctx: NodeContext):
         if isinstance(instance := CaptureLogs.instance.get(None), CaptureLogs):
@@ -36,13 +50,24 @@ def extractor(x, ctx: NodeContext):
                 instance.builder("source_file").append(ctx.fn_context.source_file)
                 instance.builder("source").append(ctx.source)
 
-                assert isinstance(ctx.fn_context.name, str)
                 instance.builder("fn").append(ctx.fn_context.name)
                 instance.builder("fn_call_count").append(
                     ctx.fn_context.call_count.get()
                 )
-                instance.builder("struct_repr").append(structured_repr(x))
-                instance.builder("value_type").append(repr(type(x)))
+
+                struct_repr = structured_repr(x)
+                instance.builder("struct_repr").append(struct_repr)
+
+                value_type = repr(type(x))
+                instance.builder("value_type").append(value_type)
+
+                type_name = type(x).__name__
+                if type_name in CaptureLogs.LSP_NO_SHOW_TYPES:
+                    instance.builder("lsp_repr").append(None)
+                else:
+                    instance.builder("lsp_repr").append(
+                        f" {struct_repr}"
+                    )  # space inserted for formatting
 
                 instance.builder("timestamp").append(time.perf_counter())
 
@@ -52,7 +77,9 @@ def extractor(x, ctx: NodeContext):
 
         return x
 
-    def __init__(self):
+    def __init__(self, from_script_path=None):
+        # TODO: support in-memory mode
+
         # Maps function UUIDs (_MAXRAY_TRANSFORM_ID) to FnContext instances
         self.fn_sources = {}
 
@@ -74,11 +101,16 @@ def __init__(self):
             "source": pa.string(),
             # Extracted data
             "struct_repr": pa.string(),
+            "lsp_repr": pa.string(),
             "value_type": pa.string(),
             "timestamp": pa.float64(),
         }
 
-        self.save_to = "/tmp/maxray-logs.arrow"
+        log_file_name = ".maxray-logs.arrow"
+        if from_script_path is not None:
+            self.save_to = Path(from_script_path).resolve(True).parent / log_file_name
+        else:
+            self.save_to = Path("/tmp") / log_file_name
 
         self.flush_every_records = 10_000
 
@@ -94,6 +126,10 @@ def builder(self, name: str):
         return self.builders[name]
 
     def flush(self):
+        if not self.builders:
+            logger.warning("Nothing to flush")
+            return
+
         arrays, names = [], []
         for col_name, col_type in self.type_schema.items():
             builder = self.builders[col_name]
@@ -111,14 +147,16 @@ def __enter__(self):
 
         # TODO: forbid re-entry
 
-        self.sink = self.write_context.enter_context(pa.OSFile(self.save_to, "wb"))
+        self.sink = self.write_context.enter_context(pa.OSFile(str(self.save_to), "wb"))
         self.writer = self.write_context.enter_context(
             pa.ipc.new_file(self.sink, self.schema())
         )
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         try:
+            self.flush()
+
             if exc_type is not None:
                 return
 

maxray/capture/logs_cli.py

@@ -0,0 +1,124 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from maxray import xray, _set_logging, NodeContext
+from maxray.capture.logs import CaptureLogs
+
+import ast
+from textwrap import indent
+import sys
+import tempfile
+from typing import Any
+from copy import copy
+from types import ModuleType
+
+import numpy as np
+
+
+# assign some dummy module that doesn't conflict
+sys.modules["override_mod"] = np
+
+
+@dataclass
+class WrapScriptIntoMain:
+    """
+    Convenience utility: Sometimes you just want to profile/annotate a script once without modifying any source code.
+    This makes a copy of the script file, wrapping its contents in a callable "main" function to which `maxray` transforms can then be applied.
+    """
+
+    script_path: str
+    temp_sourcefile: Any = field(
+        default_factory=lambda: tempfile.NamedTemporaryFile("w", delete=False)
+    )
+
+    def build(self):
+        with open(self.script_path, "r") as file:
+            source = file.read()
+
+        new_source = f"""def main():
+{indent(source, '    ')}
+"""
+
+        # inspect.getsource relies on the file actually existing
+        with open(self.temp_sourcefile.name, "w") as f:
+            f.write(new_source)
+
+        tree = ast.parse(new_source, filename=self.temp_sourcefile.name)
+        main_func = ast.fix_missing_locations(tree)
+
+        # BUG: will error if `import *` is used (doesn't work inside a fn `def`)
+        compiled_code = compile(
+            main_func, filename=self.temp_sourcefile.name, mode="exec"
+        )
+
+        namespace = {}
+        exec(compiled_code, namespace)
+        # TODO: use non-conflicting name other than "main"
+        main = namespace["main"]
+        main.__module__ = "override_mod"
+        return main
+
+    def run(self):
+        fn = xray(self.rewrite_node, initial_scope={"__name__": "__main__"})(
+            self.build()
+        )
+        with CaptureLogs(self.script_path) as cl:
+            fn()
+
+    def rewrite_node(self, x, ctx: NodeContext):
+        # functions and contextvars can't be deepcopied
+        ctx = copy(ctx)
+        ctx.fn_context = copy(ctx.fn_context)
+
+        # source_file should never be self.script_path because we've copied
+        if ctx.fn_context.source_file == self.temp_sourcefile.name:
+            ctx.fn_context.source_file = self.script_path
+
+            # subtract the "def" line and new indentation
+            ctx.location = (
+                ctx.location[0] - 1,
+                ctx.location[1] - 1,
+                ctx.location[2] - 4,
+                ctx.location[3] - 4,
+            )
+        CaptureLogs.extractor(x, ctx)
+        return x
+
+    @staticmethod
+    def empty_module(module_name: str):
+        module_code = ""
+
+        module_ast = ast.parse(module_code, filename=f"{module_name}.py")
+
+        module_code_object = compile(
+            module_ast, filename=f"{module_name}.py", mode="exec"
+        )
+
+        module = ModuleType(module_name)
+
+        exec(module_code_object, module.__dict__)
+
+        return module
+
+
+def run_script():
+    _set_logging(True)
+
+    match sys.argv:
+        case (
+            _,
+            script_path,
+        ) if (
+            path := Path(script_path).resolve(True)
+        ).exists() and path.suffix == ".py":
+            WrapScriptIntoMain(str(path)).run()
+
+        case (_0, script_path, "--", *args):
+            sys.argv = [_0, *args]
+            path = Path(script_path).resolve(True)
+            if path.exists() and path.suffix == ".py":
+                WrapScriptIntoMain(str(path)).run()
+
+        case _:
+            raise RuntimeError(
+                f"Incorrect argument usage - expected `capture-logs <script_path> -- script_args...` (got {sys.argv[1:]})"
+            )

maxray/transforms.py

@@ -58,6 +58,9 @@ class NodeContext:
     """
 
     location: tuple[int, int, int, int]
+    """
+    (start_line, end_line, start_col, end_col)
+    """
 
     local_scope: Any = None
 
@@ -458,7 +461,7 @@ def recompile_fn_with_transform(
     transform_fn,
     ast_pre_callback=None,
     ast_post_callback=None,
-    initial_scope={},
+    override_scope={},
     pass_scope=False,
     special_use_instance_type=None,
     is_maxray_root=False,
@@ -555,7 +558,7 @@ def recompile_fn_with_transform(
     fn_call_counter = ContextVar("maxray_call_counter", default=0)
     fn_context = FnContext(
         source_fn,
-        source_fn.__name__,
+        source_fn.__qualname__,
         module.__name__,
         source,
         sourcefile,
@@ -595,7 +598,7 @@ def patch_mro(super_type: super):
             "_MAXRAY_BUILTINS_LOCALS": locals,
             "_MAXRAY_PATCH_MRO": patch_mro,
         },
-        "override": initial_scope,
+        "override": override_scope,
         "class_local": {},
         "module": {},
         "closure": {},

pyproject.toml

@@ -15,6 +15,8 @@ python = "^3.11"
 result = "^0.16.1"
 loguru = "^0.7.2"
 
+[tool.poetry.scripts]
+capture-logs = "maxray.capture.logs_cli:run_script"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.1.1"