diff --git a/langfun/core/structured/schema/base.py b/langfun/core/structured/schema/base.py index 62d82430..cafc3b4d 100644 --- a/langfun/core/structured/schema/base.py +++ b/langfun/core/structured/schema/base.py @@ -349,7 +349,8 @@ def _add_dependency(cls_or_classes): def _fill_dependencies(vs: pg.typing.ValueSpec, include_subclasses: bool): if isinstance(vs, pg.typing.Object): cls = vs.cls - if cls.__module__ == 'builtins': + if (not include_generated_subclasses + and getattr(cls, '__module__', 'builtins') == 'builtins'): return if cls not in seen: @@ -376,7 +377,8 @@ def _fill_dependencies(vs: pg.typing.ValueSpec, include_subclasses: bool): # polluting the generation space, classes dynamically created by # 'eval' (which have __module__ == 'builtins') are excluded from # dependencies by default. - if ((include_generated_subclasses or subcls.__module__ != 'builtins') + if ((include_generated_subclasses + or getattr(subcls, '__module__', 'builtins') != 'builtins') and subcls not in dependencies): _fill_dependencies( pg.typing.Object(subcls), include_subclasses=True diff --git a/langfun/core/structured/schema/base_test.py b/langfun/core/structured/schema/base_test.py index 6c31e178..9bdee71e 100644 --- a/langfun/core/structured/schema/base_test.py +++ b/langfun/core/structured/schema/base_test.py @@ -261,6 +261,27 @@ class C(A): self.assertEqual(base.class_dependencies(1), []) + def test_class_dependencies_missing_module(self): + class Meta(type): + @property + def __module__(cls): + raise AttributeError('Testing AttributeError') + + class Base: + pass + + class Sub(Base, metaclass=Meta): + pass + + self.assertEqual( + base.class_dependencies(Base, include_generated_subclasses=True), + [Base, Sub], + ) + self.assertEqual( + base.class_dependencies(Base, include_generated_subclasses=False), + [Base], + ) + class AnnotationTest(unittest.TestCase):