Add a simplification for ast.parse

This was failing on the wrapped lines in Dataset.map (and all subclasses) because textwrap.dedent didn't remove the leading space.

PiperOrigin-RevId: 438447039

Author: MarkDaoust

tools/tensorflow_docs/api_generator/get_source.py

@@ -13,12 +13,27 @@
 # limitations under the License.
 # ==============================================================================
 """Simple get_source."""
+import ast
 import inspect
 import textwrap
 
 from typing import Any, Optional, Sequence, Tuple
 
 
+def get_ast(py_object) -> Optional[ast.AST]:
+  if isinstance(py_object, str):
+    source = textwrap.dedent(py_object)
+  else:
+    source = get_source(py_object)
+  if source is None:
+    return None
+
+  try:
+    return ast.parse(source)
+  except Exception:  # pylint: disable=broad-except
+    return None
+
+
 def get_source(py_object: Any) -> Optional[str]:
   if py_object is not None:
     try:

tools/tensorflow_docs/api_generator/public_api.py

@@ -194,14 +194,10 @@ def visit_Import(self, node):  # pylint: disable=invalid-name
     def visit_ImportFrom(self, node):  # pylint: disable=invalid-name
       self._add_imported_symbol(node)
 
-  if isinstance(obj, str):
-    source = textwrap.dedent(obj)
-  else:
-    source = get_source.get_source(obj)
-  if source is None:
+  tree = get_source.get_ast(obj)
+  if tree is None:
     return []
 
-  tree = ast.parse(source)
   visitor = ImportNodeVisitor()
   visitor.visit(tree)
   return visitor.imported_symbols

tools/tensorflow_docs/api_generator/report/linter.py

@@ -176,12 +176,13 @@ def lint_returns(
   Returns:
     A filled `ReturnLint` proto object.
   """
-  source = get_source.get_source(page_info.py_object)
-
   return_visitor = ReturnVisitor()
-  if source is not None:
+
+  source = get_source.get_source(page_info.py_object)
+  obj_ast = get_source.get_ast(page_info.py_object)
+  if obj_ast is not None:
     try:
-      return_visitor.visit(ast.parse(source))
+      return_visitor.visit(obj_ast)
     except Exception:  # pylint: disable=broad-except
       pass
 
@@ -236,12 +237,13 @@ def lint_raises(page_info: base_page.PageInfo) -> api_report_pb2.RaisesLint:
 
   # Extract the raises from the source code.
   raise_visitor = RaiseVisitor()
-  source = get_source.get_source(page_info.py_object)
-  if source is not None:
+  obj_ast = get_source.get_ast(page_info.py_object)
+  if obj_ast is not None:
     try:
-      raise_visitor.visit(ast.parse(source))
+      raise_visitor.visit(obj_ast)
     except Exception:  # pylint: disable=broad-except
       pass
+
   raises_lint.total_raises_in_code = len(raise_visitor.total_raises)
 
   # Extract the raises defined in the docstring.

tools/tensorflow_docs/api_generator/signature.py

@@ -56,9 +56,8 @@ def _preprocess_default(self, val: ast.AST) -> str:
     return text_default_val
 
   def extract(self, obj: Any):
-    obj_source = get_source.get_source(obj)
-    if obj_source is not None:
-      obj_ast = ast.parse(obj_source)
+    obj_ast = get_source.get_ast(obj)
+    if obj_ast is not None:
       self.visit(obj_ast)
 
 
@@ -674,11 +673,10 @@ def visit_FunctionDef(self, node):  # pylint: disable=invalid-name
 
   visitor = ASTDecoratorExtractor()
 
-  # Note: inspect.getsource doesn't include the decorator lines on classes,
+  # Note: get_source doesn't include the decorator lines on classes,
   # this won't work for classes until that's fixed.
-  func_source = get_source.get_source(func)
-  if func_source is not None:
-    func_ast = ast.parse(func_source)
+  func_ast = get_source.get_ast(func)
+  if func_ast is not None:
     visitor.visit(func_ast)
 
   return visitor.decorator_list