Fix AttributeError when __module__ raises exception in Schema class dependencies resolution.

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

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):