Add support for using the module dir as base_dir.

Currently, if we try to document `/path/module/__init__.py` using `base_dir=['/path/module']`, everything under `/path/module` is trimmed for being outside of the base dir.

This is in prep for adding more modules to the TF Lite Support library, where submodules come from disparate parts of the repo.

PiperOrigin-RevId: 446334715

Author: markmcd

tools/tensorflow_docs/api_generator/public_api.py

@@ -331,8 +331,9 @@ def __call__(self, path: Sequence[str], parent: Any,
       # contents will get filtered when the submodules are checked.
       if len(mod_base_dirs) == 1:
         mod_base_dir = mod_base_dirs[0]
-        # Check that module is in one of the `self._base_dir`s
-        if not any(base in mod_base_dir.parents for base in self.base_dirs):
+        # Check that module is, or is in one of the `self._base_dir`s
+        if not (any(base in mod_base_dir.parents for base in self.base_dirs) or
+                mod_base_dir in self.base_dirs):
           continue
       yield name, child
 

tools/tensorflow_docs/api_generator/public_api_test.py

@@ -210,6 +210,21 @@ def test_filter_base_dirs(self):
                   ('sub2', module.sub2)])
     self.assertEqual([('a', module.a), ('sub1', module.sub1)], list(result))
 
+  def test_filter_base_dir_pointing_to_submodule_dir(self):
+    module = types.ModuleType('module')
+    module.__file__ = '/1/2/3/module'
+    module.submodule = types.ModuleType('submodule')
+    module.submodule.__file__ = '/1/2/3/submodule/__init__.py'
+
+    test_filter = public_api.FilterBaseDirs(
+        base_dirs=[pathlib.Path('/1/2/3/submodule')])
+    result = test_filter(
+        path=('module',),
+        parent=module,
+        children=[('submodule', module.submodule)])
+
+    self.assertEqual([('submodule', module.submodule)], list(result))
+
 
 if __name__ == '__main__':
   absltest.main()