Support dataclass inheritance

[gautiervarjo]

🚀 Feature request

The documentation describes how dataclass-like classes are "not expected to have subclasses", and therefore do not accept specifying a class_path. However dataclasses.dataclass does support inheritance: python docs

Would it be feasible to add support for this? I can't say whether this makes sense for all dataclass-like objects, or just the plain python ones.

My use-case for dataclass+inheritance is pretty basic, my main reason for using dataclasses is reducing boilerplate of writing __init__/__str__/etc functions.

To reproduce

A working example with plain classes

import dataclasses
import jsonargparse

class Base:
    def __init__(self, name: str) -> None:
        self.name = name

class Derived(Base):
    def __init__(self) -> None:
        super().__init__(name="derived plain")

parser = jsonargparse.ArgumentParser()
parser.add_subclass_arguments(Base, "thing")
print(parser.parse_args(["--thing.class_path=Derived"]))

A failing example with dataclasses

import dataclasses
import jsonargparse

@dataclasses.dataclass
class DataclassBase:
    name: str

@dataclasses.dataclass
class DataclassDerived(DataclassBase):
    def __init__(self) -> None:
        super().__init__(name="derived dataclass")

# Does not work, jsonargparse assumes no inheritance for dataclass-like types.
parser = jsonargparse.ArgumentParser()

parser.add_subclass_arguments(DataclassBase, "thing")
# ^ this raises: "ValueError: Not allowed for dataclass-like classes."

print(parser.parse_args(["--thing.class_path=DataclassDerived"]))

A failing example with CLI + config

# Config
thing:
  class_path: Derived

# ... snip ...: same definition of DataclassBase and DataclassDerived.

def main(thing: DataclassBase) -> None:
    pass

jsonargparse.CLI(main)
# Passing the config above to this CLI produces:
# usage: parser.py [-h] [--config CONFIG] [--print_config[=flags]] [--thing CONFIG] --thing.name NAME
# parser.py: error: Configuration check failed :: Key "thing.name" is required but not included in config object or its value is None.

Expected behavior

The failing dataclass snippet should behave exactly like the working plain-class one.

Environment

• jsonargparse version: 4.21.1
• Python version: 3.8
• How jsonargparse was installed: installed via Pants/PEX, using the requirement line jsonargparse[signatures]==4.21.1.
• OS: Linux

[mauvilsa]

@gautiervarjo thank you for the proposal!

I do have in my list of things to improve, a way for the developer to specify which classes to consider or not as subclasses or as dataclass-like. Though, I haven't decided yet how this will be.

A couple of minor notes until this kind of support is implemented.

1. A use case could be to accept only a limited number of dataclasses, instead of full subclass support. In this case it can be done as:

   parser.add_argument("--thing", type=Union[DataclassBase, DataclassDerived])

   or with CLI as:

   def main(thing: Union[DataclassBase, DataclassDerived]) -> None:

2. A class that is a mixture of a dataclass and a normal class is not considered dataclass-like. Thus, the subclass behavior can be enabled by doing:

   class NormalClassBase:
       pass
   
   @dataclasses.dataclass
   class DataclassBase(NormalClassBase):
       name: str

[gautiervarjo]

Thank you for the tips! These seem very low-effort to implement so they're definitely good enough for my use-case 🙂

[CM-BF]

@mauvilsa Thank you for your workarounds! I would like to share one more user case: using Equinox in JAX, since every equinox Module is a dataclass.

[mauvilsa]

This feature request also applies to pydantic classes as mentioned in #503.

[rusmux]

@mauvilsa Are there going to be any updates in the near future? I'm particularly interested in supporting pydantic.BaseModel, as above workarounds do not work for it

[mauvilsa]

@rusmux thanks for letting me know about your interest. I do consider interest to decide on which new features to work on. Though I can't really say if soon I would work on this one since it is not simple.

[perepichka]

A similar use-case, which used to work as of version 4.20.0, was hierarchical Optional dataclasses:

With the following setup:

from dataclasses import dataclass

@dataclass
class B:
    param_1: int = 10
    param_2: str = "b"

@dataclass
class Bb(B):
    pass

and

import jsonargparse

from hierarchy_test import B
from typing import Optional


def test(
    b: Optional[B],
):
    print("success")


if __name__ == "__main__":
    jsonargparse.CLI(test)

with the following YAML:

b:
  class_path: hierarchy_test.Bb
  init_args:
    param_1: 1
    param_2: "dd"

This runs with no error on 4.20.0, and fails on 4.33.2

This issue has been preventing me from upgrading to a newer Pytorch Lightning. I've got a quasi-solution by explicitly defining a Union of types as suggested by @mauvilsa . In my production env this involves dynamically creating a Union type based on all subclasses found through a search of a section of the codebase, not the most elegant thing.