Как в Python написать декоратор, чтобы объединить логику двух функций?

Это отвлеченный пример, который описывает задачу, над которой работаю. У меня есть функция func_one с определенной логикой. Я написал вторую функцию func_two, которая часть кода берет из первой функции, чтобы затем на его основе продолжить свои вычисления.

Я бы не хотел повторять код первой функции во второй. Мне показалось, что с помощью декоратора проблему можно решить: логику первой функции реализовать во второй, не прибегая к дублированию скрипта. Однако правильно написать декоратор не вышло: два аргумента, как в приводимом ниже случае, излишни (интерпретатор ожидает один). Но как еще "сшить" обе функции через обертку, не знаю.

Код:

import numpy as np

# первая функция
def func_one(order):
    a = np.sin(x)/2
    return a

# вариант обертки (ошибочный)
def wrap_polynomial(func_one, func_two):
    def wrapper(order):
        request = func_one()
        res = func_two()
        print('Переменная b:\n', res)
    return wrapper

# вторая функция 
@wrap_polynomial
def func_two(order):
    b = a + np.cos(x) 
    return b

Ответы (1 шт):

Автор решения: Павел

Для решения вашей задачи декоратор не нужен.

Можно ограничиться обыкновенным вызовом первой функции внутри второй и затем использовать ее результат, либо, как написали в комментариях, использовать композицию при помощи вложенного вызова.

Но в образовательных целях расскажу вам о том, почему у вас не получилось создать декоратор:

  1. Декоратор должен принимать один аргумент, только ту функцию, которую он декорирует, а не две.

    Если вы дополнительно хотите опционально принимать ту функцию, которая будет предварительно выполняться, то для этого нужно сделать еще одну обертку, это называется декоратор с пробросом аргумента.

  2. Обычно в декораторе не указывают конкретные аргументы, а пишут *args, **kwrags, чтобы пробросить все аргументы в декорируемую функцию независимо от их количества.

  3. Внутри функции wrapper вы должны возвращать результат выполнения той функции, которая декорируется, иначе return в функции func_two не будет работать и результатом всегда будет None.

Выглядеть это должно примерно следующим образом:

import numpy as np

def func_one(x):
    a = np.sin(x) / 2
    return a

def wrap_polynomial(func_one=None):
    def decorator(func_two):
        def wrapper(*args, **kwargs):
            intermediate_result = func_one(*args, **kwargs)
            print('Предварительное значение: ', intermediate_result)
            result = func_two(arg=intermediate_result, **kwargs)
            return result
        return wrapper
    return decorator


@wrap_polynomial(func_one=func_one)
def func_two(arg=None):
    print('Получен аргумент: ', arg)
    b = np.cos(arg)
    return b

print(func_two(1))

Результат:

Предварительное значение: 0.42073549240394825
Получен аргумент: 0.42073549240394825
0.9127887886561931
→ Ссылка