decorator with change of docstring

Decorator that checks return value and changes docstring.py

import functools
import unittest
from pydoc import render_doc

def returns(t):
    def checktype(func):
        def checked_func(*args, **kwargs):
            rval = func(*args, **kwargs)
            if not isinstance(rval, t):
                raise TypeError("Return type is not of type: int but of type "
                                + str(type(rval)))
            return rval
        checked_func = functools.update_wrapper(checked_func, func)
        checked_func.func_doc += "\nReturn type: " + str(t)
        return checked_func
    return checktype


@returns((int, float))
def foo(a):
    """Computes the square of the input argument"""
    return a * a

print(render_doc(foo))


class DecoTest(unittest.TestCase):
    def test_int(self):
        self.assertEqual( 9, foo(3) )
    def test_float(self):
        self.assertEqual( 9.0, foo(3.0) )
    def test_complex(self):
        self.assertRaises( TypeError, foo, 3.0j)

unittest.main()