1 files changed, 43 insertions, 20 deletions

diff --git a/becommands/comment.py b/becommands/comment.py
index ec93262..045b331 100644
--- a/becommands/comment.py
+++ b/becommands/comment.py
@@ -15,40 +15,66 @@
 #    along with this program; if not, write to the Free Software
 #    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 """Add a comment to a bug"""
-from libbe import cmdutil, utility
+from libbe import cmdutil, bugdir, utility
 import os
 __desc__ = __doc__
 
 def execute(args):
     """
-    >>> from libbe import bugdir
-    >>> import os, time
-    >>> dir = bugdir.simple_bug_dir()
-    >>> os.chdir(dir.dir)
+    >>> import time
+    >>> bd = bugdir.simple_bug_dir()
+    >>> os.chdir(bd.root)
     >>> execute(["a", "This is a comment about a"])
-    >>> comment = dir.get_bug("a").list_comments()[0]
-    >>> comment.body
-    u'This is a comment about a\\n'
-    >>> comment.From == dir.rcs.get_user_id()
+    >>> bd.load()
+    >>> comment = bd.bug_from_shortname("a").comment_root[0]
+    >>> print comment.body
+    This is a comment about a
+    <BLANKLINE>
+    >>> comment.From == bd.rcs.get_user_id()
     True
     >>> comment.time <= int(time.time())
     True
     >>> comment.in_reply_to is None
     True
+
     >>> if 'EDITOR' in os.environ:
     ...     del os.environ["EDITOR"]
     >>> execute(["b"])
     Traceback (most recent call last):
     UserError: No comment supplied, and EDITOR not specified.
+
     >>> os.environ["EDITOR"] = "echo 'I like cheese' > "
     >>> execute(["b"])
-    >>> dir.get_bug("b").list_comments()[0].body
-    u'I like cheese\\n'
+    >>> bd.load()
+    >>> print bd.bug_from_shortname("b").comment_root[0].body
+    I like cheese
+    <BLANKLINE>
     """
     options, args = get_parser().parse_args(args)
-    if len(args) < 1:
-        raise cmdutil.UsageError()
-    bug, parent_comment = cmdutil.get_bug_and_comment(args[0])
+    if len(args) == 0:
+        raise cmdutil.UserError("Please specify a bug or comment id.")
+    if len(args) > 2:
+        help()
+        raise cmdutil.UserError("Too many arguments.")
+    
+    shortname = args[0]
+    if shortname.count(':') > 1:
+        raise cmdutil.UserError("Invalid id '%s'." % shortname)        
+    elif shortname.count(':') == 1:
+        # Split shortname generated by Comment.comment_shortnames()
+        bugname = shortname.split(':')[0]
+        is_reply = True
+    else:
+        bugname = shortname
+        is_reply = False
+    
+    bd = bugdir.BugDir(loadNow=True)
+    bug = bd.bug_from_shortname(bugname)
+    if is_reply:
+        parent = bug.comment_root.comment_from_shortname(shortname, bug_shortname=bugname)
+    else:
+        parent = bug.comment_root
+    
     if len(args) == 1:
         try:
             body = utility.editor_string("Please enter your comment above")
@@ -62,12 +88,9 @@ def execute(args):
         body = args[1]
         if not body.endswith('\n'):
             body+='\n'
-
-    comment = bug.new_comment(body)
-    if parent_comment is not None:
-        comment.in_reply_to = parent_comment.uuid
-    comment.save()
-
+    
+    comment = parent.new_reply(body=body)
+    bd.save()
 
 def get_parser():
     parser = cmdutil.CmdOptionParser("be comment ID COMMENT")