def testInPlaceLookupExternalClassesByFullName(self):
src1 = textwrap.dedent("""
def f1() -> bar.Bar
class Foo(object):
pass
""")
src2 = textwrap.dedent("""
def f2() -> foo.Foo
class Bar(object):
pass
""")
ast1 = self.Parse(src1).Replace(name="foo").Visit(
visitors.AddNamePrefix())
ast2 = self.Parse(src2).Replace(name="bar").Visit(
visitors.AddNamePrefix())
ast1 = ast1.Visit(
visitors.LookupExternalTypes(dict(foo=ast1, bar=ast2),
full_names=True))
ast2 = ast2.Visit(
visitors.LookupExternalTypes(dict(foo=ast1, bar=ast2),
full_names=True))
f1, = ast1.Lookup("foo.f1").signatures
f2, = ast2.Lookup("bar.f2").signatures
self.assertIs(ast2.Lookup("bar.Bar"), f1.return_type.cls)
self.assertIs(ast1.Lookup("foo.Foo"), f2.return_type.cls)
def _resolve_external_types(self, ast):
try:
ast = ast.Visit(visitors.LookupExternalTypes(
self._get_module_map(), full_names=True, self_name=ast.name))
except KeyError as e:
raise BadDependencyError(e.message, ast.name)
return ast
def testLookupTwoStarAliasesWithDefaultPyi(self):
src1 = "def __getattr__(name) -> ?"
src2 = "def __getattr__(name) -> ?"
src3 = textwrap.dedent("""
from foo import *
from bar import *
""")
ast1 = self.Parse(src1).Replace(name="foo").Visit(
visitors.AddNamePrefix())
ast2 = self.Parse(src2).Replace(name="bar").Visit(
visitors.AddNamePrefix())
ast3 = self.Parse(src3).Replace(name="baz").Visit(
visitors.AddNamePrefix())
ast3 = ast3.Visit(
visitors.LookupExternalTypes(
{
"foo": ast1,
"bar": ast2,
"baz": ast3
},
full_names=True,
self_name="baz"))
self.assertMultiLineEqual(
pytd.Print(ast3),
textwrap.dedent("""\
from typing import Any
def baz.__getattr__(name) -> Any: ..."""))
def testInPlaceFillInExternalClasses(self):
src1 = textwrap.dedent("""
def f1() -> bar.Bar
class Foo(object):
pass
""")
src2 = textwrap.dedent("""
def f2() -> foo.Foo
class Bar(object):
pass
""")
ast1 = self.Parse(src1)
ast2 = self.Parse(src2)
ast1 = ast1.Visit(visitors.LookupExternalTypes(dict(foo=ast1, bar=ast2)))
ast2 = ast2.Visit(visitors.LookupExternalTypes(dict(foo=ast1, bar=ast2)))
f1, = ast1.Lookup("f1").signatures
f2, = ast2.Lookup("f2").signatures
self.assertIs(ast2.Lookup("Bar"), f1.return_type.cls)
self.assertIs(ast1.Lookup("Foo"), f2.return_type.cls)
def ParseWithBuiltins(self, src):
ast = parser.TypeDeclParser().Parse(textwrap.dedent(src))
b, t = builtins.GetBuiltinsAndTyping()
ast = ast.Visit(visitors.LookupExternalTypes(
{"__builtin__": b, "typing": t}, full_names=True))
ast = ast.Visit(visitors.NamedTypeToClassType())
ast = visitors.AdjustTypeParameters(ast)
ast.Visit(visitors.FillInModuleClasses({"": ast}))
ast.Visit(visitors.VerifyVisitor())
return ast
def testGetBasesInMRO(self):
ast = self.parser.Parse(textwrap.dedent("""
T = TypeVar("T")
class Foo(Generic[T]): pass
class Bar(Foo[int]): pass
"""))
b, t = builtins.GetBuiltinsAndTyping()
ast = ast.Visit(visitors.LookupExternalTypes(
{"__builtin__": b, "typing": t}, full_names=True))
ast = ast.Visit(visitors.NamedTypeToClassType())
mro = utils.GetBasesInMRO(ast.Lookup("Bar"), lookup_ast=ast)
self.assertListEqual(["Foo", "typing.Generic", "__builtin__.object"],
[t.name for t in mro])
def ParseWithBuiltins(self, src):
ast = parser.parse_string(textwrap.dedent(src))
b, t = builtins.GetBuiltinsAndTyping()
ast = ast.Visit(
visitors.LookupExternalTypes({
"__builtin__": b,
"typing": t
},
full_names=True))
ast = ast.Visit(visitors.NamedTypeToClassType())
ast = ast.Visit(visitors.AdjustTypeParameters())
ast.Visit(visitors.FillInLocalPointers({"": ast, "__builtin__": b}))
ast.Visit(visitors.VerifyVisitor())
return ast
def ProcessAst(serializable_ast, module_map):
"""Postprocess a pickled ast.
Postprocessing will either just fill the ClassType references from module_map
or if module_name changed between pickling and loading rename the module
internal references to the new module_name.
Renaming is more expensive than filling references, as the whole AST needs to
be rebuild.
Args:
serializable_ast: A SerializableAst instance.
module_map: Used to resolve ClassType.cls links to already loaded modules.
The loaded module will be added to the dict.
Returns:
A pytd.TypeDeclUnit, this is either the input raw_ast with the references
set or a newly created AST with the new module_name and the references set.
Raises:
AssertionError: If module_name is already in module_map, which means that
module_name is already loaded.
UnrestorableDependencyError: If no concrete module exists in module_map for
one of the references from the pickled ast.
"""
raw_ast = serializable_ast.ast
module_map[raw_ast.name] = raw_ast
# Notice that this is also resolving local ClassType references.
class_lookup = visitors.LookupExternalTypes(module_map,
full_names=True,
self_name=None)
if serializable_ast.class_type_nodes:
for node in serializable_ast.class_type_nodes:
try:
if node is not class_lookup.VisitClassType(node):
serializable_ast = serializable_ast.Replace(
class_type_nodes=None)
break
except KeyError as e:
raise UnrestorableDependencyError("Unresolved class: %r." %
e.message)
if serializable_ast.class_type_nodes is None:
try:
raw_ast = raw_ast.Visit(class_lookup)
except KeyError as e:
raise UnrestorableDependencyError("Unresolved class: %r." %
e.message)
return raw_ast
def testLookupConstant(self):
src1 = textwrap.dedent("""
Foo = ... # type: type
""")
src2 = textwrap.dedent("""
class Bar(object):
bar = ... # type: foo.Foo
""")
ast1 = self.Parse(src1)
ast2 = self.Parse(src2)
ast2 = ast2.Visit(
visitors.LookupExternalTypes({
"foo": ast1,
"bar": ast2
}))
self.assertEqual(
ast2.Lookup("Bar").constants[0],
pytd.Constant(name="bar", type=pytd.AnythingType()))
def GetBuiltinsAndTyping(python_version):
"""Get __builtin__.pytd and typing.pytd."""
assert python_version
global _cached_builtins_pytd
if _cached_builtins_pytd.cache:
assert _cached_builtins_pytd.version == python_version
else:
t = parser.parse_string(_FindBuiltinFile("typing", python_version),
name="typing",
python_version=python_version)
b = parser.parse_string(_FindBuiltinFile("__builtin__",
python_version),
name="__builtin__",
python_version=python_version)
b = b.Visit(
visitors.LookupExternalTypes({"typing": t},
full_names=True,
self_name="__builtin__"))
t = t.Visit(visitors.LookupBuiltins(b))
b = b.Visit(visitors.NamedTypeToClassType())
t = t.Visit(visitors.NamedTypeToClassType())
b = b.Visit(visitors.AdjustTypeParameters())
t = t.Visit(visitors.AdjustTypeParameters())
b = b.Visit(visitors.CanonicalOrderingVisitor())
t = t.Visit(visitors.CanonicalOrderingVisitor())
b.Visit(
visitors.FillInLocalPointers({
"": b,
"typing": t,
"__builtin__": b
}))
t.Visit(
visitors.FillInLocalPointers({
"": t,
"typing": t,
"__builtin__": b
}))
b.Visit(visitors.VerifyLookup())
t.Visit(visitors.VerifyLookup())
b.Visit(visitors.VerifyContainers())
t.Visit(visitors.VerifyContainers())
_cached_builtins_pytd = Cache(python_version, (b, t))
return _cached_builtins_pytd.cache
def testGetBasesInMRO(self):
ast = parser.parse_string(textwrap.dedent("""
from typing import Generic, TypeVar
T = TypeVar("T")
class Foo(Generic[T]): pass
class Bar(Foo[int]): pass
"""),
python_version=self.PYTHON_VERSION)
b, t = builtins.GetBuiltinsAndTyping(self.PYTHON_VERSION)
ast = ast.Visit(
visitors.LookupExternalTypes({
"__builtin__": b,
"typing": t
},
full_names=True))
ast = ast.Visit(visitors.NamedTypeToClassType())
bases = mro.GetBasesInMRO(ast.Lookup("Bar"), lookup_ast=ast)
self.assertListEqual(["Foo", "typing.Generic", "__builtin__.object"],
[t.name for t in bases])
def ParseWithBuiltins(self, src):
ast = parser.parse_string(textwrap.dedent(src),
python_version=self.PYTHON_VERSION)
ast = ast.Visit(
visitors.LookupExternalTypes(
{
"__builtin__": self.loader.builtins,
"typing": self.loader.typing
},
full_names=True))
ast = ast.Visit(visitors.NamedTypeToClassType())
ast = ast.Visit(visitors.AdjustTypeParameters())
ast.Visit(
visitors.FillInLocalPointers({
"": ast,
"__builtin__": self.loader.builtins
}))
ast.Visit(visitors.VerifyVisitor())
return ast
def testLookupStarAliasInUnnamedModule(self):
src1 = textwrap.dedent("""
class A(object): ...
""")
src2 = "from foo import *"
ast1 = self.Parse(src1).Replace(name="foo").Visit(
visitors.AddNamePrefix())
ast2 = self.Parse(src2)
name = ast2.name
ast2 = ast2.Visit(
visitors.LookupExternalTypes({"foo": ast1},
full_names=True,
self_name=None))
self.assertEqual(name, ast2.name)
self.assertMultiLineEqual(
pytd.Print(ast2),
textwrap.dedent("""\
import foo
A = foo.A"""))
def _load_and_resolve_ast_dependencies(self, ast, ast_name=None):
"""Fill in all ClassType.cls pointers."""
deps = visitors.CollectDependencies()
ast.Visit(deps)
if deps.modules:
for name in deps.modules:
if name not in self._modules:
other_ast = self._import_name(name)
if other_ast is None:
raise DependencyNotFoundError(name, ast_name
or ast.name)
module_map = {
name: module.ast
for name, module in self._modules.items()
}
ast = ast.Visit(
visitors.LookupExternalTypes(module_map,
full_names=True,
self_name=ast_name))
return ast
def _LookupClassReferences(serializable_ast, module_map, self_name):
"""Fills .cls references in serializable_ast.ast with ones from module_map.
Already filled references are not changed. References to the module self._name
are not filled. Setting self_name=None will fill all references.
Args:
serializable_ast: A SerializableAst instance.
module_map: Used to resolve ClassType.cls links to already loaded modules.
The loaded module will be added to the dict.
self_name: A string representation of a module which should not be resolved,
for example: "foo.bar.module1" or None to resolve all modules.
Returns:
A SerializableAst with an updated .ast. .class_type_nodes is set to None
if any of the Nodes needed to be regenerated.
"""
class_lookup = visitors.LookupExternalTypes(module_map,
full_names=True,
self_name=self_name)
raw_ast = serializable_ast.ast
if serializable_ast.class_type_nodes:
for node in serializable_ast.class_type_nodes:
try:
if node is not class_lookup.VisitClassType(node):
serializable_ast = serializable_ast.Replace(
class_type_nodes=None)
break
except KeyError as e:
raise UnrestorableDependencyError("Unresolved class: %r." %
e.message)
if serializable_ast.class_type_nodes is None:
try:
raw_ast = raw_ast.Visit(class_lookup)
except KeyError as e:
raise UnrestorableDependencyError("Unresolved class: %r." %
e.message)
serializable_ast = serializable_ast.Replace(ast=raw_ast)
return serializable_ast
def testLookupStarAliasWithDifferentGetAttr(self):
src1 = "def __getattr__(name) -> int"
src2 = textwrap.dedent("""
from foo import *
def __getattr__(name) -> str
""")
ast1 = self.Parse(src1).Replace(name="foo").Visit(
visitors.AddNamePrefix())
ast2 = self.Parse(src2).Replace(name="bar").Visit(
visitors.AddNamePrefix())
ast2 = ast2.Visit(
visitors.LookupExternalTypes({
"foo": ast1,
"bar": ast2
},
full_names=True,
self_name="bar"))
self.assertMultiLineEqual(
pytd.Print(ast2),
textwrap.dedent("""\
def bar.__getattr__(name) -> str: ..."""))
def testLookupTwoStarAliases(self):
src1 = "class A(object): ..."
src2 = "class B(object): ..."
src3 = textwrap.dedent("""
from foo import *
from bar import *
""")
ast1 = self.Parse(src1).Replace(name="foo").Visit(
visitors.AddNamePrefix())
ast2 = self.Parse(src2).Replace(name="bar").Visit(
visitors.AddNamePrefix())
ast3 = self.Parse(src3).Replace(name="baz").Visit(
visitors.AddNamePrefix())
ast3 = ast3.Visit(
visitors.LookupExternalTypes(
{
"foo": ast1,
"bar": ast2,
"baz": ast3
},
full_names=True,
self_name="baz"))
self.assertSetEqual({a.name for a in ast3.aliases}, {"baz.A", "baz.B"})
def testLookupTwoStarAliasesWithSameClass(self):
src1 = "class A(object): ..."
src2 = "class A(object): ..."
src3 = textwrap.dedent("""
from foo import *
from bar import *
""")
ast1 = self.Parse(src1).Replace(name="foo").Visit(
visitors.AddNamePrefix())
ast2 = self.Parse(src2).Replace(name="bar").Visit(
visitors.AddNamePrefix())
ast3 = self.Parse(src3).Replace(name="baz").Visit(
visitors.AddNamePrefix())
self.assertRaises(
KeyError, ast3.Visit,
visitors.LookupExternalTypes(
{
"foo": ast1,
"bar": ast2,
"baz": ast3
},
full_names=True,
self_name="baz"))
def _load_and_resolve_ast_dependencies(self, ast, ast_name=None):
"""Fill in all ClassType.cls pointers."""
deps = visitors.CollectDependencies()
ast.Visit(deps)
if deps.modules:
for name in deps.modules:
if name not in self._modules:
other_ast = self._import_name(name)
if other_ast is None:
error = "Can't find pyi for %r" % name
raise BadDependencyError(error, ast_name or ast.name)
module_map = {
name: module.ast
for name, module in self._modules.items()
}
try:
ast = ast.Visit(
visitors.LookupExternalTypes(module_map,
full_names=True,
self_name=ast_name))
except KeyError as e:
raise BadDependencyError(e.message, ast_name or ast.name)
return ast
def testLookupTwoStarAliasesWithDifferentGetAttrs(self):
src1 = "def __getattr__(name) -> int"
src2 = "def __getattr__(name) -> str"
src3 = textwrap.dedent("""
from foo import *
from bar import *
""")
ast1 = self.Parse(src1).Replace(name="foo").Visit(
visitors.AddNamePrefix())
ast2 = self.Parse(src2).Replace(name="bar").Visit(
visitors.AddNamePrefix())
ast3 = self.Parse(src3).Replace(name="baz").Visit(
visitors.AddNamePrefix())
self.assertRaises(
KeyError, ast3.Visit,
visitors.LookupExternalTypes(
{
"foo": ast1,
"bar": ast2,
"baz": ast3
},
full_names="True",
self_name="baz"))
def GetBuiltinsAndTyping():
"""Get __builtin__.pytd and typing.pytd."""
global _cached_builtins_pytd
if not _cached_builtins_pytd:
t = parser.TypeDeclParser().Parse(_FindBuiltinFile("typing"),
name="typing")
t = t.Visit(visitors.AddNamePrefix())
b = parser.TypeDeclParser().Parse(_FindBuiltinFile("__builtin__"),
name="__builtin__")
b = b.Visit(visitors.AddNamePrefix())
b = b.Visit(visitors.NamedTypeToClassType())
b = b.Visit(
visitors.LookupExternalTypes({"typing": t},
full_names=True,
self_name="__builtin__"))
t = t.Visit(visitors.LookupBuiltins(b))
t = t.Visit(visitors.NamedTypeToClassType())
b = visitors.AdjustTypeParameters(b)
t = visitors.AdjustTypeParameters(t)
b.Visit(
visitors.FillInModuleClasses({
"": b,
"typing": t,
"__builtin__": b
}))
t.Visit(
visitors.FillInModuleClasses({
"": t,
"typing": t,
"__builtin__": b
}))
b.Visit(visitors.VerifyLookup())
t.Visit(visitors.VerifyLookup())
b.Visit(visitors.VerifyContainers())
t.Visit(visitors.VerifyContainers())
_cached_builtins_pytd = b, t
return _cached_builtins_pytd
def testLookupStarAliasWithDuplicateClass(self):
src1 = "class A(object): ..."
src2 = textwrap.dedent("""
from foo import *
class A(object):
x = ... # type: int
""")
ast1 = self.Parse(src1).Replace(name="foo").Visit(
visitors.AddNamePrefix())
ast2 = self.Parse(src2).Replace(name="bar").Visit(
visitors.AddNamePrefix())
ast2 = ast2.Visit(
visitors.LookupExternalTypes({
"foo": ast1,
"bar": ast2
},
full_names=True,
self_name="bar"))
self.assertMultiLineEqual(
pytd.Print(ast2),
textwrap.dedent("""\
class bar.A(object):
x = ... # type: int
"""))
def GetBuiltinsAndTyping():
"""Get __builtin__.pytd and typing.pytd."""
global _cached_builtins_pytd
if not _cached_builtins_pytd:
t = parser.parse_string(_FindBuiltinFile("typing"), name="typing")
b = parser.parse_string(_FindBuiltinFile("__builtin__"),
name="__builtin__")
b = b.Visit(
visitors.LookupExternalTypes({"typing": t},
full_names=True,
self_name="__builtin__"))
t = t.Visit(visitors.LookupBuiltins(b))
b = b.Visit(visitors.NamedTypeToClassType())
t = t.Visit(visitors.NamedTypeToClassType())
b = b.Visit(visitors.AdjustTypeParameters())
t = t.Visit(visitors.AdjustTypeParameters())
b = b.Visit(visitors.CanonicalOrderingVisitor())
t = t.Visit(visitors.CanonicalOrderingVisitor())
b.Visit(
visitors.FillInLocalPointers({
"": b,
"typing": t,
"__builtin__": b
}))
t.Visit(
visitors.FillInLocalPointers({
"": t,
"typing": t,
"__builtin__": b
}))
b.Visit(visitors.VerifyLookup())
t.Visit(visitors.VerifyLookup())
b.Visit(visitors.VerifyContainers())
t.Visit(visitors.VerifyContainers())
_cached_builtins_pytd = b, t
return _cached_builtins_pytd
def testLookupStarAlias(self):
src1 = textwrap.dedent("""
x = ... # type: int
T = TypeVar("T")
class A(object): ...
def f(x: T) -> T: ...
B = A
""")
src2 = "from foo import *"
ast1 = self.Parse(src1).Replace(name="foo").Visit(
visitors.AddNamePrefix())
ast2 = self.Parse(src2).Replace(name="bar").Visit(
visitors.AddNamePrefix())
ast2 = ast2.Visit(
visitors.LookupExternalTypes({
"foo": ast1,
"bar": ast2
},
full_names=True,
self_name="bar"))
self.assertEqual("bar", ast2.name)
self.assertSetEqual({a.name
for a in ast2.aliases},
{"bar.x", "bar.T", "bar.A", "bar.f", "bar.B"})
