Merge pull request #76 from KeerthiVasudevan/add_hooks_for_decorator

Add hooks for decorators

src/dynapyt/analyses/TraceAll.py

@@ -1396,6 +1396,63 @@ def exit_with(self, dyn_ast: str, iid: int, is_suppressed: bool, exc_value):
         """
         self.log(iid, "Exited with")
 
+    def enter_decorator(self, dyn_ast: str, iid: int, decorator_name, args, kwargs):
+        """Hook for entering a decorator.
+
+
+        Parameters
+        ----------
+        dyn_ast : str
+            The path to the original code. Can be used to extract the syntax tree.
+
+        iid : int
+            Unique ID of the syntax tree node.
+
+        decorator_name : str
+            The name of the decorator function.
+
+        args : List
+            The positional arguments passed to the decorator.
+
+        kwargs : Dict
+            The keyword arguments passed to the decorator.
+
+        """
+        self.log(iid, "Entered decorator", decorator_name)
+
+    def exit_decorator(self, dyn_ast: str, iid: int, decorator_name, result, args, kwargs) -> Any:
+        """Hook for exiting a decorator.
+
+
+        Parameters
+        ----------
+        dyn_ast : str
+            The path to the original code. Can be used to extract the syntax tree.
+
+        iid : int
+            Unique ID of the syntax tree node.
+
+        decorator_name : str
+            The name of the decorator function.
+
+        result : Any
+            The result of the decorator.
+
+        args : List
+            The positional arguments passed to the decorator.
+
+        kwargs : Dict
+            The keyword arguments passed to the decorator.
+
+        
+        Returns
+        -------
+        Any
+            If provided, overwrites the result returned by the function 
+
+        """
+        self.log(iid, "Exited decorator", decorator_name)
+
     # Top level
 
     def runtime_event(self, dyn_ast: str, iid: int) -> None:

src/dynapyt/instrument/CodeInstrumenter.py

@@ -1930,3 +1930,31 @@ def leave_WithItem(self, original_node, updated_node):
         self.to_import.add("_enter_with_")
         call = cst.Call(func=callee_name, args=[ast_arg, iid_arg, ctx_manager_arg])
         return updated_node.with_changes(item=call)
+
+
+    def leave_Decorator(self, original_node, updated_node):
+        print("decorator node: ", original_node)
+        if ("enter_decorator" not in self.selected_hooks) and (
+            "exit_decorator" not in self.selected_hooks
+        ):
+            return updated_node
+
+        iid = self.__create_iid(original_node)
+        ast_arg = cst.Arg(value=cst.Name("_dynapyt_ast_"))
+        iid_arg = cst.Arg(value=cst.Integer(value=str(iid)))
+        dynapyt_decorator_attr = cst.Attribute(
+            value=cst.Name("_rt"), attr=cst.Name("dynapyt_decorator"),
+        )
+        dynapyt_decorator_call = cst.Call(
+            func=dynapyt_decorator_attr,
+            args=[ast_arg, iid_arg],
+        )
+        dynapyt_decorator = cst.Decorator(
+            decorator=dynapyt_decorator_call
+        )
+
+        self.to_import.add("dynapyt_decorator")
+
+        return cst.FlattenSentinel([dynapyt_decorator, updated_node])
+
+

src/dynapyt/runtime.py

@@ -13,7 +13,7 @@
 import json
 import sys
 import os
-from functools import lru_cache
+from functools import lru_cache, wraps
 from tempfile import gettempdir
 from .utils.hooks import snake, get_name
 from .instrument.IIDs import IIDs
@@ -830,3 +830,29 @@ def _enter_with_(self, dyn_ast, iid, ctx_manager_arg):
         self.call_if_exists("runtime_event", dyn_ast, iid)
         self.call_if_exists("control_flow_event", dyn_ast, iid)
         self.call_if_exists("exit_with", dyn_ast, iid, is_suppressed, None)
+
+
+    def _enter_decorator_(self, dyn_ast, iid, decorator, args, kwargs):
+        self.call_if_exists("runtime_event", dyn_ast, iid)
+        self.call_if_exists("control_flow_event", dyn_ast, iid)
+        self.call_if_exists("enter_decorator", dyn_ast, iid, decorator, args, kwargs)
+
+    def _exit_decorator_(self, dyn_ast, iid, decorator, result, args, kwargs):
+        self.call_if_exists("runtime_event", dyn_ast, iid)
+        self.call_if_exists("control_flow_event", dyn_ast, iid)
+        result = self.call_if_exists("exit_decorator", dyn_ast, iid, decorator, result, args, kwargs)
+        return result
+
+    def dynapyt_decorator(self, iid_arg, ast_arg):
+        def dynapyt_decorator_wrapper(func):
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                self._enter_decorator_(ast_arg, iid_arg, func.__name__, args, kwargs)
+                result = func(*args, **kwargs)
+                r = self._exit_decorator_(ast_arg, iid_arg, func.__name__, result, args, kwargs)
+                if (r is not None):
+                    return r
+                return result
+
+            return wrapper
+        return dynapyt_decorator_wrapper

src/dynapyt/utils/hierarchy.json

@@ -94,7 +94,9 @@
                 "_return": {},
                 "_yield": {},
                 "implicit_return": {}
-            }
+            },
+            "enter_decorator": {},
+            "exit_decorator": {}
         },
         "memory_access": {
             "read": {

tests/manipulate_single_hook/decorator/analysis.py

@@ -0,0 +1,21 @@
+from dynapyt.analyses.BaseAnalysis import BaseAnalysis
+
+
+class TestAnalysis(BaseAnalysis):
+    def begin_execution(self) -> None:
+        print("begin execution")
+
+    def enter_decorator(self, dyn_ast: str, iid: int, decorator: str, args, kwargs) -> None:
+        print("enter decorator: ", decorator)
+
+    def exit_decorator(self, dyn_ast: str, iid: int, decorator: str, result, args, kwargs) -> None:
+        print("exit decorator: ", decorator)
+
+    def runtime_event(self, dyn_ast: str, iid: int) -> None:
+        print("runtime event")
+
+    def control_flow_event(self, dyn_ast: str, iid: int) -> None:
+        print("control flow event")
+
+    def end_execution(self) -> None:
+        print("end execution")

tests/manipulate_single_hook/decorator/expected.txt

@@ -0,0 +1,39 @@
+begin execution
+runtime event
+control flow event
+runtime event
+runtime event
+control flow event
+runtime event
+runtime event
+control flow event
+runtime event
+control flow event
+enter decorator:  wrapper
+runtime event
+control flow event
+runtime event
+runtime event
+control flow event
+Decorator function before
+runtime event
+runtime event
+control flow event
+runtime event
+control flow event
+runtime event
+runtime event
+control flow event
+Simple function
+runtime event
+control flow event
+runtime event
+runtime event
+control flow event
+Decorator function after
+runtime event
+control flow event
+runtime event
+control flow event
+exit decorator:  wrapper
+end execution

tests/manipulate_single_hook/decorator/program.py

@@ -0,0 +1,13 @@
+def decorator_function_one(func):
+    def wrapper():
+        print("Decorator function before")
+        func()
+        print("Decorator function after")
+    return wrapper
+
+@decorator_function_one
+def simple_function():
+    print("Simple function")
+
+simple_function()
+

tests/trace_single_hook/decorators/builtin_decorator/analysis.py

@@ -0,0 +1,14 @@
+from dynapyt.analyses.BaseAnalysis import BaseAnalysis
+
+class TestAnalysis(BaseAnalysis):
+    def begin_execution(self) -> None:
+        print("begin execution")
+
+    def enter_decorator(self, dyn_ast: str, iid: int, decorator: str, args, kwargs) -> None:
+        print("enter decorator: ", decorator)
+
+    def exit_decorator(self, dyn_ast: str, iid: int, decorator: str, result, args, kwargs) -> None:
+        print("exit decorator: ", decorator)
+
+    def end_execution(self) -> None:
+        print("end execution")

tests/trace_single_hook/decorators/builtin_decorator/expected.txt

@@ -0,0 +1,7 @@
+begin execution
+enter decorator:  my_context_manager
+exit decorator:  my_context_manager
+Entering context
+Inside context
+Exiting context
+end execution

tests/trace_single_hook/decorators/builtin_decorator/program.py

@@ -0,0 +1,13 @@
+from contextlib import contextmanager
+
+@contextmanager
+def my_context_manager():
+    print('Entering context')
+    yield
+    print('Exiting context')
+
+with my_context_manager():
+    print('Inside context')
+
+
+

tests/trace_single_hook/decorators/decorator_class_attribute/analysis.py

@@ -0,0 +1,14 @@
+from dynapyt.analyses.BaseAnalysis import BaseAnalysis
+
+class TestAnalysis(BaseAnalysis):
+    def begin_execution(self) -> None:
+        print("begin execution")
+
+    def enter_decorator(self, dyn_ast: str, iid: int, decorator: str, args, kwargs) -> None:
+        print("enter decorator: ", decorator)
+
+    def exit_decorator(self, dyn_ast: str, iid: int, decorator: str, result, args, kwargs) -> None:
+        print("exit decorator: ", decorator)
+
+    def end_execution(self) -> None:
+        print("end execution")

tests/trace_single_hook/decorators/decorator_class_attribute/expected.txt

@@ -0,0 +1,7 @@
+begin execution
+enter decorator:  wrapper
+Decorator function before
+Simple function
+Decorator function after
+exit decorator:  wrapper
+end execution

tests/trace_single_hook/decorators/decorator_class_attribute/program.py

@@ -0,0 +1,15 @@
+class foo:
+    def decorator_func(self, func):
+        def wrapper():
+            print("Decorator function before")
+            func()
+            print("Decorator function after")
+        return wrapper
+
+f = foo()
+
+@f.decorator_func
+def simple_function():
+    print("Simple function")
+
+simple_function()

tests/trace_single_hook/decorators/decorator_func_with_args/analysis.py

@@ -0,0 +1,15 @@
+from dynapyt.analyses.BaseAnalysis import BaseAnalysis
+
+class TestAnalysis(BaseAnalysis):
+    def begin_execution(self) -> None:
+        print("begin execution")
+
+    def enter_decorator(self, dyn_ast: str, iid: int, decorator: str, args, kwargs) -> None:
+        print("enter decorator: ", decorator)
+        print("args: ", args)
+
+    def exit_decorator(self, dyn_ast: str, iid: int, decorator: str, result, args, kwargs) -> None:
+        print("exit decorator: ", decorator)
+
+    def end_execution(self) -> None:
+        print("end execution")

tests/trace_single_hook/decorators/decorator_func_with_args/expected.txt

@@ -0,0 +1,10 @@
+begin execution
+enter decorator:  wrapper
+args:  ('foo', 'bar')
+Decorator function before
+Simple function
+parameter 1:  foo
+parameter 2:  bar
+Decorator function after
+exit decorator:  wrapper
+end execution

tests/trace_single_hook/decorators/decorator_func_with_args/program.py

@@ -0,0 +1,15 @@
+def decorator_function_one(func):
+    def wrapper(*args):
+        print("Decorator function before")
+        func(*args)
+        print("Decorator function after")
+    return wrapper
+
+@decorator_function_one
+def simple_function(arg1, arg2):
+    print("Simple function")
+    print("parameter 1: ", arg1)
+    print("parameter 2: ", arg2)
+
+simple_function("foo", "bar")
+

tests/trace_single_hook/decorators/decorator_result_modification/analysis.py

@@ -0,0 +1,17 @@
+from dynapyt.analyses.BaseAnalysis import BaseAnalysis
+
+class TestAnalysis(BaseAnalysis):
+    def begin_execution(self) -> None:
+        print("begin execution")
+
+    def enter_decorator(self, dyn_ast: str, iid: int, decorator: str, args, kwargs) -> None:
+        print("enter decorator: ", decorator)
+
+    def exit_decorator(self, dyn_ast: str, iid: int, decorator: str, result, args, kwargs) -> None:
+        print("exit decorator: ", decorator)
+        number = 10
+        print("Number returned from exit_decorator: ", number)
+        return number
+
+    def end_execution(self) -> None:
+        print("end execution")

tests/trace_single_hook/decorators/decorator_result_modification/expected.txt

@@ -0,0 +1,9 @@
+begin execution
+enter decorator:  wrapper
+Decorator function before
+Number returned from function:  2
+Decorator function after
+exit decorator:  wrapper
+Number returned from exit_decorator:  10
+Number returned from decorated function:  10
+end execution

tests/trace_single_hook/decorators/decorator_result_modification/program.py

@@ -0,0 +1,17 @@
+def decorator_function_one(func):
+    def wrapper():
+        print("Decorator function before")
+        func()
+        print("Decorator function after")
+    return wrapper
+
+@decorator_function_one
+def return_number():
+    number = 2
+    print("Number returned from function: ", number)
+    return number
+
+result = return_number()
+print("Number returned from decorated function: ", result)
+
+