Function for restricting the type of attributes in an object (Python, mypy)

user_13184550 0
Viewed 545

Example setup:

from typing import Optional


class A(object):
    def __init__(self):
        self.a: Optional[int] = None
        
    def check_a(self) -> bool:
        return self.a is not None
        
a = A()
if a.check_a():
    print(a.a + 1)  # error: Unsupported operand types for + ("None" and "int")

The check_a method checks what type of variable a is, but mypy does not see this and writes an error. TypeGuard will not help, because it can create a function to check the type, and not a function to check the type of an object variable

Is it possible to somehow make mypy notice this in order to use the function to check the type of the variable self.a without explicitly referring to it in the check? (Use if a.a_check instead of if a.a is not None)?

python-typing mypy type-hinting python
2 Answers

There are two ways in which this can be done with the new TypeGuard feature, which can be imported from typing in Python 3.10, and is available from the PyPI typing_extensions package in earlier Python versions. Note that typing_extensions is already a dependency of Mypy, so if you're using Mypy, you likely already have it.

The first option is to change your check_a method to a staticmethod that takes in a variable that might be int or None as an input, and verifies whether or not it is an int. (Apologies, I have changed the names of some of your variables, as I found it quite confusing to have a class A that also had an a attribute.)

from typing import TypeGuard, Optional

class Foo1:
    def __init__(self, bar: Optional[int] = None) -> None:
        self.bar = bar
    
    @staticmethod
    def check_bar(bar: Optional[int]) -> TypeGuard[int]:
        return bar is not None
        

f1 = Foo1()
if f1.check_bar(f1.bar):
    print(f1.bar + 1)

The second option is to use structural subtyping to assert that an instance of class Foo (or class A in your original question) has certain properties at a certain point in time. This requires altering the test method so that it becomes a classmethod, and is a little more complex to set up, but leads to a nicer check once you have it set up.

from typing import TypeGuard, Optional, Protocol, TypeVar

class HasIntBar(Protocol):
    bar: int


F = TypeVar('F', bound='Foo2')


class Foo2:
    def __init__(self, bar: Optional[int] = None) -> None:
        self.bar = bar
    
    @classmethod
    def check_bar(cls: type[F], instance: F) -> TypeGuard[HasIntBar]:
        return instance.bar is not None
        

f2 = Foo2()
if Foo2.check_bar(f2): # could also write this as `if f2.check_bar(f2)`
    print(f2.bar + 1)

You can try both of these options out on Mypy playground here.

user_139900160

The problem is that, as far as the type checker is concerned, your bool says nothing about the type of your .a. For example, you could have written

class A:
    ...

    def check_a(self) -> bool:
        return True

However, as you probably know, mypy can rule out None within an if x is not None:. Your problem is that you've split your if x is not None: between the function definition and call site, so mypy can't use it to infer the value's not None.

What you can do to fix this is pass the operation you want to do to the A, in which case you'd have something like what's called foreach in other languages. In the context of Optional (and A is a wrapper around a mutable Optional) this can apply a function to values if they exist.

class A:
    ...

    def foreach(self, f: Callable[[int], None]) -> None:
        if self.a is not None:
            f(self.a)

a = A()
a.foreach(lambda x: print(x + 1))

Note I don't have a check_a here.

user_59869070
Related