« Back to Index

Function Timer Decorator

View original Gist on GitHub

Tags: #python

1. Python Function Timer Decorator Sync.py

import functools
import time


def timer(func):
    """calculate run time of decorated function."""

    @functools.wraps(func)
    def wrap_timer(*args, **kwargs):
        start_time = time.perf_counter()  # could also use perf_counter_ns()
        result = func(*args, **kwargs)

        end_time = time.perf_counter()
        run_time = end_time - start_time
        print(run_time)
        return result

    return wrap_timer


@timer
def slow_op():
    print("start")
    time.sleep(5)
    print("end")
    return "done"


print(slow_op())

2. Python Function Timer Decorator Async.py

import asyncio  # NEW: import asyncio module
import functools
import time


def timer(func):
    """calculate run time of decorated function."""

    @functools.wraps(func)
    async def wrap_timer(*args, **kwargs):  # NEW: async
        start_time = time.perf_counter_ns()
        result = await func(*args, **kwargs)  # NEW: await

        end_time = time.perf_counter_ns()
        run_time = end_time - start_time
        print(run_time)
        return result

    return wrap_timer


@timer
async def slow_op():  # NEW: async
    print("start")
    time.sleep(5)
    print("end")


asyncio.run(slow_op())  # NEW: run async function on event loop

3. Python Function Timer Decorator Async with Arguments.py

import asyncio
import functools
import time


def timer(metric_name):  # NEW: we're now accepting an argument
    """calculate run time of decorated function."""

    def nested_timer(func):  # NEW: basically the wrapping parent function is new and all other code was indented
        @functools.wraps(func)
        async def wrap_timer(*args, **kwargs):
            print(f"metric_name: {metric_name}")

            start_time = time.perf_counter_ns()
            result = await func(*args, **kwargs)
            end_time = time.perf_counter_ns()

            run_time = end_time - start_time
            print(run_time)

            return result

        return wrap_timer

    return nested_timer


@timer("foo.bar")  # NEW: passing in an argument
async def slow_op():
    print("start")
    time.sleep(5)
    print("end")


asyncio.run(slow_op())

4. Python Function Timer Decorator with check for Async and Non-Async.py

import asyncio
import functools
import time


def timer(metric_name):
    """calculate run time of decorated function."""

    def nested_timer(func):
        if asyncio.iscoroutinefunction(func):
            @functools.wraps(func)
            async def wrap_timer(*args, **kwargs):
                print("async call")
                print(f"metric_name: {metric_name}")

                start_time = time.perf_counter_ns()
                result = await func(*args, **kwargs)
                end_time = time.perf_counter_ns()

                run_time = end_time - start_time
                print(run_time)

                return result
        else:
            @functools.wraps(func)
            def wrap_timer(*args, **kwargs):
                print("non-async call")
                print(f"metric_name: {metric_name}")

                start_time = time.perf_counter_ns()
                result = func(*args, **kwargs)
                end_time = time.perf_counter_ns()

                run_time = end_time - start_time
                print(run_time)

                return result

        return wrap_timer

    return nested_timer


@timer("foo.bar")
async def slow_op(a, b, c):
    print(a, b, c)
    print("start")
    time.sleep(5)
    print("end")


asyncio.run(slow_op("A", "B", "C"))


@timer("foo.bar")
def another_slow_op():
    print("start another")
    time.sleep(3)
    print("end another")


another_slow_op()

5. Python Function Timer Decorator with check for Async and Non-Async and works as Context Manager.py

import asyncio
import functools
import time


class Timer():
    def __init__(self, metric_name):
        self.metric_name = metric_name

    def __call__(self, func):
        """decorator implementation."""

        if asyncio.iscoroutinefunction(func):
            @functools.wraps(func)
            async def wrap_timer(*args, **kwargs):
                print("async call")
                print(f"metric_name: {self.metric_name}")

                start_time = time.perf_counter_ns()
                result = await func(*args, **kwargs)
                end_time = time.perf_counter_ns()

                run_time = end_time - start_time
                print(run_time)

                return result
        else:
            @functools.wraps(func)
            def wrap_timer(*args, **kwargs):
                print("non-async call")
                print(f"metric_name: {self.metric_name}")

                start_time = time.perf_counter_ns()
                result = func(*args, **kwargs)
                end_time = time.perf_counter_ns()

                run_time = end_time - start_time
                print(run_time)

                return result

        return wrap_timer

    def __enter__(self):
        self.start_time = time.perf_counter_ns()
        return self

    def __exit__(self, *args):
        end_time = time.perf_counter_ns()
        run_time = end_time - self.start_time
        print(run_time)


def timer(metric_name):
    return Timer(metric_name)


@timer("foo.bar")
async def slow_op(a, b, c):
    print(a, b, c)
    print("start")
    time.sleep(5)
    print("end")


asyncio.run(slow_op("A", "B", "C"))


@timer("foo.bar")
def another_slow_op():
    print("start another")
    time.sleep(3)
    print("end another")


another_slow_op()

with timer("foo.bar"):
    print("context manager start")
    time.sleep(10)
    print("context manager finished")