Conditional typing for mapping input to correct method in Protocol-Based Implementation

[BrandonNav]

Hello everyone,

I'm currently working on a Python project where I have various implementations (Impl1 and Impl2) of a base class that uses Protocols for type hinting. The base class has a 'run' method which should accept an input and process it accordingly based on the specific implementation (Impl1 or Impl2).

Here's my challenge: I've stored the classes in a registry as Enums, and I'm trying to use the Enum to create an instance and run the associated method. The line impl.run(input) is causing an error:

Argument of type "Impl1Input | Impl2Input" cannot be assigned to parameter "input" of type "Impl1Input" in function "run"
  Type "Impl1Input | Impl2Input" cannot be assigned to type "Impl1Input"
    "Impl2Input" is incompatible with "Impl1Input"

What I am trying to do is to have conditional typing such that the 'run' method can correctly map my input to the appropriate 'run' method based on the instance's class. However, I'm not sure how to achieve this, or if I'm going about this the correct way.

Is there a way in Python to do conditional typing to solve this issue? Or is there a better design pattern or strategy to achieve the same effect? Any help or pointers would be greatly appreciated!

Thanks in advance.

Below is the code to reproduce the issue:

from enum import Enum
from typing import Generic, NewType, Protocol, TypeVar


# --------Type Defs-------------

InputT = NewType("InputT", str)
OutputT = NewType("OutputT", str)

_Input = TypeVar("_Input", bound=InputT, contravariant=True)
_Output = TypeVar("_Output", bound=OutputT, covariant=True)


class Template(Protocol[_Input, _Output]):
    @staticmethod
    def render(input: _Input) -> _Output:
        ...


class Base(Generic[_Input, _Output]):
    template: Template[_Input, _Output]
    input_type: type[_Input]
    
    @classmethod
    def create_input(cls: type[Self], input: str):
        return cls.input_type(input)

    def run(self, input: _Input) -> _Output:
        ...
        
        
# ----------Impl1---------------

class Impl1Input(InputT):
    pass

class Impl1Output(OutputT):
    pass


class Impl1Template(Template[Impl1Input, Impl1Output]):
    @staticmethod
    def render(input: Impl1Input) -> Impl1Output:
        return Impl1Output(input)


class Impl1(Base[Impl1Input, Impl1Output]):
    template = Impl1Template
    input_type = Impl1Input

    def run(self, input: Impl1Input) -> Impl1Output:
        return self.template.render(input)
    
# ----------Impl2---------------

class Impl2Input(InputT):
    pass

class Impl2Output(OutputT):
    pass

class Impl2Template(Template[Impl2Input, Impl2Output]):
    @staticmethod
    def render(input: Impl2Input) -> Impl2Output:
        return Impl2Output(input)

class Impl2(Base[Impl2Input, Impl2Output]):
    template = Impl2Template
    input_type = Impl2Input

    def run(self, input: Impl2Input) -> Impl2Output:
        return self.template.render(input)

# ------------------------------

class Registry(Enum):
    Impl1 = Impl1
    Impl2 = Impl2


def do(registry_entry: Registry):
    impl = registry_entry.value()  # type is Impl1 | Impl2
    input = impl.input_type("input") # type is Impl1Input | Impl2Input

    impl.run(input)
    # Argument of type "Impl1Input | Impl2Input" cannot be assigned to parameter "input" of type "Impl1Input" in function "run"
    #  Type "Impl1Input | Impl2Input" cannot be assigned to type "Impl1Input"
    #    "Impl2Input" is incompatible with "Impl1Input"

[BrandonNav]

Progress update:

I can manually use type narrowing to specialise the types but it requires duplicating every derived Impl

def do(registry_entry: Registry):
    match (impl:= registry_entry.value()):
        case Impl1():
            input = impl.input_type("input")
            impl.run(input)
            impl.template.render(input)
        case Impl2():
            input = impl.input_type("input")
            impl.run(input)
            impl.template.render(input)

I think TypeGuards might help here (so long as this approach is still valid)

[BrandonNav]

Searching for alternatives I found: #1273
Which gives me the impression that what I want is not currently possible.