[pypy-commit] pypy py3k: add support for more complex types in enforceargs, like [int] or {str:int}; also move the imports as late as possible, to prevent circular imports

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: py3k
Changeset: r56772:f5c7bfc60abd
Date: 2012-08-21 11:48 +0200
http://bitbucket.org/pypy/pypy/changeset/f5c7bfc60abd/

Log:	add support for more complex types in enforceargs, like [int] or
	{str:int}; also move the imports as late as possible, to prevent
	circular imports

diff --git a/pypy/rlib/objectmodel.py b/pypy/rlib/objectmodel.py
--- a/pypy/rlib/objectmodel.py
+++ b/pypy/rlib/objectmodel.py
@@ -125,19 +125,39 @@
             return f
         return decorator
     #
-    from pypy.annotation.signature import annotationoftype
-    from pypy.annotation.model import SomeObject
-    def decorator(f):
+    def decorator(f): 
         def get_annotation(t):
+            from pypy.annotation.signature import annotation
+            from pypy.annotation.model import SomeObject
             if isinstance(t, SomeObject):
                 return t
-            return annotationoftype(t)
+            return annotation(t)
+        def get_type_descr_of_argument(arg):
+            # we don't want to check *all* the items in list/dict: we assume
+            # they are already homogeneous, so we only check the first
+            # item. The case of empty list/dict is handled inside typecheck()
+            if isinstance(arg, list):
+                item = arg[0]
+                return [get_type_descr_of_argument(item)]
+            elif isinstance(arg, dict):
+                key, value = next(arg.iteritems())
+                return {get_type_descr_of_argument(key): get_type_descr_of_argument(value)}
+            else:
+                return type(arg)
         def typecheck(*args):
+            from pypy.annotation.model import SomeList, SomeDict
             for i, (expected_type, arg) in enumerate(zip(types, args)):
                 if expected_type is None:
                     continue
                 s_expected = get_annotation(expected_type)
-                s_argtype = get_annotation(type(arg))
+                # special case: if we expect a list or dict and the argument
+                # is an empty list/dict, the typecheck always pass
+                if isinstance(s_expected, SomeList) and arg == []:
+                    continue
+                if isinstance(s_expected, SomeDict) and arg == {}:
+                    continue
+                #
+                s_argtype = get_annotation(get_type_descr_of_argument(arg))
                 if not s_expected.contains(s_argtype):
                     msg = "%s argument number %d must be of type %s" % (
                         f.func_name, i+1, expected_type)
diff --git a/pypy/rlib/test/test_objectmodel.py b/pypy/rlib/test/test_objectmodel.py
--- a/pypy/rlib/test/test_objectmodel.py
+++ b/pypy/rlib/test/test_objectmodel.py
@@ -444,6 +444,19 @@
     # in RPython there is an implicit int->float promotion
     assert f(42) == 42
 
+def test_enforceargs_complex_types():
+    @enforceargs([int], {str: int})
+    def f(a, b):
+        return a, b
+    x = [0, 1, 2]
+    y = {'a': 1, 'b': 2}
+    assert f(x, y) == (x, y)
+    assert f([], {}) == ([], {})
+    assert f(None, None) == (None, None)
+    py.test.raises(TypeError, "f(['hello'], y)")
+    py.test.raises(TypeError, "f(x, {'a': 'hello'})")
+    py.test.raises(TypeError, "f(x, {0: 42})")
+
 def test_enforceargs_no_typecheck():
     @enforceargs(int, str, None, typecheck=False)
     def f(a, b, c):