2 files changed, 16 insertions, 7 deletions
diff --git a/libbe/command/base.py b/libbe/command/base.py
index 357940f..2f0ccc6 100644
--- a/libbe/command/base.py
+++ b/libbe/command/base.py
@@ -294,7 +294,7 @@ class Command (object):
         if argument == None:
             ret = ['--%s' % o.name for o in self.options]
             if len(self.args) > 0 and self.args[0].completion_callback != None:
-                ret.extend(self.args[0].completion_callback(self, argument))
+                ret.extend(self.args[0].completion_callback(self, argument, fragment))
             return ret
         elif argument.completion_callback != None:
             # finish a particular argument
@@ -486,12 +486,18 @@ class UserInterface (object):
         raise NotImplementedError
 
     def run(self, command, options=None, args=None):
-        command.ui = self
-        self.io.setup_command(command)
-        self.storage_callbacks.setup_command(command)
+        self.setup_command(command)
+        return command.run(options, args)
+
+    def setup_command(self, command):
+        if command.ui == None:
+            command.ui = self
+        if self.io != None:
+            self.io.setup_command(command)
+        if self.storage_callbacks != None:
+            self.storage_callbacks.setup_command(command)        
         command.restrict_file_access = self.restrict_file_access
         command._get_user_id = self._get_user_id
-        return command.run(options, args)
 
     def _get_user_id(self):
         """Callback for use by commands that need it."""
diff --git a/libbe/ui/command_line.py b/libbe/ui/command_line.py
index b36d251..1c7399d 100755
--- a/libbe/ui/command_line.py
+++ b/libbe/ui/command_line.py
@@ -277,24 +277,27 @@ def main():
     ui.restrict_file_access = False
     ui.storage_callbacks = None
     be = BE(ui=ui)
+    ui.setup_command(be)
+
     parser = CmdOptionParser(be)
     try:
         options,args = parser.parse_args()
     except CallbackExit:
         return 0
     except libbe.command.UserError, e:
-        print >> be.stdout, 'ERROR:\n', e
+        print >> ui.io.stdout, 'ERROR:\n', e
         return 1
 
     command_name = args.pop(0)
     try:
         Class = libbe.command.get_command_class(command_name=command_name)
     except libbe.command.UnknownCommand, e:
-        print >> be.stdout, e
+        print >> ui.io.stdout, e
         return 1
 
     ui.storage_callbacks = libbe.command.StorageCallbacks(options['repo'])
     command = Class(ui=ui)
+    ui.setup_command(command)
 
     if command.name in ['comment']:
         paginate = 'never'