List subclasses of a class

Question

Description

• Given a class, return all its subclasses (recursively).
• As you can see I've eliminated recursion using a stack.

What I want reviewed

• Is there a better way to do this?
• How can I make this code more generic and easier to use?
• Is it pythonic?
• Better way to eliminate recursion?

Code

def all_subclasses(cls):
 if cls == type:
 raise ValueError("Invalid class - 'type' is not a class")
 subclasses = set()
 stack = []
 try:
 immediate_subclasses = cls.__subclasses__()
 except (TypeError, AttributeError) as ex:
 raise ValueError("Invalid class" + repr(cls)) from ex
 for subclass in immediate_subclasses:
 stack.append(subclass) 
 while stack:
 sub = stack.pop()
 subclasses.add(sub)
 try:
 sub_subclasses = sub.__subclasses__()
 except (TypeError, AttributeError) as _:
 continue
 if sub_subclasses:
 stack.extend(sub_subclasses)
 return list(subclasses)

Tests

import unittest
from class_util import all_subclasses
def names(classes): 
 return sorted([cls.__name__ for cls in classes])
class A:
 @classmethod
 def all_subclasses(cls):
 return all_subclasses(cls)
class B(A):
 pass
class C(B):
 pass
class D(C):
 pass
class E(C):
 pass
class F(E, C):
 pass
class AllSublassesTestCase(unittest.TestCase):
 def test_nested_classes(self):
 self.assertEqual(names(A.all_subclasses()), ["B", "C", "D", "E", "F"])
 def test_work_with_buitins(self):
 self.assertTrue(names(all_subclasses(dict)))
 self.assertTrue(names(all_subclasses(tuple))) 
 self.assertTrue(names(all_subclasses(list)))
 def test_value_error_is_raised_on_invalid_classes(self):
 self.assertRaises(ValueError, all_subclasses, type)
 self.assertRaises(ValueError, all_subclasses, "")
 self.assertRaises(ValueError, all_subclasses, None)
 self.assertRaises(ValueError, all_subclasses, [])
if __name__ == "__main__":
 unittest.main()

Answer 1

• While working on the stack you are using stack.extend. You can also use this in the part where you add the immediate subclasses.
• There is no need to check if a list is empty before using extend, if it is empty it will just do nothing.
• If you don't need the exception, just don't catch it with as _.
• Not sure if you should be doing if cls is type instead of if cls == type.

def all_subclasses(cls):
 if cls == type:
 raise ValueError("Invalid class - 'type' is not a class")
 subclasses = set()
 stack = []
 try:
 stack.extend(cls.__subclasses__())
 except (TypeError, AttributeError) as ex:
 raise ValueError("Invalid class" + repr(cls)) from ex 
 while stack:
 sub = stack.pop()
 subclasses.add(sub)
 try:
 stack.extend(sub.__subclasses__())
 except (TypeError, AttributeError):
 continue
 return list(subclasses)

One way to optimize this further is to make sure you don't visit a class multiple times:

 while stack:
 sub = stack.pop()
 subclasses.add(sub)
 try:
 stack.extend(s for s in sub.__subclasses__() if s not in subclasses)
 except (TypeError, AttributeError):
 continue

This should prevent having to visit (almost) every class twice with convoluted hierarchies like this:

 A
 / \
 B C
 \ /
 D
 / | | \
E F G ...

• JaDogg: Nice. Thanks for this. I'll post more questions when I can.