diff --git a/src/dynapyt/analyses/TraceAll.py b/src/dynapyt/analyses/TraceAll.py index e7f738a..d1e8a7b 100644 --- a/src/dynapyt/analyses/TraceAll.py +++ b/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: diff --git a/src/dynapyt/instrument/CodeInstrumenter.py b/src/dynapyt/instrument/CodeInstrumenter.py index 84c315f..871fb53 100644 --- a/src/dynapyt/instrument/CodeInstrumenter.py +++ b/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]) + + \ No newline at end of file diff --git a/src/dynapyt/runtime.py b/src/dynapyt/runtime.py index 38fb56a..a14ef59 100644 --- a/src/dynapyt/runtime.py +++ b/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 \ No newline at end of file diff --git a/src/dynapyt/utils/hierarchy.json b/src/dynapyt/utils/hierarchy.json index 6b2f99e..91ccf90 100644 --- a/src/dynapyt/utils/hierarchy.json +++ b/src/dynapyt/utils/hierarchy.json @@ -94,7 +94,9 @@ "_return": {}, "_yield": {}, "implicit_return": {} - } + }, + "enter_decorator": {}, + "exit_decorator": {} }, "memory_access": { "read": { diff --git a/tests/manipulate_single_hook/decorator/analysis.py b/tests/manipulate_single_hook/decorator/analysis.py new file mode 100644 index 0000000..8b1578d --- /dev/null +++ b/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") \ No newline at end of file diff --git a/tests/manipulate_single_hook/decorator/expected.txt b/tests/manipulate_single_hook/decorator/expected.txt new file mode 100644 index 0000000..5825f43 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/manipulate_single_hook/decorator/program.py b/tests/manipulate_single_hook/decorator/program.py new file mode 100644 index 0000000..964106f --- /dev/null +++ b/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() + diff --git a/tests/trace_single_hook/decorators/builtin_decorator/analysis.py b/tests/trace_single_hook/decorators/builtin_decorator/analysis.py new file mode 100644 index 0000000..cc08bca --- /dev/null +++ b/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") \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/builtin_decorator/expected.txt b/tests/trace_single_hook/decorators/builtin_decorator/expected.txt new file mode 100644 index 0000000..53289f8 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/builtin_decorator/program.py b/tests/trace_single_hook/decorators/builtin_decorator/program.py new file mode 100644 index 0000000..d6505c4 --- /dev/null +++ b/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') + + + diff --git a/tests/trace_single_hook/decorators/decorator_class_attribute/analysis.py b/tests/trace_single_hook/decorators/decorator_class_attribute/analysis.py new file mode 100644 index 0000000..cc08bca --- /dev/null +++ b/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") \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/decorator_class_attribute/expected.txt b/tests/trace_single_hook/decorators/decorator_class_attribute/expected.txt new file mode 100644 index 0000000..41d805c --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/decorator_class_attribute/program.py b/tests/trace_single_hook/decorators/decorator_class_attribute/program.py new file mode 100644 index 0000000..efcd28e --- /dev/null +++ b/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() \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/decorator_func_with_args/analysis.py b/tests/trace_single_hook/decorators/decorator_func_with_args/analysis.py new file mode 100644 index 0000000..762cf03 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/decorator_func_with_args/expected.txt b/tests/trace_single_hook/decorators/decorator_func_with_args/expected.txt new file mode 100644 index 0000000..0587cf0 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/decorator_func_with_args/program.py b/tests/trace_single_hook/decorators/decorator_func_with_args/program.py new file mode 100644 index 0000000..5244530 --- /dev/null +++ b/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") + diff --git a/tests/trace_single_hook/decorators/decorator_result_modification/analysis.py b/tests/trace_single_hook/decorators/decorator_result_modification/analysis.py new file mode 100644 index 0000000..4de3d88 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/decorator_result_modification/expected.txt b/tests/trace_single_hook/decorators/decorator_result_modification/expected.txt new file mode 100644 index 0000000..ba1e0eb --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/decorator_result_modification/program.py b/tests/trace_single_hook/decorators/decorator_result_modification/program.py new file mode 100644 index 0000000..cc99138 --- /dev/null +++ b/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) + + diff --git a/tests/trace_single_hook/decorators/decorator_simple/analysis.py b/tests/trace_single_hook/decorators/decorator_simple/analysis.py new file mode 100644 index 0000000..cc08bca --- /dev/null +++ b/tests/trace_single_hook/decorators/decorator_simple/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") \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/decorator_simple/expected.txt b/tests/trace_single_hook/decorators/decorator_simple/expected.txt new file mode 100644 index 0000000..41d805c --- /dev/null +++ b/tests/trace_single_hook/decorators/decorator_simple/expected.txt @@ -0,0 +1,7 @@ +begin execution +enter decorator: wrapper +Decorator function before +Simple function +Decorator function after +exit decorator: wrapper +end execution \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/decorator_simple/program.py b/tests/trace_single_hook/decorators/decorator_simple/program.py new file mode 100644 index 0000000..964106f --- /dev/null +++ b/tests/trace_single_hook/decorators/decorator_simple/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() + diff --git a/tests/trace_single_hook/decorators/decorator_with_args/analysis.py b/tests/trace_single_hook/decorators/decorator_with_args/analysis.py new file mode 100644 index 0000000..06059a6 --- /dev/null +++ b/tests/trace_single_hook/decorators/decorator_with_args/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") \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/decorator_with_args/expected.txt b/tests/trace_single_hook/decorators/decorator_with_args/expected.txt new file mode 100644 index 0000000..41d805c --- /dev/null +++ b/tests/trace_single_hook/decorators/decorator_with_args/expected.txt @@ -0,0 +1,7 @@ +begin execution +enter decorator: wrapper +Decorator function before +Simple function +Decorator function after +exit decorator: wrapper +end execution \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/decorator_with_args/program.py b/tests/trace_single_hook/decorators/decorator_with_args/program.py new file mode 100644 index 0000000..655afde --- /dev/null +++ b/tests/trace_single_hook/decorators/decorator_with_args/program.py @@ -0,0 +1,15 @@ +def decorator_function_with_args(arg1, arg2): + def decorator_function(func): + def wrapper(): + print("Decorator function before") + func() + print("Decorator function after") + return wrapper + return decorator_function + +@decorator_function_with_args("arg1", "arg2") +def simple_function(): + print("Simple function") + +simple_function() + diff --git a/tests/trace_single_hook/decorators/nested_decorator/analysis.py b/tests/trace_single_hook/decorators/nested_decorator/analysis.py new file mode 100644 index 0000000..cc08bca --- /dev/null +++ b/tests/trace_single_hook/decorators/nested_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") \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/nested_decorator/expected.txt b/tests/trace_single_hook/decorators/nested_decorator/expected.txt new file mode 100644 index 0000000..71e7268 --- /dev/null +++ b/tests/trace_single_hook/decorators/nested_decorator/expected.txt @@ -0,0 +1,12 @@ +begin execution +enter decorator: wrapper_two +Decorator function two: before +enter decorator: wrapper_one +Decorator function one: before +Inside simple function +Decorator function one: after +exit decorator: wrapper_one +Decorator function two: after +exit decorator: wrapper_two +2 +end execution \ No newline at end of file diff --git a/tests/trace_single_hook/decorators/nested_decorator/program.py b/tests/trace_single_hook/decorators/nested_decorator/program.py new file mode 100644 index 0000000..ff64d20 --- /dev/null +++ b/tests/trace_single_hook/decorators/nested_decorator/program.py @@ -0,0 +1,24 @@ +def decorator_function_one(func): + def wrapper_one(): + print("Decorator function one: before") + result = func() + print("Decorator function one: after") + return result + 1 + return wrapper_one + +def decorator_function_two(func): + def wrapper_two(): + print("Decorator function two: before") + result = func() + print("Decorator function two: after") + return result + return wrapper_two + +@decorator_function_two +@decorator_function_one +def simple_function(): + print("Inside simple function") + return 1 + +print(simple_function()) +