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.