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)