Skip to content

Commit

Permalink
Treat NewTypes like normal subclasses
Browse files Browse the repository at this point in the history
NewTypes are assumed not to inherit any members from their base classes.
This results in incorrect inference results. Avoid this by changing the
transformation for NewTypes to treat them like any other subclass.

pylint-dev/pylint#3162
pylint-dev/pylint#2296
  • Loading branch information
colatkinson committed Dec 17, 2021
1 parent a9f5f97 commit 4cf70d3
Show file tree
Hide file tree
Showing 2 changed files with 72 additions and 9 deletions.
53 changes: 44 additions & 9 deletions astroid/brain/brain_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,6 @@
from astroid.util import Uninferable

TYPING_NAMEDTUPLE_BASENAMES = {"NamedTuple", "typing.NamedTuple"}
TYPING_TYPEVARS = {"TypeVar", "NewType"}
TYPING_TYPEVARS_QUALIFIED = {"typing.TypeVar", "typing.NewType"}
TYPING_TYPE_TEMPLATE = """
class Meta(type):
def __getitem__(self, item):
Expand All @@ -52,6 +50,11 @@ def __args__(self):
class {0}(metaclass=Meta):
pass
"""
# PEP484 suggests NewType is equivalent to this for typing purposes
TYPING_NEWTYPE_TEMPLATE = """
class {derived}({base}):
pass
"""
TYPING_MEMBERS = set(getattr(typing, "__all__", []))

TYPING_ALIAS = frozenset(
Expand Down Expand Up @@ -106,23 +109,32 @@ def __class_getitem__(cls, item):
"""


def looks_like_typing_typevar_or_newtype(node):
def looks_like_typing_typevar(node):
func = node.func
if isinstance(func, Attribute):
return func.attrname == "TypeVar"
if isinstance(func, Name):
return func.name == "TypeVar"
return False


def looks_like_typing_newtype(node):
func = node.func
if isinstance(func, Attribute):
return func.attrname in TYPING_TYPEVARS
return func.attrname == "NewType"
if isinstance(func, Name):
return func.name in TYPING_TYPEVARS
return func.name == "NewType"
return False


def infer_typing_typevar_or_newtype(node, context_itton=None):
def infer_typing_typevar(node, context_itton=None):
"""Infer a typing.TypeVar(...) or typing.NewType(...) call"""
try:
func = next(node.func.infer(context=context_itton))
except (InferenceError, StopIteration) as exc:
raise UseInferenceDefault from exc

if func.qname() not in TYPING_TYPEVARS_QUALIFIED:
if func.qname() != "typing.TypeVar":
raise UseInferenceDefault
if not node.args:
raise UseInferenceDefault
Expand All @@ -132,6 +144,24 @@ def infer_typing_typevar_or_newtype(node, context_itton=None):
return node.infer(context=context_itton)


def infer_typing_newtype(node, context_itton=None):
"""Infer a typing.TypeVar(...) or typing.NewType(...) call"""
try:
func = next(node.func.infer(context=context_itton))
except (InferenceError, StopIteration) as exc:
raise UseInferenceDefault from exc

if func.qname() != "typing.NewType":
raise UseInferenceDefault
if not node.args:
raise UseInferenceDefault

derived = node.args[0].as_string().strip("'")
base = node.args[1].as_string().strip("'")
node = extract_node(TYPING_NEWTYPE_TEMPLATE.format(derived=derived, base=base))
return node.infer(context=context_itton)


def _looks_like_typing_subscript(node):
"""Try to figure out if a Subscript node *might* be a typing-related subscript"""
if isinstance(node, Name):
Expand Down Expand Up @@ -409,8 +439,13 @@ def infer_typing_cast(

AstroidManager().register_transform(
Call,
inference_tip(infer_typing_typevar_or_newtype),
looks_like_typing_typevar_or_newtype,
inference_tip(infer_typing_typevar),
looks_like_typing_typevar,
)
AstroidManager().register_transform(
Call,
inference_tip(infer_typing_newtype),
looks_like_typing_newtype,
)
AstroidManager().register_transform(
Subscript, inference_tip(infer_typing_attr), _looks_like_typing_subscript
Expand Down
28 changes: 28 additions & 0 deletions tests/unittest_brain.py
Original file line number Diff line number Diff line change
Expand Up @@ -1659,6 +1659,34 @@ def test_typing_types(self) -> None:
inferred = next(node.infer())
self.assertIsInstance(inferred, nodes.ClassDef, node.as_string())

def test_typing_newtype_attrs(self) -> None:
ast_nodes = builder.extract_node(
"""
from typing import List, NewType
import typing
NewType("Foo", str) #@
NewType("Bar", "int") #@
"""
)
assert isinstance(ast_nodes, list)

# Base type given by reference
foo_node = ast_nodes[0]
foo_inferred = next(foo_node.infer())
self.assertIsInstance(foo_inferred, astroid.ClassDef)

# Check base type method is inferred
foo_base_class_method = foo_inferred.getattr("endswith")[0]
self.assertIsInstance(foo_base_class_method, astroid.FunctionDef)

# Base type given by string
bar_node = ast_nodes[1]
bar_inferred = next(bar_node.infer())
self.assertIsInstance(bar_inferred, astroid.ClassDef)

bar_base_class_method = bar_inferred.getattr("bit_count")[0]
self.assertIsInstance(bar_base_class_method, astroid.FunctionDef)

def test_namedtuple_nested_class(self):
result = builder.extract_node(
"""
Expand Down

0 comments on commit 4cf70d3

Please sign in to comment.