How to recover types from serialized OmegaConf?

[rhaps0dy]

Suppose I have a config with classes that have subtypes, for example:

from dataclasses import dataclass, field
import pytest
from omegaconf import OmegaConf, SCMode, ValidationError

@dataclass
class AbstractLetter:
    pass

@dataclass
class A(AbstractLetter):
    a: str = "a"

@dataclass
class B(AbstractLetter):
    b: int = 2

@dataclass
class Config:
    l: AbstractLetter = field(default_factory=A)

cfg = OmegaConf.structured(Config(l=B()))

Now with this we get nice type checking at runtime if we assign the wrong thing to subfields of cfg.

# Attempting to set the wrong type here does raise a ValidationError
with pytest.raises(ValidationError):
    cfg.l.b = "hello"

However, this is not the case if we serialize cfg to yaml and then de-serialize it. This is because it becomes a regular DictConfig from that.

serialized = OmegaConf.to_yaml(cfg)
cfg2 = OmegaConf.create(serialized)
# This has no type information, so attempting to set the wrong type is completely fine.
cfg2.l.b = "hello"

If cfg.l was a concrete type, we could just merge the yaml with a OmegaConf.structured(Config) and get type information that way. However, because it's abstract and the default constructor is A, that fails:

OmegaConf.merge(OmegaConf.structured(Config), cfg2)

How do I serialize in such a way that we can get back type information?

[Jasha10]

One strategy is to use OmegaConf.merge to combine the deserialized config with the target type. Something like this:

typed_cfg2 = OmegaConf.merge(Config, cfg2)

[rhaps0dy]

Thank you for answering! As I wrote at the very end of my question, this doesn't work. The reason is that cfg2 has different attributes (because it was created from a non-default AbstractLetter) and the merging thus correctly fails.

If cfg2 was typed then the merge would realize it has to replace the entire cfg2.l instead of just cfg2.l.b, and it would succeed; but it's not typed.

[Jasha10]

Sorry, I should have read the end of your question before answering.

This is a use case that would be covered by #144, which is not yet fully implemented.
The usage pattern would look something like this:

@dataclass
class ConfigUnion:
    l: A | B = field(default_factory=A)

That A | B is a union type. Currently, running OmegaConf.merge(ConfigUnion, cfg2) gives the following error message:

omegaconf.errors.ConfigValueError: Unions of containers are not supported:
l: Union[A, B]
    full_key:
    object_type=None

Unfortunately there's no good workaround for this as far as I'm aware.

[rhaps0dy]

That makes sense, thank you! I'm probably going to write my own thing that supports this use case then.

[rhaps0dy]

It'd still be annoying to have to explicitly write out A | B | ... (if there are more sub-classes) under #144.

[Jasha10]

It'd still be annoying to have to explicitly write out A | B | ... (if there are more sub-classes) under #144.

Hmm, interesting.

In terms of OmegaConf's internals, I believe implementing merge support for explicit unions l: A | B | ... would be more straightforward than implementing merge support for arbitrary subclasses l: AbstractLetter. This is because, when you call, typed_cfg2 = OmegaConf.merge(Config, cfg2), OmegaConf needs to have a way to know what "candidate subclasses" of AbstractLetter are available. In the "explicit unions" case, OmegaConf can parse the type annotations on the Config classes (as is done in the omegaconf._utils.get_dataclass_data function) to learn that Config.l must be either A or B. In the "arbitrary subclasses" case, however, there is no readily-accessible list of candidate subclasses of AbstractLetter that OmegaConf can get access to. When calling OmegaConf.merge on a union type, OmegaConf iterates over the components of the Union annotation until it finds a component such that the merge can succeed (this is done in the omegaconf.base.UnionNode._set_value_impl method).

In theory something like python's abstract base class facility could be used to register and keep track of candidate subclasses of AbstractLetter (i.e. by making AbstractLetter an abstract base class and then later calling AbstractLetter.register(A) and AbstractLetter.register(B)). This would likely take a bit more work to implement, however.

[rhaps0dy]

Thanks! I agree with this, except it's actually rather easy to get subclasses explicitly. Here's a recursive function that does so:

def all_subclasses(cls) -> set[type]:
    a = set(cls.__subclasses__())
    for c in a:
        a = a | all_subclasses(c)
    return a

One problem with this is that dataclasses with slots=True create two classes with the same name and module, and only one of them should be used. I have some code with battle-scars for this which you might find useful; but also maybe it's too much complexity to add. In that case you need to keep track of the order in which subclasses have been added and only use the last one which has the same (cls.__module__, cls.__qualname__) tuple.

EDIT: I no longer believe the code I pasted does the right thing.