diff --git a/schemas/.utils.py.swp b/schemas/.utils.py.swp new file mode 100644 index 0000000..fb8d39c Binary files /dev/null and b/schemas/.utils.py.swp differ diff --git a/schemas/blocks.py b/schemas/blocks.py index fb71d97..18e41b1 100644 --- a/schemas/blocks.py +++ b/schemas/blocks.py @@ -2,11 +2,9 @@ import json from enum import Enum from typing import Dict, List, Optional -from hexbytes import HexBytes from pydantic import BaseModel -from web3.datastructures import AttributeDict -from .utils import CamelModel +from .utils import CamelModel, Web3Model class BlockCallType(Enum): @@ -30,7 +28,7 @@ class BlockCall(CamelModel): error: Optional[str] -class Block(BaseModel): +class Block(Web3Model): block_number: int calls: List[BlockCall] data: dict @@ -39,12 +37,6 @@ class Block(BaseModel): transaction_hashes: List[str] txs_gas_data: Dict[str, dict] - class Config: - json_encoders = { - AttributeDict: dict, - HexBytes: lambda h: h.hex(), - } - def get_filtered_calls(self, hash: str) -> List[BlockCall]: return [ call for call in self.calls diff --git a/schemas/utils.py b/schemas/utils.py index 8a57059..a3cb04b 100644 --- a/schemas/utils.py +++ b/schemas/utils.py @@ -1,6 +1,8 @@ import json +from hexbytes import HexBytes from pydantic import BaseModel +from web3.datastructures import AttributeDict def to_camel(string: str) -> str: @@ -10,13 +12,23 @@ def to_camel(string: str) -> str: ) +def to_original_json_dict(model: BaseModel) -> dict: + return json.loads(model.json(by_alias=True, exclude_unset=True)) + + +class Web3Model(BaseModel): + """BaseModel that handles web3's unserializable objects""" + + class Config: + json_encoders = { + AttributeDict: dict, + HexBytes: lambda h: h.hex(), + } + + class CamelModel(BaseModel): """BaseModel that translates from camelCase to snake_case""" - class Config: + class Config(Web3Model.Config): alias_generator = to_camel allow_population_by_field_name = True - - -def to_original_json_dict(model: BaseModel) -> dict: - return json.loads(model.json(by_alias=True, exclude_unset=True))