Improve matching of typing names in the pyi parser.

Adds support for aliasing names like TypeVar and simplifies name matching by
requiring the caller to pass in only one variant of a name (e.g.,
"typing.Callable") and having the matcher automatically try all variants
("Callable", "typing.Callable", "collections.abc.Callable").

See python/typeshed#10201 for motivation.

Fixes #1430.

PiperOrigin-RevId: 534721075

Author: rchen152

pytype/pyi/CMakeLists.txt

@@ -66,6 +66,7 @@ py_library(
     .classdef
     .metadata
     .types
+    pytype.utils
     pytype.pytd.pytd_for_parser
     pytype.pytd.parse.parse
 )

pytype/pyi/classdef.py

@@ -3,11 +3,10 @@
 import collections
 import sys
 
-from typing import cast, Dict, List
+from typing import cast, Callable, Dict, List
 
 from pytype.pyi import types
 from pytype.pytd import pytd
-from pytype.pytd import pytd_utils
 from pytype.pytd.parse import node as pytd_node
 
 # pylint: disable=g-import-not-at-top
@@ -17,18 +16,17 @@
   from typed_ast import ast3
 # pylint: enable=g-import-not-at-top
 
-
-_PROTOCOL_ALIASES = ("typing.Protocol", "typing_extensions.Protocol")
 ParseError = types.ParseError
 
 
-def get_bases(bases: List[pytd.Type]) -> List[pytd.Type]:
+def get_bases(
+    bases: List[pytd.Type], type_match: Callable[..., bool]) -> List[pytd.Type]:
   """Collect base classes."""
 
   bases_out = []
   namedtuple_index = None
   for i, p in enumerate(bases):
-    if p.name and pytd_utils.MatchesFullName(p, _PROTOCOL_ALIASES):
+    if p.name and type_match(p.name, "typing.Protocol"):
       if isinstance(p, pytd.GenericType):
         # From PEP 544: "`Protocol[T, S, ...]` is allowed as a shorthand for
         # `Protocol, Generic[T, S, ...]`."

pytype/pyi/definitions.py

@@ -4,8 +4,9 @@
 import dataclasses
 import sys
 
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
+from pytype import utils
 from pytype.pyi import classdef
 from pytype.pyi import metadata
 from pytype.pyi import types
@@ -19,6 +20,7 @@
 from pytype.pytd.codegen import namedtuple
 from pytype.pytd.codegen import pytdgen
 from pytype.pytd.parse import node as pytd_node
+from pytype.pytd.parse import parser_constants
 
 # pylint: disable=g-import-not-at-top
 if sys.version_info >= (3, 8):
@@ -31,13 +33,6 @@
 # Typing members that represent sets of types.
 _TYPING_SETS = ("typing.Intersection", "typing.Optional", "typing.Union")
 
-# Aliases for some typing.X types
-_ANNOTATED_TYPES = ("typing.Annotated", "typing_extensions.Annotated")
-_CALLABLE_TYPES = ("typing.Callable", "collections.abc.Callable")
-_CONCATENATE_TYPES = ("typing.Concatenate", "typing_extensions.Concatenate")
-_LITERAL_TYPES = ("typing.Literal", "typing_extensions.Literal")
-_TUPLE_TYPES = ("tuple", "builtins.tuple", "typing.Tuple")
-
 
 class StringParseError(ParseError):
   pass
@@ -493,23 +488,43 @@ def add_import(self, from_package, import_list):
         if t:
           self.aliases[t.new_name] = t.pytd_alias()
 
-  def _matches_full_name(self, t, full_name):
-    """Whether t.name matches full_name in format {module}.{member}."""
-    return pytd_utils.MatchesFullName(
-        t, full_name, self.module_info.module_name, self.aliases)
+  def _resolve_alias(self, name: str) -> str:
+    if name in self.aliases:
+      alias = self.aliases[name].type
+      if isinstance(alias, pytd.NamedType):
+        name = alias.name
+      elif isinstance(alias, pytd.Module):
+        name = alias.module_name
+    return name
+
+  def matches_type(self, name: str, target: Union[str, Tuple[str, ...]]):
+    """Checks whether 'name' matches the 'target' type."""
+    if isinstance(target, tuple):
+      return any(self.matches_type(name, t) for t in target)
+    assert "." in target, "'target' must be a fully qualified type name"
+    if "." in name:
+      prefix, name_base = name.rsplit(".", 1)
+      name = f"{self._resolve_alias(prefix)}.{name_base}"
+    else:
+      name = self._resolve_alias(name)
+    name = utils.strip_prefix(name, parser_constants.EXTERNAL_NAME_PREFIX)
+    if name == target:
+      return True
+    module, target_base = target.rsplit(".", 1)
+    if name == target_base:
+      return True
+    if module == "builtins":
+      return self.matches_type(name, f"typing.{target_base.title()}")
+    equivalent_modules = {"typing", "collections.abc", "typing_extensions"}
+    if module not in equivalent_modules:
+      return False
+    return any(name == f"{mod}.{target_base}" for mod in equivalent_modules)
 
   def _matches_named_type(self, t, names):
     """Whether t is a NamedType matching any of names."""
     if not isinstance(t, pytd.NamedType):
       return False
-    for name in names:
-      if "." in name:
-        if self._matches_full_name(t, name):
-          return True
-      else:
-        if t.name == name:
-          return True
-    return False
+    return self.matches_type(t.name, names)
 
   def _is_empty_tuple(self, t):
     return isinstance(t, pytd.TupleType) and not t.parameters
@@ -551,22 +566,22 @@ def _remove_unsupported_features(self, parameters):
 
   def _parameterized_type(self, base_type: Any, parameters):
     """Return a parameterized type."""
-    if self._matches_named_type(base_type, _LITERAL_TYPES):
+    if self._matches_named_type(base_type, "typing.Literal"):
       return pytd_literal(parameters, self.aliases)
-    elif self._matches_named_type(base_type, _ANNOTATED_TYPES):
+    elif self._matches_named_type(base_type, "typing.Annotated"):
       return pytd_annotated(parameters)
     self._verify_no_literal_parameters(base_type, parameters)
     arg_is_paramspec = False
-    if self._matches_named_type(base_type, _TUPLE_TYPES):
+    if self._matches_named_type(base_type, "builtins.tuple"):
       if len(parameters) == 2 and parameters[1] is self.ELLIPSIS:
         parameters = parameters[:1]
         builder = pytd.GenericType
       else:
         builder = pytdgen.heterogeneous_tuple
-    elif self._matches_named_type(base_type, _CONCATENATE_TYPES):
+    elif self._matches_named_type(base_type, "typing.Concatenate"):
       assert parameters
       builder = pytd.Concatenate
-    elif self._matches_named_type(base_type, _CALLABLE_TYPES):
+    elif self._matches_named_type(base_type, "typing.Callable"):
       if parameters[0] is self.ELLIPSIS:
         parameters = (pytd.AnythingType(),) + parameters[1:]
       if parameters and isinstance(parameters[0], pytd.NamedType):
@@ -661,7 +676,7 @@ def build_class(
       self, class_name, bases, keywords, decorators, defs
   ) -> pytd.Class:
     """Build a pytd.Class from definitions collected from an ast node."""
-    bases = classdef.get_bases(bases)
+    bases = classdef.get_bases(bases, self.matches_type)
     keywords = classdef.get_keywords(keywords)
     constants, methods, aliases, slots, classes = _split_definitions(defs)
 

pytype/pyi/parser.py

@@ -30,20 +30,6 @@
 # reexport as parser.ParseError
 ParseError = types.ParseError
 
-_TYPEVAR_IDS = ("TypeVar", "typing.TypeVar")
-_PARAMSPEC_IDS = (
-    "ParamSpec", "typing.ParamSpec", "typing_extensions.ParamSpec")
-_TYPING_NAMEDTUPLE_IDS = ("NamedTuple", "typing.NamedTuple")
-_COLL_NAMEDTUPLE_IDS = ("namedtuple", "collections.namedtuple")
-_TYPEDDICT_IDS = (
-    "TypedDict", "typing.TypedDict", "typing_extensions.TypedDict")
-_NEWTYPE_IDS = ("NewType", "typing.NewType")
-_ANNOTATED_IDS = (
-    "Annotated", "typing.Annotated", "typing_extensions.Annotated")
-_FINAL_IDS = ("typing.Final", "typing_extensions.Final")
-_TYPE_ALIAS_IDS = ("typing.TypeAlias", "typing_extensions.TypeAlias")
-_TYPING_LITERAL_IDS = ("Literal", "typing.Literal", "typing_extensions.Literal")
-
 #------------------------------------------------------
 # imports
 
@@ -218,7 +204,8 @@ def _convert_typing_annotated(self, node):
   def enter_Subscript(self, node):
     if isinstance(node.value, ast3.Attribute):
       node.value = _attribute_to_name(node.value).id
-    if getattr(node.value, "id", None) in _ANNOTATED_IDS:
+    if self.defs.matches_type(getattr(node.value, "id", ""),
+                              "typing.Annotated"):
       self._convert_typing_annotated(node)
     self.subscripted.append(node.value)
 
@@ -244,8 +231,13 @@ def visit_Attribute(self, node):
   def visit_BinOp(self, node):
     if self.subscripted:
       last = self.subscripted[-1]
-      if ((isinstance(last, ast3.Name) and last.id in _TYPING_LITERAL_IDS) or
-          isinstance(last, str) and last in _TYPING_LITERAL_IDS):
+      if isinstance(last, ast3.Name):
+        last_id = last.id
+      elif isinstance(last, str):
+        last_id = last
+      else:
+        last_id = ""
+      if self.defs.matches_type(last_id, "typing.Literal"):
         raise ParseError("Expressions are not allowed in typing.Literal.")
     if isinstance(node.op, ast3.BitOr):
       return self.defs.new_type("typing.Union", [node.left, node.right])
@@ -446,7 +438,7 @@ def visit_AnnAssign(self, node):
       typ = pytd.NamedType("tuple")
       val = None
     elif typ.name:
-      if pytd_utils.MatchesFullName(typ, _FINAL_IDS):
+      if self.defs.matches_type(typ.name, "typing.Final"):
         if isinstance(node.value, types.Pyval):
           # to_pytd_literal raises an exception if the value is a float, but
           # checking upfront allows us to generate a nicer error message.
@@ -460,7 +452,7 @@ def visit_AnnAssign(self, node):
         elif isinstance(val, pytd.NamedType):
           typ = pytd.Literal(val)
           val = None
-      elif pytd_utils.MatchesFullName(typ, _TYPE_ALIAS_IDS):
+      elif self.defs.matches_type(typ.name, "typing.TypeAlias"):
         typ = val
         val = None
         is_alias = True
@@ -668,31 +660,32 @@ def enter_Call(self, node):
     # passing them to internal functions directly in visit_Call.
     if isinstance(node.func, ast3.Attribute):
       node.func = _attribute_to_name(node.func)
-    if node.func.id in _TYPEVAR_IDS:
+    if self.defs.matches_type(node.func.id, "typing.TypeVar"):
       self._convert_typevar_args(node)
-    elif node.func.id in _PARAMSPEC_IDS:
+    elif self.defs.matches_type(node.func.id, "typing.ParamSpec"):
       self._convert_paramspec_args(node)
-    elif node.func.id in _TYPING_NAMEDTUPLE_IDS:
+    elif self.defs.matches_type(node.func.id, "typing.NamedTuple"):
       self._convert_typing_namedtuple_args(node)
-    elif node.func.id in _COLL_NAMEDTUPLE_IDS:
+    elif self.defs.matches_type(node.func.id, "collections.namedtuple"):
       self._convert_collections_namedtuple_args(node)
-    elif node.func.id in _TYPEDDICT_IDS:
+    elif self.defs.matches_type(node.func.id, "typing.TypedDict"):
       self._convert_typed_dict_args(node)
-    elif node.func.id in _NEWTYPE_IDS:
+    elif self.defs.matches_type(node.func.id, "typing.NewType"):
       return self._convert_newtype_args(node)
 
   def visit_Call(self, node):
-    if node.func.id in _TYPEVAR_IDS:
+    if self.defs.matches_type(node.func.id, "typing.TypeVar"):
       if self.level > 0:
         raise ParseError("TypeVars need to be defined at module level")
       return _TypeVar.from_call(node)
-    elif node.func.id in _PARAMSPEC_IDS:
+    elif self.defs.matches_type(node.func.id, "typing.ParamSpec"):
       return _ParamSpec.from_call(node)
-    elif node.func.id in _TYPING_NAMEDTUPLE_IDS + _COLL_NAMEDTUPLE_IDS:
+    elif self.defs.matches_type(
+        node.func.id, ("typing.NamedTuple", "collections.namedtuple")):
       return self.defs.new_named_tuple(*node.args)
-    elif node.func.id in _TYPEDDICT_IDS:
+    elif self.defs.matches_type(node.func.id, "typing.TypedDict"):
       return self.defs.new_typed_dict(*node.args, node.keywords)
-    elif node.func.id in _NEWTYPE_IDS:
+    elif self.defs.matches_type(node.func.id, "typing.NewType"):
       return self.defs.new_new_type(*node.args)
     # Convert all other calls to NamedTypes; for example:
     # * typing.pyi uses things like

pytype/pyi/parser_test.py

@@ -559,6 +559,19 @@ class B:
           __match_args__: tuple
     """)
 
+  def test_typevar_alias(self):
+    self.check("""
+      from typing import TypeVar as _TypeVar
+      T = _TypeVar('T')
+      def f(x: T) -> T: ...
+    """, """
+      from typing import TypeVar, TypeVar as _TypeVar
+
+      T = TypeVar('T')
+
+      def f(x: T) -> T: ...
+    """)
+
 
 class QuotedTypeTest(parser_test_base.ParserTestBase):
 

pytype/pytd/CMakeLists.txt

@@ -158,7 +158,6 @@ py_library(
     .pytd_visitors
     pytype.utils
     pytype.platform_utils.platform_utils
-    pytype.pytd.parse.parse
 )
 
 py_library(

pytype/pytd/pytd_utils.py

@@ -22,7 +22,6 @@
 from pytype.pytd import printer
 from pytype.pytd import pytd
 from pytype.pytd import pytd_visitors
-from pytype.pytd.parse import parser_constants
 
 
 ANON_PARAM = re.compile(r"_[0-9]+")
@@ -439,30 +438,3 @@ def MergeBaseClass(cls, base):
                     decorators=decorators,
                     slots=slots,
                     template=cls.template or base.template)
-
-
-def MatchesFullName(t, full_name, current_module_name=None, aliases=None):
-  """Whether t.name matches full_name in format {module}.{member}."""
-  if isinstance(full_name, tuple):
-    return any(MatchesFullName(t, name, current_module_name, aliases)
-               for name in full_name)
-  expected_module_name, expected_name = full_name.rsplit(".", 1)
-  if current_module_name == expected_module_name:
-    # full_name is inside the current module, so check for the name without
-    # the module prefix.
-    return t.name == expected_name
-  elif "." not in t.name:
-    # full_name is not inside the current module, so a local type can't match.
-    return False
-  else:
-    module_name, name = t.name.rsplit(".", 1)
-    if aliases and module_name in aliases:
-      # Adjust the module name if it has been aliased with `import x as y`.
-      # See test_pyi.PYITest.testTypingAlias.
-      module = aliases[module_name].type
-      if isinstance(module, pytd.Module):
-        module_name = module.module_name
-    expected_module_names = {
-        expected_module_name,
-        parser_constants.EXTERNAL_NAME_PREFIX + expected_module_name}
-    return module_name in expected_module_names and name == expected_name