How to use wrapt - 10 common examples

def test_wrap_function_module_name(self):

        _args = (1, 2)
        _kwargs = { 'one': 1, 'two': 2 }

        called = []

        def wrapper(wrapped, instance, args, kwargs):
            called.append((args, kwargs))
            self.assertEqual(instance, None)
            self.assertEqual(args, _args)
            self.assertEqual(kwargs, _kwargs)
            return wrapped(*args, **kwargs)

        wrapt.wrap_function_wrapper(__name__, 'global_function_1', wrapper)

        result = global_function_1(*_args, **_kwargs)

        self.assertEqual(result, (_args, _kwargs))
        self.assertEqual(called[0], (_args, _kwargs))

def test_wrap_class_method_module_name(self):

        _args = (1, 2)
        _kwargs = { 'one': 1, 'two': 2 }

        called = []

        def wrapper(wrapped, instance, args, kwargs):
            called.append((args, kwargs))
            self.assertEqual(instance, Class_2)
            self.assertEqual(args, _args)
            self.assertEqual(kwargs, _kwargs)
            return wrapped(*args, **kwargs)

        wrapt.wrap_function_wrapper(__name__, 'Class_2.method',
                wrapper)

        result = Class_2.method(*_args, **_kwargs)

        self.assertEqual(result, (_args, _kwargs))
        self.assertEqual(called[0], (_args, _kwargs))

def test_wrap_static_method_module_name(self):

        _args = (1, 2)
        _kwargs = { 'one': 1, 'two': 2 }

        called = []

        def wrapper(wrapped, instance, args, kwargs):
            called.append((args, kwargs))
            self.assertEqual(instance, None)
            self.assertEqual(args, _args)
            self.assertEqual(kwargs, _kwargs)
            return wrapped(*args, **kwargs)

        wrapt.wrap_function_wrapper(__name__, 'Class_3.method',
                wrapper)

        result = Class_3.method(*_args, **_kwargs)

        self.assertEqual(result, (_args, _kwargs))
        self.assertEqual(called[0], (_args, _kwargs))

        @wrapt.decorator
        def _decorator(wrapped, instance, args, kwargs):
            return wrapped(*args, **kwargs)

        class Class(object):
            @_decorator
            @classmethod
            def function(cls, *args, **kwargs):
                return args, kwargs

            self.assertTrue(isinstance(function, wrapt.FunctionWrapper))
            self.assertTrue(isinstance(function, wrapt.ObjectProxy))

        instance = Class()

        self.assertFalse(isinstance(instance.function, wrapt.FunctionWrapper))
        self.assertTrue(isinstance(instance.function, wrapt.ObjectProxy))

def test_update_qualname_modified_on_original(self):
        def function():
            pass

        def wrapper(wrapped, instance, args, kwargs):
            return wrapped(*args, **kwargs)

        instance = wrapt.FunctionWrapper(function, wrapper)

        if six.PY3:
            method = self.test_update_qualname_modified_on_original
            self.assertEqual(instance.__qualname__,
                    (method.__qualname__ + '..function'))

        instance.__qualname__ = 'override_qualname'

        self.assertEqual(function.__qualname__, 'override_qualname')
        self.assertEqual(instance.__qualname__, 'override_qualname')

def test_re_bind_after_none(self):

        def function():
            pass

        def wrapper(wrapped, instance, args, kwargs):
            return wrapped(*args, **kwargs)

        _wrapper = wrapt.FunctionWrapper(function, wrapper)

        self.assertTrue(isinstance(_wrapper, wrapt.FunctionWrapper))

        instance = object()

        _bound_wrapper_1 = _wrapper.__get__(None, type(instance))

        self.assertTrue(_bound_wrapper_1._self_parent is _wrapper)

        self.assertTrue(isinstance(_bound_wrapper_1,
                wrapt.BoundFunctionWrapper))
        self.assertEqual(_bound_wrapper_1._self_instance, None)

        _bound_wrapper_2 = _bound_wrapper_1.__get__(instance, type(instance))

        self.assertTrue(_bound_wrapper_2._self_parent is _wrapper)

def test_update_doc_modified_on_original(self):
        def function():
            """documentation"""
            pass

        def wrapper(wrapped, instance, args, kwargs):
            return wrapped(*args, **kwargs)

        instance = wrapt.FunctionWrapper(function, wrapper)

        self.assertEqual(instance.__doc__, "documentation")

        instance.__doc__ = 'override_doc'

        self.assertEqual(function.__doc__, 'override_doc')
        self.assertEqual(instance.__doc__, 'override_doc')

def test_update_annotations_modified_on_original(self):
        def function():
            pass

        def wrapper(wrapped, instance, args, kwargs):
            return wrapped(*args, **kwargs)

        instance = wrapt.FunctionWrapper(function, wrapper)

        if six.PY3:
            self.assertEqual(instance.__annotations__, {})

        else:
            def run(*args):
                instance.__annotations__

            self.assertRaises(AttributeError, run, ())

        override_annotations = { 'override_annotations': '' }
        instance.__annotations__ = override_annotations

        self.assertEqual(function.__annotations__, override_annotations)
        self.assertEqual(instance.__annotations__, override_annotations)

def test_unpatch(self):
        gql = graphql.graphql
        unpatch()
        assert gql == graphql.graphql
        assert not isinstance(graphql.graphql, FunctionWrapper)
        patch()
        assert isinstance(graphql.graphql, FunctionWrapper)

        tracer, schema = get_traced_schema()
        graphql.graphql(schema, '{ hello }')
        span = tracer.writer.pop()[0]

        unpatch()
        assert gql == graphql.graphql

        cb_args = {}
        def test_cb(**kwargs):
            cb_args.update(kwargs)
        patch(span_callback=test_cb)
        assert isinstance(graphql.graphql, FunctionWrapper)

    @classmethod
    def function2(cls):
        print('function2')

    @wrapt.synchronized
    @staticmethod
    def function3():
        print('function3')

c1 = C1()

@wrapt.synchronized
class C2(object):
    pass

@wrapt.synchronized
class C3:
    pass

class C4(object):

    # XXX This yields undesirable results due to how class method is
    # implemented. The classmethod doesn't bind the method to the class
    # before calling. As a consequence, the decorator wrapper function
    # sees the instance as None with the class being explicitly passed
    # as the first argument. It isn't possible to detect and correct
    # this.

    @classmethod
    @wrapt.synchronized
    def function2(cls):
        print('function2')