git » pymisc » master » tree

[master] / contract.py

"""
Design by contract decorators for Python
Alberto Bertogli (albertito@blitiri.com.ar)
8/Jul/2005
-------------------------------------------

This module provides a set of three decorators to implement design by
contract, using pre and post functions embedded inside the function body,
like:

@contract
def f(x):
	@precond(f)
	def pre(x):
		assert x > 0

	@postcond(f)
	def post(ret, x):
		assert ret > 0

	return (2 * x) / x


The precondition will run as soon as it gets defined, so it's wise to put it
on top of the function definition.
The postcondition can be anywhere.
You can omit either or both, at your discresion.

The precondition function takes the same parameters as the function; the
postcondition function takes the returned value first, and then all the
function parameters.


Expressing the conditions as independant functions is trivial, and this module
doesn't implement it. Putting them inside the function have some drawbacks,
most notably it's harder to reuse them, and it's not really that different
from open-coding them directly in the function. However, it can be more
comfortable if you prefer looking at everything inside the same function body.
"""


def contract(f):
	"Decorator for functions that will support DBC."

	class Wrapper:
		"This class will wrap around the function"
		def __init__(self):
			# function in question
			self.f = f

			# post-condition function to run
			self.__atexit = None

			# arguments passed to f()
			self.args = []
			self.kwargs = {}

		def put_post(self, postfun):
			"Register a post-condition function"
			self.__atexit = postfun

		def run_pre(self, prefun):
			"Run a pre-condition function"
			prefun(*self.args, **self.kwargs)

		def __call__(self, *args, **kwargs):
			"Emulate a call to the function."
			# save the arguments first
			self.args = args
			self.kwargs = kwargs

			# run the function; here the precondition will be ran
			# by its decorator calling run_pre(), and put_post()
			# will be called by the postcondition decorator to
			# register the postcondition function
			r = self.f(*args, **kwargs)

			# if we have a postcondition, run it
			if self.__atexit:
				self.__atexit(r, *args, **kwargs)
			return r
	return Wrapper()


def precond(pfun):
	"""Decorator for the precondition function, takes the parent function
	as the first parameter, and returns the decorator to use."""
	def mkdeco(prefun):
		# run the precondition function
		pfun.run_pre(prefun)
		return prefun
	return mkdeco


def postcond(pfun):
	"""Decorator for the postcondition function, takes the parent function
	as the first parameter, and returns the decorator to use."""
	def mkdeco(postfun):
		# register the postcondition function
		pfun.put_post(postfun)
		return postfun
	return mkdeco


#
# Example!
#

if __name__ == '__main__':

	@contract
	def f(a):
		"My beautiful function."

		@precond(f)
		def pre(a):
			"Precondition."
			print 'pre', a

		@postcond(f)
		def post(ret, a):
			"Postcondition."
			print 'post', ret, a

		print 'f', a
		return 2 * a


	print '* run f'
	r = f(26)
	print '* f returned', r
	print '* exit f'