[PATCH 2 of 3] convert/mtn: use mtn "automate stdio" when available

[daniel.atallah at gmail.com]

# HG changeset patch
# User Daniel Atallah <daniel.atallah at gmail.com>
# Date 1300904816 14400
# Branch stable
# Node ID d2ef3495b1a62785683fc2a57212e8f2d2500fc6
# Parent  0fb617760b931cd25537e3362af52b2cfbd1a329
convert/mtn: use mtn "automate stdio" when available

This causes a single long-running mtn process to be used instead of spawning a
new process per operation.

diff --git a/hgext/convert/monotone.py b/hgext/convert/monotone.py
--- a/hgext/convert/monotone.py
+++ b/hgext/convert/monotone.py
@@ -19,6 +19,7 @@
 
         self.ui = ui
         self.path = path
+        self.automatestdio = False
 
         norepo = NoRepo(_("%s does not look like a monotone repository")
                         % path)
@@ -73,9 +74,105 @@
         self.rev = rev
 
     def mtnrun(self, *args, **kwargs):
+        if self.automatestdio:
+            return self.mtnrunstdio(*args, **kwargs)
+        else:
+            return self.mtnrunsingle(*args, **kwargs)
+
+    def mtnrunsingle(self, *args, **kwargs):
         kwargs['d'] = self.path
         return self.run0('automate', *args, **kwargs)
 
+    def mtnrunstdio(self, *args, **kwargs):
+        #Prepare the command in automate stdio format
+        command = []
+        for k, v in kwargs.iteritems():
+            command.append("%s:%s" % (len(k), k))
+            try:
+                command.append("%s:%s" % (len(v), v))
+            except TypeError:
+                pass
+        if len(command) > 0:
+            command.insert(0, 'o')
+            command.append('e')
+
+        command.append('l')
+        for arg in args:
+            command += "%s:%s" % (len(arg), arg)
+        command.append('e')
+        commandstr = ''.join(command)
+
+        self.ui.debug("mtn: Sending '%s'\n" % commandstr)
+        self.mtnwritefp.write(commandstr)
+        self.mtnwritefp.flush()
+
+        return self.mtnstdioreadcommandoutput(commandstr)
+
+    def mtnstdioreadpacket(self):
+        read = None
+        commandnbr = ''
+        while not read == ':':
+            read = self.mtnreadfp.read(1)
+            if read == '':
+                raise util.Abort('bad mtn packet - no end of commandnbr')
+            commandnbr += read
+        commandnbr = commandnbr[:-1]
+
+        stream = self.mtnreadfp.read(1)
+        if not stream in ['m', 'e', 'w', 'p', 't', 'l']:
+            raise util.Abort('bad mtn packet - bad stream type %s' % stream)
+
+        read = self.mtnreadfp.read(1)
+        if not read == ':':
+            raise util.Abort('bad mtn packet - no divider before size')
+
+        read = None
+        lengthstr = ''
+        while not read == ':':
+            read = self.mtnreadfp.read(1)
+            if read == '':
+                raise util.Abort('bad mtn packet - no end of packet size')
+            lengthstr += read
+        try:
+            length = long(lengthstr[:-1])
+        except TypeError:
+            raise util.Abort('bad mtn packet - bad packet size %s' % lengthstr)
+
+        read = self.mtnreadfp.read(length)
+        if not len(read) == length:
+            raise util.Abort('bad mtn packet - unable to read full packet " \
+            " read %s of %s' % (len(read), length))
+
+        return (commandnbr, stream, length, read)
+
+    def mtnstdioreadcommandoutput(self, command):
+        retval = ''
+        done = False
+        while not done:
+            commandnbr, stream, length, output = self.mtnstdioreadpacket()
+            self.ui.debug('mtn: read packet %s:%s:%s\n' %
+                (commandnbr, stream, length))
+
+            if stream == 'l':
+                #End of command
+                if not output == '0':
+                    raise util.Abort('mtn command \'%s\' returned %s' %
+                        (command, output))
+                done = True
+            elif stream in ['e', 'w']:
+                #Error, warning output
+                self.ui.warn(_('%s error:\n') % self.command)
+                self.ui.warn(output)
+            elif stream == 'p':
+                #Progress messages
+                self.ui.debug('mtn: ' + output)
+            elif stream == 'm':
+                #Main stream - command output
+                retval = output
+
+        return retval
+
+
     def mtnloadmanifest(self, rev):
         if self.manifest_rev == rev:
             return
@@ -225,3 +322,43 @@
         # This function is only needed to support --filemap
         # ... and we don't support that
         raise NotImplementedError()
+
+    def before(self):
+        #Check if we have a new enough version to use automate stdio
+        version = 0.0
+        try:
+            versionstr = self.mtnrunsingle("interface_version")
+            version = float(versionstr)
+        except:
+            raise util.Abort("unable to determine mtn automate interface "
+                "version")
+
+        if version >= 12.0:
+            self.automatestdio = True
+            self.ui.debug("mtn automate version %s - using automate stdio\n" %
+                (version))
+
+            #Mark commandline not to redirect stdin
+            self.redirectstdin = False
+            #launch the long-running automate stdio process
+            self.mtnwritefp, self.mtnreadfp = self._run2('automate', 'stdio',
+                '-d', self.path)
+            #read the headers
+            read = self.mtnreadfp.readline()
+            if not read == 'format-version: 2\n':
+                raise util.Abort('mtn automate stdio header unexpected: %s' %
+                    (read))
+            while not read == '\n':
+                read = self.mtnreadfp.readline()
+                if read == '':
+                    raise util.Abort("failed to reach end of mtn automate stdio"
+                        " headers")
+        else:
+            self.ui.debug("mtn automate version %s - not using automate stdio "
+                "(automate >= 12.0 - mtn >= 0.46 is needed)\n" %(version))
+
+    def after(self):
+        if self.automatestdio:
+            self.mtnwritefp.close()
+            self.mtnreadfp.close()
+