Added a new tensorflow_docs public filter and removed the _allowed_symbols variables from the TFF packages.

PiperOrigin-RevId: 306446724

Author: michaelreneer

tools/tensorflow_docs/api_generator/public_api.py

@@ -14,12 +14,13 @@
 # limitations under the License.
 # ==============================================================================
 """Visitor restricting traversal to only the public tensorflow API."""
-import inspect
 
+import ast
+import inspect
+import typing
 
 from tensorflow_docs.api_generator import doc_controls
 
-import typing
 _TYPING = frozenset(id(value) for value in typing.__dict__.values())
 
 
@@ -121,6 +122,77 @@ def util_2
  return filtered_children
 
 
+def _get_imported_symbols(obj):
+ """Returns a list of symbol names imported by the given `obj`."""
+
+ class ImportNodeVisitor(ast.NodeVisitor):
+ """An `ast.Visitor` that collects the names of imported symbols."""
+
+ def __init__(self):
+ self.imported_symbols = []
+
+ def _add_imported_symbol(self, node):
+ self.imported_symbols.extend([alias.name for alias in node.names])
+
+ def visit_Import(self, node): # pylint: disable=invalid-name
+ self._add_imported_symbol(node)
+
+ def visit_ImportFrom(self, node): # pylint: disable=invalid-name
+ self._add_imported_symbol(node)
+
+ source = inspect.getsource(obj)
+ tree = ast.parse(source)
+ visitor = ImportNodeVisitor()
+ visitor.visit(tree)
+ return visitor.imported_symbols
+
+
+def explicit_package_contents_filter(path, parent, children):
+ """Filter modules to only include explicit contents.
+
+ This function returns the children explicitly included by this module, meaning
+ that it will exclude:
+
+ * Modules in a package not explicitly imported by the package (submodules
+ are implicitly injected into their parent's namespace).
+ * Modules imported by a module that is not a package.
+
+ This filter is useful if you explicitly define your API in the packages of
+ your library, but do not expliticly define that API in the `__all__` variable
+ of each module. The purpose is to make it easier to maintain that API.
+
+ Note: This filter does work with wildcard imports, however it is generally not
+ recommended to use wildcard imports.
+
+ Args:
+ path: A tuple of names forming the path to the object.
+ parent: The parent object.
+ children: A list of (name, value) tuples describing the attributes of the
+ patent.
+
+ Returns:
+ A filtered list of children `(name, value)` pairs.
+ """
+ del path # Unused
+ is_parent_module = inspect.ismodule(parent)
+ is_parent_package = is_parent_module and hasattr(parent, '__path__')
+ if is_parent_package:
+ imported_symbols = _get_imported_symbols(parent)
+ filtered_children = []
+ for child in children:
+ name, obj = child
+ if inspect.ismodule(obj):
+ # Do not include modules in a package not explicitly imported by the
+ # package.
+ if is_parent_package and name not in imported_symbols:
+ continue
+ # Do not include modules imported by a module that is not a package.
+ if is_parent_module and not is_parent_package:
+ continue
+ filtered_children.append(child)
+ return filtered_children
+
+
 class PublicAPIFilter(object):
  """Visitor to use with `traverse` to filter just the public API."""
 

tools/tensorflow_docs/api_generator/public_api_test.py

@@ -21,6 +21,8 @@
 import typing
 
 from absl.testing import absltest
+# This import is using to test
+from tensorflow_docs import api_generator
 from tensorflow_docs.api_generator import public_api
 
 
@@ -150,6 +152,45 @@ def public_members(obj):
 
  self.assertCountEqual([], filtered_names)
 
+ def test_explicit_package_contents_filter_removes_modules_not_explicitly_imported(
+ self):
+ path = ('tensorflow_docs', 'api_generator')
+ parent = api_generator
+ members = inspect.getmembers(parent)
+ members.append(('inspect', inspect))
+
+ # Assert that parent is a module and is a package, and that the members of
+ # parent include a module named `inspect`.
+ self.assertTrue(inspect.ismodule(parent))
+ self.assertTrue(hasattr(parent, '__path__'))
+ self.assertIn('inspect', [name for name, _ in members])
+ self.assertTrue(inspect.ismodule(inspect))
+
+ filtered_members = public_api.explicit_package_contents_filter(
+ path, parent, members)
+
+ # Assert that the filtered_members do not include a module named `inspect`.
+ self.assertNotIn('inspect', [name for name, _ in filtered_members])
+
+ def test_explicit_package_contents_filter_removes_modules_imported_by_modules(
+ self):
+ path = ('tensorflow_docs', 'api_generator', 'public_api')
+ parent = public_api
+ members = inspect.getmembers(parent)
+
+ # Assert that parent is a module and not a package, and that the members of
+ # parent include a module named `inspect`.
+ self.assertTrue(inspect.ismodule(parent))
+ self.assertFalse(hasattr(parent, '__path__'))
+ self.assertIn('inspect', [name for name, _ in members])
+ self.assertTrue(inspect.ismodule(inspect))
+
+ filtered_members = public_api.explicit_package_contents_filter(
+ path, parent, members)
+
+ # Assert that the filtered_members do not include a module named `inspect`.
+ self.assertNotIn('inspect', [name for name, _ in filtered_members])
+
  def test_ignore_typing(self):
  children_before = [('a', 1), ('b', 3), ('c', typing.List)]
  children_after = public_api.ignore_typing('ignored', 'ignored',