Fix type hints and move the get annotations to another file

mythril/analysis/analysis_module_helpers.py

@@ -0,0 +1,18 @@
+from typing import Callable
+from mythril.laser.ethereum.state.annotation import StateAnnotation
+
+
+def get_annotation(state, annotation_type: Callable) -> StateAnnotation:
+    """
+    Annotation is searched, if not found then a new annotation is created
+    :param state: Get's annotation from state
+    :param annotation_type: The type of the annotation
+    :return: The annotation
+    """
+    annotations = list(state.get_annotations(annotation_type))
+    if len(annotations) == 0:
+        state.annotate(annotation_type())
+        annotations = list(state.get_annotations(annotation_type))
+    assert len(annotations) == 1
+    return annotations[0]
+

mythril/analysis/modules/arbitrary_write.py

@@ -2,21 +2,17 @@
 withdrawal."""
 import logging
 from copy import copy
-from typing import List, cast, Tuple
+from typing import List, Tuple
 from mythril.analysis import solver
+from mythril.analysis.analysis_module_helpers import get_annotation
 from mythril.analysis.modules.base import DetectionModule
 
 from mythril.analysis.report import Issue
-from mythril.laser.ethereum.transaction.symbolic import (
-    ATTACKER_ADDRESS,
-    CREATOR_ADDRESS,
-)
 from mythril.analysis.swc_data import WRITE_TO_ARBITRARY_STORAGE
 from mythril.exceptions import UnsatError
-from mythril.laser.ethereum.transaction import ContractCreationTransaction
 from mythril.laser.ethereum.state.annotation import StateAnnotation
 from mythril.laser.ethereum.state.global_state import GlobalState
-from mythril.laser.smt import UGT, Sum, symbol_factory, BVAddNoOverflow, If, BitVecFunc
+from mythril.laser.smt import symbol_factory, BitVecFunc
 
 log = logging.getLogger(__name__)
 
@@ -44,24 +40,6 @@ class AribtraryWriteAnnotation(StateAnnotation):
         return result
 
 
-def get_annotation(state, AnnotationType) -> StateAnnotation:
-    """
-    Annotation is searched
-    :param state:
-    :return:
-    """
-    annotations = cast(
-        List[AnnotationType], list(state.get_annotations(AnnotationType))
-    )
-    if len(annotations) == 0:
-        state.annotate(AnnotationType())
-        annotations = cast(
-            List[AnnotationType], list(state.get_annotations(AnnotationType))
-        )
-    assert len(annotations) == 1
-    return annotations[0]
-
-
 class ArbitraryStorage(DetectionModule):
     """This module search for cases where Ether can be withdrawn to a user-
     specified address."""