diff --git a/core/google/api/core/gapic_v1/method.py b/core/google/api/core/gapic_v1/method.py index d402c1e220ae..48accdbb3d51 100644 --- a/core/google/api/core/gapic_v1/method.py +++ b/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/core/google/api/core/helpers/general_helpers.py b/core/google/api/core/helpers/general_helpers.py new file mode 100644 index 000000000000..0c8e40893252 --- /dev/null +++ b/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/core/google/api/core/helpers/grpc_helpers.py b/core/google/api/core/helpers/grpc_helpers.py index 28e22fde8888..da9eaf53bc66 100644 --- a/core/google/api/core/helpers/grpc_helpers.py +++ b/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/core/google/api/core/retry.py b/core/google/api/core/retry.py index d772b333556c..625749e8da0d 100644 --- a/core/google/api/core/retry.py +++ b/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/core/google/api/core/timeout.py b/core/google/api/core/timeout.py index 7f8f6b6d8ff8..1a4d3bfdccef 100644 --- a/core/google/api/core/timeout.py +++ b/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/core/tests/unit/api_core/helpers/test_general_helpers.py b/core/tests/unit/api_core/helpers/test_general_helpers.py new file mode 100644 index 000000000000..9b39ea39d1c9 --- /dev/null +++ b/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