5
\$\begingroup\$

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()
\$\endgroup\$
2
\$\begingroup\$
  • 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 ...
\$\endgroup\$
  • 1
    \$\begingroup\$ Nice. Thanks for this. I'll post more questions when I can. \$\endgroup\$ – 422_unprocessable_entity Apr 9 at 18:28

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.