mypy "Incompatible return value type" with dict comprehension and isinstance-CodePudding

I have an issue with mypy giving the following error:

Incompatible return value type (got "Dict[str, Pin]", expected "Union[Dict[str, Input], Dict[str, Output]]") [return-value]mypy

in the following example code:

from typing import Dict, List, Union, Type

class Pin():
    def __init__(self):
        pass

class Input(Pin):
    def __init__(self):
        pass

class Output(Pin):
    def __init__(self):
        pass

INPUT_STRINGS: List[str] = ["in", "input", "i"]
OUTPUT_STRINGS: List[str] = ["out", "output", "o"]

def filter_pins(pins: Dict[str, Pin], which) -> Union[Dict[str, Input], Dict[str, Output]]:
        direction: Union[Type[Output], Type[Input]]
        if which.lower() in OUTPUT_STRINGS:
            direction = Output
        elif which.lower() in INPUT_STRINGS:
            direction = Input
        filtered = {name: pin for name, pin in pins.items() if isinstance(pin, direction)}
        return filtered

The goal of the filter_pins function is to take a dict of Pin instances and only return either Input or Output instances in a dict. So in the dict comprehension, I'm using isinstance to check which subclass of Pin, each pin is. It seems like mypy doesn't recognize that pin will always be either Input or Output. I think this is a case of 'type narrowing' and should in principle work. Any ideas whether I can and should fix this or it's a mypy issue?

Thanks!

CodePudding user response：

Since Input and Output are both subclasses of Pin, would simply using Pin in the output type work for you?

def filter_pins(pins: Dict[str, Pin], which: str) -> Dict[str, Pin]:

If not, here's an alternative: cast the output to the Union type (note that cast does nothing at runtime, it is only for static type checking)

return cast(Union[Dict[str, Input], Dict[str, Output]], filtered)

And if you can accept changing the signature of filter_pins:

from typing import Dict, List, TypeVar, Type

T = TypeVar('T', Input, Output)

def filter_pins(pins: Dict[str, Pin], filter_type: Type[T]) -> Dict[str, T]:
    filtered = {name: pin for name, pin in pins.items() if isinstance(pin, filter_type)}
    return filtered

class PinCollection:
    pins: Dict[str, Pin] = {}

    @property
    def inputs(self) -> Dict[str, Input]:
        """Returns only the input pins."""
        return filter_pins(self.pins, Input)