A way to subclass NamedTuple for purposes of typechecking

user_5208398 0
Viewed 9124

I have several namedtuples that share some fields. I have a function that accepts these tuples and is guaranteed to only interact with the shared fields. I want to typecheck such code in mypy.

An example of the code would be:

from typing import NamedTuple

class Base(NamedTuple):
    x: int
    y: int


class BaseExtended(NamedTuple):
    x: int
    y: int
    z: str

def DoSomething(tuple: Base):
    return tuple.x + tuple.y

base = Base(3, 4)
base_extended = BaseExtended(5, 6, 'foo')

DoSomething(base)
DoSomething(base_extended)

When I run mypy on this code, I get a predictable error:

mypy_example.py:20: error: Argument 1 to "DoSomething" has incompatible type "BaseExtended"; expected "Base"

Is there no way to structure my code and keep mypy typechecking? I cannot inherit BaseExtended from Base, since there's a bug in the NamedTuple inheritance implementation.

I don't want to use an ugly Union[Base, BaseExtended] either, since this breaks when I try to typecheck a List, since List[Union[Base, BaseExtended]] is not equal to List[BaseExtended] due to some mypy magic about variant/covariant types.

Should I just abandon the idea?

python-typing mypy type-hinting namedtuple python
2 Answers

The way named tuples are constructed make inheritance from typing.NamedTuple classes as yet not possible. You'd have to write your own metaclass to extend the typing.NamedTupleMeta class to make subclassing work, and even then the class generated by collections.namedtuple() is just not built to extend.

Instead, you want to use the new dataclasses module to define your classes and achieve inheritance:

from dataclasses import dataclass

@dataclass(frozen=True)
class Base:
    x: int
    y: int

@dataclass(frozen=True)
class BaseExtended(Base):
    z: str

The module is new in Python 3.7 but you can pip install dataclasses the backport on Python 3.6.

The above defines two immutable classes with x and y attributes, with the BaseExtended class adding one more attribute. BaseExtended is a full subclass of Base, so for typing purposes fits the requirements for the DoSomething() function.

The classes are not full named tuples, as they don't have a length or support indexing, but that's trivially added by creating a baseclass that inherits from collections.abc.Sequence, adding two methods to access fields by index. If you add order=True to the @dataclass() decorator then your instances become fully orderable the same way (named) tuples are:

from collections.abc import Sequence
from dataclasses import dataclass, fields

class DataclassSequence(Sequence):
    # make a dataclass tuple-like by accessing fields by index
    def __getitem__(self, i):
        return getattr(self, fields(self)[i].name)
    def __len__(self):
        return len(fields(self))

@dataclass(frozen=True, order=True)
class Base(DataclassSequence):
    x: int
    y: int

MyPy will soon support dataclasses explicitly; in version 0.600 you'll get errors still as it doesn't recognise the dataclasses module import or that a __new__ method is generated.

In Python 3.6 and earlier, you can also install the attrs project to achieve the same effects; the above sequence base class looks like this using attrs:

from collections.abc import Sequence
import attr

class AttrsSequence(Sequence):
    # make a dataclass tuple-like by accessing fields by index
    def __getitem__(self, i):
        return getattr(self, attr.fields(type(self))[i].name)
    def __len__(self):
        return len(attr.fields(type(self)))

@attr.s(frozen=True, auto_attribs=True)
class Base(AttrsSequence):
    x: int
    y: int

dataclasses is directly based on attrs, with attrs providing more functionality; mypy fully supports classes generated with attrs.

user_1002970
Related