Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 7576f65

Browse files
tkukushkintkukushkin
authored
Add __dataclass_fields__ and __attrs_attrs__ to dataclasses (#8578)
Fixes #6568.
1 parent fbedea5 commit 7576f65

15 files changed

Lines changed: 228 additions & 80 deletions

‎mypy/plugin.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,18 @@ def named_type(self, qualified_name: str, args: Optional[List[Type]] = None) ->
254254
"""Construct an instance of a builtin type with given type arguments."""
255255
raise NotImplementedError
256256

257+
@abstractmethod
258+
def named_type_or_none(self,
259+
qualified_name: str,
260+
args: Optional[List[Type]] = None) -> Optional[Instance]:
261+
"""Construct an instance of a type with given type arguments.
262+
263+
Return None if a type could not be constructed for the qualified
264+
type name. This is possible when the qualified name includes a
265+
module name and the module has not been imported.
266+
"""
267+
raise NotImplementedError
268+
257269
@abstractmethod
258270
def parse_bool(self, expr: Expression) -> Optional[bool]:
259271
"""Parse True/False literals."""

‎mypy/plugins/attrs.py‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@
2121
deserialize_and_fixup_type
2222
)
2323
from mypy.types import (
24-
Type, AnyType, TypeOfAny, CallableType, NoneType, TypeVarType,
25-
Overloaded, UnionType, FunctionLike, get_proper_type
24+
TupleType, Type, AnyType, TypeOfAny, CallableType, NoneType, TypeVarType,
25+
Overloaded, UnionType, FunctionLike, get_proper_type,
2626
)
2727
from mypy.typeops import make_simplified_union, map_type_from_supertype
2828
from mypy.typevars import fill_typevars
@@ -300,6 +300,8 @@ def attr_class_maker_callback(ctx: 'mypy.plugin.ClassDefContext',
300300
ctx.api.defer()
301301
return
302302

303+
_add_attrs_magic_attribute(ctx, raw_attr_types=[info[attr.name].type for attr in attributes])
304+
303305
# Save the attributes so that subclasses can reuse them.
304306
ctx.cls.info.metadata['attrs'] = {
305307
'attributes': [attr.serialize() for attr in attributes],
@@ -703,6 +705,27 @@ def _add_init(ctx: 'mypy.plugin.ClassDefContext', attributes: List[Attribute],
703705
adder.add_method('__init__', args, NoneType())
704706

705707

708+
def _add_attrs_magic_attribute(ctx: 'mypy.plugin.ClassDefContext',
709+
raw_attr_types: 'List[Optional[Type]]') -> None:
710+
attr_name = '__attrs_attrs__'
711+
any_type = AnyType(TypeOfAny.explicit)
712+
attributes_types: 'List[Type]' = [
713+
ctx.api.named_type_or_none('attr.Attribute', [attr_type or any_type]) or any_type
714+
for attr_type in raw_attr_types
715+
]
716+
fallback_type = ctx.api.named_type('__builtins__.tuple', [
717+
ctx.api.named_type_or_none('attr.Attribute', [any_type]) or any_type,
718+
])
719+
var = Var(name=attr_name, type=TupleType(attributes_types, fallback=fallback_type))
720+
var.info = ctx.cls.info
721+
var._fullname = ctx.cls.info.fullname + '.' + attr_name
722+
ctx.cls.info.names[attr_name] = SymbolTableNode(
723+
kind=MDEF,
724+
node=var,
725+
plugin_generated=True,
726+
)
727+
728+
706729
class MethodAdder:
707730
"""Helper to add methods to a TypeInfo.
708731

‎mypy/plugins/dataclasses.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,10 @@
1313
add_method, _get_decorator_bool_argument, deserialize_and_fixup_type,
1414
)
1515
from mypy.typeops import map_type_from_supertype
16-
from mypy.types import Type, Instance, NoneType, TypeVarType, CallableType, get_proper_type
16+
from mypy.types import (
17+
Type, Instance, NoneType, TypeVarType, CallableType, get_proper_type,
18+
AnyType, TypeOfAny,
19+
)
1720
from mypy.server.trigger import make_wildcard_trigger
1821

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

188191
self.reset_init_only_vars(info, attributes)
189192

193+
self._add_dataclass_fields_magic_attribute()
194+
190195
info.metadata['dataclass'] = {
191196
'attributes': [attr.serialize() for attr in attributes],
192197
'frozen': decorator_arguments['frozen'],
@@ -417,6 +422,23 @@ def _is_kw_only_type(self, node: Optional[Type]) -> bool:
417422
return False
418423
return node_type.type.fullname == 'dataclasses.KW_ONLY'
419424

425+
def _add_dataclass_fields_magic_attribute(self) -> None:
426+
attr_name = '__dataclass_fields__'
427+
any_type = AnyType(TypeOfAny.explicit)
428+
field_type = self._ctx.api.named_type_or_none('dataclasses.Field', [any_type]) or any_type
429+
attr_type = self._ctx.api.named_type('__builtins__.dict', [
430+
self._ctx.api.named_type('__builtins__.str'),
431+
field_type,
432+
])
433+
var = Var(name=attr_name, type=attr_type)
434+
var.info = self._ctx.cls.info
435+
var._fullname = self._ctx.cls.info.fullname + '.' + attr_name
436+
self._ctx.cls.info.names[attr_name] = SymbolTableNode(
437+
kind=MDEF,
438+
node=var,
439+
plugin_generated=True,
440+
)
441+
420442

421443
def dataclass_class_maker_callback(ctx: ClassDefContext) -> None:
422444
"""Hooks into the class typechecking process to add support for dataclasses.

‎test-data/unit/check-attr.test‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1390,3 +1390,15 @@ class B(A):
13901390

13911391
reveal_type(B) # N: Revealed type is "def (foo: builtins.int) -> __main__.B"
13921392
[builtins fixtures/bool.pyi]
1393+
1394+
[case testAttrsClassHasAttributeWithAttributes]
1395+
import attr
1396+
1397+
@attr.s
1398+
class A:
1399+
b: int = attr.ib()
1400+
c: str = attr.ib()
1401+
1402+
reveal_type(A.__attrs_attrs__) # N: Revealed type is "Tuple[attr.Attribute[builtins.int], attr.Attribute[builtins.str]]"
1403+
1404+
[builtins fixtures/attr.pyi]

0 commit comments

Comments
 (0)