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())
self.assertItemsEqual(["classobj"],
[t.name for t in data[ast.Lookup("A")]])
self.assertItemsEqual(["classobj"],
[t.name for t in data[ast.Lookup("B")]])
self.assertItemsEqual(["A"], [t.name for t in data[ast.Lookup("C")]])
self.assertItemsEqual(["A", "B"],
[t.name for t in data[ast.Lookup("D")]])
self.assertItemsEqual(["C", "D", "A"],
[t.name for t in data[ast.Lookup("E")]])
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)
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()))
hierarchy = {cls: [superclass for superclass in superclasses
if (hasattr(superclass, "name") and
is_complete(superclass))]
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 abc_hierarchy.Invert(hierarchy)