Patch the use of inspect.findsource with an accompanying linecache.update, which seems to resolve the issue of incorrect source code being resolved when running from a compressed file.

PiperOrigin-RevId: 244704628

tensorflow/python/autograph/pyct/BUILD

@@ -100,6 +100,15 @@ py_test(
     ],
 )
 
+sh_test(
+    name = "inspect_utils_test_par",
+    srcs = ["inspect_utils_test.sh"],
+    data = [
+        ":inspect_utils_test.par",
+    ],
+    tags = ["no_oss"],
+)
+
 py_test(
     name = "origin_info_test",
     srcs = ["origin_info_test.py"],

tensorflow/python/autograph/pyct/inspect_utils.py

@@ -21,13 +21,20 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import inspect
 import itertools
+import linecache
+import sys
+import threading
 import types
 
 import six
 
 from tensorflow.python.util import tf_inspect
 
+# This lock seems to help avoid linecache concurrency errors.
+_linecache_lock = threading.Lock()
+
 
 # These functions test negative for isinstance(*, types.BuiltinFunctionType)
 # and inspect.isbuiltin, and are generally not visible in globals().
@@ -83,6 +90,41 @@ def isbuiltin(f):
   return False
 
 
+def _fix_linecache_record(obj):
+  """Fixes potential corruption of linecache in the presence of functools.wraps.
+
+  functools.wraps modifies the target object's __module__ field, which seems
+  to confuse linecache in special instances, for example when the source is
+  loaded from a .par file (see https://google.github.io/subpar/subpar.html).
+
+  This function simply triggers a call to linecache.updatecache when a mismatch
+  was detected between the object's __module__ property and the object's source
+  file.
+
+  Args:
+    obj: Any
+  """
+  if hasattr(obj, '__module__'):
+    obj_file = inspect.getfile(obj)
+    obj_module = obj.__module__
+
+    # A snapshot of the loaded modules helps avoid "dict changed size during
+    # iteration" errors.
+    loaded_modules = tuple(sys.modules.values())
+    for m in loaded_modules:
+      if hasattr(m, '__file__') and m.__file__ == obj_file:
+        if obj_module is not m:
+          linecache.updatecache(obj_file, m.__dict__)
+
+
+def getimmediatesource(obj):
+  """A variant of inspect.getsource that ignores the __wrapped__ property."""
+  with _linecache_lock:
+    _fix_linecache_record(obj)
+    lines, lnum = inspect.findsource(obj)
+    return ''.join(inspect.getblock(lines[lnum:]))
+
+
 def getnamespace(f):
   """Returns the complete namespace of a function.
 

tensorflow/python/autograph/pyct/inspect_utils_test.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests for unspect_utils module."""
+"""Tests for inspect_utils module."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -21,6 +21,7 @@ from __future__ import print_function
 import collections
 import functools
 import imp
+import textwrap
 import types
 import weakref
 
@@ -28,6 +29,7 @@ import six
 
 from tensorflow.python import lib
 from tensorflow.python.autograph.pyct import inspect_utils
+from tensorflow.python.autograph.pyct.testing import decorators
 from tensorflow.python.autograph.pyct.testing import future_import_module
 from tensorflow.python.eager import function
 from tensorflow.python.framework import constant_op
@@ -133,6 +135,125 @@ class InspectUtilsTest(test.TestCase):
 
     self.assertTrue(inspect_utils.isnamedtuple(NamedTupleSubclass))
 
+  def assertSourceIdentical(self, actual, expected):
+    self.assertEqual(
+        textwrap.dedent(actual).strip(),
+        textwrap.dedent(expected).strip()
+    )
+
+  def test_getimmediatesource_basic(self):
+
+    def test_decorator(f):
+
+      def f_wrapper(*args, **kwargs):
+        return f(*args, **kwargs)
+
+      return f_wrapper
+
+    expected = """
+      def f_wrapper(*args, **kwargs):
+        return f(*args, **kwargs)
+    """
+
+    @test_decorator
+    def test_fn(a):
+      """Test docstring."""
+      return [a]
+
+    self.assertSourceIdentical(
+        inspect_utils.getimmediatesource(test_fn), expected)
+
+  def test_getimmediatesource_noop_decorator(self):
+
+    def test_decorator(f):
+      return f
+
+    expected = '''
+      @test_decorator
+      def test_fn(a):
+        """Test docstring."""
+        return [a]
+    '''
+
+    @test_decorator
+    def test_fn(a):
+      """Test docstring."""
+      return [a]
+
+    self.assertSourceIdentical(
+        inspect_utils.getimmediatesource(test_fn), expected)
+
+  def test_getimmediatesource_functools_wrapper(self):
+
+    def wrapper_decorator(f):
+
+      @functools.wraps(f)
+      def wrapper(*args, **kwargs):
+        return f(*args, **kwargs)
+
+      return wrapper
+
+    expected = textwrap.dedent("""
+      @functools.wraps(f)
+      def wrapper(*args, **kwargs):
+        return f(*args, **kwargs)
+    """)
+
+    @wrapper_decorator
+    def test_fn(a):
+      """Test docstring."""
+      return [a]
+
+    self.assertSourceIdentical(
+        inspect_utils.getimmediatesource(test_fn), expected)
+
+  def test_getimmediatesource_functools_wrapper_different_module(self):
+
+    expected = textwrap.dedent("""
+      @functools.wraps(f)
+      def wrapper(*args, **kwargs):
+        return f(*args, **kwargs)
+    """)
+
+    @decorators.wrapping_decorator
+    def test_fn(a):
+      """Test docstring."""
+      return [a]
+
+    self.assertSourceIdentical(
+        inspect_utils.getimmediatesource(test_fn), expected)
+
+  def test_getimmediatesource_normal_decorator_different_module(self):
+
+    expected = textwrap.dedent("""
+      def standalone_wrapper(*args, **kwargs):
+        return f(*args, **kwargs)
+    """)
+
+    @decorators.standalone_decorator
+    def test_fn(a):
+      """Test docstring."""
+      return [a]
+
+    self.assertSourceIdentical(
+        inspect_utils.getimmediatesource(test_fn), expected)
+
+  def test_getimmediatesource_normal_functional_decorator_different_module(
+      self):
+
+    expected = textwrap.dedent("""
+      def functional_wrapper(*args, **kwargs):
+        return f(*args, **kwargs)
+    """)
+
+    @decorators.functional_decorator()
+    def test_fn(a):
+      """Test docstring."""
+      return [a]
+
+    self.assertSourceIdentical(
+        inspect_utils.getimmediatesource(test_fn), expected)
+
   def test_getnamespace_globals(self):
     ns = inspect_utils.getnamespace(factory)
     self.assertEqual(ns['free_function'], free_function)

tensorflow/python/autograph/pyct/inspect_utils_test.sh

@@ -0,0 +1,19 @@
+# Copyright 2017 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.
+# ==============================================================================
+# Test that runs inspect_utils_test as a .par file.
+
+
+SCRIPT_DIR="$(dirname ${BASH_SOURCE[0]})"
+${SCRIPT_DIR}/inspect_utils_test.par

tensorflow/python/autograph/pyct/parser.py

@@ -22,16 +22,13 @@ from __future__ import division
 from __future__ import print_function
 
 import textwrap
-import threading
 
 import gast
 
+from tensorflow.python.autograph.pyct import inspect_utils
 from tensorflow.python.util import tf_inspect
 
 
-_parse_lock = threading.Lock()  # Prevents linecache concurrency errors.
-
-
 STANDARD_PREAMBLE = textwrap.dedent("""
     from __future__ import division
     from __future__ import print_function
@@ -53,8 +50,7 @@ def parse_entity(entity, future_features):
     generate the AST (including any prefixes that this function may have added).
   """
   try:
-    with _parse_lock:
+    source = inspect_utils.getimmediatesource(entity)
-      source = tf_inspect.getsource_no_unwrap(entity)
   except (IOError, OSError) as e:
     raise ValueError(
         'Unable to locate the source code of {}. Note that functions defined'

tensorflow/python/autograph/pyct/testing/BUILD

@@ -17,6 +17,7 @@ filegroup(
 py_library(
     name = "test_modules",
     srcs = [
+        "decorators.py",
         "future_import_module.py",
     ],
     srcs_version = "PY2AND3",

tensorflow/python/autograph/pyct/testing/decorators.py

@@ -0,0 +1,50 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Module with test decorators."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+
+
+def wrapping_decorator(f):
+
+  @functools.wraps(f)
+  def wrapper(*args, **kwargs):
+    return f(*args, **kwargs)
+
+  return wrapper
+
+
+def standalone_decorator(f):
+
+  def standalone_wrapper(*args, **kwargs):
+    return f(*args, **kwargs)
+
+  return standalone_wrapper
+
+
+def functional_decorator():
+
+  def decorator(f):
+
+    def functional_wrapper(*args, **kwargs):
+      return f(*args, **kwargs)
+
+    return functional_wrapper
+
+  return decorator

tensorflow/python/util/tf_inspect.py

@@ -399,22 +399,3 @@ def isroutine(object):  # pylint: disable=redefined-builtin
 def stack(context=1):
   """TFDecorator-aware replacement for inspect.stack."""
   return _inspect.stack(context)[1:]
-
-
-def getsource_no_unwrap(obj):
-  """Return source code for an object. Does not unwrap TFDecorators.
-
-  The source code is returned literally, including indentation for functions not
-  at the top level. This function is analogous to inspect.getsource, with one
-  key difference - it doesn't unwrap decorators. For simplicity, support for
-  some Python object types is dropped (tracebacks, frames, code objects).
-
-  Args:
-      obj: a class, method, or function object.
-
-  Returns:
-      source code as a string
-
-  """
-  lines, lnum = _inspect.findsource(obj)
-  return ''.join(_inspect.getblock(lines[lnum:]))

tensorflow/python/util/tf_inspect_test.py

@@ -741,73 +741,6 @@ class TfInspectGetCallArgsTest(test.TestCase):
         'c': 'goodbye'
     }, tf_inspect.getcallargs(decorated, 4, c='goodbye'))
 
-  def testGetSourceNoUnwrapHandlesPlainDecorator(self):
-    def dec(f):
-      def wrapper(*args, **kwargs):
-        return f(*args, **kwargs)
-      return wrapper
-
-    @dec
-    def f():
-      return 1
-
-    source = tf_inspect.getsource_no_unwrap(f)
-    self.assertNotIn('dec', source)
-    self.assertIn('wrapper', source)
-    self.assertNotIn('return 1', source)
-
-  def testGetSourceNoUnwrapHandlesFunctoolsDecorator(self):
-    def dec(f):
-      @functools.wraps(f)
-      def wrapper(*args, **kwargs):
-        return f(*args, **kwargs)
-      return wrapper
-
-    @dec
-    def f():
-      return 1
-
-    source = tf_inspect.getsource_no_unwrap(f)
-    self.assertNotIn('dec', source)
-    self.assertIn('wrapper', source)
-    self.assertNotIn('return 1', source)
-
-  def testGetSourceNoUnwrapHandlesPlainDecoratorFactory(self):
-    def dec_factory():
-      def dec(f):
-        def wrapper(*args, **kwargs):
-          return f(*args, **kwargs)
-        return wrapper
-      return dec
-
-    @dec_factory()
-    def f():
-      return 1
-
-    source = tf_inspect.getsource_no_unwrap(f)
-    self.assertNotIn('factory', source)
-    self.assertNotIn('dec', source)
-    self.assertIn('wrapper', source)
-    self.assertNotIn('return 1', source)
-
-  def testGetSourceNoUnwrapHandlesFunctoolsDecoratorFactory(self):
-    def dec_factory():
-      def dec(f):
-        @functools.wraps(f)
-        def wrapper(*args, **kwargs):
-          return f(*args, **kwargs)
-        return wrapper
-      return dec
-
-    @dec_factory()
-    def f():
-      return 1
-
-    source = tf_inspect.getsource_no_unwrap(f)
-    self.assertNotIn('factory', source)
-    self.assertNotIn('dec', source)
-    self.assertIn('wrapper', source)
-    self.assertNotIn('return 1', source)
 
 if __name__ == '__main__':
   test.main()