Add types for ABIs

mev_inspect/abi.py

@@ -0,0 +1,20 @@
+import json
+from typing import Optional
+
+from pydantic import parse_obj_as
+
+from mev_inspect.config import load_config
+from mev_inspect.schemas import ABI
+
+
+ABI_CONFIG_KEY = "ABI"
+
+config = load_config()
+
+
+def get_abi(abi_name: str) -> Optional[ABI]:
+    if abi_name in config[ABI_CONFIG_KEY]:
+        abi_json = json.loads(config[ABI_CONFIG_KEY][abi_name])
+        return parse_obj_as(ABI, abi_json)
+
+    return None

mev_inspect/schemas/__init__.py

@@ -1 +1,2 @@
+from .abi import ABI
 from .blocks import Block, BlockCall, BlockCallType

mev_inspect/schemas/abi.py

@@ -0,0 +1,52 @@
+from enum import Enum
+from typing import List, Union
+from typing_extensions import Literal
+
+from hexbytes import HexBytes
+from pydantic import BaseModel
+from web3 import Web3
+
+from .utils import CamelModel
+
+
+class ABIDescriptionType(str, Enum):
+    function = "function"
+    constructor = "constructor"
+    fallback = "fallback"
+    event = "event"
+    receive = "receive"
+
+
+NON_FUNCTION_DESCRIPTION_TYPES = Union[
+    Literal[ABIDescriptionType.constructor],
+    Literal[ABIDescriptionType.fallback],
+    Literal[ABIDescriptionType.event],
+    Literal[ABIDescriptionType.receive],
+]
+
+
+class ABIDescriptionInput(CamelModel):
+    internal_type: str
+
+
+class ABIGenericDescription(BaseModel):
+    type: NON_FUNCTION_DESCRIPTION_TYPES
+
+
+class ABIFunctionDescription(BaseModel):
+    type: Literal[ABIDescriptionType.function]
+    name: str
+    inputs: List[ABIDescriptionInput]
+
+    def get_selector(self) -> HexBytes:
+        signature = self.get_signature()
+        return Web3.sha3(text=signature)[0:4]
+
+    def get_signature(self) -> str:
+        joined_input_types = ",".join(input.internal_type for input in self.inputs)
+
+        return f"{self.name}({joined_input_types})"
+
+
+ABIDescription = Union[ABIFunctionDescription, ABIGenericDescription]
+ABI = List[ABIDescription]