Is it possible to add/overwrite a type hint in case of the following example? The example is just to get an idea of what I mean, by no means is this something that I would use in this way.

from dataclasses import dataclass


def wrapper(f):
    def deco(instance):
        if not instance.user:
            instance.user = data(name="test")
        return f(instance)

    return deco


@dataclass
class data:
    name: str


class test_class:
    def __init__(self):
        self.user: None | data = None

    @wrapper
    def test(self):
        print(self.user.name)


x = test_class()
x.test()

The issue is that type hinting does not understand that the decorated method's user attribute is not None, thus showing a linting error that name is not a known member of none.

Of course this code could be altered so that instead of using a decorator it would just do something like this:

def test(self):
   if not self.user:
      ...
   print(self.user.name)

But that is not the point. I just want to know if it is possible to let the type hinter know that the attribute is not None. I could also just suppress the warning but that is not what I am looking for.

CodePudding user response：

I would use the good ol' assert and be done with it:

    ...
    @wrapper
    def test(self):
        assert isinstance(self.user, data)
        print(self.user.name)

I realize this is a crude way as opposed to some annotation magic you might have expected for the decorator, but in my opinion this is the most practical approach.

There are countless other situations that can be constructed, where the type of some instance attribute may be altered externally. In those cases the use of such a simple assertion is not only for the benefit of the static type checker, but can also save you from shooting yourself in the foot, if you decide to alter that external behavior.

Alternative - Getter

Another possibility is to make the user attribute private and add a function (or property) to get it, which ensures that it is not None. Here is a working example:

from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
from typing import TypeVar


T = TypeVar("T")


@dataclass
class Data:
    name: str


def wrapper(f: Callable[[TestClass], T]) -> Callable[[TestClass], T]:
    def deco(self: TestClass) -> T:
        try:
            _ = self.user
        except RuntimeError:
            self.user = Data(name="test")
        return f(self)
    return deco


class TestClass:
    def __init__(self) -> None:
        self._user: None | Data = None

    @property
    def user(self) -> Data:
        if self._user is None:
            raise RuntimeError
        return self._user

    @user.setter
    def user(self, data: Data) -> None:
        self._user = data

    @wrapper
    def test(self) -> None:
        print(self.user.name)


if __name__ == '__main__':
    x = TestClass()
    x.test()

Depending on the use case, this might actually be preferred because otherwise, user being a public attribute, all outside code that wants to use TestClass will face the same problem of never being sure if user is None or not, thus being forced to do the same checks again and again.

CodePudding user response：

Sadly there isn't really a satisfactory answer to your question. The problem is that no type-checkers execute any code - that means that any dynamic type generation doesn't work. For that reason, if you want to tell the type-checker that the self.user is not None you need to create a class where user is not Optional.

I don't think it's a good idea but here is how you could achieve what you want to achieve. Note though that that way you need to keep the two classes in sync and some type-checkers have trouble with decorators...

from typing import ParamSpec, TypeVar, Concatenate, Callable, cast
from dataclasses import dataclass


T = TypeVar("T") # generic return value
P = ParamSpec("P") # all other params after self

def wrapper( # this wrapper works on any functions in 'test_class'
    f: Callable[Concatenate["test_class", P], T]
) -> Callable[Concatenate["__non_optional_user_test_class", P], T]:
    def deco(instance: "test_class", *args: P.args, **kwargs: P.kwargs):
        if not instance.user:
            instance.user = data(name="test")
        return f(cast("__non_optional_user_test_class", instance), *args, **kwargs)

    return deco

@dataclass
class data:
    name: str


class __non_optional_user_test_class:
    user: data


class test_class:
    def __init__(self):
        self.user: None | data = None

    @wrapper
    def test(self):
        print(self.user.name)


x = test_class()
x.test()

You sadly cannot generate the __non_optional_user_test_class dynamically in such a way that type-checkers understand them... And you would need to write a new wrapper for all classes where you want to apply this @wrapper.