Rename classification to classified trace

This commit is contained in:
Luke Van Seters 2021-07-25 21:14:46 -04:00
parent c2dc095c7d
commit 99d321f4e7
2 changed files with 16 additions and 16 deletions

View File

@ -5,7 +5,7 @@ from mev_inspect.decode import ABIDecoder
from mev_inspect.schemas.blocks import Block, Trace, TraceType from mev_inspect.schemas.blocks import Block, Trace, TraceType
from mev_inspect.schemas.classifications import ( from mev_inspect.schemas.classifications import (
Classification, Classification,
ClassificationType, ClassifiedTrace,
DecodeSpec, DecodeSpec,
) )
@ -28,22 +28,22 @@ class Processor:
def process( def process(
self, self,
block: Block, block: Block,
) -> List[Classification]: ) -> List[ClassifiedTrace]:
return [ return [
self._classify(trace) self._classify(trace)
for trace in block.traces for trace in block.traces
if trace.type != TraceType.reward if trace.type != TraceType.reward
] ]
def _classify(self, trace: Trace) -> Classification: def _classify(self, trace: Trace) -> ClassifiedTrace:
if trace.type == TraceType.call: if trace.type == TraceType.call:
classification = self._classify_call(trace) classified_trace = self._classify_call(trace)
if classification is not None: if classified_trace is not None:
return classification return classified_trace
return self._build_unknown_classification(trace) return self._build_unknown_classified_trace(trace)
def _classify_call(self, trace) -> Optional[Classification]: def _classify_call(self, trace) -> Optional[ClassifiedTrace]:
to_address = trace.action["to"] to_address = trace.action["to"]
for spec in self._decode_specs: for spec in self._decode_specs:
@ -58,12 +58,12 @@ class Processor:
call_data = decoder.decode(trace.action["input"]) call_data = decoder.decode(trace.action["input"])
if call_data is not None: if call_data is not None:
return Classification( return ClassifiedTrace(
transaction_hash=trace.transaction_hash, transaction_hash=trace.transaction_hash,
block_number=trace.block_number, block_number=trace.block_number,
trace_type=trace.type, trace_type=trace.type,
trace_address=trace.trace_address, trace_address=trace.trace_address,
classification_type=ClassificationType.unknown, classification=Classification.unknown,
protocol=spec.protocol, protocol=spec.protocol,
function_name=call_data.function_name, function_name=call_data.function_name,
function_signature=call_data.function_signature, function_signature=call_data.function_signature,
@ -73,13 +73,13 @@ class Processor:
return None return None
@staticmethod @staticmethod
def _build_unknown_classification(trace): def _build_unknown_classified_trace(trace):
return Classification( return ClassifiedTrace(
transaction_hash=trace.transaction_hash, transaction_hash=trace.transaction_hash,
block_number=trace.block_number, block_number=trace.block_number,
trace_type=trace.type, trace_type=trace.type,
trace_address=trace.trace_address, trace_address=trace.trace_address,
classification_type=ClassificationType.unknown, classification=Classification.unknown,
protocol=None, protocol=None,
function_name=None, function_name=None,
function_signature=None, function_signature=None,

View File

@ -6,7 +6,7 @@ from pydantic import BaseModel
from .blocks import TraceType from .blocks import TraceType
class ClassificationType(Enum): class Classification(Enum):
unknown = "unknown" unknown = "unknown"
swap = "swap" swap = "swap"
@ -16,12 +16,12 @@ class Protocol(Enum):
sushiswap = "sushiswap" sushiswap = "sushiswap"
class Classification(BaseModel): class ClassifiedTrace(BaseModel):
transaction_hash: str transaction_hash: str
block_number: int block_number: int
trace_type: TraceType trace_type: TraceType
trace_address: List[int] trace_address: List[int]
classification_type: ClassificationType classification: Classification
protocol: Optional[Protocol] protocol: Optional[Protocol]
function_name: Optional[str] function_name: Optional[str]
function_signature: Optional[str] function_signature: Optional[str]