Decorator to Find the Time of Execution of a Function in Python

Question

from datetime import datetime
import time

def decorating_func(own_func):
    print(datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
    return own_func

@decorating_func
def addition(num1,num2):
    return num1 + num2

@decorating_func
def multiply(num1,num2):
    return num1 * num2

print(addition(10,5))
time.sleep(71)
print(multiply(10,5))

I am new to using decorators, and I want the above code to produce

output:
2021-05-15 19:46:01
15
2021-05-15 19:47:11
50

instead of

2021-05-15 19:46:01
2021-05-15 19:46:01
15
50

Can someone explain how above code should be modified to achieve needed output?

Answer 1

A decorator takes a function and returns a new function to replace the original. It is evaluated immediately.

So:

def decorate(func):
  print("Decorate!")
  return func

@decorate
def add(x, y):
  return x + y

will print "Decorate!" even without invoking add.

What you want to do is "wrap" the original function in a newer function which does some work and then invokes the original.

def decorate(func):
  def _wrap(*args, **kwargs):
    print("Start")
    return func(*args, **kwargs)

  return _wrap

@decorate
def add(x, y):
  return x + y

Now "add" will be replaced by _wrap which when invoked will run your print statement first and then invoke the original function with the same args and return the result. If you want to run some "post" logic as well, you can leverage a try...finally:

def decorate(func):
  def _wrap(*args, **kwargs):
    print("Start")
    try:
       return func(*args, **kwargs)
    finally:
      print("End")

  return _wrap

@decorate
def add(x, y):
  return x + y

The finally block is guaranteed to execute. If you don't care about printing "End" in error scenarios, you could also store the return result:

def decorate(func):
  def _wrap(*args, **kwargs):
    print("Start")
    result = func(*args, **kwargs)
    print("End")
    return result

  return _wrap

@decorate
def add(x, y):
  return x + y

If you need to pass args into the decorator, you just create a decorator factory, i.e. a function which takes some args and returns a decorate function.

def makeDecorator(text):
  def decorate(func):
    def _wrap(*args, **kwargs):
      print(f"Start: {text}")
      try:
        return func(*args, **kwargs)
      finally:
        print("End")
    
    return _wrap
  return decorate

@makeDecorator(text="Hi")
def add(x, y):
  return x + y

The key is that decorators "replace" the decorated object. So if you're decorating a function, you need to return a function. The decorate function itself is not the function that gets invoked when you call add().

Additional Tip: When you return the _wrap function, its name will be "_wrap" rather than the name of the original function. The functools standard library provides some helpers such as the wraps decorator which you can use on _wrap to give it the same name as the original function. Useful for debugging.