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

Infer the type of the result of calling typing.cast() #1076

Merged
merged 2 commits into from
Jul 30, 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
1 change: 1 addition & 0 deletions ChangeLog
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ What's New in astroid 2.6.6?
============================
Release date: TBA

* Added support to infer return type of ``typing.cast()``


What's New in astroid 2.6.5?
Expand Down
34 changes: 34 additions & 0 deletions astroid/brain/brain_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
Call,
Const,
Name,
NodeNG,
Subscript,
)
from astroid.scoped_nodes import ClassDef, FunctionDef
Expand Down Expand Up @@ -356,6 +357,36 @@ def infer_tuple_alias(
return iter([class_def])


def _looks_like_typing_cast(node: Call) -> bool:
return isinstance(node, Call) and (
isinstance(node.func, Name)
and node.func.name == "cast"
or isinstance(node.func, Attribute)
and node.func.attrname == "cast"
)


def infer_typing_cast(
node: Call, ctx: context.InferenceContext = None
) -> typing.Iterator[NodeNG]:
"""Infer call to cast() returning same type as casted-from var"""
if not isinstance(node.func, (Name, Attribute)):
raise UseInferenceDefault
Comment on lines +373 to +374
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Technically this is covered by _looks_like_typing_cast but type checker don't know about that.


try:
func = next(node.func.infer(context=ctx))
except InferenceError as exc:
raise UseInferenceDefault from exc
if (
not isinstance(func, FunctionDef)
or func.qname() != "typing.cast"
or len(node.args) != 2
):
raise UseInferenceDefault

return node.args[1].infer(context=ctx)


AstroidManager().register_transform(
Call,
inference_tip(infer_typing_typevar_or_newtype),
Expand All @@ -364,6 +395,9 @@ def infer_tuple_alias(
AstroidManager().register_transform(
Subscript, inference_tip(infer_typing_attr), _looks_like_typing_subscript
)
AstroidManager().register_transform(
Call, inference_tip(infer_typing_cast), _looks_like_typing_cast
)

if PY39_PLUS:
AstroidManager().register_transform(
Expand Down
32 changes: 32 additions & 0 deletions tests/unittest_brain.py
Original file line number Diff line number Diff line change
Expand Up @@ -1810,6 +1810,38 @@ def test_typing_object_builtin_subscriptable(self):
self.assertIsInstance(inferred, nodes.ClassDef)
self.assertIsInstance(inferred.getattr("__iter__")[0], nodes.FunctionDef)

def test_typing_cast(self):
node = builder.extract_node(
"""
from typing import cast
class A:
pass

b = 42
a = cast(A, b)
a
"""
)
inferred = next(node.infer())
assert isinstance(inferred, nodes.Const)
assert inferred.value == 42
Comment on lines +1826 to +1827
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Usually it's a bit easier to test for a const value.


def test_typing_cast_attribute(self):
node = builder.extract_node(
"""
import typing
class A:
pass

b = 42
a = typing.cast(A, b)
a
"""
)
inferred = next(node.infer())
assert isinstance(inferred, nodes.Const)
assert inferred.value == 42


class ReBrainTest(unittest.TestCase):
def test_regex_flags(self):
Expand Down