(Python) Парочка полезных декораторов

import should
  
@should.give((5,2),7)
@should.give(("aa","bbb"),"aabbb")
@should.give(([1],[2,3]), [1,2,3])
@should.give((1,1),1) # test
def add(a,b):
  return a+b
  
@should.throw((1,0), Exception)
@should.throw((5,0), ZeroDivisionError)
@should.throw((5,0), TypeError) # test
@should.throw((5,1), TypeError) # test
def div(a,b):
  return a/b


>pythonw -u "should.py"
[!] add(1, 1) should give 1, but got 2.
[!] div(5, 1) should raise TypeError, but raised nothing.
[!] div(5, 0) should raise TypeError, but raised ZeroDivisionError.
>Exit code: 0


Идея в том, чтоб описать тест-кейсы к функции в качестве декоратора. Тогда они будут выполнять двойную функцию:

• тестирование (автоматическая проверка при запуске кода после изменения)
• документация (примеры использования функции)

Реализация:

# file: should.py

def render_func_args(f, args):
  return '%s(%s)' % (f.__name__, ', '.join(map(repr,args)))

def _inform(msg):
  print '[!]', msg

def should(args, expected):
  def check(f):
    val = f(*args)
    if val != expected:
      _inform('%s should give %s, but got %s.' % \
        (render_func_args(f, args), repr(expected), repr(val)))    
  def decor(f):
    check(f)
    return f
    
  return decor
  
def should_raise(args, ExcptClass):
  def check(f):
    excpt = None
    try:
      f(*args)
    except Exception, ex:
      excpt = ex
      
    if excpt == None:
      _inform("%s should raise %s, but raised nothing." % \
        (render_func_args(f, args), ExcptClass.__name__))
    elif not isinstance(excpt, ExcptClass):
      _inform("%s should raise %s, but raised %s." % \
        (render_func_args(f, args), ExcptClass.__name__, excpt.__class__.__name__))
  def decor(f):
    check(f)
    return f
  
  return decor
  
throw = should_raise
give = should
  
if __name__=='__main__':
  import should
  
  @should.give((5,2),7)
  @should.give(("aa","bbb"),"aabbb")
  @should.give(([1],[2,3]), [1,2,3])
  @should.give((1,1),1) # test
  def add(a,b):
    return a+b
    
  @should.throw((1,0), Exception)
  @should.throw((5,0), ZeroDivisionError)
  @should.throw((5,0), TypeError) # test
  @should.throw((5,1), TypeError) # test
  def div(a,b):
    return a/b

* This source code was highlighted with Source Code Highlighter.

UPD. Тем кто не знаком с декораторами — ссылка для ознакомления.