Create a typing from a custom dataclass

Question

I'm currently working with a dataclass (defined in one of the dependencies I'm currently working with inside of my project), which is defined similar to this snippet:

from dataclasses import dataclass
from typing import Union

class X:
    ...

class Y:
    ...

@dataclass
class Container:
    data: Union[X, Y, dict]

Then, within my project, I wanted to specify some cases where I already know the type of data. With Container being a dataclass, I know it is not possible to hint directly over the attribute as following:

my_container = Container(X())
my_container.data: X  # IDE complains

I was wondering if it is possible then to specify it anyhow similar to Container[X] to define that the only attribute it is from the type I already know of.

If not, is there any dunder method to add at the dataclass to allow this typing?

Note: I am not allowed to modify the original Container code.

Answer 1

Since Container is not a generic, I do not believe there is a way to directly parameterize it like Container[x]. However, to accomplish something similar to what you attempted in your example, you can use typing.cast() to tell the type checker what type you think the data is. For example, you could do this:

def only_accept_X(x: X) -> None:
    ...


my_container = Container(X())
my_container.data = cast(X, my_container.data)

only_accept_X(my_container.data)

Your type checker not complain about the last line because it will recognize my_container.data as being of type X.
Depending on your use case, you might also prefer to write a Generic wrapper class:

from typing import cast, Generic, TypeVar

T = TypeVar("T", X, Y, dict)

class TypedContainer(Generic[T]):
    def __init__(self, data: T):
        self.container = Container(data)

    @property
    def data(self) -> T:
        return cast(T, self.container.data)

Now you can do stuff like this:

def only_accept_X_container(x: TypedContainer[X]) -> None:
    ...

def only_accept_X(x: X) -> None:
    ...
    

my_container = TypedContainer(X())

only_accept_X_container(my_container)
only_accept_X(my_container.data)

Answer 2

Note: I am not allowed to modify the original Container code.

You're always allowed to get creative with what "modification" means.

For example, one way to solve this is using a type-check-only generic subclass which overrides data; the following snippet should work across type-checkers and work with reflection-heavy code using lots of isinstance or issubclass checks.

Demo (mypy Playground, pyright Playground, pyre Playground)

from <dependency> import Container, X

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing import TypeVar, Generic
    from dataclasses import dataclass

    T = TypeVar("T")

    @dataclass
    class MyContainer(Container, Generic[T]):
        data: T  # type: ignore[assignment]
else:
    MyContainer = Container


reveal_type(MyContainer(X()).data)  # Revealed type is "X"