Copy over get_all_subclasses (#39)

src/lightning_utilities/core/inheritance.py

@@ -0,0 +1,14 @@
+from typing import Iterator, Set, Type
+
+
+def get_all_subclasses_iterator(cls: Type) -> Iterator[Type]:
+    def recurse(cl: Type) -> Iterator[Type]:
+        for subclass in cl.__subclasses__():
+            yield subclass
+            yield from recurse(subclass)
+
+    yield from recurse(cls)
+
+
+def get_all_subclasses(cls: Type) -> Set[Type]:
+    return set(get_all_subclasses_iterator(cls))

tests/unittests/core/test_inheritance.py

@@ -0,0 +1,24 @@
+from lightning_utilities.core.inheritance import get_all_subclasses
+
+
+def test_get_all_subclasses():
+    class A1:
+        ...
+
+    class A2(A1):
+        ...
+
+    class B1:
+        ...
+
+    class B2(B1):
+        ...
+
+    class C(A2, B2):
+        ...
+
+    assert get_all_subclasses(A1) == {A2, C}
+    assert get_all_subclasses(A2) == {C}
+    assert get_all_subclasses(B1) == {B2, C}
+    assert get_all_subclasses(B2) == {C}
+    assert get_all_subclasses(C) == set()