Add functools.wraps decorator to TAC handler

ucsb-seclab · Nov 4, 2023 · 77c19dc · 77c19dc
1 parent 659b5d5
commit 77c19dc
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/greed/TAC/base.py b/greed/TAC/base.py
@@ -1,3 +1,4 @@
+import functools
 import logging
 from typing import List, Mapping, Callable
 
@@ -118,6 +119,7 @@ def handler_without_side_effects(func: Callable[["TAC_Statement", SymbolicEVMSta
         (can just read and return statically computed results).
         """
 
+        @functools.wraps(func)
         def wrap(self, state: SymbolicEVMState):
             self.set_arg_val(state)
             state.trace.append(state.curr_stmt)
@@ -150,6 +152,7 @@ def handler_with_side_effects(func: Callable[["TAC_Statement", SymbolicEVMState]
         (can't just read and return statically computed results).
         """
 
+        @functools.wraps(func)
         def wrap(self, state: SymbolicEVMState):
             self.set_arg_val(state)
             state.trace.append(state.curr_stmt)