Как украсить метод, который принадлежит другому классу в моем текущем классе?

Question

У нас есть внутренняя библиотека, которая возвращает объект-потребитель kafka. Я хочу реализовать декоратор, который использует этот объект для запуска потребителя (метод для запуска потребителя также происходит из этой внутренней библиотеки). Цель состоит в том, чтобы импортировать этот декоратор и использовать его для любого метода, который требует запуска потребителя перед выполнением. Вот фрагмент из файла, скажем, utils.py:

from internal_library import KafkaMsgConsumer
class KafkaMessageConsumer(object):

    def __init__(self, topic_name, kafka_host, kafka_port='some_port'):
        self.topic_name = topic_name
        self.kafka_host = kafka_host
        self.kafka_port = kafka_port
        self.consumer = KafkaMsgConsumer(kafka_host_name=(str(self.kafka_host) + ":" + str(self.kafka_port)),
                         topic_name=self.topic_name)

    def consumer_required(self):
        def decorator():
            consumer = self.consumer.start_consumer()
            return consumer
        return decorator

Тогда есть мой основной скрипт, где я хочу использовать этот декоратор. Содержание из script.py:

mes = KafkaMessageConsumer(topic_name='some_topic', kafka_host='some_host',
                           kafka_port='some_port')


@mes.consumer_required()
def post_act():
    ''' some processing goes here before which I require the consumer to be started'''

Answer 1

Ваш декоратор должен принять декорированную функцию в качестве аргумента, и вам нужно вернуть соответствующую оболочку:

from functools import wraps

def consumer_required(self):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            self.consumer.start_consumer()
            return func(*args, **kwargs)
        return wrapper
    return decorator

В качестве альтернативы вы также можете вставить consumer в декорированную функцию через

kwargs.update(consumer=self.consumer.start_consumer())

Answer 2

import functools


def consumer_required(self):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            self.consumer.start_consumer()
            return func(*args, **kwargs)

        return wrapper

    return decorator