pythonpython-typingmypy

How Can I Type a Method That Accepts Any Subclass of a Base Class?


I'm having some trouble typing subclasses of a base class.

I have a base class called Table, which defines shared behavior for all table subclasses. Then, I have a Database class that manages collections of these tables.

The code below works fine at runtime. However, when I try to type the methods in the Database class using Table as the type annotation, I run into issues when passing a subclass (e.g., Ducks). The types break, and MyPy complains.

I think I need to use a bounded generic or a protocol, but I couldn't figure out how to make it work properly.

Here’s the relevant code:

#!/usr/bin/env python3
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Self, TypeVar

# _Table = TypeVar("_Table", bound="Table") <- Tried this but did not work

@dataclass
class Table:
    @classmethod
    def get_name(cls) -> str:
        return cls.__name__.lower()


@dataclass
class Database:
    records: Mapping[str, Sequence[Table]]

    def filter(
        self, table: type[Table], _filter: Callable[[Table], bool]
    ) -> Iterable[Table]:
        yield from (i for i in self.records[table.get_name()] if _filter(i))
        # USING cast(T) here kind of works but i think is kind a cheat.

@dataclass
class Ducks(Table):
    name: str
    age: int


if __name__ == "__main__":

    records = {
        "ducks": [
            Ducks(name="Patinhas", age=104),
            Ducks(name="Donald", age=30),
            Ducks(name="Huguinho", age=12),
            Ducks(name="Zezinho", age=12),
            Ducks(name="Luizinho", age=12),
        ]
    }

    db = Database(records)
    f: Callable[[Ducks], bool] = lambda t: t.age < 100 # <- BONUS QUESTION: is there a way to not use this long typing lambda and just use the untyped notation ?
    print(*db.filter(Ducks, _filter=f), sep="\n") # <- PROBLEM HERE
$ python3 ex.py
Ducks(name='Donald', age=30)
Ducks(name='Huguinho', age=12)
Ducks(name='Zezinho', age=12)
Ducks(name='Luizinho', age=12)

Mypy lint: ex.py 46 37 error arg-type Argument "_filter" to "filter" of "Database" has incompatible type "Callable[[Ducks], bool]"; expected "Callable[[Table], bool]" (lsp))

This makes sense because Callable[[Table], bool] doesn't guarantee compatibility with Ducks, even though it's a subclass. I want the method to work with any subclass of Table and have the type checker understand that. I need a way to tell the type checker that the Database functions should work with any subclass of Table, and that specific operations may require a more specific type.


Solution

  • As stated in the existing answer by @user2357112, you can't make this fully type safe: mapping from a name (string) to some sequence of instances can't be encoded in python's static type system, you can't map a name to a type.

    However, your filter method is probably the best possible approach. You didn't try to throw in more stringly-typed magic, and that's great: you only need one cast in the implementation to make the interface (how others interact with your class) statically correct.

    Here it is:

    @dataclass
    class Database:
        records: Mapping[str, Sequence[Table]]
    
        def filter(
            self, table: type[_Table], _filter: Callable[[_Table], bool]
        ) -> Iterable[_Table]:
            yield from (
                i 
                for i in cast("Sequence[_Table]", self.records[table.get_name()]) 
                if _filter(i)
            )
    

    Note that I cast the whole sequence. It's the only invariant you have to verify yourself (namely that a key of table.get_name() always corresponds to instances of that table - please fix the naming, tables are not rows nor records, it's mind-blowing), the rest is on type checker shoulders.

    And yes, this supports inline lambdas as you wanted:

    db = Database(records)
    print(
        *db.filter(Ducks, _filter=lambda t: t.age < 100),
        sep="\n"
    )
    

    And here's a playground with your complete example modified for this answer.