Plugins/VcsPlugins/vcsMercurial/HgClient.py

# -*- coding: utf-8 -*-

# Copyright (c) 2011 Detlev Offenbach <detlev@die-offenbachs.de>
#

"""
Module implementing an interface to the Mercurial command server.
"""

import struct
import io

from PyQt4.QtCore import QProcess, QProcessEnvironment, QObject, QByteArray, \
    QCoreApplication

import Preferences


class HgClient(QObject):
    """
    Class implementing the Mercurial command server interface.
    """
    InputFormat = ">I"
    OutputFormat = ">cI"
    OutputFormatSize = struct.calcsize(OutputFormat)
    ReturnFormat = ">i"
    
    def __init__(self, repoPath, encoding, parent=None):
        """
        Constructor
        
        @param repoPath root directory of the repository (string)
        @param encoding encoding to be used by the command server (string)
        @param parent reference to the parent object (QObject)
        """
        super().__init__(parent)
        
        self.__server = QProcess()
        self.__started = False
        self.__version = None
        self.__encoding = Preferences.getSystem("IOEncoding")
        
        # connect signals
        self.__server.finished.connect(self.__serverFinished)
        
        # generate command line and environment
        self.__serverArgs = []
        self.__serverArgs.append("serve")
        self.__serverArgs.append("--cmdserver")
        self.__serverArgs.append("pipe")
        self.__serverArgs.append("--config")
        self.__serverArgs.append("ui.interactive=True")
        if repoPath:
            self.__serverArgs.append("--repository")
            self.__serverArgs.append(repoPath)
        
        if encoding:
            env = QProcessEnvironment.systemEnvironment()
            env.insert("HGENCODING", encoding)
            self.__server.setProcessEnvironment(env)
            self.__encoding = encoding
    
    def startServer(self):
        """
        Public method to start the command server.
        
        @return tuple of flag indicating a successful start (boolean) and
            an error message (string) in case of failure
        """
        self.__server.start('hg', self.__serverArgs)
        serverStarted = self.__server.waitForStarted()
        if not serverStarted:
            return False, self.trUtf8(
                    'The process {0} could not be started. '
                    'Ensure, that it is in the search path.'
                ).format('hg')
        
        self.__server.setReadChannel(QProcess.StandardOutput)
        ok, error = self.__readHello()
        self.__started = ok
        return ok, error
    
    def stopServer(self):
        """
        Public method to stop the command server.
        """
        self.__server.closeWriteChannel()
        self.__server.waitForFinished()
    
    def restartServer(self):
        """
        Public method to restart the command server.
        
        @return tuple of flag indicating a successful start (boolean) and
            an error message (string) in case of failure
        """
        self.stopServer()
        return self.startServer()
    
    def __readHello(self):
        """
        Private method to read the hello message sent by the command server.
        
        @return tuple of flag indicating success (boolean) and an error message
            in case of failure (string)
        """
        ch, msg = self.__readChannel()
        if not ch:
            return False, self.trUtf8("Did not receive the 'hello' message.")
        elif ch != "o":
            return False, self.trUtf8("Received data on unexpected channel.")
        
        msg = msg.split("\n")
        
        if not msg[0].startswith("capabilities: "):
            return False, self.trUtf8("Bad 'hello' message, expected 'capabilities: '"
                                      " but got '{0}'.").format(msg[0])
        self.__capabilities = msg[0][len('capabilities: '):]
        if not self.__capabilities:
            return False, self.trUtf8("'capabilities' message did not contain"
                                      " any capability.")
        
        self.__capabilities = set(self.__capabilities.split())
        if "runcommand" not in self.__capabilities:
            return False, "'capabilities' did not contain 'runcommand'."
        
        if not msg[1].startswith("encoding: "):
            return False, self.trUtf8("Bad 'hello' message, expected 'encoding: '"
                                      " but got '{0}'.").format(msg[1])
        encoding = msg[1][len('encoding: '):]
        if not encoding:
            return False, self.trUtf8("'encoding' message did not contain"
                                      " any encoding.")
        self.__encoding = encoding
        
        return True, ""
    
    def __serverFinished(self, exitCode, exitStatus):
        """
        Private slot connected to the finished signal.
        
        @param exitCode exit code of the process (integer)
        @param exitStatus exit status of the process (QProcess.ExitStatus)
        """
        self.__started = False
    
    def __readChannel(self):
        """
        Private method to read data from the command server.
        
        @return tuple of channel designator and channel data (string, integer or string)
        """
        if self.__server.bytesAvailable() > 0 or \
           self.__server.waitForReadyRead(10000):
            data = bytes(self.__server.read(HgClient.OutputFormatSize))
            if not data:
                return "", ""
            
            channel, length = struct.unpack(HgClient.OutputFormat, data)
            channel = channel.decode(self.__encoding)
            if channel in "IL":
                return channel, length
            else:
                return (channel,
                        str(self.__server.read(length), self.__encoding, "replace"))
        else:
            return "", ""
    
    def __writeDataBlock(self, data):
        """
        Private slot to write some data to the command server.
        
        @param data data to be sent (string)
        """
        if not isinstance(data, bytes):
            data = data.encode(self.__encoding)
        self.__server.write(QByteArray(struct.pack(HgClient.InputFormat, len(data))))
        self.__server.write(QByteArray(data))
        self.__server.waitForBytesWritten()
    
    def __runcommand(self, args, inputChannels, outputChannels):
        """
        Private method to run a command in the server (low level).
        
        @param args list of arguments for the command (list of string)
        @param inputChannels dictionary of input channels. The dictionary must
            have the keys 'I' and 'L' and each entry must be a function receiving
            the number of bytes to write.
        @param outputChannels dictionary of output channels. The dictionary must
            have the keys 'o' and 'e' and each entry must be a function receiving
            the data.
        """
        if not self.__started:
            return -1
        
        self.__server.write(QByteArray(b'runcommand\n'))
        self.__writeDataBlock('\0'.join(args))
        
        while True:
            QCoreApplication.processEvents()
            if self.__server.bytesAvailable() == 0:
                continue
            channel, data = self.__readChannel()
            
            # input channels
            if channel in inputChannels:
                self.__writeDataBlock(inputChannels[channel](data))
            
            # output channels
            elif channel in outputChannels:
                outputChannels[channel](data)
            
            # result channel, command is finished
            elif channel == "r":
                return struct.unpack(HgClient.ReturnFormat, 
                                     data.encode(self.__encoding))[0]
            
            # unexpected but required channel
            elif channel.isupper():
                raise RuntimeError(
                    "Unexpected but required channel '{0}'.".format(channel))
            
            # optional channels
            else:
                pass
    
    def runcommand(self, args, prompt=None, input=None):
        """
        Public method to execute a command via the command server.
        
        @param args list of arguments for the command (list of string)
        @param prompt function to reply to prompts by the server. It
            receives the max number of bytes to return and the contents
            of the output channel received so far.
        @param input function to reply to bulk data requests by the server.
            It receives the max number of bytes to return.
        @return output and errors of the command server (string)
        """
        output = io.StringIO()
        error = io.StringIO()
        outputChannels = {
            "o": output.write,
            "e": error.write
        }
        
        inputChannels = {}
        if prompt is not None:
            def func(size):
                reply = prompt(size, output.getvalue())
                return reply
            inputChannels["L"] = func
        if input is not None:
            inputChannels["I"] = input
        
        self.__runcommand(args, inputChannels, outputChannels)
        out = output.getvalue()
        err = error.getvalue()
        
        return out, err