Execute globals() and locals() in the correct context in converted code.

PiperOrigin-RevId: 320716956
Change-Id: If1eae7e65edbfb965edf73fa27d21bfc7bea8f18

tensorflow/python/autograph/impl/api.py

@@ -377,6 +377,10 @@ def converted_call(f,
       return py_builtins.eval_in_original_context(f, args, caller_fn_scope)
     if f is super:
       return py_builtins.super_in_original_context(f, args, caller_fn_scope)
+    if f is globals:
+      return py_builtins.globals_in_original_context(caller_fn_scope)
+    if f is locals:
+      return py_builtins.locals_in_original_context(caller_fn_scope)
     if kwargs:
       return py_builtins.overload_of(f)(*args, **kwargs)
     else:

tensorflow/python/autograph/operators/BUILD

@@ -120,8 +120,7 @@ py_test(
     name = "py_builtins_test",
     srcs = ["py_builtins_test.py"],
     python_version = "PY3",
-    srcs_version = "PY2AND3",
+    srcs_version = "PY3",
-    tags = ["no_windows"],
     deps = [
         ":operators",
         "//tensorflow/python:client_testlib",
@@ -133,23 +132,6 @@ py_test(
     ],
 )
 
-py_test(
-    name = "py_builtins_py3_test",
-    srcs = ["py_builtins_py3_test.py"],
-    python_version = "PY3",
-    srcs_version = "PY3",
-    tags = [
-        "no_windows",
-        # TODO(kkb): Temporay workaround since KokoroPresubmit was failing.
-        #                 cl/259400943 for more context.
-        "no_oss_py2",
-    ],
-    deps = [
-        ":operators",
-        "//tensorflow/python:client_testlib",
-    ],
-)
-
 py_test(
     name = "slices_test",
     srcs = ["slices_test.py"],

tensorflow/python/autograph/operators/py_builtins.py

@@ -89,6 +89,16 @@ def _find_originating_frame(caller_fn_scope, innermost=True):
   return result
 
 
+def locals_in_original_context(caller_fn_scope):
+  """Executes the locals function in the context of a specified function."""
+  return _find_originating_frame(caller_fn_scope, innermost=True).f_locals
+
+
+def globals_in_original_context(caller_fn_scope):
+  """Executes the locals function in the context of a specified function."""
+  return _find_originating_frame(caller_fn_scope, innermost=True).f_globals
+
+
 def eval_in_original_context(f, args, caller_fn_scope):
   """Executes the eval function in the context of a specified function."""
   # When control flow is rewritten using functions, eval should use the

tensorflow/python/autograph/operators/py_builtins_py3_test.py

@@ -1,123 +0,0 @@
-# Copyright 2019 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.
-# ==============================================================================
-"""Tests for py_builtins_py3 module."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.python.autograph.core import converter
-from tensorflow.python.autograph.core import function_wrappers
-from tensorflow.python.autograph.operators import py_builtins
-from tensorflow.python.platform import test
-
-
-class TestBaseClass(object):
-
-  def overridden_method(self, x):
-    return x + 20
-
-
-class PyBuiltinsTest(test.TestCase):
-
-  def _basic_function_scope(self):
-    return function_wrappers.FunctionScope(
-        'test_function_name',
-        'test_scope',  # Note: this must match the name in the `with` statement.
-        converter.ConversionOptions())
-
-  def test_super_in_original_context_niladic_call(self):
-    test_case_self = self
-
-    class TestSubclass(TestBaseClass):
-
-      def overridden_method(self, x):
-        test_case_self.fail('This should never be called.')
-
-      def test_method(self):
-        with test_case_self._basic_function_scope() as test_scope:
-          b = py_builtins.super_in_original_context(super, (), test_scope)
-          return b.overridden_method(1)
-
-    tc = TestSubclass()
-    self.assertEqual(tc.test_method(), 21)
-
-  def test_super_in_original_context_caller_with_locals(self):
-    test_case_self = self
-
-    class TestSubclass(TestBaseClass):
-
-      def overridden_method(self, x):
-        test_case_self.fail('This should never be called.')
-
-      def test_method(self, x):
-        y = 7
-        with test_case_self._basic_function_scope() as test_scope:
-          z = 7
-          return py_builtins.super_in_original_context(
-              super, (), test_scope).overridden_method(x + y - z)
-
-    tc = TestSubclass()
-    self.assertEqual(tc.test_method(1), 21)
-
-  def test_super_in_original_context_inner_function(self):
-    test_case_self = self
-
-    class TestSubclass(TestBaseClass):
-
-      def overridden_method(self, x):
-        test_case_self.fail('This should never be called.')
-
-      def test_method(self, x):
-        with test_case_self._basic_function_scope() as test_scope:
-          # Oddly, it's sufficient to use `self` in an inner function
-          # to gain access to __class__ in this scope.
-          # TODO(mdan): Is this true across implementations?
-          # Note: normally, it's illegal to use super() in inner functions (it
-          # throws an error), but the generated code may create them.
-          def inner_fn():
-            return py_builtins.super_in_original_context(
-                super, (), test_scope).overridden_method(x)
-
-          return inner_fn()
-
-    tc = TestSubclass()
-    self.assertEqual(tc.test_method(1), 21)
-
-  def test_super_in_original_context_inner_lambda(self):
-    test_case_self = self
-
-    class TestSubclass(TestBaseClass):
-
-      def overridden_method(self, x):
-        test_case_self.fail('This should never be called.')
-
-      def test_method(self, x):
-        with test_case_self._basic_function_scope() as test_scope:
-          # Oddly, it's sufficient to use `self` in an inner function
-          # to gain access to __class__ in this scope.
-          # TODO(mdan): Is this true across implementations?
-          # Note: normally, it's illegal to use super() in inner functions (it
-          # throws an error), but the generated code may create them.
-          l = lambda: py_builtins.super_in_original_context(  # pylint:disable=g-long-lambda
-              super, (), test_scope).overridden_method(x)
-          return l()
-
-    tc = TestSubclass()
-    self.assertEqual(tc.test_method(1), 21)
-
-
-if __name__ == '__main__':
-  test.main()

tensorflow/python/autograph/operators/py_builtins_test.py

@@ -40,10 +40,11 @@ from tensorflow.python.platform import test
 
 class TestBase(object):
 
-  def plus_twenty(self, x):
+  def overridden_method(self, x):
     return x + 20
 
 
+@test_util.run_all_in_graph_and_eager_modes
 class PyBuiltinsTest(test.TestCase):
 
   def test_abs(self):
@@ -400,12 +401,67 @@ class PyBuiltinsTest(test.TestCase):
 
     self.assertEqual(test_fn(), 2)
 
+  def test_locals_in_original_context(self):
+
+    def test_fn():
+      l = 1  # pylint:disable=unused-variable
+      with self._basic_function_scope() as test_scope:
+        return py_builtins.locals_in_original_context(test_scope)
+
+    locs = test_fn()
+
+    self.assertEqual(locs['l'], 1)
+
+  def test_locals_in_original_context_inner_function(self):
+
+    def test_fn():
+      l = 1  # pylint:disable=unused-variable
+      with self._basic_function_scope() as test_scope:
+
+        def inner_fn():
+          # Note: a user function without a top-level function scope should
+          # never be found in user code; it's only possible in generated code.
+          l = 2  # pylint:disable=unused-variable
+          return py_builtins.locals_in_original_context(test_scope)
+
+        return inner_fn()
+
+    locs = test_fn()
+
+    self.assertEqual(locs['l'], 2)
+
+  def test_globals_in_original_context(self):
+
+    def test_fn():
+      with self._basic_function_scope() as test_scope:
+        return py_builtins.globals_in_original_context(test_scope)
+
+    globs = test_fn()
+
+    self.assertIs(globs['TestBase'], TestBase)
+
+  def test_globals_in_original_context_inner_function(self):
+
+    def test_fn():
+      with self._basic_function_scope() as test_scope:
+
+        def inner_fn():
+          # Note: a user function without a top-level function scope should
+          # never be found in user code; it's only possible in generated code.
+          return py_builtins.globals_in_original_context(test_scope)
+
+        return inner_fn()
+
+    globs = test_fn()
+
+    self.assertIs(globs['TestBase'], TestBase)
+
   def test_super_in_original_context_unary_call(self):
     test_case_self = self
 
     class TestSubclass(TestBase):
 
-      def plus_twenty(self, x):
+      def overridden_method(self, x):
         test_case_self.fail('This should never be called.')
 
       def test_method(self):
@@ -413,7 +469,7 @@ class PyBuiltinsTest(test.TestCase):
           test_base_unbound = py_builtins.super_in_original_context(
               super, (TestSubclass,), test_scope)
           test_base = test_base_unbound.__get__(self, TestSubclass)
-          return test_base.plus_twenty(1)
+          return test_base.overridden_method(1)
 
     tc = TestSubclass()
     self.assertEqual(tc.test_method(), 21)
@@ -423,18 +479,98 @@ class PyBuiltinsTest(test.TestCase):
 
     class TestSubclass(TestBase):
 
-      def plus_twenty(self, x):
+      def overridden_method(self, x):
         test_case_self.fail('This should never be called.')
 
       def test_method(self):
         with test_case_self._basic_function_scope() as test_scope:
           test_base = py_builtins.super_in_original_context(
               super, (TestSubclass, self), test_scope)
-          return test_base.plus_twenty(1)
+          return test_base.overridden_method(1)
 
     tc = TestSubclass()
     self.assertEqual(tc.test_method(), 21)
 
+  def test_super_in_original_context_niladic_call(self):
+    test_case_self = self
+
+    class TestSubclass(TestBase):
+
+      def overridden_method(self, x):
+        test_case_self.fail('This should never be called.')
+
+      def test_method(self):
+        with test_case_self._basic_function_scope() as test_scope:
+          b = py_builtins.super_in_original_context(super, (), test_scope)
+          return b.overridden_method(1)
+
+    tc = TestSubclass()
+    self.assertEqual(tc.test_method(), 21)
+
+  def test_super_in_original_context_caller_with_locals(self):
+    test_case_self = self
+
+    class TestSubclass(TestBase):
+
+      def overridden_method(self, x):
+        test_case_self.fail('This should never be called.')
+
+      def test_method(self, x):
+        y = 7
+        with test_case_self._basic_function_scope() as test_scope:
+          z = 7
+          return py_builtins.super_in_original_context(
+              super, (), test_scope).overridden_method(x + y - z)
+
+    tc = TestSubclass()
+    self.assertEqual(tc.test_method(1), 21)
+
+  def test_super_in_original_context_inner_function(self):
+    test_case_self = self
+
+    class TestSubclass(TestBase):
+
+      def overridden_method(self, x):
+        test_case_self.fail('This should never be called.')
+
+      def test_method(self, x):
+        with test_case_self._basic_function_scope() as test_scope:
+          # Oddly, it's sufficient to use `self` in an inner function
+          # to gain access to __class__ in this scope.
+          # TODO(mdan): Is this true across implementations?
+          # Note: normally, it's illegal to use super() in inner functions (it
+          # throws an error), but the generated code may create them.
+          def inner_fn():
+            return py_builtins.super_in_original_context(
+                super, (), test_scope).overridden_method(x)
+
+          return inner_fn()
+
+    tc = TestSubclass()
+    self.assertEqual(tc.test_method(1), 21)
+
+  def test_super_in_original_context_inner_lambda(self):
+    test_case_self = self
+
+    class TestSubclass(TestBase):
+
+      def overridden_method(self, x):
+        test_case_self.fail('This should never be called.')
+
+      def test_method(self, x):
+        with test_case_self._basic_function_scope() as test_scope:
+          # Oddly, it's sufficient to use `self` in an inner function
+          # to gain access to __class__ in this scope.
+          # TODO(mdan): Is this true across implementations?
+          # Note: normally, it's illegal to use super() in inner functions (it
+          # throws an error), but the generated code may create them.
+          l = lambda: py_builtins.super_in_original_context(  # pylint:disable=g-long-lambda
+              super, (), test_scope).overridden_method(x)
+          return l()
+
+    tc = TestSubclass()
+    self.assertEqual(tc.test_method(1), 21)
+
   def test_filter(self):
     self.assertListEqual(
         list(py_builtins.filter_(lambda x: x == 'b', ['a', 'b', 'c'])), ['b'])