Skip to content

Commit

Permalink
Add __dataclass_fields__ and __attrs_attrs__ to dataclasses (#8578)
Browse files Browse the repository at this point in the history
Fixes #6568.
  • Loading branch information
tkukushkin committed Aug 20, 2021
1 parent fbedea5 commit 7576f65
Show file tree
Hide file tree
Showing 15 changed files with 228 additions and 80 deletions.
12 changes: 12 additions & 0 deletions mypy/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,18 @@ def named_type(self, qualified_name: str, args: Optional[List[Type]] = None) ->
"""Construct an instance of a builtin type with given type arguments."""
raise NotImplementedError

@abstractmethod
def named_type_or_none(self,
qualified_name: str,
args: Optional[List[Type]] = None) -> Optional[Instance]:
"""Construct an instance of a type with given type arguments.
Return None if a type could not be constructed for the qualified
type name. This is possible when the qualified name includes a
module name and the module has not been imported.
"""
raise NotImplementedError

@abstractmethod
def parse_bool(self, expr: Expression) -> Optional[bool]:
"""Parse True/False literals."""
Expand Down
27 changes: 25 additions & 2 deletions mypy/plugins/attrs.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,8 @@
deserialize_and_fixup_type
)
from mypy.types import (
Type, AnyType, TypeOfAny, CallableType, NoneType, TypeVarType,
Overloaded, UnionType, FunctionLike, get_proper_type
TupleType, Type, AnyType, TypeOfAny, CallableType, NoneType, TypeVarType,
Overloaded, UnionType, FunctionLike, get_proper_type,
)
from mypy.typeops import make_simplified_union, map_type_from_supertype
from mypy.typevars import fill_typevars
Expand Down Expand Up @@ -300,6 +300,8 @@ def attr_class_maker_callback(ctx: 'mypy.plugin.ClassDefContext',
ctx.api.defer()
return

_add_attrs_magic_attribute(ctx, raw_attr_types=[info[attr.name].type for attr in attributes])

# Save the attributes so that subclasses can reuse them.
ctx.cls.info.metadata['attrs'] = {
'attributes': [attr.serialize() for attr in attributes],
Expand Down Expand Up @@ -703,6 +705,27 @@ def _add_init(ctx: 'mypy.plugin.ClassDefContext', attributes: List[Attribute],
adder.add_method('__init__', args, NoneType())


def _add_attrs_magic_attribute(ctx: 'mypy.plugin.ClassDefContext',
raw_attr_types: 'List[Optional[Type]]') -> None:
attr_name = '__attrs_attrs__'
any_type = AnyType(TypeOfAny.explicit)
attributes_types: 'List[Type]' = [
ctx.api.named_type_or_none('attr.Attribute', [attr_type or any_type]) or any_type
for attr_type in raw_attr_types
]
fallback_type = ctx.api.named_type('__builtins__.tuple', [
ctx.api.named_type_or_none('attr.Attribute', [any_type]) or any_type,
])
var = Var(name=attr_name, type=TupleType(attributes_types, fallback=fallback_type))
var.info = ctx.cls.info
var._fullname = ctx.cls.info.fullname + '.' + attr_name
ctx.cls.info.names[attr_name] = SymbolTableNode(
kind=MDEF,
node=var,
plugin_generated=True,
)


class MethodAdder:
"""Helper to add methods to a TypeInfo.
Expand Down
24 changes: 23 additions & 1 deletion mypy/plugins/dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,10 @@
add_method, _get_decorator_bool_argument, deserialize_and_fixup_type,
)
from mypy.typeops import map_type_from_supertype
from mypy.types import Type, Instance, NoneType, TypeVarType, CallableType, get_proper_type
from mypy.types import (
Type, Instance, NoneType, TypeVarType, CallableType, get_proper_type,
AnyType, TypeOfAny,
)
from mypy.server.trigger import make_wildcard_trigger

# The set of decorators that generate dataclasses.
Expand Down Expand Up @@ -187,6 +190,8 @@ def transform(self) -> None:

self.reset_init_only_vars(info, attributes)

self._add_dataclass_fields_magic_attribute()

info.metadata['dataclass'] = {
'attributes': [attr.serialize() for attr in attributes],
'frozen': decorator_arguments['frozen'],
Expand Down Expand Up @@ -417,6 +422,23 @@ def _is_kw_only_type(self, node: Optional[Type]) -> bool:
return False
return node_type.type.fullname == 'dataclasses.KW_ONLY'

def _add_dataclass_fields_magic_attribute(self) -> None:
attr_name = '__dataclass_fields__'
any_type = AnyType(TypeOfAny.explicit)
field_type = self._ctx.api.named_type_or_none('dataclasses.Field', [any_type]) or any_type
attr_type = self._ctx.api.named_type('__builtins__.dict', [
self._ctx.api.named_type('__builtins__.str'),
field_type,
])
var = Var(name=attr_name, type=attr_type)
var.info = self._ctx.cls.info
var._fullname = self._ctx.cls.info.fullname + '.' + attr_name
self._ctx.cls.info.names[attr_name] = SymbolTableNode(
kind=MDEF,
node=var,
plugin_generated=True,
)


def dataclass_class_maker_callback(ctx: ClassDefContext) -> None:
"""Hooks into the class typechecking process to add support for dataclasses.
Expand Down
12 changes: 12 additions & 0 deletions test-data/unit/check-attr.test
Original file line number Diff line number Diff line change
Expand Up @@ -1390,3 +1390,15 @@ class B(A):

reveal_type(B) # N: Revealed type is "def (foo: builtins.int) -> __main__.B"
[builtins fixtures/bool.pyi]

[case testAttrsClassHasAttributeWithAttributes]
import attr

@attr.s
class A:
b: int = attr.ib()
c: str = attr.ib()

reveal_type(A.__attrs_attrs__) # N: Revealed type is "Tuple[attr.Attribute[builtins.int], attr.Attribute[builtins.str]]"

[builtins fixtures/attr.pyi]

0 comments on commit 7576f65

Please sign in to comment.