29 lines
885 B
Python
29 lines
885 B
Python
|
from collections.abc import Callable
|
||
|
|
||
|
import pandas as pd
|
||
|
from pydantic import BaseModel, ValidationError
|
||
|
|
||
|
|
||
|
class Interseptor:
|
||
|
def __init__(self, model: BaseModel):
|
||
|
self.model = model
|
||
|
self.not_valid_rows = []
|
||
|
|
||
|
def __call__(self, func: Callable[..., pd.DataFrame]):
|
||
|
def wrapped(*args, **kwrds):
|
||
|
res = func(*args, **kwrds)
|
||
|
df_dict = res.to_dict(orient="records")
|
||
|
valid_rows = []
|
||
|
for i, r in enumerate(df_dict):
|
||
|
try:
|
||
|
self.model(**r)
|
||
|
except ValidationError:
|
||
|
r["InterseptorOrigin"] = func.__name__
|
||
|
r["InterseptorIndex"] = i
|
||
|
self.not_valid_rows.append(r)
|
||
|
else:
|
||
|
valid_rows.append(r)
|
||
|
return pd.DataFrame.from_records(valid_rows)
|
||
|
|
||
|
return wrapped
|