Skip to content

Commit

Permalink
Restrict type inference of dataclass attributes. (#1130)
Browse files Browse the repository at this point in the history
For now, from the typing module only generic collection types are inferred:
Dict, FrozenSet, List, Set, Tuple. Astroid proxies these to the built-in
collection types (e.g., dict).

Other type annotations from typing like Callable and Union yield Uninferable;
these would need to be handled on a case by case basis.

Co-authored-by: Pierre Sassoulas <[email protected]>
  • Loading branch information
david-yz-liu and Pierre-Sassoulas authored Aug 16, 2021
1 parent 02a4c26 commit 2129907
Show file tree
Hide file tree
Showing 3 changed files with 104 additions and 10 deletions.
5 changes: 5 additions & 0 deletions ChangeLog
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,13 @@ What's New in astroid 2.7.1?
============================
Release date: TBA

* When processing dataclass attributes, only do typing inference on collection types.
Support for instantiating other typing types is left for the future, if desired.


* Fixed LookupMixIn missing from ``astroid.node_classes``.

Closes #1129


What's New in astroid 2.7.0?
Expand Down
49 changes: 39 additions & 10 deletions astroid/brain/brain_dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,16 +97,7 @@ def infer_dataclass_attribute(
if value is not None:
yield from value.infer(context=ctx)
if annotation is not None:
klass = None
try:
klass = next(annotation.infer())
except (InferenceError, StopIteration):
yield Uninferable

if not isinstance(klass, ClassDef):
yield Uninferable
else:
yield klass.instantiate_class()
yield from _infer_instance_from_annotation(annotation, ctx=ctx)
else:
yield Uninferable

Expand Down Expand Up @@ -229,6 +220,44 @@ def _is_init_var(node: NodeNG) -> bool:
return getattr(inferred, "name", "") == "InitVar"


# Allowed typing classes for which we support inferring instances
_INFERABLE_TYPING_TYPES = frozenset(
(
"Dict",
"FrozenSet",
"List",
"Set",
"Tuple",
)
)


def _infer_instance_from_annotation(
node: NodeNG, ctx: context.InferenceContext = None
) -> Generator:
"""Infer an instance corresponding to the type annotation represented by node.
Currently has limited support for the typing module.
"""
klass = None
try:
klass = next(node.infer(context=ctx))
except (InferenceError, StopIteration):
yield Uninferable
if not isinstance(klass, ClassDef):
yield Uninferable
elif klass.root().name in (
"typing",
"",
): # "" because of synthetic nodes in brain_typing.py
if klass.name in _INFERABLE_TYPING_TYPES:
yield klass.instantiate_class()
else:
yield Uninferable
else:
yield klass.instantiate_class()


if PY37_PLUS:
AstroidManager().register_transform(
ClassDef, dataclass_transform, is_decorated_with_dataclass
Expand Down
60 changes: 60 additions & 0 deletions tests/unittest_brain_dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from astroid import bases, nodes
from astroid.const import PY37_PLUS
from astroid.exceptions import InferenceError
from astroid.util import Uninferable

if not PY37_PLUS:
pytest.skip("Dataclasses were added in 3.7", allow_module_level=True)
Expand Down Expand Up @@ -235,3 +236,62 @@ class A:
assert len(inferred) == 1
assert isinstance(inferred[0], nodes.Const)
assert inferred[0].value == "hi"


def test_inference_generic_collection_attribute():
"""Test that an attribute with a generic collection type from the
typing module is inferred correctly.
"""
attr_nodes = astroid.extract_node(
"""
from dataclasses import dataclass, field
import typing
@dataclass
class A:
dict_prop: typing.Dict[str, str]
frozenset_prop: typing.FrozenSet[str]
list_prop: typing.List[str]
set_prop: typing.Set[str]
tuple_prop: typing.Tuple[int, str]
a = A({}, frozenset(), [], set(), (1, 'hi'))
a.dict_prop #@
a.frozenset_prop #@
a.list_prop #@
a.set_prop #@
a.tuple_prop #@
"""
)
names = (
"Dict",
"FrozenSet",
"List",
"Set",
"Tuple",
)
for node, name in zip(attr_nodes, names):
inferred = next(node.infer())
assert isinstance(inferred, bases.Instance)
assert inferred.name == name


def test_inference_callable_attribute():
"""Test that an attribute with a Callable annotation is inferred as Uninferable.
See issue#1129.
"""
instance = astroid.extract_node(
"""
from dataclasses import dataclass
from typing import Any, Callable
@dataclass
class A:
enabled: Callable[[Any], bool]
A(lambda x: x == 42).enabled #@
"""
)
inferred = next(instance.infer())
assert inferred is Uninferable

0 comments on commit 2129907

Please sign in to comment.