Add type annotations to plugin

ucsb-seclab · Feb 13, 2024 · f9c406a · f9c406a
1 parent 8af5263
commit f9c406a
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/greed/state_plugins/plugin.py b/greed/state_plugins/plugin.py
@@ -1,13 +1,20 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from greed.state import SymbolicEVMState
+
 
 class SimStatePlugin:
     """
     Interface for SimState plugins.
     """
+    state: "SymbolicEVMState"
+
     def __init__(self):
         self.state = None
         return
 
-    def set_state(self, state):
+    def set_state(self, state: "SymbolicEVMState"):
         self.state = state
 
     def copy(self):