CODEWITHSUNDEEP
pythonmypypython-typing
Paulo Costa
Paulo Costa

Reputation: 1635

Type hints when only a subset of attribute types are accepted

Say I have a class with multiple possible types for an attribute and a function that only works on some of these types.

How would I type-hint that function so that it specifies it only accepts some of the types accepted by the original class?

import dataclasses

@dataclasses.dataclass
class Foo:
    a: int | str

def plus_1_a(foo: Foo) -> None:
    """This will only work if foo.a is an int, not a string."""
    foo.a += 1

Upvotes: 1

Views: 150

Answers (1)

Paweł Rubin
Paweł Rubin

Reputation: 3439

Consider making Foo generic over the type of a.

import dataclasses
from typing import Generic, TypeVar

T = TypeVar("T", int, str)


@dataclasses.dataclass
class Foo(Generic[T]):
    a: T


def plus_1_a(foo: Foo[int]) -> None:
    """This will only work if foo.a is an int, not a string."""
    foo.a += 1

Upvotes: 0

Related Questions