Add type hints to all classes in mythril module

mythril/mythril/mythril_analyzer.py

@@ -7,6 +7,9 @@
 
 import logging
 import traceback
+from typing import Optional, List
+
+from . import MythrilDisassembler
 from mythril.support.source_support import Source
 from mythril.support.loader import DynLoader
 from mythril.analysis.symbolic import SymExecWrapper
@@ -14,17 +17,23 @@ from mythril.analysis.callgraph import generate_graph
 from mythril.analysis.traceexplore import get_serializable_statespace
 from mythril.analysis.security import fire_lasers, retrieve_callback_issues
 from mythril.analysis.report import Report
+from mythril.ethereum.evmcontract import EVMContract
 
 log = logging.getLogger(__name__)
 
 
-class MythrilAnalyzer(object):
+class MythrilAnalyzer:
     """
     The Mythril Analyzer class
     Responsible for the analysis of the smart contracts
     """
 
-    def __init__(self, disassembler, requires_dynld=False, onchain_storage_access=True):
+    def __init__(
+        self,
+        disassembler: MythrilDisassembler,
+        requires_dynld: bool = False,
+        onchain_storage_access: bool = True,
+    ):
         """
 
         :param disassembler: The MythrilDisassembler class
@@ -32,21 +41,21 @@ class MythrilAnalyzer(object):
         :param onchain_storage_access: Whether onchain access should be done or not
         """
         self.eth = disassembler.eth
-        self.contracts = disassembler.contracts or []
+        self.contracts = disassembler.contracts or []  # type: List[EVMContract]
         self.enable_online_lookup = disassembler.enable_online_lookup
         self.dynld = requires_dynld
         self.onchain_storage_access = onchain_storage_access
 
     def dump_statespace(
         self,
-        strategy,
+        strategy: str,
-        contract,
+        contract: List[EVMContract],
-        address=None,
+        address: Optional[str] = None,
-        max_depth=None,
+        max_depth: Optional[int] = None,
-        execution_timeout=None,
+        execution_timeout: Optional[int] = None,
-        create_timeout=None,
+        create_timeout: Optional[int] = None,
-        enable_iprof=False,
+        enable_iprof: bool = False,
-    ):
+    ) -> str:
         """
         Returns serializable statespace of the contract
         :param strategy: The search strategy to go through the CFG
@@ -77,16 +86,16 @@ class MythrilAnalyzer(object):
 
     def graph_html(
         self,
-        strategy,
+        strategy: str,
-        contract,
+        contract: List[EVMContract],
-        address,
+        address: str,
-        max_depth=None,
+        max_depth: Optional[int] = None,
-        enable_physics=False,
+        enable_physics: bool = False,
-        phrackify=False,
+        phrackify: bool = False,
-        execution_timeout=None,
+        execution_timeout: Optional[int] = None,
-        create_timeout=None,
+        create_timeout: Optional[int] = None,
-        enable_iprof=False,
+        enable_iprof: bool = False,
-    ):
+    ) -> str:
         """
 
         :param strategy: The search strategy to go through the CFG
@@ -118,17 +127,17 @@ class MythrilAnalyzer(object):
 
     def fire_lasers(
         self,
-        strategy,
+        strategy: str,
-        contracts=None,
+        contracts: Optional[List[EVMContract]] = None,
-        address=None,
+        address: Optional[str] = None,
-        modules=None,
+        modules: Optional[List[str]] = None,
-        verbose_report=False,
+        verbose_report: bool = False,
-        max_depth=None,
+        max_depth: Optional[int] = None,
-        execution_timeout=None,
+        execution_timeout: Optional[int] = None,
-        create_timeout=None,
+        create_timeout: Optional[int] = None,
-        transaction_count=None,
+        transaction_count: Optional[int] = None,
-        enable_iprof=False,
+        enable_iprof: bool = False,
-    ):
+    ) -> Report:
         """
         :param strategy: The search strategy to go through the CFG
         :param contracts: The Contracts list on which the analysis should be done

mythril/mythril/mythril_config.py

@@ -7,6 +7,7 @@ import re
 from pathlib import Path
 from shutil import copyfile
 from configparser import ConfigParser
+from typing import Optional
 
 from mythril.exceptions import CriticalError
 from mythril.ethereum.interface.rpc.client import EthJsonRpc
@@ -15,16 +16,16 @@ from mythril.ethereum.interface.leveldb.client import EthLevelDB
 log = logging.getLogger(__name__)
 
 
-class MythrilConfig(object):
+class MythrilConfig:
     def __init__(self):
         self.mythril_dir = self._init_mythril_dir()
         self.config_path = os.path.join(self.mythril_dir, "config.ini")
         self.leveldb_dir = self._init_config()
-        self.eth = None
+        self.eth = None  # type: Optional[EthJsonRpc]
-        self.eth_db = None
+        self.eth_db = None  # type: Optional[EthLevelDB]
 
     @staticmethod
-    def _init_mythril_dir():
+    def _init_mythril_dir() -> str:
         """
         Initializes the mythril dir and config.ini file
         :return: The mythril dir's path
@@ -48,7 +49,7 @@ class MythrilConfig(object):
 
         return mythril_dir
 
-    def _init_config(self):
+    def _init_config(self) -> str:
         """If no config file exists, create it and add default options.
 
         Default LevelDB path is specified based on OS
@@ -83,7 +84,7 @@ class MythrilConfig(object):
         return os.path.expanduser(leveldb_dir)
 
     @staticmethod
-    def _get_fallback_dir():
+    def _get_fallback_dir() -> str:
         """
         Returns the LevelDB path
         :return: The LevelDB path
@@ -103,7 +104,7 @@ class MythrilConfig(object):
         return os.path.join(leveldb_fallback_dir, "geth", "chaindata")
 
     @staticmethod
-    def _add_default_options(config):
+    def _add_default_options(config: ConfigParser) -> None:
         """
         Adds defaults option to config.ini
         :param config: The config file object
@@ -112,7 +113,7 @@ class MythrilConfig(object):
         config.add_section("defaults")
 
     @staticmethod
-    def _add_leveldb_option(config, leveldb_fallback_dir):
+    def _add_leveldb_option(config: ConfigParser, leveldb_fallback_dir: str) -> None:
         """
         Sets a default leveldb path in .mythril/config.ini file
         :param config: The config file object
@@ -129,7 +130,7 @@ class MythrilConfig(object):
         config.set("defaults", "leveldb_dir", leveldb_fallback_dir)
 
     @staticmethod
-    def _add_dynamic_loading_option(config):
+    def _add_dynamic_loading_option(config: ConfigParser) -> None:
         """
         Sets the dynamic loading config option in .mythril/config.ini file
         :param config: The config file object
@@ -146,17 +147,17 @@ class MythrilConfig(object):
         )
         config.set("defaults", "dynamic_loading", "infura")
 
-    def set_api_leveldb(self, leveldb_path):
+    def set_api_leveldb(self, leveldb_path: str) -> None:
         """
         """
         self.eth_db = EthLevelDB(leveldb_path)
 
-    def set_api_rpc_infura(self):
+    def set_api_rpc_infura(self) -> None:
         """Set the RPC mode to INFURA on Mainnet."""
         log.info("Using INFURA Main Net for RPC queries")
         self.eth = EthJsonRpc("mainnet.infura.io", 443, True)
 
-    def set_api_rpc(self, rpc=None, rpctls=False):
+    def set_api_rpc(self, rpc: str = None, rpctls: bool = False) -> None:
         """
         Sets the RPC mode to either ganache or infura
         """
@@ -181,12 +182,12 @@ class MythrilConfig(object):
         else:
             raise CriticalError("Invalid RPC settings, check help for details.")
 
-    def set_api_rpc_localhost(self):
+    def set_api_rpc_localhost(self) -> None:
         """Set the RPC mode to a local instance."""
         log.info("Using default RPC settings: http://localhost:8545")
         self.eth = EthJsonRpc("localhost", 8545)
 
-    def set_api_from_config_path(self):
+    def set_api_from_config_path(self) -> None:
         """Set the RPC mode based on a given config file."""
         config = ConfigParser(allow_no_value=False)
         config.optionxform = str

mythril/mythril/mythril_disassembler.py

@@ -5,8 +5,9 @@ import os
 
 from ethereum import utils
 from solc.exceptions import SolcError
-from typing import List, Optional
+from typing import List, Tuple, Optional
 from mythril.ethereum import util
+from mythril.ethereum.interface.rpc.client import EthJsonRpc
 from mythril.exceptions import CriticalError, CompilerError, NoContractFoundError
 from mythril.support import signatures
 from mythril.support.truffle import analyze_truffle_project
@@ -19,17 +20,21 @@ log = logging.getLogger(__name__)
 
 class MythrilDisassembler:
     def __init__(
-        self, eth, solc_version=None, solc_args=None, enable_online_lookup=False
+        self,
-    ):
+        eth: Optional[EthJsonRpc],
+        solc_version: str = None,
+        solc_args: str = None,
+        enable_online_lookup: bool = False,
+    ) -> None:
         self.solc_binary = self._init_solc_binary(solc_version)
         self.solc_args = solc_args
         self.eth = eth
         self.enable_online_lookup = enable_online_lookup
         self.sigs = signatures.SignatureDB(enable_online_lookup=enable_online_lookup)
-        self.contracts = []
+        self.contracts = []  # type: List[EVMContract]
 
     @staticmethod
-    def _init_solc_binary(version):
+    def _init_solc_binary(version: str) -> str:
         """
         Only proper versions are supported. No nightlies, commits etc (such as available in remix).
         :param version: Version of the solc binary required
@@ -68,7 +73,9 @@ class MythrilDisassembler:
 
         return solc_binary
 
-    def load_from_bytecode(self, code, bin_runtime=False, address=None):
+    def load_from_bytecode(
+        self, code: str, bin_runtime: bool = False, address: Optional[str] = None
+    ) -> Tuple[str, EVMContract]:
         """
         Returns the address and the contract class for the given bytecode
         :param code: Bytecode
@@ -96,7 +103,7 @@ class MythrilDisassembler:
             )
         return address, self.contracts[-1]  # return address and contract object
 
-    def load_from_address(self, address):
+    def load_from_address(self, address: str) -> Tuple[str, EVMContract]:
         """
         Returns the contract given it's on chain address
         :param address: The on chain address of a contract
@@ -130,7 +137,9 @@ class MythrilDisassembler:
                 )
         return address, self.contracts[-1]  # return address and contract object
 
-    def load_from_solidity(self, solidity_files):
+    def load_from_solidity(
+        self, solidity_files: List[str]
+    ) -> Tuple[str, List[SolidityContract]]:
         """
 
         :param solidity_files: List of solidity_files
@@ -180,18 +189,18 @@ class MythrilDisassembler:
 
         return address, contracts
 
-    def analyze_truffle_project(self, *args, **kwargs):
+    def analyze_truffle_project(self, *args, **kwargs) -> None:
         """
         :param args:
         :param kwargs:
         :return:
         """
-        return analyze_truffle_project(
+        analyze_truffle_project(
             self.sigs, *args, **kwargs
         )  # just passthru by passing signatures for now
 
     @staticmethod
-    def hash_for_function_signature(func):
+    def hash_for_function_signature(func: str) -> str:
         """
         Return function name's corresponding signature hash
         :param func: function name

mythril/mythril/mythril_leveldb.py

@@ -2,7 +2,7 @@ import re
 from mythril.exceptions import CriticalError
 
 
-class MythrilLevelDB(object):
+class MythrilLevelDB:
     """
     Class which does search operations on leveldb
     There are two DBs

tests/mythril/mythril_analyzer_test.py

@@ -1,6 +1,5 @@
-import pytest
 from pathlib import Path
-from mythril.mythril import *
+from mythril.mythril import MythrilDisassembler, MythrilAnalyzer
 
 
 def test_fire_lasers():

tests/mythril/mythril_config_test.py

@@ -0,0 +1,2 @@
+from mythril.mythril import MythrilConfig
+