Python type checker decorator.

typecheck.py

#!/usr/bin/env python
# encoding: utf-8

# Copyright (c) 2014, Ivan Alejandro
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
# The views and conclusions contained in the software and documentation are
# those of the authors and should not be interpreted as representing official
# policies, either expressed or implied, of the FreeBSD Project.
"""
Decorator that allows you to define the valid types of the parameters in a
function.
All the existing solutions that I've found had some problem, so I've made my
own.
This helper:
    - doesn't modify `args` or `kwargs`
    - doesn't move arguments from one place to another
    - support both `args` and `kwargs`
    - allows you to use all the python 2.x supported parameters usage
    - use named arguments type definition
    - allows you to enforce a subset of all the available arguments
    - rise a standard TypeError if the arguments type are wrong
"""

import functools
import inspect


def decorator_with_args(decorator):
    """
    Helper to decorate a decorator and be able to use arguments in a 'simple'
    decorator without too much hassle.
    """
    def new(*args, **kwargs):
        def new2(func):
            return decorator(func, *args, **kwargs)
        return new2
    return new


def check_type(arg, valid):
    """
    Check that the `arg` corresponds with type/types of `valid`.

    :param arg: the arg to check
    :type arg: any object
    :param valid: the type/types that is/are valid for `arg`
    :type valid: type or tuple(type, type, ...)
    """
    if isinstance(valid, tuple):
        if sum(isinstance(arg, type_) for type_ in valid) == 0:
            msg = "{0} is not a valid type.\nValid types: {1}"
            msg = msg.format(arg, ', '.join(str(x) for x in valid))
            raise TypeError(msg)

    if not isinstance(arg, valid):
        raise TypeError('%s is not a %s type' % (arg, valid))


@decorator_with_args
def typecheck(fn, **argument_types):
    """
    Return a decorator function that requires specified types.

    :param argument_types: each element of which is a type or class or a tuple
                           of several types or classes.
    :type argument_types: tuple (type or tuple of types)

    Example to require a string then a numeric argument

    @require(str, (int, long, float))
    def (a, b): ...
    """
    argnames, _, keywords_arg, _ = inspect.getargspec(fn)

    @functools.wraps(fn)
    def check_call(*args, **kwargs):
        """
        This wrapper check the defined types and call the `fn` callable if the
        types are ok.
        """
        callargs = inspect.getcallargs(fn, *args, **kwargs)

        # add the kwargs to the named list
        kwargs = callargs.get('kwargs', {})
        callargs.update(kwargs)

        # not named args are not checked
        callargs.pop('args', None)

        # print "Callargs: ", callargs

        for name, type_ in argument_types.iteritems():
            check_type(callargs[name], type_)

        # types are ok, do the call
        fn(*args, **kwargs)

    return check_call

# Examples:


class Empty():
    pass


@typecheck(a=int, b=float)
def foo(a, b):
    print "foo:", a+b


@typecheck(a=str, b=(tuple, list), c=bool, d=int)
def bar(a, b, c, d=1):
    print "bar:", c, a, ','.join(b)


@typecheck(a=int)
def baz(a, b):
    print "baz:", a, b


@typecheck(a=str, b=int)
def test(a, b=2, *args, **kwargs):
    print "test:", a, b, args, kwargs


@typecheck(a=str, b=int, c=Empty)
def demo(a, b, c):
    print "demo:", a, b, c


@typecheck(name=str, age=(int, float, long))
def name_age(name, age):
    print "name_age:", name, age


test('asdf')
test('asdf', 123, 'lkajsdf', 345, z='B')

e = Empty()
demo('ivan', 24, e)

foo(1, 2.)
bar('abc', ('bar', 'quz'), True)
baz(1, 2)
baz(1, 'asdf')

name_age("John Doe", 42)
name_age("John Doe", 'asdf')  # Error

Output:

test: asdf 2 () {}
test: asdf 123 ('lkajsdf', 345) {'z': 'B'}
demo: ivan 24 <__main__.Empty instance at 0x7fc072f70c20>
foo: 3.0
bar: True abc bar,quz
baz: 1 2
baz: 1 asdf
name_age: John Doe 42
Traceback (most recent call last):
  File "/home/user/tests/type-check-python/type-check-custom.py", line 171, in <module>
    name_age("John Doe", 'asdf')  # Error
  File "/home/user/tests/type-check-python/type-check-custom.py", line 115, in check_call
    check_type(callargs[name], type_)
  File "/home/user/tests/type-check-python/type-check-custom.py", line 75, in check_type
    raise TypeError(msg)
TypeError: asdf is not a valid type.
Valid types: <type 'int'>, <type 'float'>, <type 'long'>

Packaged and moved to its own repo, see https://github.com/ivanalejandro0/RequireType