def testSuperClasses(self):
src = textwrap.dedent("""
class classobj:
pass
class A():
pass
class B():
pass
class C(A):
pass
class D(A,B):
pass
class E(C,D,A):
pass
""")
ast = visitors.LookupClasses(self.Parse(src))
data = ast.Visit(visitors.ExtractSuperClasses())
six.assertCountEqual(self,
["classobj"], [t.name for t in data[ast.Lookup("A")]])
six.assertCountEqual(self,
["classobj"], [t.name for t in data[ast.Lookup("B")]])
six.assertCountEqual(self, ["A"], [t.name for t in data[ast.Lookup("C")]])
six.assertCountEqual(self,
["A", "B"], [t.name for t in data[ast.Lookup("D")]])
six.assertCountEqual(self, ["C", "D", "A"],
[t.name for t in data[ast.Lookup("E")]])
def get_all_subclasses(asts):
"""Compute a class->subclasses mapping.
Args:
asts: A list of ASTs.
Returns:
A dictionary, mapping instances of pytd.Type (types) to lists of
pytd.Class (the derived classes).
"""
hierarchy = {}
for ast in asts:
hierarchy.update(ast.Visit(visitors.ExtractSuperClasses()))
def filter_superclasses(superclasses):
return [
superclass for superclass in superclasses
if hasattr(superclass, "name") and is_complete(superclass)
]
hierarchy = {
cls: filter_superclasses(superclasses)
for cls, superclasses in hierarchy.items() if is_complete(cls)
}
# typically this is a fairly short list, e.g.:
# [ClassType(basestring), ClassType(int), ClassType(object)]
return utils.invert_dict(hierarchy)
def GetAllSubClasses(ast):
"""Compute a class->subclasses mapping.
Args:
ast: Parsed PYTD.
Returns:
A dictionary, mapping instances of pytd.TYPE (types) to lists of
pytd.Class (the derived classes).
"""
hierarchy = ast.Visit(visitors.ExtractSuperClasses())
hierarchy = {cls: [superclass for superclass in superclasses]
for cls, superclasses in hierarchy.items()}
return utils.invert_dict(hierarchy)
