mypy allow redefinition at a specific line

Question

I'm trying to lint my python code with mypy. A utility that I depend on uses a callback to which arguments can take multiple types. Type is typically defined by an accompanying topic string.

I'm having an issue in a situation where a callback takes multiple topics, each with a unique type. In this case, I'm unsure how to communicate to mypy that type is more restricted than the union of the two.

def callback(self, topic: str, msg: Union[A, B]):
    if topic == "a":
        msg: A
        # Do something specific to topic "a"
    if topic == "b":
        msg: B
        # Do something specific to topic "b"

I've tried writing msg: A # type: ignore instead, but then mypy just ignores this line. I also cannot use isinstance due to the that utility's structure.

Potential solutions

There are two things that I believe should work, but I'm uncertain if mypy has such capabilities:

1. Allow redefinition on specific line. When running with --allow-redefinition the error is suppressed. Unfortuantely, running it for the whole package is too loose in my view, so I would prefer to avoid it.
2. Bind type of variable to value of another variable. I'm unsure if it exists, but it can be similar to TypedDict where keys can bind to different types of values, instead on two separate variables.

Answer 1

Overload signatures will make your function externally correct, even if you have to cast on the inside.

class Example:

    @overload
    def callback(self, topic: Literal["a"], msg: A) -> None: ...
    @overload
    def callback(self, topic: Literal["b"], msg: B) -> None: ...
    def callback(self, topic: str, msg: Union[A, B]):
        if topic == "a":
            a_msg = cast(A, msg)
            reveal_locals()
        if topic == "b":
            b_msg = cast(B, msg)
            reveal_locals()

(Note: Literal was added to typing in Python 3.8. If you're using an earlier version, you can get it from typing_extensions)

Inside the function, we still have to cast explicitly, but externally, the caller will always be required to call the function with either (1) the literal "a" and an A, or (2) the literal "b" and a B.

Answer 2

It takes some annoying boilerplate and a bit of a hack, but you can do this with a PEP 647 TypeGuard.

TypeGuards exist to tell MyPy that some custom function is able to do type narrowing in a similar way to things like isinstance. For example, you might have a function which takes a list and returns True if everything in that list is an integer:

def are_all_ints(probe: List[Any]) -> TypeGuard[List[int]]:
    return all(isinstance(int, p) for p in probe)

xs: List[Any] = ...

if are_all_ints(xs):
    # MyPy trusts the check was done correctly and therefore
    reveal_type(xs) # this gives List[int] within this section
    

As you can see, TypeGuards operate on the variable that is passed to them. You could get narrowing of your topic string if you pass it to a suitable typeguard function, but not of msg. The workaround for that is to package them together into a temporary variable.

def callback(self, topic: str, msg: Union[A, B]):
    
    def is_callback_from_a(p: Tuple[str, Union[A, B]]) -> TypeGuard[Tuple[Literal["a"], A]]:
        return p[0] == "a"
    
    def is_callback_from_b(p: Tuple[str, Union[A, B]]) -> TypeGuard[Tuple[Literal["b"], B]]:
        return p[0] == "b"
    
    # Package the two parameters into one variable so that 
    # the guard clears them simulataneously
    topic_msg_pair = (topic, msg)
    
    if is_callback_from_a(topic_msg_pair):
        ...
        # topic_msg_pair[1] is known to be an A here
    elif is_callback_from_b(topic_msg_pair):
        ...
        # topic_msg_pair[1] is known to be a B here