Internal change

PiperOrigin-RevId: 351734884
Change-Id: Ieb09826fe0759fb5831ed33c7064d8c35e2b3de4

tensorflow/python/util/tf_inspect.py

@@ -25,6 +25,20 @@ import six
 
 from tensorflow.python.util import tf_decorator
 
+
+# inspect.signature() is preferred over inspect.getfullargspec() in PY3.
+# Note that while it can handle TFDecorators, it will ignore a TFDecorator's
+# provided ArgSpec/FullArgSpec and instead return the signature of the
+# inner-most function.
+def signature(obj, *, follow_wrapped=True):
+  """TFDecorator-aware replacement for inspect.signature."""
+  return _inspect.signature(
+      tf_decorator.unwrap(obj)[1], follow_wrapped=follow_wrapped)
+
+
+Parameter = _inspect.Parameter
+Signature = _inspect.Signature
+
 ArgSpec = _inspect.ArgSpec
 
 

tensorflow/python/util/tf_inspect_test.py

@@ -492,6 +492,24 @@ class TfInspectTest(test.TestCase):
 
     self.assertEqual(argspec, tf_inspect.getfullargspec(NewClass))
 
+  def testSignatureOnDecoratorsThatDontProvideFullArgSpec(self):
+    signature = tf_inspect.signature(test_decorated_function_with_defaults)
+
+    self.assertEqual([
+        tf_inspect.Parameter('a', tf_inspect.Parameter.POSITIONAL_OR_KEYWORD),
+        tf_inspect.Parameter(
+            'b', tf_inspect.Parameter.POSITIONAL_OR_KEYWORD, default=2),
+        tf_inspect.Parameter(
+            'c', tf_inspect.Parameter.POSITIONAL_OR_KEYWORD, default='Hello')
+    ], list(signature.parameters.values()))
+
+  def testSignatureFollowsNestedDecorators(self):
+    signature = tf_inspect.signature(test_decorated_function)
+
+    self.assertEqual(
+        [tf_inspect.Parameter('x', tf_inspect.Parameter.POSITIONAL_OR_KEYWORD)],
+        list(signature.parameters.values()))
+
   def testGetDoc(self):
     self.assertEqual('Test Decorated Function With Defaults Docstring.',
                      tf_inspect.getdoc(test_decorated_function_with_defaults))