Return type in a class method based on a class variable in Python

Bram Vanroy 0
Viewed 61

I am working with a codebase where you have couples of classes, always one dataclass and another execution class. The dataclass serves as a data collector (as the name suggests).

To "connect" the dataclass to the other class, I set a class variable in the other class to make clear what the relevant dataclass is. This works fine - I can use this class variable to instantiate the data class as I please. However, it is not clear to me how I can use this to specify for a given method that it will return an instance of the linked data class.

Take this example (executable):

from abc import ABC
from dataclasses import dataclass
from typing import ClassVar

@dataclass
class Name(ABC):
    name: str


class RelatedName(ABC):
    _INDIVIDAL: ClassVar[Name]

    def return_name(self, **properties) -> Name:
        # There is a typing issue here too but you can ignore that for now
        return self._INDIVIDAL(**properties)  


@dataclass
class BiggerName(Name):
    other_name: str

class RelatedBiggerName(RelatedName):
    _INDIVIDAL: ClassVar[Name] = BiggerName


if __name__ == "__main__":
    biggie = RelatedBiggerName()
    biggiename = biggie.return_name(name="Alfred", other_name="Biggie").other_name
    print(biggiename)

The script works fine, but there is a typing problem. In the last but one line, you'll see the issue that the attribute other_name is undefined for the Name class. This is to be expected, but I am not sure how I can change the output type of return_name so that it will use the class that is defined in _INDIVIDUAL.

I tried def return_name(self, **properties) -> _INDIVIDAL but that naturally leads to name '_INDIVIDAL' is not defined.

Perhaps it is not possible what I am after. Is it at all possible to have typing within a class that depends on class variables? I'm interested in Python 3.8 and higher.

python-typing class typing python
2 Answers

Can you use generics?

from abc import ABC
from dataclasses import dataclass
from typing import ClassVar, TypeVar, Generic, Type

T = TypeVar("T", bound="Name")

@dataclass
class Name(ABC):
    name: str


class RelatedName(ABC, Generic[T]):
    # This would resolve what juanpa.arrivillaga pointed out, but mypy says:
    # ClassVar cannot contain type variables, so I guess your use-case is unsupported
    # _INDIVIDAL: ClassVar[Type[T]]
    # One option:
    # _INDIVIDAL: ClassVar
    # Second option to demonstrate Type[T]
    _INDIVIDAL: Type[T]

    def return_name(self, **properties) -> T:
        return self._INDIVIDAL(**properties)  


@dataclass
class BiggerName(Name):
    other_name: str

class RelatedBiggerName(RelatedName[BiggerName]):
    # see above
    _INDIVIDAL: Type[BiggerName] = BiggerName


if __name__ == "__main__":
    biggie = RelatedBiggerName()
    biggiename = biggie.return_name(name="Alfred", other_name="Biggie").other_name
    print(biggiename)

mypy reports no errors on this and I think conceptually this is what you want. I tested on python 3.10.

cherrywoods0

I agree with @cherrywoods that a custom generic base class seems like the way to go here.

I would like to add my own variation that should do what you want:

from abc import ABC
from dataclasses import dataclass
from typing import Any, Generic, Optional, Type, TypeVar, get_args, get_origin


T = TypeVar("T", bound="Name")


@dataclass
class Name(ABC):
    name: str


class RelatedName(ABC, Generic[T]):
    _INDIVIDUAL: Optional[Type[T]] = None

    @classmethod
    def __init_subclass__(cls, **kwargs: Any) -> None:
        """Identifies and saves the type argument"""
        super().__init_subclass__(**kwargs)
        for base in cls.__orig_bases__:  # type: ignore[attr-defined]
            origin = get_origin(base)
            if origin is None or not issubclass(origin, RelatedName):
                continue
            type_arg = get_args(base)[0]
            # Do not set the attribute for GENERIC subclasses!
            if not isinstance(type_arg, TypeVar):
                cls._INDIVIDUAL = type_arg
                return

    @classmethod
    def get_individual(cls) -> Type[T]:
        """Getter ensuring that we are not dealing with a generic subclass"""
        if cls._INDIVIDUAL is None:
            raise AttributeError(
                f"{cls.__name__} is generic; type argument unspecified"
            )
        return cls._INDIVIDUAL

    def __setattr__(self, name: str, value: Any) -> None:
        """Prevent instances from overwriting `_INDIVIDUAL`"""
        if name == "_INDIVIDUAL":
            raise AttributeError("Instances cannot modify `_INDIVIDUAL`")
        super().__setattr__(name, value)

    def return_name(self, **properties: Any) -> T:
        return self.get_individual()(**properties)


@dataclass
class BiggerName(Name):
    other_name: str


class RelatedBiggerName(RelatedName[BiggerName]):
    pass


if __name__ == "__main__":
    biggie = RelatedBiggerName()
    biggiename = biggie.return_name(name="Alfred", other_name="Biggie").other_name
    print(biggiename)

Works without problems or complaints from mypy --strict.

Differences

  • The _INDIVIDUAL attribute is no longer marked as a ClassVar because that (for no good reason) disallows type variables.
  • To protect it from being changed by instances, we use a simple customization of the __setattr__ method.
  • You no longer need to explicitly set _INDIVIDUAL on any specific subclass of RelatedName. This is taken care of automatically during subclassing by __init_subclass__. (If you are interested in details, I explain them in this post.)
  • Direct access to the _INDIVIDUAL attribute is discouraged. Instead there is the get_individual getter. If the additional parentheses annoy you, I suppose you can play around with discriptors to construct a property-like situation for _INDIVIDUAL. (Note: You can still just use cls._INDIVIDUAL or self._INDIVIDUAL, it's just that there will be the possible None-type issue.)
  • The base class is obviously a bit more complicated this way, but on the other hand the creation of specific subclasses is much nicer in my opinion.

Hope this helps.

Daniil Fajnberg0
Related