Add google.api.core.helpers.general_helpers.wraps (#4166)

Jon Wayne Parrott · web-flow · commit 0720a44e3c30 · 2017-10-12T10:11:04.000-07:00
diff --git a/packages/google-cloud-core/google/api/core/gapic_v1/method.py b/packages/google-cloud-core/google/api/core/gapic_v1/method.py
@@ -26,6 +26,7 @@
 
 from google.api.core import page_iterator
 from google.api.core import timeout
+from google.api.core.helpers import general_helpers
 from google.api.core.helpers import grpc_helpers
 
 _PY_VERSION = platform.python_version()
@@ -252,7 +253,7 @@ def get_topic(name, timeout=None):
     if metadata is not None:
         metadata = _prepare_metadata(metadata)
 
-    return six.wraps(func)(
+    return general_helpers.wraps(func)(
         _GapicCallable(func, default_retry, default_timeout, metadata))
 
 
diff --git a/packages/google-cloud-core/google/api/core/helpers/general_helpers.py b/packages/google-cloud-core/google/api/core/helpers/general_helpers.py
@@ -0,0 +1,32 @@
+# Copyright 2017 Google Inc.
+#
+# 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.
+
+"""Helpers for general Python functionality."""
+
+import functools
+
+import six
+
+
+# functools.partial objects lack several attributes present on real function
+# objects. In Python 2 wraps fails on this so use a restricted set instead.
+_PARTIAL_VALID_ASSIGNMENTS = ('__doc__',)
+
+
+def wraps(wrapped):
+    """A functools.wraps helper that handles partial objects on Python 2."""
+    if isinstance(wrapped, functools.partial):
+        return six.wraps(wrapped, assigned=_PARTIAL_VALID_ASSIGNMENTS)
+    else:
+        return six.wraps(wrapped)
diff --git a/packages/google-cloud-core/google/api/core/helpers/grpc_helpers.py b/packages/google-cloud-core/google/api/core/helpers/grpc_helpers.py
@@ -18,6 +18,7 @@
 import six
 
 from google.api.core import exceptions
+from google.api.core.helpers import general_helpers
 import google.auth
 import google.auth.transport.grpc
 import google.auth.transport.requests
@@ -63,7 +64,7 @@ def _wrap_stream_errors(callable_):
     """
     _patch_callable_name(callable_)
 
-    @six.wraps(callable_)
+    @general_helpers.wraps(callable_)
     def error_remapped_callable(*args, **kwargs):
         try:
             result = callable_(*args, **kwargs)
diff --git a/packages/google-cloud-core/google/api/core/retry.py b/packages/google-cloud-core/google/api/core/retry.py
@@ -66,6 +66,7 @@ def check_if_exists():
 
 from google.api.core import exceptions
 from google.api.core.helpers import datetime_helpers
+from google.api.core.helpers import general_helpers
 
 _LOGGER = logging.getLogger(__name__)
 _DEFAULT_INITIAL_DELAY = 1.0  # seconds
@@ -244,7 +245,7 @@ def __call__(self, func, on_error=None):
             Callable: A callable that will invoke ``func`` with retry
                 behavior.
         """
-        @six.wraps(func)
+        @general_helpers.wraps(func)
         def retry_wrapped_func(*args, **kwargs):
             """A wrapper that calls target function with retry."""
             target = functools.partial(func, *args, **kwargs)
diff --git a/packages/google-cloud-core/google/api/core/timeout.py b/packages/google-cloud-core/google/api/core/timeout.py
@@ -58,6 +58,7 @@ def is_thing_ready(timeout=None):
 import six
 
 from google.api.core.helpers import datetime_helpers
+from google.api.core.helpers import general_helpers
 
 _DEFAULT_INITIAL_TIMEOUT = 5.0  # seconds
 _DEFAULT_MAXIMUM_TIMEOUT = 30.0  # seconds
@@ -92,7 +93,7 @@ def __call__(self, func):
         Returns:
             Callable: The wrapped function.
         """
-        @six.wraps(func)
+        @general_helpers.wraps(func)
         def func_with_timeout(*args, **kwargs):
             """Wrapped function that adds timeout."""
             kwargs['timeout'] = self._timeout
@@ -198,7 +199,7 @@ def __call__(self, func):
         timeouts = _exponential_timeout_generator(
             self._initial, self._maximum, self._multiplier, self._deadline)
 
-        @six.wraps(func)
+        @general_helpers.wraps(func)
         def func_with_timeout(*args, **kwargs):
             """Wrapped function that adds timeout."""
             kwargs['timeout'] = next(timeouts)
diff --git a/packages/google-cloud-core/tests/unit/api_core/helpers/test_general_helpers.py b/packages/google-cloud-core/tests/unit/api_core/helpers/test_general_helpers.py
@@ -0,0 +1,43 @@
+# Copyright 2017, Google Inc.
+#
+# 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.
+
+import functools
+
+from google.api.core.helpers import general_helpers
+
+
+def test_wraps_normal_func():
+
+    def func():
+        return 42
+
+    @general_helpers.wraps(func)
+    def replacement():
+        return func()
+
+    assert replacement() == 42
+
+
+def test_wraps_partial():
+
+    def func():
+        return 42
+
+    partial = functools.partial(func)
+
+    @general_helpers.wraps(partial)
+    def replacement():
+        return func()
+
+    assert replacement() == 42
