Write a decorator with parameter in Python

Question

Suppose a decorator changes the behavior of a function to_lower():

def transition(fn):
    def to_upper(text: str):
        print('the original behavior:' + fn(text))
        return text.upper()
    
    def to_capital(text: str):
        print('the original behavior:' + fn(text))
        return text.capitalize()

    return to_upper

@transition
def to_lower(text: str):
    return text.lower()

print(to_lower('AaBb'))
>>>
the original behavior:aabb
AABB

This works fine, and if we change the return statement of transition from return to_upper to return to_capital, it changes the behavior from to_lower to to_capital, which makes AaBb to Aabb

Instead of manually modifying the return of the decorator, can we modify the decorator with sort of parameter like mode, if we call @transition(mode='to_upper'), it works as return to_upper and when we call @transition(mode='to_capital'), it works as return to_capital for the decorator?

This answer is very comprehensive and might help: https://stackoverflow.com/a/5929165/7013263 — dejanualex, Jan 05 '21 at 09:34

score 2 · Accepted Answer · answered Jan 05 '21 at 09:28

A very basic implementation:


def transition(mode):
    def deco(f):
        if mode == "to_upper":
            def wrapper(text):
                return text.upper()
        elif mode == "to_capital":
            def wrapper(text):
                return text.capitalize()
        else:
            # implement different behavior for other use cases  
            pass
        return wrapper
    return deco

    
    
        
    
@transition(mode="to_capital")
def to_lower(text):
    return text.lower()


@transition(mode="to_upper")
def to_lower(text):
    return text.lower()

abc · Answer 2 · 2021-01-05T09:26:41.983

You just need to add an additional layer for the decorator argument. Example:

import functools

def transition(mode):
    def wrapped(fn):
        @functools.wraps(fn)
        def inner(*args, **kwargs):
            res = fn(*args, **kwargs)
            if mode == 'to_capital':
                return res.capitalize()
            elif mode == 'to_upper':
                return res.upper()
            else:
                raise ValueError("invalid mode")

        return inner

    return wrapped

which can be applied as

@transition(mode='to_capital')
def to_lower(text: str):
    return text.lower()

score 0 · Answer 3 · answered Jan 08 '21 at 03:48

You could make parameters optional and also code is cleaner.

from functools import wraps

def process(fn, mode):
    @wraps(fn)
    def newfn(text):
        if mode == "to_upper":
            return text.upper()
        else:
            return text.capitalize()
    return newfn

def transition(fn=None, /, *, mode="to_upper"):
    def wrap(fn):
        return process(fn, mode)
    if fn is None:
        return wrap
    return wrap(fn)

@transition(mode="upper")
def to_lower(text: str):
    return text.lower()

@transition(mode="capitalize")
def to_lower(text: str):
    return text.lower()

@transition
def to_lower(text: str):
    return text.lower()

Write a decorator with parameter in Python

3 Answers3