Skip to content

Commit

Permalink
Add support for next-gen attrs API (#9396)
Browse files Browse the repository at this point in the history
These include the attr class makers: define, mutable, frozen
And the attrib maker: field

Also includes support for auto_attribs=None which means auto_detect which method of attributes are being used.
  • Loading branch information
euresti authored Nov 18, 2020
1 parent 849a7f7 commit 6e99a2d
Show file tree
Hide file tree
Showing 4 changed files with 219 additions and 7 deletions.
52 changes: 46 additions & 6 deletions mypy/plugins/attrs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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]:
Expand All @@ -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):
Expand Down Expand Up @@ -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]:
Expand Down
13 changes: 12 additions & 1 deletion mypy/plugins/default.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
40 changes: 40 additions & 0 deletions test-data/unit/check-attr.test
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
121 changes: 121 additions & 0 deletions test-data/unit/lib-stub/attr.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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: ...

0 comments on commit 6e99a2d

Please sign in to comment.