A potential solution for improving enum-like classes

[bruno-f-cruz]

After a brainstorming session with @jtyoung84 , we came up with this idea:

from pydantic import BaseModel, RootModel, Field
from typing import Annotated, List, Literal, Union, ClassVar, Dict

class _Base(BaseModel):
    key: str
    a: int

    def __hash__(self) -> int:
        return hash(self.key)


class _Foo(_Base):
    key: Literal["Foo"] = "Foo"
    a: int = 1

class _Bar(_Base):
    key: Literal["Bar"] = "Bar"
    a: int = 2 


class MyClass(RootModel):
    root: Annotated[Union[_Foo, _Bar], Field(..., discriminator="key")]
    ONE_OF: ClassVar[Annotated[Union[_Foo, _Bar], Field(..., discriminator="key")]] = Annotated[Union[_Foo, _Bar], Field(..., discriminator="key")]
    Foo: ClassVar[_Foo] = _Foo()
    Bar: ClassVar[_Bar] = _Bar()


class Container(BaseModel):
    prop_to_be_deprecated: MyClass.ONE_OF
    prop: MyClass
    dict_of_prop: Dict[MyClass.ONE_OF, int] # This doesn't work :(

new = Container(
        prop_to_be_deprecated=MyClass.Bar,
        prop=MyClass.Foo,
        dict_of_prop={MyClass.Foo: 1})

# However this works :D
a = {}
a[_Foo()] = 1
a[_Bar()] = 2

print(a[_Bar()])

This would solve a few problems. Namely:

• The ability to be able to use discriminated unions using the RootModel class (Current Enum-like Union pattern does not provide type hinting aind-data-schema#843, Refactor discriminated unions to ensure references to original discriminator are preserved aind-data-schema#794)
• Ability to provide type hinting and enumeration with the same base class (instead of using ONE_OF)
• Fully backwards compatible (including the ONE_OF)
• Refactor can be made via Automatically generate literal models via codegen #64
• Ability to use model instances as keys (However, we cant yet use them to type dictionaries signatures as demonstrated above...) (Enum-like classes cant be used as dictionary keys aind-data-schema#960)