How to speed up a Python function (#python #cache #dev)

How to speed up a Python function

Overview

This post shows a way to speed up the execution of an indeppontent Python function. It’s a simple way that can easily overlooked.

So, imagine the following scenario: you wrote an amazing Python application, and it’s working perfectly, but it takes some time to run. After debugging you found out that it’s because of an expensive function that is called multiple times. It’s a pretty optimized function, so there’s nothing much you can do to improve performance there, however, this function is indeponent, meaning that it will always return the same output for the same input.

There’s a one-liner solution that can help you speed up things: @functools.cache

Let’s see an exmaple using the Fibonacci sequence, since calculating it is an expensive operation.

n_map = {
    0: 0,
    1: 1,
    2: 1,
}

def fibonacci(n):
    if n <= 0:
        return 0

    return n_map.get(n, fibonacci(n - 1) + fibonacci(n - 2))

In the example above, the function is called multiple times (recursive) with the same input, and we expect the same result, so we can use this decorator to cache the results.

To test it, we will:

  1. Create two versions of it: cached and uncached;
  2. Use timeit to measure the performance of each one.
  3. Create a wrapper to capture the result, so we can make sure we improved performance, but didn’t mess up the results.

Here’s the cached version:

import functools

cached_result = -1

n_map = {
    0: 0,
    1: 1,
    2: 1,
}

@functools.cache
def cached_fibonacci(n):
    if n <= 0:
        return 0

    return n_map.get(n, cached_fibonacci(n - 1) + cached_fibonacci(n - 2))

and here’s the uncached version:

uncached_result = -1

n_map = {
    0: 0,
    1: 1,
    2: 1,
}

def uncached_fibonacci(n):
    if n <= 0:
        return 0

    return n_map.get(n, uncached_fibonacci(n - 1) + uncached_fibonacci(n - 2))

def uncached_result_wrapper(n):
    global uncached_result
    uncached_result = uncached_fibonacci(n)

(Please, ignore the code duplication, it’s it’s to make things easier for this example)

Lastly, this is the code to measure the performance:

import timeit

n = 40  # Desired Fibonacci number to test
cached_time_taken = timeit.timeit('cached_result_wrapper(n)', globals=globals(), number=1)
uncached_time_taken = timeit.timeit('uncached_result_wrapper(n)', globals=globals(), number=1)

print(f"Time taken to calculate the {n}th Fibonacci ({cached_result}) 1 times cached: {cached_time_taken:.6f} seconds")
print(f"Time taken to calculate the {n}th Fibonacci ({uncached_result}) 1 times uncached: {uncached_time_taken:.6f} seconds")

Running this (on my machine) will output something like:

Time taken to calculate the 40th Fibonacci (102334155) 1 times cached: 0.000073 seconds
Time taken to calculate the 40th Fibonacci (102334155) 1 times uncached: 45.812993 seconds

Process finished with exit code 0

So, you can see that the same exact code is way more efficient when using the @functools.cache decorator. Without the caching, the code takes almost 46 seconds to run, but with it, it takes less than a millisecond.

It works by creating a wrapper around your function and memoizing the results, so the first time an input is presented, it will calculate the result and store it in a dictionary. The next time the same input is presented, it will return the result from the dictionary, instead of recalculating it.

The @functools.cache decorator is available since Python 3.9, so make sure you are using a compatible version. If you’re a bit behind and can’t use it, you can use the functools.lru_cache instead, which is available since Python 3.2.

  • Documentation to the @functools.cache can be found here.
  • Documentation to the functools.lru_cache can be found here.

If you want this working example, you can get it on my GithHub

Hope that helps. 🙂

Translations: