CODEWITHSUNDEEP
pythonlistpydantic
Karl
Karl

Reputation: 45

Define custom list for Pydantic

I created a "special" list by subclassing the standard list and adding some extra functions.

class SpecialList(list):

    def double(self):
        for a in self:
            a.value_2 *= 2

Now I would like to use this special list in my pydantic models:

class A(BaseModel):
    value_1: str
    value_2: int


class B(BaseModel):
    a: SpecialList[A] = SpecialList()
        
    class Config:
        arbitrary_types_allowed=True

Generally, this works and I can use my special function as expected:

t = B()
t.a.append(A(value_1='X', value_2=3))
t.a.double()
# returns: B(a=[A(value_1='X', value_2=6)])

But, here comes the issue, I cannot initialize the model from existing data, e.g. a dict:

B.model_validate({'a': [{'value_1': 'X', 'value_2': 3}]})

The following error is then raised:

"""
ValidationError: 1 validation error for B
a
  Input should be an instance of SpecialList [type=is_instance_of, input_value=[{'value_1': 'X', 'value_2': 3}], input_type=list]
    For further information visit https://errors.pydantic.dev/2.4/v/is_instance_of

"""

How can I make the subclassed SpecialList used like a normal list in Pydantic?

I tried to use the @field_validator to convert the list to my SpecialList, but then I had a list of dicts, not a list of As as intended.

Upvotes: 0

Views: 330

Answers (1)

Michael H
Michael H

Reputation: 572

My Answer

You can use pydantic.RootModel.

doc: https://docs.pydantic.dev/latest/concepts/models/#rootmodel-and-custom-root-types

My Example

from pydantic import VERSION, BaseModel, RootModel


print(VERSION)
# > 2.0.3


class SpecialList(RootModel[list]):
    def double(self):
        for a in self.root:  # < use `self.root`
            a.value_2 *= 2

    def append(self, element):  # < if you don't want call `.root` for all instances
        self.root.append(element)


class A(BaseModel):
    value_1: str
    value_2: int


class B(BaseModel):
    a: SpecialList = SpecialList(root=[])

    class Config:
        arbitrary_types_allowed = True


t = B()
t.a.append(A(value_1='X', value_2=3))
print(t.a)
# > root=[A(value_1='X', value_2=3)]


# same as above but without defining method `append`:
t.a.root.append(A(value_1='X', value_2=3))
print(t.a)
# > root=[A(value_1='X', value_2=3), A(value_1='X', value_2=3)]


t.a.double()
print(t.a)
# > root=[A(value_1='X', value_2=6), A(value_1='X', value_2=6)]


b = B.model_validate({'a': [{'value_1': 'X', 'value_2': 3}]})
print(b)
# > a=SpecialList(root=[{'value_1': 'X', 'value_2': 3}])

Upvotes: 1

Related Questions