Fix usage of docs_controls.should_skip.

`docs_controls` was recently added to allow users to tag objects with `do_not_generate_docs` (it can be used as a decorator). This allows you to skip objects by _identity_, instead of path.

In my initial implementation I checked for the tag in the doc generation step.

This change fixes it to skip objects during the api-crawling step. This is necessary to avoid crawling the internals of all the excluded objects.

To enable this, the `_is_private` method of PublicAPIVisitor needs access to the object, not just the path to it. The changes in `public_api` allow this.

`generate_lib` uses the `_is_private` signature change to inspect the object.

PiperOrigin-RevId: 209587765

tensorflow/tools/common/public_api.py

@@ -102,9 +102,10 @@ class PublicAPIVisitor(object):
     """Override the default root name of 'tf'."""
     self._root_name = root_name
 
-  def _is_private(self, path, name):
+  def _is_private(self, path, name, obj=None):
     """Return whether a name is private."""
     # TODO(wicke): Find out what names to exclude.
+    del obj  # Unused.
     return ((path in self._private_map and
              name in self._private_map[path]) or
             (name.startswith('_') and not re.match('__.*__$', name) or
@@ -129,7 +130,7 @@ class PublicAPIVisitor(object):
 
     # Remove things that are not visible.
     for name, child in list(children):
-      if self._is_private(full_path, name):
+      if self._is_private(full_path, name, child):
         children.remove((name, child))
 
     self._visitor(path, parent, children)

tensorflow/tools/docs/generate_lib.py

@@ -112,9 +112,6 @@ def write_docs(output_dir,
             _is_free_function(py_object, full_name, parser_config.index)):
       continue
 
-    if doc_controls.should_skip(py_object):
-      continue
-
     sitepath = os.path.join('api_docs/python',
                             parser.documentation_path(full_name)[:-3])
 
@@ -298,6 +295,15 @@ def _get_default_do_not_descend_map():
   }
 
 
+class DocControlsAwareCrawler(public_api.PublicAPIVisitor):
+  """A `docs_controls` aware API-crawler."""
+
+  def _is_private(self, path, name, obj):
+    if doc_controls.should_skip(obj):
+      return True
+    return super(DocControlsAwareCrawler, self)._is_private(path, name, obj)
+
+
 def extract(py_modules,
             private_map,
             do_not_descend_map,
@@ -305,7 +311,7 @@ def extract(py_modules,
   """Extract docs from tf namespace and write them to disk."""
   # Traverse the first module.
   visitor = visitor_cls(py_modules[0][0])
-  api_visitor = public_api.PublicAPIVisitor(visitor)
+  api_visitor = DocControlsAwareCrawler(visitor)
   api_visitor.set_root_name(py_modules[0][0])
   add_dict_to_dict(private_map, api_visitor.private_map)
   add_dict_to_dict(do_not_descend_map, api_visitor.do_not_descend_map)