Standardize the use of inspect.getsource

PiperOrigin-RevId: 438269328

Author: MarkDaoust

tools/tensorflow_docs/api_generator/get_source.py

@@ -0,0 +1,40 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Simple get_source."""
+import inspect
+import textwrap
+
+from typing import Any, Optional, Sequence, Tuple
+
+
+def get_source(py_object: Any) -> Optional[str]:
+  if py_object is not None:
+    try:
+      return textwrap.dedent(inspect.getsource(py_object))
+    except Exception:  # pylint: disable=broad-except
+      # A wide-variety of errors can be thrown here.
+      pass
+  return None
+
+
+def get_source_lines(
+    py_object: Any) -> Tuple[Optional[Sequence[str]], Optional[int]]:
+  if py_object is not None:
+    try:
+      return inspect.getsourcelines(py_object)
+    except Exception:  # pylint: disable=broad-except
+      # A wide-variety of errors can be thrown here.
+      pass
+  return None, None

tools/tensorflow_docs/api_generator/parser.py

@@ -28,6 +28,7 @@
 
 from tensorflow_docs.api_generator import config
 from tensorflow_docs.api_generator import doc_generator_visitor
+from tensorflow_docs.api_generator import get_source
 from tensorflow_docs.api_generator import obj_type as obj_type_lib
 from tensorflow_docs.api_generator import signature as signature_lib
 
@@ -653,15 +654,14 @@ def get_defined_in(
   if code_url_prefix is None:
     return None
 
-  try:
-    lines, start_line = inspect.getsourcelines(py_object)
+  lines, start_line = get_source.get_source_lines(py_object)
+  if start_line is None:
+    end_line = None
+  else:
     end_line = start_line + len(lines) - 1
     if 'MACHINE GENERATED' in lines[0]:
       # don't link to files generated by tf_export
       return None
-  except (IOError, TypeError, IndexError):
-    start_line = None
-    end_line = None
 
   # In case this is compiled, point to the original
   if posix_rel_path_str.endswith('.pyc'):

tools/tensorflow_docs/api_generator/public_api.py

@@ -23,6 +23,7 @@
 
 from tensorflow_docs.api_generator import doc_controls
 from tensorflow_docs.api_generator import doc_generator_visitor
+from tensorflow_docs.api_generator import get_source
 
 _TYPING_IDS = frozenset(
     id(obj)
@@ -184,11 +185,10 @@ def visit_Import(self, node):  # pylint: disable=invalid-name
     def visit_ImportFrom(self, node):  # pylint: disable=invalid-name
       self._add_imported_symbol(node)
 
-  try:
-    source = inspect.getsource(obj)
-  except OSError:
-    # inspect raises an OSError for empty module files.
+  source = get_source.get_source(obj)
+  if source is None:
     return []
+
   tree = ast.parse(source)
   visitor = ImportNodeVisitor()
   visitor.visit(tree)

tools/tensorflow_docs/api_generator/report/linter.py

@@ -23,19 +23,12 @@
 
 import astor
 
+from tensorflow_docs.api_generator import get_source
 from tensorflow_docs.api_generator import parser
 from tensorflow_docs.api_generator.pretty_docs import base_page
 from tensorflow_docs.api_generator.report.schema import api_report_generated_pb2 as api_report_pb2
 
 
-def _get_source(py_object: Any) -> Optional[str]:
-  if py_object is not None:
-    try:
-      source = textwrap.dedent(inspect.getsource(py_object))
-      return source
-    except Exception:  # pylint: disable=broad-except
-      return None
-  return None
 
 
 def _count_empty_param(items: List[Tuple[str, Optional[str]]]) -> int:
@@ -183,7 +176,7 @@ def lint_returns(
   Returns:
     A filled `ReturnLint` proto object.
   """
-  source = _get_source(page_info.py_object)
+  source = get_source.get_source(page_info.py_object)
 
   return_visitor = ReturnVisitor()
   if source is not None:
@@ -243,7 +236,7 @@ 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(page_info.py_object)
+  source = get_source.get_source(page_info.py_object)
   if source is not None:
     try:
       raise_visitor.visit(ast.parse(source))

tools/tensorflow_docs/api_generator/signature.py

@@ -30,6 +30,7 @@
 import astor
 
 from tensorflow_docs.api_generator import config
+from tensorflow_docs.api_generator import get_source
 from tensorflow_docs.api_generator import public_api
 
 EMPTY = inspect.Signature.empty
@@ -55,9 +56,10 @@ def _preprocess_default(self, val: ast.AST) -> str:
     return text_default_val
 
   def extract(self, obj: Any):
-    obj_source = textwrap.dedent(inspect.getsource(obj))
-    obj_ast = ast.parse(obj_source)
-    self.visit(obj_ast)
+    obj_source = get_source.get_source(obj)
+    if obj_source is not None:
+      obj_ast = ast.parse(obj_source)
+      self.visit(obj_ast)
 
 
 class _ArgDefaultAndAnnotationExtractor(_BaseDefaultAndAnnotationExtractor):
@@ -672,14 +674,11 @@ def visit_FunctionDef(self, node):  # pylint: disable=invalid-name
 
   visitor = ASTDecoratorExtractor()
 
-  try:
-    # Note: inspect.getsource doesn't include the decorator lines on classes,
-    # this won't work for classes until that's fixed.
-    func_source = textwrap.dedent(inspect.getsource(func))
+  # Note: inspect.getsource 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)
     visitor.visit(func_ast)
-  except Exception:  # pylint: disable=broad-except
-    # A wide-variety of errors can be thrown here.
-    pass
 
   return visitor.decorator_list