Go like defer functionality in Python

Go has a nice keyword defer, which is used to defer the execution of some code until the surrounding function returns. It is useful for registering cleanup work to be done before the function returns. The prototypical example is closing a file handle:

package main

import (
  "os"
)

func main() {
  // Open a file   
  file, err := os.Open("test.txt")
  // I'm leaving out error handling and other stuff you would have in reality.

  // Lots of code here

  // Close the file handle when done
  file.Close()
}

You must remember to close the file handle when done with it, which means calling file.Close() when you return from the function. That can be tricky if there you write a lot of code in between opening the file and closing it. There can also be multiple early returns, and you need to remember to close the file in each. Even worse, the function could panic, and then you would still wish to close the file.

As a solution to this Go offer the defer statement. defer essentially tells Go to run some code when the function returns, no matter how it returns, even if it is a panic. So, you would instead do something like this:

package main

import (
  "os"
)

func main() {
  // Open a file   
  file, err := os.Open("test.txt")
  defer file.Close()
  // I'm leaving out error handling and other stuff you would have in reality.

  // Lots of code here
}

Here Go will defer running file.Close() until the function returns. This is one of the niceties I really like about Go. It lets me schedule cleanup at the same time as I create the mess, and whatever happens later it will be taken care of.

Sadly, in my DayJob, I don't program in Go, but in Python, and that left me feeling bereft of defer. Python has the with construct, which sort of does this, but it is not as neat once you have a lot of file handles and stuff to clean up. I set out to whip up a replacement in Python. Using it looks like this:

@defer
def main(defer):
  file = open('test.txt')
  defer(lambda: file.close())

I am using a decorator defer to inject a defer function as an argument into the python function. I then pass the function a lambda with the code I want to write on return. I need the lambda in order to delay running the code, otherwise the file.close() would run immediately. This is not as slick as the Go version, but I can't add language built-ins.

The implementation of defer is the following:

from functools import wraps

def defer(func):
  @wraps(func)
  def func_wrapper(*args, **kwargs):
    deferred = []
    defer = lambda f: deferred.append(f)
    try:
      return func(*args, defer=defer, **kwargs)
    finally:
      deferred.reverse()
      for f in deferred:
        f()
  return func_wrapper

Aside from the standard boilerplate for creating a python decorator the implementation is straight forward. I create a deferred list to store all the lambdas that are deferred. I then create a defer lambda that just appends whatever it receives to the deferred list. I pass the defer lambda to the function so that the function can use it to add things to the deferred list. Finally, after the function has run I reverse the order of deferred list do that things are cleaned up in reverse order they were created. Lastly, it’s just looping through all the deferred lambdas stored in deferred and running them.

I use try/finally so that the cleanup is run even if the function raises and exception instead of returning normally.