Repository: cthoyt/class-resolver

Issue:
ClassResolver.from_subclasses does not collect __main__ classes
Since `ClassResolver.from_subclasses` excludes classes from private modules by default, and the name of the module from which a Python script is invoked is replaced by `__main__`, classes defined there are excluded, since they are considered private as their module name starts with `_`.

cf. https://github.com/pykeen/pykeen/commit/442e586cb4f2d67b644fde867db91a9819aa507e

Comments:

Patch:
diff --git a/setup.cfg b/setup.cfg
index 35f325e..5b5dd0f 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -69,6 +69,7 @@ docs =
     sphinx-autodoc-typehints
     sphinx_automodapi
 tests =
+    docdata
     coverage
     pytest
 docdata =
diff --git a/src/class_resolver/utils.py b/src/class_resolver/utils.py
index ca53bf7..c2ae7e4 100644
--- a/src/class_resolver/utils.py
+++ b/src/class_resolver/utils.py
@@ -64,10 +64,35 @@ OneOrManyHintOrType = Optional[OneOrSequence[HintOrType[X]]]
 OneOrManyOptionalKwargs = Optional[OneOrSequence[OptionalKwargs]]
 
 
+def is_private(class_name: str, module_name: str, main_is_private: bool = True) -> bool:
+    """
+    Decide whether a class in a module is considered private.
+
+    :param class_name:
+        the class name, i.e., `cls.__name__`
+    :param module_name:
+        the module name, i.e., `cls.__module__`
+    :param main_is_private:
+        whether the `__main__` module is considered private
+
+    :return:
+        whether the class should be considered private
+    """
+    # note: this method has been separated for better testability
+    if class_name.startswith("_"):
+        return True
+    if not main_is_private and module_name.startswith("__main__"):
+        return False
+    if any(part.startswith("_") for part in module_name.split(".")):
+        return True
+    return False
+
+
 def get_subclasses(
     cls: Type[X],
     exclude_private: bool = True,
     exclude_external: bool = True,
+    main_is_private: bool = True,
 ) -> Iterable[Type[X]]:
     """Get all subclasses.
 
@@ -77,15 +102,17 @@ def get_subclasses(
         done when having shadow duplicate classes implemented in C
     :param exclude_external: If true, will exclude any class that does not originate
         from the same package as the base class.
+    :param main_is_private: If true, __main__ is considered a private module.
     :yields: Descendant classes of the ancestor class
     """
     for subclass in cls.__subclasses__():
         yield from get_subclasses(subclass)
-        if exclude_private:
-            if any(part.startswith("_") for part in subclass.__module__.split(".")):
-                continue
-            if subclass.__name__.startswith("_"):
-                continue
+        if exclude_private and is_private(
+            class_name=subclass.__name__,
+            module_name=subclass.__module__,
+            main_is_private=main_is_private,
+        ):
+            continue
         if exclude_external and not same_module(cls, subclass):
             continue
         yield subclass


Test patch:
diff --git a/tests/test_utils.py b/tests/test_utils.py
index e0b7278..3dc141f 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -6,13 +6,30 @@ import enum
 import unittest
 from collections import Counter, defaultdict
 
-from class_resolver.utils import get_subclasses, normalize_with_default, same_module
+from class_resolver.utils import (
+    get_subclasses,
+    is_private,
+    normalize_with_default,
+    same_module,
+)
 from tests._private_extras import PrivateDict
 
 
 class TestUtilities(unittest.TestCase):
     """Test utilities."""
 
+    def test_is_private(self):
+        """Test whether a module should be considered private."""
+        self.assertFalse(is_private("A", "", main_is_private=False))
+        self.assertFalse(is_private("A", "", main_is_private=True))
+        self.assertTrue(is_private("_module", "", main_is_private=True))
+        self.assertTrue(is_private("_module", "", main_is_private=True))
+        self.assertTrue(is_private("A", "B._module", main_is_private=True))
+        self.assertTrue(is_private("A", "__main__", main_is_private=True))
+        self.assertFalse(is_private("A", "__main__", main_is_private=False))
+        self.assertTrue(is_private("_A", "__main__", main_is_private=True))
+        self.assertTrue(is_private("_A", "__main__", main_is_private=False))
+
     def test_same_module(self):
         """Test getting subclasses."""
         self.assertFalse(same_module(Counter, dict))