Commits

Antonio Cuni  committed d84e1c8

add the possibility of doing @enforceargs(foo=int) in case we want to enforce only one specific argument

  • Participants
  • Parent commits 6dab408

Comments (0)

Files changed (2)

File pypy/rlib/objectmodel.py

 
 specialize = _Specialize()
 
-def enforceargs(*types, **kwds):
+def enforceargs(*types_, **kwds):
     """ Decorate a function with forcing of RPython-level types on arguments.
     None means no enforcing.
 
     typechecking by passing ``typecheck=False`` to @enforceargs.
     """
     typecheck = kwds.pop('typecheck', True)
-    if kwds:
-        raise TypeError, 'got an unexpected keyword argument: %s' % kwds.keys()
+    if types_ and kwds:
+        raise TypeError, 'Cannot mix positional arguments and keywords'
+
     if not typecheck:
         def decorator(f):
-            f._annenforceargs_ = types
+            f._annenforceargs_ = types_
             return f
         return decorator
     #
-    def decorator(f): 
+    def decorator(f):
         def get_annotation(t):
             from pypy.annotation.signature import annotation
             from pypy.annotation.model import SomeObject
         # not RPython. Instead, we generate a function with exactly the same
         # argument list
         srcargs, srcvarargs, srckeywords, defaults = inspect.getargspec(f)
+        if kwds:
+            types = tuple([kwds.get(arg) for arg in srcargs])
+        else:
+            types = types_
         assert len(srcargs) == len(types), (
             'not enough types provided: expected %d, got %d' %
             (len(types), len(srcargs)))

File pypy/rlib/test/test_objectmodel.py

         return a+b
     assert f(2) == 42
 
+def test_enforceargs_keywords():
+    @enforceargs(b=int)
+    def f(a, b, c):
+        return a+b
+    assert f._annenforceargs_ == (None, int, None)
+
 def test_enforceargs_int_float_promotion():
     @enforceargs(float)
     def f(x):