diff --git a/pubsub/google/cloud/pubsub_v1/_gapic.py b/pubsub/google/cloud/pubsub_v1/_gapic.py index 25cb3e5fa33c..da755dfbca09 100644 --- a/pubsub/google/cloud/pubsub_v1/_gapic.py +++ b/pubsub/google/cloud/pubsub_v1/_gapic.py @@ -1,4 +1,4 @@ -# Copyright 2017, Google LLC All rights reserved. +# Copyright 2019, Google LLC 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. @@ -25,29 +25,21 @@ def add_methods(source_class, blacklist=()): not added. """ - def wrap(wrapped_fx): + def wrap(wrapped_fx, lookup_fx): """Wrap a GAPIC method; preserve its name and docstring.""" - # If this is a static or class method, then we need to *not* + # If this is a static or class method, then we do *not* # send self as the first argument. # - # Similarly, for instance methods, we need to send self.api rather + # For instance methods, we need to send self.api rather # than self, since that is where the actual methods were declared. - instance_method = True - # If this is a bound method it's a classmethod. - self = getattr(wrapped_fx, "__self__", None) - if issubclass(type(self), type): - instance_method = False - - # Okay, we have figured out what kind of method this is; send - # down the correct wrapper function. - if instance_method: + if isinstance(lookup_fx, (classmethod, staticmethod)): + fx = lambda *a, **kw: wrapped_fx(*a, **kw) # noqa + return staticmethod(functools.wraps(wrapped_fx)(fx)) + else: fx = lambda self, *a, **kw: wrapped_fx(self.api, *a, **kw) # noqa return functools.wraps(wrapped_fx)(fx) - fx = lambda *a, **kw: wrapped_fx(*a, **kw) # noqa - return staticmethod(functools.wraps(wrapped_fx)(fx)) - def actual_decorator(cls): # Reflectively iterate over most of the methods on the source class # (the GAPIC) and make wrapped versions available on this client. @@ -66,7 +58,9 @@ def actual_decorator(cls): continue # Add a wrapper method to this object. - fx = wrap(getattr(source_class, name)) + lookup_fx = source_class.__dict__[name] + fx = wrap(attr, lookup_fx) + setattr(cls, name, fx) # Return the augmented class. diff --git a/pubsub/tests/unit/pubsub_v1/test__gapic.py b/pubsub/tests/unit/pubsub_v1/test__gapic.py new file mode 100644 index 000000000000..5478aee18213 --- /dev/null +++ b/pubsub/tests/unit/pubsub_v1/test__gapic.py @@ -0,0 +1,63 @@ +# Copyright 2019 Google LLC +# +# 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 +# +# https://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. + + +from google.cloud.pubsub_v1 import _gapic + + +class SourceClass(object): + def __init__(self): + self.x = "x" + + def method(self): + return "source class instance method" + + @staticmethod + def static_method(): + return "source class static method" + + @classmethod + def class_method(cls): + return "source class class method" + + @classmethod + def blacklisted_method(cls): + return "source class blacklisted method" + + +def test_add_method(): + @_gapic.add_methods(SourceClass, ("blacklisted_method",)) + class Foo(object): + def __init__(self): + self.api = SourceClass() + + def method(self): + return "foo class instance method" + + foo = Foo() + + # Any method that's callable and not blacklisted is "inherited". + assert set(["method", "static_method", "class_method"]) <= set(dir(foo)) + assert "blacklisted_method" not in dir(foo) + + # Source Class's static and class methods become static methods. + assert type(Foo.__dict__["static_method"]) == staticmethod + assert foo.static_method() == "source class static method" + assert type(Foo.__dict__["class_method"]) == staticmethod + assert foo.class_method() == "source class class method" + + # The decorator changes the behavior of instance methods of the wrapped class. + # method() is called on an instance of the Source Class (stored as an + # attribute on the wrapped class). + assert foo.method() == "source class instance method"