feat: patch __init__ and __call__

Author: blepabyte

maxray/__init__.py

@@ -80,12 +80,54 @@ def callable_allowed_for_transform(x, ctx: NodeContext):
     )
 
 
+def instance_init_allowed_for_transform(x, ctx: NodeContext):
+    """
+    Decides whether the __init__ method can be transformed.
+    """
+    return (
+        type(x) is type
+        and hasattr(x, "__init__")
+        and not hasattr(x, "_MAXRAY_TRANSFORMED")
+    )
+
+
+def instance_call_allowed_for_transform(x, ctx: NodeContext):
+    """
+    Decides whether the __call__ method can be transformed.
+    """
+    return (
+        type(x) is type
+        and hasattr(x, "__call__")
+        and not hasattr(x, "_MAXRAY_TRANSFORMED")
+    )
+
+
 def _maxray_walker_handler(x, ctx: NodeContext):
     # We ignore writer calls triggered by code execution in other writers to prevent easily getting stuck in recursive hell
     if _GLOBAL_WRITER_ACTIVE_FLAG.get():
         return x
 
-    # 1. logic to recursively patch callables
+    # 1.  logic to recursively patch callables
+    # 1a. special-case callables: __init__ and __call__
+    if instance_init_allowed_for_transform(x, ctx):
+        # TODO: should we somehow delay doing this until before an actual call?
+        match recompile_fn_with_transform(
+            x.__init__, _maxray_walker_handler, special_use_instance_type=x
+        ):
+            case Ok(init_patch):
+                logger.info(f"Patching __init__ for class {x}")
+                setattr(x, "__init__", init_patch)
+
+    if instance_call_allowed_for_transform(x, ctx):
+        # TODO: should we somehow delay doing this until before an actual call?
+        match recompile_fn_with_transform(
+            x.__call__, _maxray_walker_handler, special_use_instance_type=x
+        ):
+            case Ok(call_patch):
+                logger.info(f"Patching __call__ for class {x}")
+                setattr(x, "__call__", call_patch)
+
+    # 1b. normal functions or bound methods or method descriptors like @classmethod and @staticmethod
     if callable_allowed_for_transform(x, ctx):
         # TODO: don't cache objects w/ __call__
         if x in _MAXRAY_FN_CACHE:
@@ -94,25 +136,24 @@ def _maxray_walker_handler(x, ctx: NodeContext):
             hook.active_call_state.get() and hook.descend_predicate(x, ctx)
             for hook in _MAXRAY_REGISTERED_HOOKS
         ):
-            # user-defined filters for which nodes to descend into
+            # user-defined filters for which nodes (not) to descend into
             pass
         else:
             match recompile_fn_with_transform(x, _maxray_walker_handler):
                 case Ok(x_trans):
                     # NOTE: x_trans now has _MAXRAY_TRANSFORMED field to True
                     if inspect.ismethod(x):
                         # Two cases: descriptor vs bound method
-                        # TODO: handle callables and .__call__ patching
                         match x.__self__:
                             case type():
                                 # Descriptor
-                                logger.warning(
+                                logger.debug(
                                     f"monkey-patching descriptor method {x.__name__} on type {x.__self__}"
                                 )
                                 parent_cls = x.__self__
                             case _:
                                 # Bound method
-                                logger.warning(
+                                logger.debug(
                                     f"monkey-patching bound method {x.__name__} on type {type(x.__self__)}"
                                 )
                                 parent_cls = type(x.__self__)

maxray/transforms.py

@@ -357,7 +357,7 @@ def get_fn_name(fn):
         try:
             name = repr(fn)
         except Exception:
-            name = "<unrepresentable function>"
+            name = f"<unrepresentable function of type {type(fn)}>"
 
     return f"{name} @ {id(fn)}"
 
@@ -369,10 +369,18 @@ def recompile_fn_with_transform(
     ast_post_callback=None,
     initial_scope={},
     pass_scope=False,
+    special_use_instance_type=None,
 ) -> Result[Callable, str]:
     """
     Recompiles `source_fn` so that essentially every node of its AST tree is wrapped by a call to `transform_fn` with the evaluated value along with context information about the source code.
     """
+    # TODO: use non-overridable __getattribute__ instead?
+    if not hasattr(source_fn, "__name__"):  # Safety check against weird functions
+        return Err(f"There is no __name__ for function {get_fn_name(source_fn)}")
+
+    if source_fn.__name__ == "<lambda>":
+        return Err("Cannot safely recompile lambda functions")
+
     # handle `functools.wraps`
     if hasattr(source_fn, "__wrapped__"):
         # SOUNDNESS: failure when decorators aren't applied at the definition site (will look for the original definition, ignoring any transformations that have been applied before the wrap but after definition)
@@ -404,10 +412,8 @@ def recompile_fn_with_transform(
         return Err(
             f"No source code for probable built-in function {get_fn_name(source_fn)}"
         )
-
-    # TODO: use non-overridable __getattribute__ instead?
-    if not hasattr(source_fn, "__name__"):  # Safety check against weird functions
-        return Err(f"There is no __name__ for function {get_fn_name(source_fn)}")
+    except SyntaxError:
+        return Err(f"Syntax error in function {get_fn_name(source_fn)}")
 
     if "super()" in source:
         # TODO: we could replace calls to super() with super(__class__, self)?
@@ -438,6 +444,11 @@ def recompile_fn_with_transform(
                 # Bound method
                 parent_cls = type(source_fn.__self__)
 
+    # yeah yeah an unbound __init__ isn't actually a method but we can basically treat it as one
+    if special_use_instance_type is not None:
+        fn_is_method = True
+        parent_cls = special_use_instance_type
+
     fn_call_counter = ContextVar("maxray_call_counter", default=0)
     fn_context = FnContext(
         source_fn,

tests/test_transforms.py

@@ -402,7 +402,7 @@ def uh():
         z = X()
         return z()
 
-    assert uh() == 3
+    assert uh() == 4
 
 
 def test_junk_annotations():