dispatch: generalize signature checking for extension command wrapping

mercurial/dispatch.py

@@ -7,7 +7,7 b''
 
 from i18n import _
 from repo import RepoError
-import os, sys, atexit, signal, pdb, traceback, socket, errno, shlex, time
+import os, sys, atexit, signal, pdb, socket, errno, shlex, time
 import util, commands, hg, lock, fancyopts, revlog, version, extensions, hook
 import cmdutil
 import ui as _ui
@@ -356,9 +356,9 b' def _dispatch(ui, args):'
                     raise RepoError(_("There is no Mercurial repository here"
                                       " (.hg not found)"))
                 raise
-        d = lambda: func(ui, repo, *args, **cmdoptions)
-    else:
-        d = lambda: func(ui, *args, **cmdoptions)
+        args.insert(0, repo)
+
+    d = lambda: util.checksignature(func)(ui, *args, **cmdoptions)
 
     # run pre-hook, and abort if it fails
     ret = hook.hook(lui, repo, "pre-%s" % cmd, False, args=" ".join(fullargs))
@@ -374,11 +374,7 b' def _runcommand(ui, options, cmd, cmdfun'
     def checkargs():
         try:
             return cmdfunc()
-        except TypeError:
-            # was this an argument error?
-            tb = traceback.extract_tb(sys.exc_info()[2])
-            if len(tb) != 2: # no
-                raise
+        except util.SignatureError:
             raise ParseError(cmd, _("invalid arguments"))
 
     if options['profile']:

mercurial/extensions.py

@@ -96,7 +96,8 b' def wrapcommand(table, command, wrapper)'
 
     origfn = entry[0]
     def wrap(*args, **kwargs):
-        return wrapper(origfn, *args, **kwargs)
+        return util.checksignature(wrapper)(
+            util.checksignature(origfn), *args, **kwargs)
 
     wrap.__doc__ = getattr(origfn, '__doc__')
     wrap.__module__ = getattr(origfn, '__module__')

mercurial/util.py

@@ -13,7 +13,7 b' platform-specific details from the core.'
 """
 
 from i18n import _
-import cStringIO, errno, getpass, re, shutil, sys, tempfile
+import cStringIO, errno, getpass, re, shutil, sys, tempfile, traceback
 import os, stat, threading, time, calendar, ConfigParser, locale, glob, osutil
 import imp
 
@@ -671,6 +671,21 b' def system(cmd, environ={}, cwd=None, on'
         if cwd is not None and oldcwd != cwd:
             os.chdir(oldcwd)
 
+class SignatureError:
+    pass
+
+def checksignature(func):
+    '''wrap a function with code to check for calling errors'''
+    def check(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except TypeError:
+            if len(traceback.extract_tb(sys.exc_info()[2])) == 1:
+                raise SignatureError
+            raise
+
+    return check
+
 # os.path.lexists is not available on python2.3
 def lexists(filename):
     "test whether a file with this name exists. does not follow symlinks"