2 years ago

#66264

test-img

Adam Kern

Typing for decorator that wraps attrs.frozen and adds a new field

I am trying to set up a class decorator in Python that acts like attr.frozen but adds an additional field before creation (as well as a few other things). While the code works fine, I'm having trouble getting mypy to realize that the new class has the new field. I've tried to do this through a combination of a custom mypy plugin (exactly as described in attr's documentation) and a Protocol that defines that the new class has the given field. In summary, the code breaks down as follows (all in a single file, although I've broken it up here).

It should be noted I'm running Python 3.7, so I'm using typing_extensions where needed, but I believe this problem persists regardless of version.

First define the Protocol that should inform mypy that the new class has the new field (called added here):

from typing_extensions import Protocol

class Proto(Protocol):
    def __init__(self, added: float, *args, **kwargs):
        ...
    @property
    def added(self) -> float:
        ...

Now define the field_transformer function that adds the new field, as per attr's documentation:

from typing import Type, List

import attr

def _field_transformer(cls: type, fields: List[attr.Attribute]) -> List[attr.Attribute]:
    return [
        # For some reason mypy has trouble with attr.Attribute's signature
        # Bonus points if someone can point out a fix that doesn't use type: ignore
        attr.Attribute (  # type: ignore
            "added",  # name
            attr.NOTHING,  # default
            None,  # validator
            True,  # repr
            None,  # cmp
            None,  # hash
            True,  # init
            False,  # inherited
            type=float,
            order=float,
        ),
        *fields,
    ]

Now, finally, set up a class decorator that does what we want:

from functools import wraps
from typing import Callable, TypeVar

_T = TypeVar("_T", bound=Proto)
_C = TypeVar("_C", bound=type)

def transform(_cls: _C = None, **kwargs):
    def transform_decorator(cls: _C) -> Callable[[], Type[_T]]:
        @wraps(cls)
        def wrapper() -> Type[_T]:
            if "field_transformer" not in kwargs:
                kwargs["field_transformer"] = _field_transformer
            return attr.frozen(cls, **kwargs)

        return wrapper()

    if _cls is None:
        return transform_decorator
    return transform_decorator(_cls)

And now for the (failing) mypy tests:

@transform
class Test:
    other_field: str

# E: Too many arguments for "Test"
# E: Argument 1 to "Test" has incompatible type "float"; expected "str"
t = Test(0.0, "hello, world")  
print(t.added)  # E: "Test" has no attribute "added"

Ideally I'd like mypy to eliminate all three of these errors. I am frankly not sure whether this is possible; it could be that the dynamic addition of an attribute is just not typeable and we may have to force users of our library to write custom typing stubs when they use the decorator. However, since we always add the same attribute(s) to the generated class, it would be great if there is a solution, even if that means writing a custom mypy plugin that supports this decorator in particular (if that's even possible).

python

mypy

python-typing

python-attrs

0 Answers

Your Answer

Accepted video resources