Search code examples
pythontypessubclasstyping

Subclass Python list with all elements of the same type


I'm looking for a way to make a subclass of list where all elements are of the same type. I don't need a fully generic implementation, as described in this answer.

What I'd like is something roughly equivalent to the following code (using typing):

class NumberParty(list[int]):
    def __str__(self) -> str:
        "🎉".join([str(x) for x in self])

    def total(self) -> int:
        return sum(self)

The reason I want this class instead of using List[int] annotations is that I use the __str__ method described above in multiple places and want to eliminate the copied lines of code.


As a bonus question, would it be wise to overload __init__ in something similar to the following code?

__init__(self, *numbers: int):
    self = numbers
__init__(self, nums: Iterable[int]):
    self = list(nums)

The first way would be for directly creating a new NumberParty instance while the second would create one from an existing list or set of integers.


Solution

  • Unless I'm misunderstanding your question, just doing:

    from typing import List
    
    class NumberParty(List[int]):
        def __str__(self) -> str:
            return "🎉".join([str(x) for x in self])
    
        def total(self) -> int:
            return sum(self)
    

    ...should just work.

    Note that inheriting from List[int] is essentially the same as inheriting from list -- the only difference is that type checkers will understand what your custom class is supposed to contain and verify you're using methods like append(...) in a type safe way.

    However, there is no difference at runtime, and you can still freely ignore your type checker and append strs or whatever, just like how you Python lets you append strs to regular lists that are annotated to be of type List[int].

    And since this class is a subtype of List[int], code snippets like the below should type check:

    def expect_list(x: List[int]) -> None: pass
    
    def expect_number_party(x: NumberParty) -> None: pass
    
    n = NumberParty()
    
    # Both type checks
    expect_list(n)
    expect_number_party(n)
    

    Regarding your question about overriding __init__ -- you'll need to call the super constructor if you decide to take that route to ensure the class is set up properly.