[Python] Using Pydantic to define AST of boolean expressions

I'd like to use Pydantic to define and validate AST of queries that will be applied on a Pandas dataframe. Here's my code:

from typing import List, Literal, Optional, Union
from pydantic import BaseModel, Field
import pandas as pd

class ColumnCondition(BaseModel):
"""
A class that represents a condition that is applied to a single column.
"""
tag: Literal["ColumnCondition"] = "ColumnCondition"
column: str = Field(..., title="The name of the column to apply the condition to.")
operator: Literal["==", "!=", "<", ">", "<=", ">="] = Field(
..., title="The operator of the condition."
)
value: Optional[str] = Field(None, title="The value to compare the column to.")


class AndCondition(BaseModel):
"""
A class that represents an 'and' condition that is applied to two or more conditions.
"""
tag: Literal["AndCondition"] = "AndCondition"
conditions: List["Condition"]


Condition = Union[ColumnCondition, AndCondition]

class ConditionModel(BaseModel):
condition: Condition = Field(discriminator="tag")

def get_column_metadata(df: pd.DataFrame) -> dict:
return {col: str(dtype) for col, dtype in df.dtypes.items()}


if __name__ == "__main__":
"""
Example
"""
condition_json = {
"tag": "AndCondition",
"conditions": [
{
"tag": "ColumnCondition",
"column": "original_amount.currency",
"operator": ">=",
"value": "100",
},
{
"tag": "ColumnCondition",
"column": "original_amount.currency",
"operator": "<=",
"value": "1000",
},
],
}

cond = ConditionModel.model_validate({"condition": condition_json})
print(cond.model_dump_json(indent=2))


This works well, but I have a few questions:

• Is there a way to remove the ConditionModel wrapper class? I couldn't work around it.
• What is the best way to handle the types of values? should I have another field in the ColumnCondition class of its type? or maybe hold a list of columns and types?
• What is the best way to convert such a condition into a string to be used in the Dataframe.query method? should I implement a __str__ in each class? or maybe write a method that traverse the AST and creates the string?

Continue reading...