Add type hints when setting state

trailofbits · Aug 4, 2020 · c11d703 · c11d703
1 parent 9dc5f7b
commit c11d703
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/manticore/native/state.py b/manticore/native/state.py
@@ -23,19 +23,19 @@ def __init__(self, *args, **kwargs):
         self._hooks: Dict[Optional[int], Set[HookCallback]] = {}
         self._after_hooks: Dict[Optional[int], Set[HookCallback]] = {}
 
-    def __getstate__(self):
+    def __getstate__(self) -> Dict[str, Any]:
         state = super().__getstate__()
         state["hooks"] = self._hooks
         state["after_hooks"] = self._after_hooks
         return state
 
-    def __setstate__(self, state):
+    def __setstate__(self, state: Dict[str, Any]) -> None:
         super().__setstate__(state)
         self._hooks = state["hooks"]
         self._after_hooks = state["after_hooks"]
         self._resub_hooks()
 
-    def __enter__(self):
+    def __enter__(self) -> "State":
         new_state = super().__enter__()
         new_state._hooks = copy.copy(self._hooks)
         new_state._after_hooks = copy.copy(self._after_hooks)