pydantic subclass support

[findmyway]

To reproduce

from pydantic import BaseModel

from jsonargparse import auto_cli


class ModelBase(BaseModel): ...


class Model(ModelBase):
    name: str = "default"


class Experiment(ModelBase):
    model: ModelBase

exp = Experiment(model=Model())

def run(x:Experiment = exp):
    print(x)

if __name__ == "__main__":
    auto_cli(run)

  ...
  File ".venv/lib/python3.12/site-packages/jsonargparse/_core.py", line 208, in set_defaults
    raise NSKeyError(f'No action for key "{dest}" to set its default.')
jsonargparse._namespace.NSKeyError: No action for key "x.model" to set its default.

I'm not very sure whether it relates to previous discussions on dataclass and subclasses. All I can see after debugging is that an ArgumentGroup is created for x.model but there's not actions added to the _actions field.

Environment

• jsonargparse version: 4.41.0

[findmyway]

After searching for a while, it seems to just be fixed in #775

And for my case here, I need to add the below lines:

from jsonargparse._common import not_subclass_type_selectors

# not_subclass_type_selectors.pop("dataclass")
not_subclass_type_selectors.pop("pydantic")

However, the expected result should return Model(), but ModelBase() is returned instead here.

model=ModelBase()

[mauvilsa]

To support pydantic subclasses indeed #775 is necessary, though this is experimental. And it is clear that when pydantic is not supporting subclasses, there should be specific error messages to guide the developer. For the code in the issue description above it could fail saying something about the default using subclasses not being supported and hint at enabling pydantic subclass support.

Regarding ModelBase() being returned, the issue is the serialization of the default which currently is done with model_dump. Note the following:

>>> repr(exp)
"Experiment(model=Model(name='default'))"
>>> exp.model_dump()
{'model': {}}

The serialization function should be improved. But the fact that the model key becomes an empty dict makes it impossible. Any idea why model_dump behaves like this? Just curious. This could be fixed by not relying on model_dump.

[findmyway]

Any idea why model_dump behaves like this?

Ah yes, I think that's due to Serializing with duck-typing 🦆

However, even I added the SerializeAsAny annotation, the output is still model=ModelBase().

---

from pydantic import BaseModel, SerializeAsAny

from jsonargparse import auto_cli

from jsonargparse._common import not_subclass_type_selectors

not_subclass_type_selectors.pop("dataclass")
not_subclass_type_selectors.pop("pydantic")


class ModelBase(BaseModel): ...


class Model(ModelBase):
    name: str = "default"


class Experiment(ModelBase):
    model: SerializeAsAny[ModelBase]


exp = Experiment(model=Model())


def run(x: Experiment = exp):
    print(x)


if __name__ == "__main__":
    auto_cli(run)

In [3]: exp
Out[3]: Experiment(model=Model(name='default'))

In [4]: repr(exp)
Out[4]: "Experiment(model=Model(name='default'))"

In [5]: exp.model_dump()
Out[5]: {'model': {'name': 'default'}}

[mauvilsa]

Thanks for the info! Still, {'model': {'name': 'default'}} is not enough to be able to reconstruct particular subclasses. And requiring people to add SerializeAsAny wouldn't be a good user experience.

From what I have seen pydantic doesn't provide out-of-the-box a serializer that includes all necessary details. So for this to work, I will need to implement the serializer. It can be done, just that it would have been nicer if pydantic provided it.

[findmyway]

Thanks for your reply. I think now I understand how it works. Below is a workable version:

from pydantic import BaseModel, SerializeAsAny, Field

from jsonargparse import auto_cli

from jsonargparse._common import not_subclass_type_selectors

not_subclass_type_selectors.pop("dataclass")
not_subclass_type_selectors.pop("pydantic")

from typing import Annotated, Literal, Union


class ModelBase(BaseModel):
    type: Literal["base"] = "base"


class Model(ModelBase):
    name: str = "default"
    type: Literal["model"] = "model"

class Experiment(ModelBase):
    model: SerializeAsAny[ModelBase | Model] = Field(discriminator="type")


exp = Experiment(model=Model())


def run(x: Experiment = exp):
    print(x)


if __name__ == "__main__":
    auto_cli(run)

For my usage, two more works are needed:

1. Update the base class to enable the SerializeAsAny feature as defalt. https://docs.pydantic.dev/latest/concepts/serialization/#overriding-the-serialize_as_any-default-false
2. Automatically create the union type somehow. Related discussions: Subclass resolution pydantic/pydantic#7366

So for this to work, I will need to implement the serializer.

I think it's out of the scope of this package. Better to leave it for end users (with proper error messages if possible).

[findmyway]

Reopen this issue since I just realized that the NSKeyError is unavoidable.

Below is what I get so far to make it work with minor fixes in #780 :

from jsonargparse._common import not_subclass_type_selectors

# not_subclass_type_selectors.pop("dataclass")
not_subclass_type_selectors.pop("pydantic")

from typing import Any
from pydantic import BaseModel, SerializeAsAny, model_serializer, model_validator


# https://github.com/pydantic/pydantic/discussions/7008#discussioncomment-6826052
class BM(BaseModel):
    __subclasses_map__ = {}

    @model_serializer(mode="wrap")
    def __serialize_with_class_type__(self, default_serializer) -> Any:
        ret = default_serializer(self)
        if isinstance(ret, dict):
            ret["klass"] = f"{self.__class__.__module__}.{self.__class__.__qualname__}"
        return ret

    @model_validator(mode="wrap")
    @classmethod
    def __convert_to_real_type__(cls, value: Any, handler):
        if isinstance(value, dict) is False:
            return handler(value)

        # it is a dict so make sure to remove the klass
        # because we don't allow extra keywords but want to ensure
        # e.g Cat.model_validate(cat.model_dump()) works
        class_full_name = value.pop("klass", None)

        # if it's not the polymorphic base we construct via default handler
        # if not cls.__is_polymorphic_base:
        if not BM in cls.__bases__:
            return handler(value)

        # otherwise we lookup the correct polymorphic type and construct that
        # instead
        if class_full_name is None:
            raise ValueError("Missing klass field")

        class_type = cls.__subclasses_map__.get(class_full_name, None)

        if class_type is None:
            # TODO could try dynamic import
            raise TypeError(
                "Trying to instantiate {class_full_name}, which has not yet been defined!"
            )

        return class_type.model_validate(value)

    @classmethod
    def __pydantic_init_subclass__(cls, **kwargs):
        cls.__subclasses_map__[f"{cls.__module__}.{cls.__qualname__}"] = cls

    def model_dump(self, **kwargs) -> dict:
        return super().model_dump(serialize_as_any=True, **kwargs)

    def model_dump_json(self, **kwargs) -> str:
        return super().model_dump_json(serialize_as_any=True, **kwargs)


class Pet(BM):
    name: str


class Cat(Pet):
    meows: int


class SpecialCat(Cat):
    number_of_tails: int


class Dog(Pet):
    barks: float
    friend: Pet


class Person(Pet):
    name: str

    pets: list[Pet]


exp = Person(
    name="jt",
    pets=[
        SpecialCat(name="sc", number_of_tails=2, meows=3),
        Dog(name="dog", barks=2, friend=Cat(name="cc", meows=2)),
    ],
)


def run(x: Person = exp):
    print(x)


from jsonargparse import auto_cli

if __name__ == "__main__":
    auto_cli(run)

The output with python debug.py --print_config is :

x:
  class_path: __main__.Person
  init_args:
    name: jt
    pets:
    - class_path: __main__.Pet
      init_args:
        name: sc
        meows: 3
        number_of_tails: 2
        klass: __main__.SpecialCat
    - class_path: __main__.Pet
      init_args:
        name: dog
        barks: 2.0
        friend:
          name: cc
          meows: 2
          klass: __main__.Cat
        klass: __main__.Dog
    klass: __main__.Person

Note that if I removed the not_subclass_type_selectors.pop("pydantic") line, I got:

x:
  name: jt
  pets:
  - name: sc
    meows: 3
    number_of_tails: 2
    klass: __main__.SpecialCat
  - name: dog
    barks: 2.0
    friend:
      name: cc
      meows: 2
      klass: __main__.Cat
    klass: __main__.Dog

My questions are:

1. https://github.com/omni-us/jsonargparse/pull/780/files#r2403679240 I think adding some more info into the model_dump result seems inevitable. But the added field will result to the NSKeyError.
   a. Is it expected?
   b. Is it possible to somehow disable the validation with some flags or a custom ArgumentParser? (like ignore the keys starting with __)
2. The result of python debug.py --print_config looks quite different with not_subclass_type_selectors.pop("pydantic"). Is it expected? Now there's a duplicated work here (Both klass and the class_path: ... are to set the deserialized class type correctly)

---

Update:

• For 1, I added an extra field of skip_validation in the print_config. And the Minor patch to support pydantic basemodel better #780 is now updated to reflect it. But it only controls the --print_config option. To make it work when execution, I still need to manually disable that exception.
• For 2, I'm wondering if it's possible to convert pydantic model into class_path style directly (not very sure to me yet). Right now model_dump() is used to convert it into a dict. And that's why the above de/serialization work is required.
• Another issue with current implementation: https://github.com/omni-us/jsonargparse/pull/780/files#r2403780254

https://github.com/omni-us/jsonargparse/blob/3eb1bdd73df9be5089b1d55dd6d35fb693aff7b3/jsonargparse/_typehints.py#L269-L273

[mauvilsa]

I think it's out of the scope of this package. Better to leave it for end users (with proper error messages if possible).

@findmyway I am not sure why you think serialization is out of scope. Just because pydantic has some limitations with regards to subclasses, it does not mean that jsonargparse should inherit such limitations. The subclass support in jsonargparse is generic, and there is no reason how it works should be any different for the case of pydantic classes. Things should work for all pydantic classes without the need for developers to figure out how implement jsonargparse-compatible serialization. From your last example, it should work with just:

class Pet(pydantic.BaseModel):
    name: str

class Cat(Pet):
    meows: int

class SpecialCat(Cat):
    number_of_tails: int

class Dog(Pet):
    barks: float
    friend: Pet

class Person(Pet):
    name: str
    pets: list[Pet]

person = Person(
    name="jt",
    pets=[
        SpecialCat(name="sc", number_of_tails=2, meows=3),
        Dog(name="dog", barks=2, friend=Cat(name="cc", meows=2)),
    ],
)

I created pull request #781 which an improved more generic serializer. The change is not specific to pydantic, since I noticed that there is a similar issue for normal dataclasses. With this, your examples work fine.

The result of python debug.py --print_config looks quite different with not_subclass_type_selectors.pop("pydantic"). Is it expected? Now there's a duplicated work here (Both klass and the class_path: ... are to set the deserialized class type correctly)

A reminder, pydantic subclass support, i.e. not_subclass_type_selectors.pop("pydantic") is experimental and the code will change. It is not recommended to use for any real project just yet. The idea is to get feedback to identify issues and have a better final implementation when it is officially supported. And you feedback up to now has been great help.

Without not_subclass_type_selectors.pop("pydantic") it means to use the current pydantic support, which does not support subclasses. If you give it subclasses, the behavior is undefined. Now I am thinking it should just fail. I need to remember to implement this later on.

Is it possible to somehow disable the validation with some flags or a custom ArgumentParser? (like ignore the keys starting with __)

If for some non-jsonargparse reason you need to add klass then I would say just accept the duplication. But also if you have parameters with name that start with a single underscore "_", they are not included in the parser.