31 lines
1.0 KiB
Python
31 lines
1.0 KiB
Python
from collections.abc import Callable
|
|
|
|
import pandas as pd
|
|
from pydantic import BaseModel, ValidationError
|
|
|
|
|
|
class ValidationInterseptor:
|
|
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["ValidationInterseptorFunc"] = func.__name__
|
|
r["ValidationInterseptorArgs"] = args
|
|
r["ValidationInterseptorKwrds"] = kwrds
|
|
r["ValidationInterseptorIndex"] = i
|
|
self.not_valid_rows.append(r)
|
|
else:
|
|
valid_rows.append(r)
|
|
return pd.DataFrame.from_records(valid_rows)
|
|
|
|
return wrapped
|