How to write correct typings for join of two python lists with elements having same parent

Question

I have pydantic models, which share same base class

class Base(BaseModel):
    shared_arg: str

class A(Base):
    some_arg1: int

class B(Base):
    some_arg2: str

I have a function, which accepts list[Base] as an argument. I want to combine list[A] and list[B] together in order to pass it inside this function and operate on their shared_arg values.

Assume, function is

def work_on_bases(values: list[Base]) -> list[str]:
    return [v.shared_arg.upper() for v in values]

and lists look like this:

list_a = [A(shared_arg='hello', some_arg1=0), 
          A(shared_arg='world', some_arg1=5)]

list_b = [B(shared_arg='green', some_arg2=7), 
          A(shared_arg='horse', some_arg2=9)] 

I've tried to make something like this

combined: list[Base] = list_a + list_b

and this

list_a_base: list[Base] = list_a
list_b_base: list[Base] = list_b
combined: list[Base] = list_a + list_b

But pylance (pyright) still issues

Operator "+" not supported for types "list[A]" and "list[B]"

Is there a way to write correct typings here or # type: ignore is the only way to go?

Answer 1

This work perfectly with mypy. If you don't specify the type on list_b the checker will assume a list of Object but if you specify the type all will work as expected.

Note that there are a few issue in the initialization of listb in your example (B.some_arg2 expects a str and A doesn't have some_arg2 attribute)

class Base(BaseModel):
    shared_arg: str

class A(Base):
    some_arg1: int

class B(Base):
    some_arg2: str

list_a: list[Base] = [A(shared_arg='hello', some_arg1=0), 
          A(shared_arg='world', some_arg1=5)]

list_b: list[Base] = [B(shared_arg='green', some_arg2='7'),
          A(shared_arg='horse', some_arg1=9)]

combined: list[Base] = list_a + list_b