Python: PyCharm IDE method chaining suggestions not working when using decorators

Question

Given below code:

def method_friendly_decorator(method_to_decorate):
    def wrapper(self, *args, **kwargs):
        print('I got args! Look:'.join(*args, **kwargs))
        method_to_decorate(self, *args, **kwargs)

        return self

    return wrapper


class DecorateMe(object):

    def __init__(self):
        self.first_name: str = ''
        self.last_name: str = ''

    @method_friendly_decorator
    def print_first_name(self, first_name: str):
        self.first_name = first_name
        print(f'My first name is {first_name}')
        return self

    @method_friendly_decorator
    def print_last_name(self, last_name: str):
        self.last_name = last_name
        print(f'My last name is {last_name}')
        return self

    @method_friendly_decorator
    def print_full_name(self):
        print(f'My full name is {self.first_name} {self.last_name}')
        return self

I'd like to be able to do:

def test_naming():
    DecorateMe().print_first_name('Test').print_last_name('User')

When I do not use the decorators on the functions in the DecorateMe class, I can see the suggested functions available, like below:

But, if I keep the method decorators, then I don't get the suggestions any longer:

I'm wondering if this is just "how it works" or if I am doing something wrong here?

Keeping in mind that, when using the decorators, the method chain still works if you know the exact method name to use, but to me, that's defeating the purpose since the whole idea is to be able to see the available methods.

Answer 1

You can define the return type for your wrapper like this:

def method_friendly_decorator(method_to_decorate):
    def wrapper(self, *args, **kwargs) -> "DecorateMe":  # Define return type here
        print('I got args! Look:'.join(*args, **kwargs))
        return method_to_decorate(self, *args, **kwargs)
    return wrapper

Put DecorateMe in quotes to not get NameError: name 'DecorateMe' is not defined . This is called a forward reference , from PEP 484

Side note: Your wrapper is hard-coded to always return self , making the return statement in each method ignored. It should return the return value of the decorated function as I did above