omry / omegaconf

Flexible Python configuration system. The last one you will ever need.
BSD 3-Clause "New" or "Revised" License
1.97k stars 111 forks source link

[Static typing] Make DictConfig generic in the dataclass that it represents (if any) #899

Open tmke8 opened 2 years ago

tmke8 commented 2 years ago

Is your feature request related to a problem? Please describe. The return type of a lot of functions in omegaconf is a union of a lot of types, which makes them not very useful. For example, to_object returns a Union with Any, so it's effectively Any.

Describe the solution you'd like Proof of concept for a generic DictConfig:

from dataclasses import dataclass
from typing import Any, Dict, Generic, Type, TypeVar

T = TypeVar("T")

class DictConfig(Generic[T]):
    pass

class OmegaConf:
    @staticmethod
    def structured(obj: Type[T]) -> DictConfig[T]:  # the DictConfig "remembers" the original class
        raise NotImplementedError

    @staticmethod
    def create(obj: Dict[Any, Any]) -> DictConfig[None]:  # not a structured config, so we just say `None`
        raise NotImplementedError

    @staticmethod
    def to_object(cfg: DictConfig[T]) -> T:  # very precise return type
        raise NotImplementedError

@dataclass
class MyConfig:
    x: int

def main() -> None:
    d = OmegaConf.structured(MyConfig)
    reveal_type(d)  # DictConfig[MyConfig]
    o = OmegaConf.to_object(d)
    reveal_type(o)  # MyConfig

Describe alternatives you've considered There might be other solutions to the problem of unspecific return types, but I haven't thought of one yet.

Additional context None.

Jasha10 commented 2 years ago

Cool idea!

A related category of possible typing improvements is more precise use of @typing.overload for the omegaconf.OmegaConf methods.

Jasha10 commented 2 years ago

I suspect that such a change would not be backward compatible. For example, currently annotations like this are legal:

cfg: DictConfig = OmegaConf.structured(...)

If we make the change you've proposed, mypy will probably require annotations like this:

cfg: DictConfig[SomeType] = OmegaConf.structured(...)

Since we don't want to break the annotations in existing code, I'll milestone this feature request for OmegaConf 3.0.

bhogan-bdai commented 3 months ago

This feature would be grate. I'm using OmegaConf with Metaflow. Being able to parse a config file into a structured (read: correct type) config would eliminate a day or two of wokarounds I need to add.