'How can I use TypeGuards to narrow types for multiple object fields in Python?

Say I've got a Foo object which has multiple fields which can either be None or some other type. Whether or not the field is None relates to whether other fields are None or not, so by checking one field, I can instantly know if the other fields are None.

I understand this is terrible class design, but I'm not able to modify it as it is other people's code that I am annotating.

It looks like TypeGuards introduced in PEP 647 are my best bet for adding this functionality, but I can't figure out how I could specifically apply them to this situation. I've attached an attempt I made with subclasses, but it fails both in MyPy and in Pyright.

from typing import Optional
from typing_extensions import TypeGuard

class Foo:
    """A class containing the parameters `value`, `values` and `other`. If 
    `values` is `None` then both the others are not, and if the others are not
    then `values` is.
    """
    def __init__(self, value: 'int | list[int]', other: Optional[int]) -> None:
        is_singular = isinstance(value, int)
        self.value = value if is_singular else None
        self.values = None if is_singular else value
        if is_singular:
            assert other is not None
        else:
            assert other is None
        self.other = other

class SingularFoo(Foo):
    """A subclass for an instance of `Foo` where `values` is `None`
    """
    def __init__(self, value: int, other: int) -> None:
        super().__init__(value, other)

class MultiFoo(Foo):
    """A subclass for an instance of `Foo` where `values` is not `None`
    """
    def __init__(self, value: list[int]) -> None:
        super().__init__(value, None)

def isFooSingular(f: Foo) -> TypeGuard[SingularFoo]:
    """A type guard that returns whether `f` is singular (meaning that `values`
    is `None` and `value` and `other` are not)
    """
    return f.value is not None

# Create a singular `Foo`
my_foo = Foo(1, 2)
# Type guard
assert isFooSingular(my_foo)
# After the type guard, both should be considered as `int`
#
# Errors from MyPy:
# * Unsupported operand types for + ("int" and "None")
# * Unsupported operand types for + ("List[int]" and "int")
# * Unsupported operand types for + ("List[int]" and "None")
# * Unsupported operand types for + ("None" and "int")
# * Unsupported left operand type for + ("None")
print(my_foo.value + my_foo.other)

How can I modify this code such that the type guard performs the desired type narrowing operation.



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source