[PATCH 10 of 12 RFC v2] hglib: a wrapper around hg's cmdserver, written in Python

[Idan Kamara]

# HG changeset patch
# User Idan Kamara <idankk86 at gmail.com>
# Date 1307111464 -10800
# Node ID 06f8e17e7af4b21df2ee211cf8b65ef348f0e323
# Parent  2aa1aab45314c38f5636360b3a9136eed6329bcb
hglib: a wrapper around hg's cmdserver, written in Python

diff -r 2aa1aab45314 -r 06f8e17e7af4 contrib/hglib/hglib.py
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/contrib/hglib/hglib.py	Fri Jun 03 17:31:04 2011 +0300
@@ -0,0 +1,216 @@
+import subprocess
+import struct
+import cStringIO
+
+HGPATH = 'hg'
+
+def connect(path, fin, out, err, extraconfigs=None):
+    return hgclient(path, fin, out, err, extraconfigs)
+
+class CommandError(Exception):
+    pass
+
+class ServerError(Exception):
+    pass
+
+class ResponseError(ServerError, ValueError):
+    pass
+
+class CapabilityError(ServerError):
+    pass
+
+class hgclient(object):
+    inputfmt = '>I'
+    outputfmt = '>cI'
+    outputfmtsize = struct.calcsize(outputfmt)
+    retfmt = '>i'
+
+    def __init__(self, path, fin, out, err, extraconfigs):
+        self.fin = fin
+        self.ochannels = {'o' : out, 'e' : err, 'd' : err}
+        cmdline = [HGPATH, '-R', path, 'serve', '--cmdserver', 'pipe']
+
+        if extraconfigs:
+            cmdline += ['--config'] + extraconfigs
+
+        self.server = subprocess.Popen(cmdline, stdin=subprocess.PIPE,
+                                       stdout=subprocess.PIPE)
+
+        self._readhello()
+
+    @property
+    def encoding(self):
+        """ get the servers encoding """
+        if not 'getencoding' in self.capabilities:
+            raise CapabilityError('getencoding')
+
+        if not self._encoding:
+            self.server.stdin.write('getencoding\n')
+            self._encoding = self._readfromchannel('r')
+
+        return self._encoding
+
+    def _readfromchannel(self, channel, discardoptinal=True):
+        ch, length = self._readheader()
+
+        # read-loop until data on the given channel arrives
+        while ch != channel:
+            if ch in self.ochannels:
+                data = self.server.stdout.read(length)
+                self.ochannels[ch].write(data)
+                self.ochannels[ch].flush()
+            elif ch.isupper(): # is this a required channel?
+                raise ResponseError("cannot discard required channel '%s'" % ch)
+            elif not discardoptinal:
+                raise ResponseError("expected data on channel '%s', "
+                                    "got '%s'" % (channel, ch))
+            else:
+                # discard this channels data
+                self.server.stdout.read(length)
+
+            ch, length = self._readheader()
+
+        return self.server.stdout.read(length)
+
+    def _readhello(self):
+        """ read the hello message the server sends when started """
+
+        self.capabilities = self._readfromchannel('o')[len('capabilities: '):]
+        if not self.capabilities:
+            raise ResponseError("expected 'capabilities: '")
+
+        self.capabilities = set(self.capabilities.split())
+
+        # at the very least the server should be able to run commands
+        assert 'runcommand' in self.capabilities
+
+        self._encoding = self._readfromchannel('o')[len('encoding: '):]
+        if not self._encoding:
+            raise ResponseError("expected 'encoding: '")
+
+    def _readheader(self):
+        data = self.server.stdout.read(hgclient.outputfmtsize)
+        return struct.unpack(hgclient.outputfmt, data)
+
+    def runcommand(self, args):
+        if not self.server or self.terminated():
+            raise ValueError("server not connected")
+
+        self.server.stdin.write('runcommand\n')
+
+        args = '\0'.join(args)
+        self.server.stdin.write(struct.pack(hgclient.inputfmt, len(args)))
+        self.server.stdin.write(args)
+        self.server.stdin.flush()
+
+        while True:
+            channel, length = self._readheader()
+
+            # input channels
+            if channel in 'IL':
+                assert length > 0
+
+                if channel == 'I':
+                    s = self.fin.read(length)
+                else:
+                    s = self.fin.readline(length)
+
+                self.server.stdin.write(struct.pack(hgclient.inputfmt, len(s)))
+
+                if s:
+                    self.server.stdin.write(s)
+            # output channels
+            elif channel in self.ochannels:
+                data = self.server.stdout.read(length)
+                self.ochannels[channel].write(data)
+                self.ochannels[channel].flush()
+            # result channel, command finished
+            elif channel == 'r':
+                # this is the return code
+                data = self.server.stdout.read(length)
+                break
+            # a channel that we don't know and can't ignore
+            elif channel.isupper():
+                raise ResponseError("unexpected data on required channel '%s'"
+                                    % channel)
+            # optional channel
+            else:
+                self.server.stdout.read(length)
+
+        ret = struct.unpack(hgclient.retfmt, data)[0]
+        return ret
+
+    def outputruncommand(self, args):
+        """ run the command specified by args, returning (ret, output, error) """
+        self.redirect()
+        try:
+            ret = self.runcommand(args)
+        finally:
+            out, err = self.restore()
+
+        return (ret, out, err)
+
+    def inputruncommand(self, args, fin):
+        old = self.fin
+        self.fin = fin
+        try:
+            return self.outputruncommand(args)
+        finally:
+            self.fin = old
+
+    def close(self):
+        self.server.stdin.close()
+        self.server.wait()
+        ret = self.server.returncode
+        self.server = None
+
+        return ret
+
+    def terminated(self):
+        if self.server.poll() is not None:
+            return self.server.returncode
+        else:
+            return None
+
+    def redirect(self):
+        self.old = self.ochannels['o'], self.ochannels['e']
+        self.ochannels['o'], self.ochannels['e'] = \
+            cStringIO.StringIO(), cStringIO.StringIO()
+
+    def restore(self):
+        out, err = self.ochannels['o'].getvalue(), self.ochannels['e'].getvalue()
+        self.ochannels['o'], self.ochannels['e'] = self.old
+        return out, err
+
+    def status(self):
+        ret, out, err = self.outputruncommand(['status', '-0'])
+
+        if ret:
+            raise CommandError(ret, out, err)
+
+        d = dict((c, []) for c in 'MARC!?I ')
+
+        for entry in out.split('\0'):
+            if entry:
+                t, f = entry.split(' ', 1)
+                d[t].append(f)
+
+        return d
+
+    def import_(self, patch):
+        if isinstance(patch, str):
+            fp = open(patch)
+        else:
+            assert hasattr(patch, 'read')
+            assert hasattr(patch, 'readline')
+
+            fp = patch
+
+        try:
+            ret, out, err = self.inputruncommand(['import', '-'], fp)
+
+            if ret:
+                raise CommandError(ret, out, err)
+        finally:
+            if fp != patch:
+                fp.close()