Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Restrict type inference of dataclass attributes. #1130

Merged
merged 2 commits into from
Aug 16, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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