Add support for next-gen attrs API

mypy/plugins/attrs.py

@@ -38,10 +38,18 @@
 attr_dataclass_makers = {
     'attr.dataclass',
 }  # type: Final
+attr_frozen_makers = {
+    'attr.frozen'
+}  # type: Final
+attr_define_makers = {
+    'attr.define',
+    'attr.mutable'
+}  # type: Final
 attr_attrib_makers = {
     'attr.ib',
     'attr.attrib',
     'attr.attr',
+    'attr.field',
 }  # type: Final
 
 SELF_TVAR_NAME = '_AT'  # type: Final
@@ -232,7 +240,8 @@ def _get_decorator_optional_bool_argument(
 
 
 def attr_class_maker_callback(ctx: 'mypy.plugin.ClassDefContext',
-                              auto_attribs_default: bool = False) -> None:
+                              auto_attribs_default: Optional[bool] = False,
+                              frozen_default: bool = False) -> None:
     """Add necessary dunder methods to classes decorated with attr.s.
 
     attrs is a package that lets you define classes without writing dull boilerplate code.
@@ -247,10 +256,10 @@ def attr_class_maker_callback(ctx: 'mypy.plugin.ClassDefContext',
     info = ctx.cls.info
 
     init = _get_decorator_bool_argument(ctx, 'init', True)
-    frozen = _get_frozen(ctx)
+    frozen = _get_frozen(ctx, frozen_default)
     order = _determine_eq_order(ctx)
 
-    auto_attribs = _get_decorator_bool_argument(ctx, 'auto_attribs', auto_attribs_default)
+    auto_attribs = _get_decorator_optional_bool_argument(ctx, 'auto_attribs', auto_attribs_default)
     kw_only = _get_decorator_bool_argument(ctx, 'kw_only', False)
 
     if ctx.api.options.python_version[0] < 3:
@@ -293,9 +302,9 @@ def attr_class_maker_callback(ctx: 'mypy.plugin.ClassDefContext',
         _make_frozen(ctx, attributes)
 
 
-def _get_frozen(ctx: 'mypy.plugin.ClassDefContext') -> bool:
+def _get_frozen(ctx: 'mypy.plugin.ClassDefContext', frozen_default: bool) -> bool:
     """Return whether this class is frozen."""
-    if _get_decorator_bool_argument(ctx, 'frozen', False):
+    if _get_decorator_bool_argument(ctx, 'frozen', frozen_default):
         return True
     # Subclasses of frozen classes are frozen so check that.
     for super_info in ctx.cls.info.mro[1:-1]:
@@ -305,14 +314,18 @@ def _get_frozen(ctx: 'mypy.plugin.ClassDefContext') -> bool:
 
 
 def _analyze_class(ctx: 'mypy.plugin.ClassDefContext',
-                   auto_attribs: bool,
+                   auto_attribs: Optional[bool],
                    kw_only: bool) -> List[Attribute]:
     """Analyze the class body of an attr maker, its parents, and return the Attributes found.
 
     auto_attribs=True means we'll generate attributes from type annotations also.
+    auto_attribs=None means we'll detect which mode to use.
     kw_only=True means that all attributes created here will be keyword only args in __init__.
     """
     own_attrs = OrderedDict()  # type: OrderedDict[str, Attribute]
+    if auto_attribs is None:
+        auto_attribs = _detect_auto_attribs(ctx)
+
     # Walk the body looking for assignments and decorators.
     for stmt in ctx.cls.defs.body:
         if isinstance(stmt, AssignmentStmt):
@@ -380,6 +393,33 @@ def _analyze_class(ctx: 'mypy.plugin.ClassDefContext',
     return attributes
 
 
+def _detect_auto_attribs(ctx: 'mypy.plugin.ClassDefContext') -> bool:
+    """Return whether auto_attribs should be enabled or disabled.
+
+    It's disabled if there are any unannotated attribs()
+    """
+    for stmt in ctx.cls.defs.body:
+        if isinstance(stmt, AssignmentStmt):
+            for lvalue in stmt.lvalues:
+                lvalues, rvalues = _parse_assignments(lvalue, stmt)
+
+                if len(lvalues) != len(rvalues):
+                    # This means we have some assignment that isn't 1 to 1.
+                    # It can't be an attrib.
+                    continue
+
+                for lhs, rvalue in zip(lvalues, rvalues):
+                    # Check if the right hand side is a call to an attribute maker.
+                    if (isinstance(rvalue, CallExpr)
+                            and isinstance(rvalue.callee, RefExpr)
+                            and rvalue.callee.fullname in attr_attrib_makers
+                            and not stmt.new_syntax):
+                        # This means we have an attrib without an annotation and so
+                        # we can't do auto_attribs=True
+                        return False
+    return True
+
+
 def _attributes_from_assignment(ctx: 'mypy.plugin.ClassDefContext',
                                 stmt: AssignmentStmt, auto_attribs: bool,
                                 kw_only: bool) -> Iterable[Attribute]:

mypy/plugins/default.py

@@ -99,7 +99,18 @@ def get_class_decorator_hook(self, fullname: str
         elif fullname in attrs.attr_dataclass_makers:
             return partial(
                 attrs.attr_class_maker_callback,
-                auto_attribs_default=True
+                auto_attribs_default=True,
+            )
+        elif fullname in attrs.attr_frozen_makers:
+            return partial(
+                attrs.attr_class_maker_callback,
+                auto_attribs_default=None,
+                frozen_default=True,
+            )
+        elif fullname in attrs.attr_define_makers:
+            return partial(
+                attrs.attr_class_maker_callback,
+                auto_attribs_default=None,
             )
         elif fullname in dataclasses.dataclass_makers:
             return dataclasses.dataclass_class_maker_callback

test-data/unit/check-attr.test

@@ -361,6 +361,46 @@ class A:
 a = A(5)
 a.a = 16  # E: Property "a" defined in "A" is read-only
 [builtins fixtures/bool.pyi]
+[case testAttrsNextGenFrozen]
+from attr import frozen, field
+
+@frozen
+class A:
+    a = field()
+
+a = A(5)
+a.a = 16  # E: Property "a" defined in "A" is read-only
+[builtins fixtures/bool.pyi]
+
+[case testAttrsNextGenDetect]
+from attr import define, field
+
+@define
+class A:
+    a = field()
+
+@define
+class B:
+    a: int
+
+@define
+class C:
+    a: int = field()
+    b = field()
+
+@define
+class D:
+    a: int
+    b = field()
+
+reveal_type(A)  # N: Revealed type is 'def (a: Any) -> __main__.A'
+reveal_type(B)  # N: Revealed type is 'def (a: builtins.int) -> __main__.B'
+reveal_type(C)  # N: Revealed type is 'def (a: builtins.int, b: Any) -> __main__.C'
+reveal_type(D)  # N: Revealed type is 'def (b: Any) -> __main__.D'
+
+[builtins fixtures/bool.pyi]
+
+
 
 [case testAttrsDataClass]
 import attr

test-data/unit/lib-stub/attr.pyi

@@ -119,3 +119,124 @@ def attrs(maybe_cls: None = ...,
 s = attributes = attrs
 ib = attr = attrib
 dataclass = attrs # Technically, partial(attrs, auto_attribs=True) ;)
+
+# Next Generation API
+@overload
+def define(
+    maybe_cls: _C,
+    *,
+    these: Optional[Mapping[str, Any]] = ...,
+    repr: bool = ...,
+    hash: Optional[bool] = ...,
+    init: bool = ...,
+    slots: bool = ...,
+    frozen: bool = ...,
+    weakref_slot: bool = ...,
+    str: bool = ...,
+    auto_attribs: bool = ...,
+    kw_only: bool = ...,
+    cache_hash: bool = ...,
+    auto_exc: bool = ...,
+    eq: Optional[bool] = ...,
+    order: Optional[bool] = ...,
+    auto_detect: bool = ...,
+    getstate_setstate: Optional[bool] = ...,
+    on_setattr: Optional[object] = ...,
+) -> _C: ...
+@overload
+def define(
+    maybe_cls: None = ...,
+    *,
+    these: Optional[Mapping[str, Any]] = ...,
+    repr: bool = ...,
+    hash: Optional[bool] = ...,
+    init: bool = ...,
+    slots: bool = ...,
+    frozen: bool = ...,
+    weakref_slot: bool = ...,
+    str: bool = ...,
+    auto_attribs: bool = ...,
+    kw_only: bool = ...,
+    cache_hash: bool = ...,
+    auto_exc: bool = ...,
+    eq: Optional[bool] = ...,
+    order: Optional[bool] = ...,
+    auto_detect: bool = ...,
+    getstate_setstate: Optional[bool] = ...,
+    on_setattr: Optional[object] = ...,
+) -> Callable[[_C], _C]: ...
+
+mutable = define
+frozen = define  # they differ only in their defaults
+
+@overload
+def field(
+    *,
+    default: None = ...,
+    validator: None = ...,
+    repr: object = ...,
+    hash: Optional[bool] = ...,
+    init: bool = ...,
+    metadata: Optional[Mapping[Any, Any]] = ...,
+    converter: None = ...,
+    factory: None = ...,
+    kw_only: bool = ...,
+    eq: Optional[bool] = ...,
+    order: Optional[bool] = ...,
+    on_setattr: Optional[_OnSetAttrArgType] = ...,
+) -> Any: ...
+
+# This form catches an explicit None or no default and infers the type from the
+# other arguments.
+@overload
+def field(
+    *,
+    default: None = ...,
+    validator: Optional[_ValidatorArgType[_T]] = ...,
+    repr: object = ...,
+    hash: Optional[bool] = ...,
+    init: bool = ...,
+    metadata: Optional[Mapping[Any, Any]] = ...,
+    converter: Optional[_ConverterType] = ...,
+    factory: Optional[Callable[[], _T]] = ...,
+    kw_only: bool = ...,
+    eq: Optional[bool] = ...,
+    order: Optional[bool] = ...,
+    on_setattr: Optional[object] = ...,
+) -> _T: ...
+
+# This form catches an explicit default argument.
+@overload
+def field(
+    *,
+    default: _T,
+    validator: Optional[_ValidatorArgType[_T]] = ...,
+    repr: object = ...,
+    hash: Optional[bool] = ...,
+    init: bool = ...,
+    metadata: Optional[Mapping[Any, Any]] = ...,
+    converter: Optional[_ConverterType] = ...,
+    factory: Optional[Callable[[], _T]] = ...,
+    kw_only: bool = ...,
+    eq: Optional[bool] = ...,
+    order: Optional[bool] = ...,
+    on_setattr: Optional[object] = ...,
+) -> _T: ...
+
+# This form covers type=non-Type: e.g. forward references (str), Any
+@overload
+def field(
+    *,
+    default: Optional[_T] = ...,
+    validator: Optional[_ValidatorArgType[_T]] = ...,
+    repr: object = ...,
+    hash: Optional[bool] = ...,
+    init: bool = ...,
+    metadata: Optional[Mapping[Any, Any]] = ...,
+    converter: Optional[_ConverterType] = ...,
+    factory: Optional[Callable[[], _T]] = ...,
+    kw_only: bool = ...,
+    eq: Optional[bool] = ...,
+    order: Optional[bool] = ...,
+    on_setattr: Optional[object] = ...,
+) -> Any: ...