casting a pydantic base model attribute to a list

I have a class that inherits from base model, and is defined as below:

class MessageInfo(CustomBaseModel):
    completion_code: int = None
    code: str = None
    message: str = None
    severity: str = None

class OrderResponseMessageInfo(CustomBaseModel):
    message_info: Union[List[MessageInfo], MessageInfo] = None

Want this message_info to be able to consume either a list or a dict, but consistently produce it as a list when i serialise it to a dict or a json.

Since this is a parser for an API response, and the API response is not consistent, i.e sometimes the message info block is an array like so :

{
    "MessageInfo": [
        {
            "CompletionCode": 517,
            "Code": "1B61",
            "Message": "some message",
            "Severity": "Warning"
        },
        {
            "CompletionCode": 9886,
            "Code": "1B61",
            "Message": "Another message",
            "Severity": "Warning"
        }
    ]
}

And sometimes its a dict like so

{
    "MessageInfo": {
        "CompletionCode": 517,
        "Code": "1B61",
        "Message": "some message",
        "Severity": "Warning"
    }
}

When message info is a dict, and I do OrderResponseMessageInfo.dict(), i get

{
    "OrderLineMessageInfo": {
        "MessageInfo": {
                "CompletionCode": 517,
                "Code": "1B61",
                "Message": "Routing completed. Confirmation required.",
                "Severity": "Warning"
        }
    }
}

What i want is that if message info is a dict, then i should be able to type cast it to a list, so that when I do OrderResponseMessageInfo.dict(), i should get the output as:

{
    "OrderLineMessageInfo": {
        "MessageInfo": [
            {
                "CompletionCode": 517,
                "Code": "1B61",
                "Message": "Routing completed. Confirmation required.",
                "Severity": "Warning"
            }
        ]
    }
}

Is there a way to achieve this?

As a general rule, you should define your models in terms of the schema you actually want, not in terms of what you might get.

If you want a field to be of a list type, then define it as such. If you then want to allow singular elements to be turned into one-item-lists as a special case during parsing/initialization, you can define a custom pre=True field validator to do that.

Either way, at the end of the parsing and validation process you should end up with your desired schema, not something you need to mangle again at some stage.

Example:

from collections.abc import Mapping

from pydantic import BaseModel, validator


class Foo(BaseModel):
    x: int
    y: str


class Bar(BaseModel):
    foos: list[Foo]

    @validator("foos", pre=True)
    def single_to_list(cls, v: object) -> object:
        if isinstance(v, (Mapping, Foo)):
            return [Foo.parse_obj(v)]
        return v

Demo:

b1 = Bar.parse_raw('{"foos": [{"x": 1, "y": "spam"}, {"x": -1, "y": "eggs"}]}')
b2 = Bar.parse_raw('{"foos": {"x": 42, "y": "ham"}}')
f = Foo(x=69, y="beans")
b3 = Bar.parse_obj({"foos": f})
b4 = Bar.parse_obj({"foos": [f, f]})
print(b1.json(indent=4))
print(b2.json(indent=4))
print(b3.json(indent=4))
print(b4.json(indent=4))

Output:

{
    "foos": [
        {
            "x": 1,
            "y": "spam"
        },
        {
            "x": -1,
            "y": "eggs"
        }
    ]
}

{
    "foos": [
        {
            "x": 42,
            "y": "ham"
        }
    ]
}

{
    "foos": [
        {
            "x": 69,
            "y": "beans"
        }
    ]
}

{
    "foos": [
        {
            "x": 69,
            "y": "beans"
        },
        {
            "x": 69,
            "y": "beans"
        }
    ]
}

If you want message_info to be a list, then define it as list[MessageInfo] and write a custom validator to try and coerce any mapping or MessageInfo instance to a list containing one MessageInfo instance.

Of course you could instead override the dumping/serialization methods (like .dict and .json), but I would not recommend it. 1) That would result in more code, 2) it would be less efficient, 3) the model schema would still be misleading because you would always dump that field as a list, but the schema would imply a single instance might appear.