Level up your code with Python decorators | by Alexander Bailey | Towards Data Science

There comes a point in every Python user’s life where you can level up from writing good code, to great code.
Once you’ve mastered core Python functionality like list comprehensions and ternary operators, you should be ready to write more readable Python. Decorators are key to upgrading your code’s readability, and here we’ll cover the basics of them, as well as 3 ways to use them.
In short, decorators are functions that wrap other functions. If there’s something that you want at the start and end of a function, then decorators are here to make your life a lot easier.
Decorators can be spotted as an @ sign before a function definition. You might have spotted them in a flask app or a click CLI but the easiest way to explain how they work, is to work through a short example.
So let’s say at the start of the function we want to print “Started’’ and at the end we want to print “Finished’’. To achieve this we could do the following:
def circle_area(radius):    print("Started: circle_area")    area = 3.142 * radius ** 2    print("Finished")    return areaarea = circle_area(2)
This method works, but now our code is cluttered with print statements. The actual meat of the function could be written in just a single line but instead we’ve used three! So let’s make this clearer by making a new function that contains the main ingredient of our operation, and another containing just the print statements:
def circle_area(radius):    return 3.142 * radius ** 2def circle_area_and_print(radius):    print("Started: circle_area_and_print")    area = circle_area(radius)    print("Ended")    return areaarea = circle_area_and_print(2)
At first glance you could argue this looks no clearer, as we have more code than before. At second glance you might notice that the code is much more readable from our additions. Our work here isn’t done however, what if we wanted to get more out of our code?
Let’s say for example that we want to print before and after other functions too. As lazy data scientists, we don’t want to write the same line of code more than once, so let’s try to generalise the code we already have. We can do this with this new printer function.
def printer(function):    def new_function(*args):        print(f"Started: {function.__name__}")        output = function(*args)        print("Finished")        return output    return new_functiondef circle_area(radius):    return 3.142 * radius ** 2area = printer(circle_area)(2)
So what does the printer function do here?
The beauty of this function is that it constructs a new function by taking our original function and adding it to the print statements before and after the passed function is executed. Now printer returns the new function handle, allowing us to take any arbitrary function and return it with print statements.
This is where the Python decorator really comes into it’s own. Now we’ve written the printer function to return a modified version of our original function, all we need to do is prepend the function definition with @printer :
@printerdef circle_area(radius):    return 3.142 * radius ** 2area = circle_area(2)
Even in this simple example, when you compare the use of the printer decorator with the code we wrote at the beginning, our code is much more legible.
There are a few nifty tricks of the trade for decorators that you can use to make your code easier to understand and read.
We’ll take a look at:
Logging (similar to the first example but somewhat more useful)Type checkingError handling
Logging (similar to the first example but somewhat more useful)
Type checking
Error handling
For the first example, we need to write a small function to set up logging in our script:
import loggingdef setup_logging(name="logger",                  filepath=None,                  stream_log_level="DEBUG",                  file_log_level="DEBUG"):        logger = logging.getLogger(name)    logger.setLevel("DEBUG")    formatter = logging.Formatter(        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'    )    ch = logging.StreamHandler()    ch.setLevel(getattr(logging, stream_log_level))    ch.setFormatter(formatter)    logger.addHandler(ch)    if filepath is not None:        fh = logging.FileHandler(filepath)        fh.setLevel(getattr(logging, file_log_level))        fh.setFormatter(formatter)        logger.addHandler(fh)    return loggerlogger = setup_logging(name="default_log",filepath="logger.log")
Now we can write our new logging decorator:
def log_decorator(log_name):    def log_this(function):        logger = logging.getLogger(log_name)        def new_function(*args,**kwargs):            logger.debug(f"{function.__name__} - {args} - {kwargs}")            output = function(*args,**kwargs)            logger.debug(f"{function.__name__} returned: {output}")            return output        return new_function    return log_this
This is slightly more complicated that the printer decorator we looked at before. This time, we wanted to build a decorator that can take an argument (the name of the logger).
However the decorator function can only take one argument, so we wrap this function in another that can take many arguments. It takes a while to wrap your head around this inception-like function wrapping, but once you do it’s an invaluable tool in the data scientist’s tool-box.
Now this is a controversial one, so hold on to your hats we could be in for a bumpy ride...
Some will say that this decorator type checks a function’s inputs in a way that many would consider to be “un-Pythonic”. Python is a dynamically-typed language, so at the risk of ruffling too many feathers, it is important to note that in general when writing Python we should beg for forgiveness rather than ask permission.
However if you’re willing to be a little bit cheeky with the laws of Python, then you’ll find this a really handy tool, especially when developing code. It makes it crystal clear when identifying wrong types that are being passed around.
With that all in mind, here’s the code:
def accepts(*types):    def check_accepts(function):        assert len(types) == function.__code__.co_argcount,\            "Number of typed inputs must match the function inputs"        def new_function(*args, **kwargs):            for (a, t) in zip(args, types):                assert isinstance(a, t), \                       "arg %r does not match %s" % (a,t)            return function(*args, **kwargs)        return new_function    return check_accepts@accepts((int,float))def circle_area(radius):    return 3.142 * radius ** 2
So now we’ve ensured that the inputs to the circle_area function will only accept type ints or floats, otherwise it will raise an AssertionError.
For this example, let’s assume we’re trying to access data from an API but the API only allows for a given number of requests per minute. We can write a decorator that wraps the API call and will continue to try and get the data until it’s successful.
import apiimport timeAPI_WAIT_TIME = 5 #minutesMAX_RETRIES = 10def error_handling(api_function):    def trial(*args, num_retries=0, **kwargs):        try:            return api_function(*args, **kwargs)        except api.error.RateLimitError:            if num_retries > MAX_RETRIES:                raise RuntimeError("Too many retries")            else:                msg = f"rate limit reached. Waiting {API_WAIT_TIME} minutes ..."                time.sleep(API_WAIT_TIME * 60)                return trial(*args, num_retries=num_retries + 1, **kwargs)    return trial
We can use this decorator each time we request data from the api module. The decorator will keep trying until it gets the data or it hits the maximum number of allowed retries, if that’s the case, a RuntimeError will be raised.
Equally, a very similar version of this function could be used to arbitrarily catch different exceptions and deal with them in different ways. I’ll leave it up to the imagine of the reader to think about new ways to use this type of decorator.
Hopefully you now have a better idea of what a Python decorator is and feel confident using it in your codebase (if not, then I’ve failed you...).
The key take away here is that if you have a function that frequently occurs at the start and end of another function, then a decorator will be your new best friend. Decorators are going to help your code to be not only more readable, but also more modular and reusable.