如何在 python class 中仅使一个方法通用

Question

I have the following base class to represent fields:

T = TypeVar("T")


@dataclass
class BaseClass(Generic[T]):
    def validate(self, value: T):
        raise NotImplementedError

I also have an enum to represent the available implementations of this class:

class Types(Enum):
    A = auto()
    B = auto()

    @staticmethod
    def from_instance(instance: BaseClass) -> "Types":
        if isinstance(instance, ClassA):
            return Types.A
        if isinstance(instance, ClassB):
            return Types.B
        raise ValueError("Not supported")

Now, from these class, I have several implementations:

@dataclass
class ClassA(BaseClass[str]):
    def validate(self, value: str):
        pass


@dataclass
class ClassB(BaseClass[int]):
    def validate(self, value: int):
        pass

After this setup, I have another class to store a list of BaseClass :

@dataclass
class Container:
    instances: List[BaseClass]

    def get_by_type(self, type: Types) -> List[BaseClass]:
        return [instance for instance in self.instances if type == Types.from_instance(instance)]

At the end I have the following code and the following error:

def function(fields_from_class_a: List[ClassA]):
    print(fields_from_class_a)


container = Container(instances=[ClassA(), ClassB()])
fields = container.get_by_type(Types.A)


# throws error:
#  Argument 1 to "function" has incompatible type "List[BaseClass[Any]]"; expected "List[ClassA]"
function(fields)

So my question is, can I modify the code in such a way that the method get_by_type is correctly typed?

Answer 1

ClassA is only mapped at runtime to Types.A in Types.from_instance() .

You would have to use BaseClass[Literal[Types.A]] instead of ClassA :

# def function(fields_from_class_a: List[ClassA]):                     # Change this
def function(fields_from_class_a: List[BaseClass[Literal[Types.A]]]):  # to this

And type Container.get_by_type() as such:

TypesT = TypeVar("TypesT", Literal[Types.A], Literal[Types.B])   # Add this

# def get_by_type(self, type: Types) -> List[BaseClass]:         # Change this
def get_by_type(self, type: TypesT) -> List[BaseClass[TypesT]]:  # to this