How can I decode a JSON string into a pydantic model with a dataframe field?

Viewed 752

I am using MongoDB to store the results of a script into a database. When I want to reload the data back into python, I need to decode the JSON (or BSON) string into a pydantic basemodel. With a pydantic model with JSON compatible types, I can just do:

base_model = BaseModelClass.parse_raw(string)

But the default json.loads decoder doesn't know how to deal with a DataFrame. I can overwrite the .parse_raw function into something like:

from pydantic import BaseModel
import pandas as pd

class BaseModelClass(BaseModel):
    df: pd.DataFrame
    
    class Config:
        arbitrary_types_allowed = True
        json_encoders = {
            pd.DataFrame: lambda df: df.to_json()
        }

    @classmethod
    def parse_raw(cls, data):
        data = json.loads(data)
        data['df'] = pd.read_json(data['df'])
        return cls(**data)

But ideally I would want to automatically decode fields of type pd.DataFrame rather than manually change the parse_raw function every time. Is there any way of doing something like:

    class Config:
        arbitrary_types_allowed = True
        json_encoders = {
            pd.DataFrame: lambda df: df.to_json()
        }
        json_decoders = {
            pd.DataFrame: lambda df: pd.read_json(data['df'])
        }

To make the detection of any field which should be a data frame, be converted to one, without having to modify the parse_raw() script?

0 Answers
Related