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)