<commit_before># ex:ts=4:sw=4:sts=4:et
# -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-

import errno
import os
import re
from subprocess import PIPE
import sys
from threading import Thread
import time
from Queue import Queue, Empty

from quickrelease.config import ConfigSpec, ConfigSpecError
from quickrelease.exception import ReleaseFrameworkError

gUsingKillableProcess = True

try:
    if bool(ConfigSpec.GetConstant('DISABLE_KILLABLEPROCESS_PY')):
        gUsingKillableProcess = False
except ConfigSpecError:
    pass

if gUsingKillableProcess:
    from quickrelease.killableprocess import Popen
else:
    from subprocess import Popen

PIPE_STDOUT = 1
PIPE_STDERR = 2

# We used to use os.linesep here, but it turns out that doesn't work on
# MSYS, where the Win32 tools output os.linesep, but the ported Unix tools
# only output \n
REMOVE_LINE_ENDING = lambda x: re.sub('\r?\n?$', '', x)

class _OutputQueueReader(Thread):
    def __init__(self, queue=None,
                       monitoredStreams=2,
                       logHandleDescriptors=(),
                       printOutput=False, bufferedOutput=False):
        Thread.__init__(self)
        self.queue = queue
        self.printOutput = printOutput
        self.bufferedOutput = bufferedOutput
        self.logHandleDescriptors = logHandleDescriptors
        self.monitoredStreams = monitoredStreams

        self.collectedOutput = {}

        self.collectedOutput[PIPE_STDOUT] = []
        self.collectedOutput[PIPE_STDERR] = []

    def run(self):
        streamDeathCount = 0

        while True:
            try:
                lineDesc = self.queue.get()
            except Empty:
                continue

            if lineDesc.content is None:
                #print "line content on type %s is none" % (lineObj['type'])
                self.queue.task_done()
                streamDeathCount += 1
                assert (streamDeathCount >= 0 and streamDeathCount <=
                 self.monitoredStreams), "Stream monitor/death count mismatch!"
                if streamDeathCount == self.monitoredStreams:
                    break
                else:
                    continue

            if self.printOutput:
                print REMOVE_LINE_ENDING(lineDesc.content)
                if not self.bufferedOutput:
                    sys.stdout.flush()

            for h in self.logHandleDescriptors:
                if h.handle is not None and h.type == lineDesc.type:
                    h.handle.write(lineDesc.content)

            ## TODO: if collectedOutput is too big, dump to file
            self.collectedOutput[lineDesc.type].append(lineDesc)

            self.queue.task_done()

        for h in self.logHandleDescriptors:
            if h.handle is not None:
                h.handle.flush()

    def GetOutput(self, outputType=PIPE_STDOUT):
        if not self.collectedOutput.has_key(outputType):
            raise ValueError("No output type %s processed by this output "
             "monitor" % (outputType))

        return list(REMOVE_LINE_ENDING(x.content) for x in
         self.collectedOutput[outputType])

class RunShellCommandError(ReleaseFrameworkError):
    STDERR_DISPLAY_CONTEXT = 5

    def __init__(self, rscObj):
        explanation = "RunShellCommand(): "
        if rscObj.processtimedout:
            explanation += "command %s timed out" % (rscObj)
        elif rscObj.processkilled:
            explanation += "command %s killed; exit value: %d" % (rscObj,
             rscObj.returncode)
        else:
            explanation += ("command %s failed; exit value: %d, partial "
             "stderr: %s" % (rscObj, rscObj.returncode, ' '.join(rscObj.stderr[
             -self.STDERR_DISPLAY_CONTEXT:])))

        ReleaseFrameworkError.__init__(self, explanation, rscObj)

    def _GetCommandObj(self): return self.details
    command = property(_GetCommandObj)

RUN_SHELL_COMMAND_DEFAULT_ARGS = {
 'appendLogfile': True,
 'appendErrorLogfile': True,
 'autoRun': True,
 'command': (),
 'combineOutput': True,
 'errorLogfile': None,
 'logfile': None,
 'printOutput': None,
 'timeout': ConfigSpec.GetConstant('RUN_SHELL_COMMAND_DEFAULT_TIMEOUT'),
 'raiseErrors': True,
 'verbose': False,
 'workdir': None,
}

# RunShellCommand may seem a bit weird, but that's because it was originally a
# function, and later converted to a class.

# TODO: output (both stdout/stderr together), rawstdout, and rawstderr
# properties; change "partial stderr" message in RunShellCommandError to use
# new "output" property

class RunShellCommand(object):
    def __init__(self, *args, **kwargs):
        object.__init__(self)

        if len(args) > 0:
             if len(kwargs.keys()) > 0:
                 raise ValueError("Can't mix initialization styles.")

             kwargs['command'] = args

        for arg in RUN_SHELL_COMMAND_DEFAULT_ARGS.keys():
            argValue = RUN_SHELL_COMMAND_DEFAULT_ARGS[arg]
            if kwargs.has_key(arg):
                argValue = kwargs[arg]

            setattr(self, "_" + arg, argValue)

        if type(self._command) not in (list, tuple):
            raise ValueError("RunShellCommand: command must be list/tuple.")
        elif len(self._command) <= 0:
            raise ValueError("RunShellCommand: Empty command.")

        self._processWasKilled = False
        self._processTimedOut = False
        self._stdout = None
        self._stderr = None
        self._startTime = None
        self._endTime = None
        self._returncode = None

        # This makes it so we can pass int, longs, and other types to our
        # RunShellCommand that are easily convertable to strings, but which
        # Popen() will barf on if they're not strings.

        self._execArray = []

        for ndx in range(len(self._command)):
            listNdx = None
            try:
                _CheckRunShellCommandArg(type(self._command[ndx]))
                commandPart = None

                if type(self._command[ndx]) is list:
                    for lstNdx in range(len(self._command[ndx])):
                        _CheckRunShellCommandArg(type(
                         self._command[ndx][lstNdx]))
                        commandPart = str(self._command[ndx][lstNdx])
                else:
                    commandPart = str(self._command[ndx])

                if self.DEFAULT__STR__SEPARATOR in commandPart:
                    self.__str__separator = '|'

                self._execArray.append(commandPart)

            except TypeError, ex:
                errorStr = str(ex) + ": index %s" % (ndx)

                if listNdx is not None:
                    errorStr += ", sub index: %s" % (listNdx)

                raise ValueError(errorStr)

        if self._workdir is None:
            self._workdir = os.getcwd()

        if not os.path.isdir(self.workdir):
            raise ReleaseFrameworkError("RunShellCommand(): Invalid working "
             "directory: %s" % (self.workdir))

        if self._printOutput is None:
            self._printOutput = self._verbose

        try:
            if self._timeout is not None:
                self._timeout = int(self._timeout)
        except ValueError:
            raise ValueError("RunShellCommand(): Invalid timeout value '%s'"
             % self.timeout)

        if self._autoRun:
            self.Run()

    def _GetCommand(self): return self._command
    def _GetStdout(self): return self._stdout
    def _GetStderr(self): return self._stderr
    def _GetStartTime(self): return self._startTime
    def _GetEndTime(self): return self._endTime
    def _GetReturnCode(self): return self._returncode
    def _GetProcessKilled(self): return self._processWasKilled
    def _GetProcessTimedOut(self): return self._processTimedOut
    def _GetWorkDir(self): return self._workdir
    def _GetTimeout(self): return self._timeout

    def _GetRunningTime(self):
        if self._startTime is None or self._endTime is None:
            return None
        return self._endTime - self._startTime

    command = property(_GetCommand)
    stdout = property(_GetStdout)
    stderr = property(_GetStderr)
    runningtime = property(_GetRunningTime)
    starttime = property(_GetStartTime)
    endtime = property(_GetEndTime)
    returncode = property(_GetReturnCode)
    processkilled = property(_GetProcessKilled)
    processtimedout = property(_GetProcessTimedOut)
    workdir = property(_GetWorkDir)
    timeout = property(_GetTimeout)

    DEFAULT__STR__SEPARATOR = ','
    __str__separator = DEFAULT__STR__SEPARATOR
    __str__decorate = True

    def __str__(self):
        strRep = self.__str__separator.join(self._execArray)
        if self.__str__decorate:
            return "[" + strRep + "]"
        else:
            return strRep

    def SetStrOpts(self, separator=DEFAULT__STR__SEPARATOR, decorate=True):
        self.__str__separator = separator
        self.__str__decorate = decorate

    def __int__(self):
        return self.returncode

    def Run(self):
        if self._verbose:
            timeoutStr = ""
            if self.timeout is not None and gUsingKillableProcess:
                secondsStr = "seconds"
                if self.timeout == 1:
                    secondsStr = "second"
                timeoutStr = " with timeout %d %s" % (self.timeout, secondsStr)

            print >> sys.stderr, ("RunShellCommand(): Running %s in directory "
             "%s%s." % (str(self), self.workdir, timeoutStr))

            # Make sure all the output streams are flushed before we start; this
            # really only ever seems to have caused a problem on Win32
            sys.stderr.flush()
            sys.stdout.flush()

        commandLaunched = False
        try:
            logDescs = []

            if self._logfile:
                if self._appendLogfile:
                    logHandle = open(self._logfile, 'a')
                else:
                    logHandle = open(self._logfile, 'w')

                logDescs.append(_LogHandleDesc(logHandle, PIPE_STDOUT))

                if self._combineOutput:
                    logDescs.append(_LogHandleDesc(logHandle, PIPE_STDERR))

            if not self._combineOutput and self._errorLogfile is not None:
                if self._appendErrorLogfile:
                    errorLogHandle = open(self._errorLogfile, 'a')
                else:
                    errorLogHandle = open(self._errorLogfile, 'w')

                logDescs.append(_LogHandleDesc(errorLogHandle, PIPE_STDERR))

            outputQueue = Queue()

            self._startTime = time.time()
            process = Popen(self._execArray, stdout=PIPE, stderr=PIPE,
             cwd=self.workdir, bufsize=0)
            commandLaunched = True

            stdoutReader = Thread(target=_EnqueueOutput,
             args=(process.stdout, outputQueue, PIPE_STDOUT))
            stderrReader = Thread(target=_EnqueueOutput,
             args=(process.stderr, outputQueue, PIPE_STDERR))
            outputMonitor = _OutputQueueReader(queue=outputQueue,
             logHandleDescriptors=logDescs, printOutput=self._printOutput)

            stdoutReader.start()
            stderrReader.start()
            outputMonitor.start()

            try:
                # If you're not using killable process, you theoretically have
                # something else (buildbot) that's implementing a timeout for
                # you; so, all timeouts here are ignored... ...
                if self.timeout is not None and gUsingKillableProcess:
                    process.wait(self.timeout)
                else:
                    process.wait()

            except KeyboardInterrupt:
                process.kill()
                self._processWasKilled = True
        except OSError, ex:
            if ex.errno == errno.ENOENT:
                raise ReleaseFrameworkError("Invalid command or working dir")
            raise ReleaseFrameworkError("OSError: %s" % str(ex), details=ex)
        #except Exception, ex:
        #    print "EX: %s" % (ex)
        finally:
            if commandLaunched:
                procEndTime = time.time()

                #print >> sys.stderr, "Joining stderrReader"
                stderrReader.join()
                #print >> sys.stderr, "Joining stdoutReader"
                stdoutReader.join()
                #print >> sys.stderr, "Joining outputMonitor"
                outputMonitor.join()
                #print >> sys.stderr, "Joining q"
                outputQueue.join()

                for h in logDescs:
                    h.handle.close()

                # Assume if the runtime was up to/beyond the timeout, that it
                # was killed, due to timeout.
                if commandLaunched and self.runningtime >= self.timeout:
                    self._processWasKilled = True
                    self._processTimedOut = True

                #for i in range(len(outputMonitor.collectedOutput)):
                #    print "Line %d content: %s" % (i, outputMonitor.collectedOutput[i]['content'])
                #    print "Line %d time: %s" % (i, outputMonitor.collectedOutput[i]['time'])

                self._stdout = outputMonitor.GetOutput(PIPE_STDOUT)
                self._stderr = outputMonitor.GetOutput(PIPE_STDERR)
                self._endTime = procEndTime
                self._returncode = process.returncode

                if self._raiseErrors and self.returncode:
                    raise RunShellCommandError(self)

def _CheckRunShellCommandArg(argType):
    if argType not in (str, unicode, int, float, list, long):
        raise TypeError("RunShellCommand(): unexpected argument type %s" %
         (argType))

class _OutputLineDesc(object):
    def __init__(self, outputType=None, content=None):
        self.time = time.time()
        object.__init__(self)
        self.type = outputType
        self.content = content

class _LogHandleDesc(object):
    def __init__(self, handle, outputType=None):
        object.__init__(self)
        self.type = outputType
        self.handle = handle

def _EnqueueOutput(outputPipe, outputQueue, pipeType):
    for line in iter(outputPipe.readline, ''):
        assert line is not None, "Line was None"
        outputQueue.put(_OutputLineDesc(pipeType, line))

    outputPipe.close()
    outputQueue.put(_OutputLineDesc(pipeType))
<commit_msg>RunShellCommand stdin input stream support.<commit_after># ex:ts=4:sw=4:sts=4:et
# -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-

import errno
import os
import re
from subprocess import PIPE
import sys
from threading import Thread
import time
from Queue import Queue, Empty

from quickrelease.config import ConfigSpec, ConfigSpecError
from quickrelease.exception import ReleaseFrameworkError

gUsingKillableProcess = True

try:
    if bool(ConfigSpec.GetConstant('DISABLE_KILLABLEPROCESS_PY')):
        gUsingKillableProcess = False
except ConfigSpecError:
    pass

if gUsingKillableProcess:
    from quickrelease.killableprocess import Popen
else:
    from subprocess import Popen

PIPE_STDOUT = 1
PIPE_STDERR = 2

# We used to use os.linesep here, but it turns out that doesn't work on
# MSYS, where the Win32 tools output os.linesep, but the ported Unix tools
# only output \n
REMOVE_LINE_ENDING = lambda x: re.sub('\r?\n?$', '', x)

class _OutputQueueReader(Thread):
    def __init__(self, queue=None,
                       monitoredStreams=2,
                       logHandleDescriptors=(),
                       printOutput=False, bufferedOutput=False):
        Thread.__init__(self)
        self.queue = queue
        self.printOutput = printOutput
        self.bufferedOutput = bufferedOutput
        self.logHandleDescriptors = logHandleDescriptors
        self.monitoredStreams = monitoredStreams

        self.collectedOutput = {}

        self.collectedOutput[PIPE_STDOUT] = []
        self.collectedOutput[PIPE_STDERR] = []

    def run(self):
        streamDeathCount = 0

        while True:
            try:
                lineDesc = self.queue.get()
            except Empty:
                continue

            if lineDesc.content is None:
                #print "line content on type %s is none" % (lineObj['type'])
                self.queue.task_done()
                streamDeathCount += 1
                assert (streamDeathCount >= 0 and streamDeathCount <=
                 self.monitoredStreams), "Stream monitor/death count mismatch!"
                if streamDeathCount == self.monitoredStreams:
                    break
                else:
                    continue

            if self.printOutput:
                print REMOVE_LINE_ENDING(lineDesc.content)
                if not self.bufferedOutput:
                    sys.stdout.flush()

            for h in self.logHandleDescriptors:
                if h.handle is not None and h.type == lineDesc.type:
                    h.handle.write(lineDesc.content)

            ## TODO: if collectedOutput is too big, dump to file
            self.collectedOutput[lineDesc.type].append(lineDesc)

            self.queue.task_done()

        for h in self.logHandleDescriptors:
            if h.handle is not None:
                h.handle.flush()

    def GetOutput(self, outputType=PIPE_STDOUT):
        if not self.collectedOutput.has_key(outputType):
            raise ValueError("No output type %s processed by this output "
             "monitor" % (outputType))

        return list(REMOVE_LINE_ENDING(x.content) for x in
         self.collectedOutput[outputType])

class RunShellCommandError(ReleaseFrameworkError):
    STDERR_DISPLAY_CONTEXT = 5

    def __init__(self, rscObj):
        explanation = "RunShellCommand(): "
        if rscObj.processtimedout:
            explanation += "command %s timed out" % (rscObj)
        elif rscObj.processkilled:
            explanation += "command %s killed; exit value: %d" % (rscObj,
             rscObj.returncode)
        else:
            explanation += ("command %s failed; exit value: %d, partial "
             "stderr: %s" % (rscObj, rscObj.returncode, ' '.join(rscObj.stderr[
             -self.STDERR_DISPLAY_CONTEXT:])))

        ReleaseFrameworkError.__init__(self, explanation, rscObj)

    def _GetCommandObj(self): return self.details
    command = property(_GetCommandObj)

RUN_SHELL_COMMAND_DEFAULT_ARGS = {
 'appendLogfile': True,
 'appendErrorLogfile': True,
 'autoRun': True,
 'command': (),
 'combineOutput': True,
 'errorLogfile': None,
 'logfile': None,
 'printOutput': None,
 'timeout': ConfigSpec.GetConstant('RUN_SHELL_COMMAND_DEFAULT_TIMEOUT'),
 'raiseErrors': True,
 'verbose': False,
 'workdir': None,
 'input': None,
}

# RunShellCommand may seem a bit weird, but that's because it was originally a
# function, and later converted to a class.

# TODO: output (both stdout/stderr together), rawstdout, and rawstderr
# properties; change "partial stderr" message in RunShellCommandError to use
# new "output" property

class RunShellCommand(object):
    def __init__(self, *args, **kwargs):
        object.__init__(self)

        if len(args) > 0:
             if len(kwargs.keys()) > 0:
                 raise ValueError("Can't mix initialization styles.")

             kwargs['command'] = args

        for arg in RUN_SHELL_COMMAND_DEFAULT_ARGS.keys():
            argValue = RUN_SHELL_COMMAND_DEFAULT_ARGS[arg]
            if kwargs.has_key(arg):
                argValue = kwargs[arg]

            setattr(self, "_" + arg, argValue)

        if type(self._command) not in (list, tuple):
            raise ValueError("RunShellCommand: command must be list/tuple.")
        elif len(self._command) <= 0:
            raise ValueError("RunShellCommand: Empty command.")

        self._processWasKilled = False
        self._processTimedOut = False
        self._stdout = None
        self._stderr = None
        self._startTime = None
        self._endTime = None
        self._returncode = None
        self._stdin = None

        if self._input is not None:
            if type(self._input) is str:
                try:
                    self._stdin = open(self._input, 'r')
                except IOError, ex:
                    if ex.errno == errno.ENOENT:
                        raise ValueError("Invalid input stream file: %s" %
                         (self._input))
                    else:
                        raise ex
            else:
                self._stdin = self._input

        # This makes it so we can pass int, longs, and other types to our
        # RunShellCommand that are easily convertable to strings, but which
        # Popen() will barf on if they're not strings.

        self._execArray = []

        for ndx in range(len(self._command)):
            listNdx = None
            try:
                _CheckRunShellCommandArg(type(self._command[ndx]))
                commandPart = None

                if type(self._command[ndx]) is list:
                    for lstNdx in range(len(self._command[ndx])):
                        _CheckRunShellCommandArg(type(
                         self._command[ndx][lstNdx]))
                        commandPart = str(self._command[ndx][lstNdx])
                else:
                    commandPart = str(self._command[ndx])

                if self.DEFAULT__STR__SEPARATOR in commandPart:
                    self.__str__separator = '|'

                self._execArray.append(commandPart)

            except TypeError, ex:
                errorStr = str(ex) + ": index %s" % (ndx)

                if listNdx is not None:
                    errorStr += ", sub index: %s" % (listNdx)

                raise ValueError(errorStr)

        if self._workdir is None:
            self._workdir = os.getcwd()

        if not os.path.isdir(self.workdir):
            raise ReleaseFrameworkError("RunShellCommand(): Invalid working "
             "directory: %s" % (self.workdir))

        if self._printOutput is None:
            self._printOutput = self._verbose

        try:
            if self._timeout is not None:
                self._timeout = int(self._timeout)
        except ValueError:
            raise ValueError("RunShellCommand(): Invalid timeout value '%s'"
             % self.timeout)

        if self._autoRun:
            self.Run()

    def _GetCommand(self): return self._command
    def _GetStdout(self): return self._stdout
    def _GetStderr(self): return self._stderr
    def _GetStartTime(self): return self._startTime
    def _GetEndTime(self): return self._endTime
    def _GetReturnCode(self): return self._returncode
    def _GetProcessKilled(self): return self._processWasKilled
    def _GetProcessTimedOut(self): return self._processTimedOut
    def _GetWorkDir(self): return self._workdir
    def _GetTimeout(self): return self._timeout

    def _GetRunningTime(self):
        if self._startTime is None or self._endTime is None:
            return None
        return self._endTime - self._startTime

    command = property(_GetCommand)
    stdout = property(_GetStdout)
    stderr = property(_GetStderr)
    runningtime = property(_GetRunningTime)
    starttime = property(_GetStartTime)
    endtime = property(_GetEndTime)
    returncode = property(_GetReturnCode)
    processkilled = property(_GetProcessKilled)
    processtimedout = property(_GetProcessTimedOut)
    workdir = property(_GetWorkDir)
    timeout = property(_GetTimeout)

    DEFAULT__STR__SEPARATOR = ','
    __str__separator = DEFAULT__STR__SEPARATOR
    __str__decorate = True

    def __str__(self):
        strRep = self.__str__separator.join(self._execArray)
        if self.__str__decorate:
            return "[" + strRep + "]"
        else:
            return strRep

    def SetStrOpts(self, separator=DEFAULT__STR__SEPARATOR, decorate=True):
        self.__str__separator = separator
        self.__str__decorate = decorate

    def __int__(self):
        return self.returncode

    def Run(self):
        if self._verbose:
            timeoutStr = ""
            if self.timeout is not None and gUsingKillableProcess:
                secondsStr = "seconds"
                if self.timeout == 1:
                    secondsStr = "second"
                timeoutStr = " with timeout %d %s" % (self.timeout, secondsStr)

            print >> sys.stderr, ("RunShellCommand(): Running %s in directory "
             "%s%s." % (str(self), self.workdir, timeoutStr))

            # Make sure all the output streams are flushed before we start; this
            # really only ever seems to have caused a problem on Win32
            sys.stderr.flush()
            sys.stdout.flush()

        commandLaunched = False
        try:
            logDescs = []

            if self._logfile:
                if self._appendLogfile:
                    logHandle = open(self._logfile, 'a')
                else:
                    logHandle = open(self._logfile, 'w')

                logDescs.append(_LogHandleDesc(logHandle, PIPE_STDOUT))

                if self._combineOutput:
                    logDescs.append(_LogHandleDesc(logHandle, PIPE_STDERR))

            if not self._combineOutput and self._errorLogfile is not None:
                if self._appendErrorLogfile:
                    errorLogHandle = open(self._errorLogfile, 'a')
                else:
                    errorLogHandle = open(self._errorLogfile, 'w')

                logDescs.append(_LogHandleDesc(errorLogHandle, PIPE_STDERR))

            outputQueue = Queue()

            stdinArg = None
            if self._input is not None:
                stdinArg = PIPE

            self._startTime = time.time()
            process = Popen(self._execArray, stdin=stdinArg, gstdout=PIPE,
             stderr=PIPE, cwd=self.workdir, bufsize=0)
            commandLaunched = True

            if self._stdin is not None:
                #print >> sys.stderr, "Starting stdinWriter"
                stdinWriter = Thread(target=_WriteInput,
                 name="RunShellCommand() stdin writer",
                 args=(self._stdin, process.stdin))
                stdinWriter.start()

            stdoutReader = Thread(target=_EnqueueOutput,
             name="RunShellCommand() stdout reader",
             args=(process.stdout, outputQueue, PIPE_STDOUT))
            stderrReader = Thread(target=_EnqueueOutput,
             name="RunShellCommand() stderr reader",
             args=(process.stderr, outputQueue, PIPE_STDERR))
            outputMonitor = _OutputQueueReader(queue=outputQueue,
             logHandleDescriptors=logDescs, printOutput=self._printOutput)

            stdoutReader.start()
            stderrReader.start()
            outputMonitor.start()

            try:
                # If you're not using killable process, you theoretically have
                # something else (buildbot) that's implementing a timeout for
                # you; so, all timeouts here are ignored... ...
                if self.timeout is not None and gUsingKillableProcess:
                    process.wait(self.timeout)
                else:
                    process.wait()

            except KeyboardInterrupt:
                process.kill()
                self._processWasKilled = True
        except OSError, ex:
            if ex.errno == errno.ENOENT:
                raise ReleaseFrameworkError("Invalid command or working dir")
            raise ReleaseFrameworkError("OSError: %s" % str(ex), details=ex)
        #except Exception, ex:
        #    print "EX: %s" % (ex)
        finally:
            if commandLaunched:
                procEndTime = time.time()

                #print >> sys.stderr, "Joining stderrReader"
                stderrReader.join()
                #print >> sys.stderr, "Joining stdoutReader"
                stdoutReader.join()
                #print >> sys.stderr, "Joining outputMonitor"
                outputMonitor.join()
                #print >> sys.stderr, "Joining q"
                outputQueue.join()
                if stdinWriter is not None:
                    print >> sys.stderr, "Joining stdinWriter"
                    stdinWriter.join()

                for h in logDescs:
                    h.handle.close()

                # Assume if the runtime was up to/beyond the timeout, that it
                # was killed, due to timeout.
                if commandLaunched and self.runningtime >= self.timeout:
                    self._processWasKilled = True
                    self._processTimedOut = True

                #for i in range(len(outputMonitor.collectedOutput)):
                #    print "Line %d content: %s" % (i, outputMonitor.collectedOutput[i]['content'])
                #    print "Line %d time: %s" % (i, outputMonitor.collectedOutput[i]['time'])

                self._stdout = outputMonitor.GetOutput(PIPE_STDOUT)
                self._stderr = outputMonitor.GetOutput(PIPE_STDERR)
                self._endTime = procEndTime
                self._returncode = process.returncode

                if self._input is not None and type(self._input) is str:
                    #print >> sys.stderr, "Closing stdin file."
                    self._stdin.close()

                if self._raiseErrors and self.returncode:
                    raise RunShellCommandError(self)

def _CheckRunShellCommandArg(argType):
    if argType not in (str, unicode, int, float, list, long):
        raise TypeError("RunShellCommand(): unexpected argument type %s" %
         (argType))

class _OutputLineDesc(object):
    def __init__(self, outputType=None, content=None):
        self.time = time.time()
        object.__init__(self)
        self.type = outputType
        self.content = content

class _LogHandleDesc(object):
    def __init__(self, handle, outputType=None):
        object.__init__(self)
        self.type = outputType
        self.handle = handle

def _WriteInput(inputStream, procStdinPipe):
    for line in iter(inputStream.readline, ''):
        assert line is not None, "Line was None"
        procStdinPipe.write(line)

    procStdinPipe.close()

def _EnqueueOutput(outputPipe, outputQueue, pipeType):
    for line in iter(outputPipe.readline, ''):
        assert line is not None, "Line was None"
        outputQueue.put(_OutputLineDesc(pipeType, line))

    outputPipe.close()
    outputQueue.put(_OutputLineDesc(pipeType))
<|endoftext|>


<commit_before>#!/usr/bin/env python

import logging
import argparse
import os.path
import itertools

from maroon import Model

from settings import settings
from explore import peek, sprawl, fixgis, graph
from predict import prep, fl
from base import gob
from base import utils

try:
    import networkx as nx
except ImportError:
    pass

try:
    import numpy
except ImportError:
    numpy = None

def parse_args(argv):
    parser = argparse.ArgumentParser(description='Run gob jobs.')
    parser.add_argument('job',nargs='*')
    parser.add_argument('-s','--single',action="store_true",
                        help='run in a single process')
    parser.add_argument('-m','--mongo')
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-i','--input')
    group.add_argument('--rm',action="store_true")
    group.add_argument('--force',action="store_true")
    group.add_argument('--dot',action="store_true")
    return parser.parse_args(argv)

def create_jobs(g):
    # the crawler
    g.add_source(utils.read_json, name='geotweets')
    g.add_map_job(sprawl.parse_geotweets,'geotweets',saver='split_save')
    g.add_map_job(sprawl.mloc_users,'parse_geotweets')
    g.add_map_job(sprawl.EdgeFinder.find_contacts,'mloc_users',
              reducer=gob.set_reduce)
    g.add_map_job(sprawl.contact_split,'find_contacts',saver='split_save')

    g.add_map_job(fixgis.gnp_gps,'mloc_users')
    g.add_cat('cat_gnp_gps','gnp_gps')
    g.add_map_job(fixgis.mdists,'cat_gnp_gps')
    g.add_map_job(sprawl.ContactLookup.lookup_contacts,'contact_split',procs=15)
    g.add_map_job(sprawl.mloc_uids,'mloc_users')
    g.add_map_job(sprawl.trash_extra_mloc,'mloc_uids')
    g.add_map_job(sprawl.MDistFixer.fix_mloc_mdists,'mloc_uids',requires=['mdists'])

    g.add_map_job(sprawl.pick_nebrs,'mloc_uids',
              requires=['lookup_contacts','mdists','trash_extra_mloc',
                        'fix_mloc_mdists'],
              reducer=gob.set_reduce,
              )
    g.add_map_job(sprawl.nebr_split, 'pick_nebrs', saver='split_save')
    g.add_map_job(sprawl.EdgeFinder.find_leafs,'nebr_split',reducer=gob.set_reduce)
    g.add_map_job(sprawl.contact_split,'find_leafs',
              name='leaf_split',saver='split_save')
    g.add_map_job(sprawl.saved_users,saver='split_save')
    g.add_map_job(sprawl.ContactLookup.lookup_contacts, 'leaf_split',
              name='lookup_leafs',requires=['saved_users'])

    # FIXME: nebr_split_2 is temporary!
    g.add_map_job(None,name='nebr_split_59',saver='split_save')
    g.add_map_job(peek.contact_blur,'nebr_split_59',reducer=gob.avg_reduce)

    # the graphs
    g.add_map_job(peek.geo_ats,saver='split_save')
    g.add_cat('cat_geo_ats','geo_ats')
    g.add_map_job(peek.at_tuples,'cat_geo_ats',saver='split_save')
    g.add_map_job(peek.geo_ated,'at_tuples')

    g.add_map_job(prep.pred_users,'mloc_uids')
    g.add_map_job(peek.EdgesDict.edges_d,'pred_users',requires=['geo_ats'],procs=4)
    g.add_map_job(peek.edge_dists,'edges_d',reducer=gob.join_reduce)
    g.add_map_job(peek.rfrd_dists,'edges_d')
    g.add_cat('cat_rfrd_dists','rfrd_dists')
    g.add_map_job(peek.edge_leaf_dists,'edges_d')
    g.add_map_job(graph.graph_edge_types_cuml,'edge_dists')
    g.add_map_job(graph.graph_edge_types_prot,'edge_dists')
    g.add_map_job(graph.graph_edge_types_norm,'edge_dists')
    g.add_map_job(graph.graph_com_types,'edge_dists')
    g.add_map_job(graph.graph_edge_count,'cat_rfrd_dists')
    g.add_map_job(graph.graph_locals,'cat_rfrd_dists')
    g.add_cat('cat_edges_d','edges_d')
    g.add_map_job(graph.graph_local_groups,'cat_edges_d')

    g.add_map_job(peek.rfr_triads,'pred_users')
    g.add_cat('cat_rfr_triads','rfr_triads')
    g.add_map_job(graph.near_triads,'cat_rfr_triads')

    # add noise to location field of geolocated users
    g.add_map_job(peek.contact_mdist,'contact_split')
    g.add_map_job(peek.contact_mdist,'mloc_uids',name='mloc_mdist')
    g.add_map_job(prep.MlocBlur.mloc_blur,
              requires=['mloc_mdist','contact_mdist'])

    # fb predictor
    g.add_map_job(peek.contact_count,'contact_split',reducer=gob.sum_reduce)
    g.add_map_job(peek.mloc_tile,'mloc_uids',reducer=gob.join_reduce)
    g.add_map_job(peek.tile_split, 'mloc_tile', saver='split_save')
    g.add_map_job(peek.nebr_dists, 'tile_split')
    g.add_map_job(peek.StrangerDists.stranger_dists, 'tile_split',
              requires=['contact_count'])
    g.add_map_job(peek.strange_nebr_bins, 'stranger_dists',
              name='strange_bins',reducer=gob.sum_reduce)
    g.add_map_job(peek.strange_nebr_bins, 'nebr_dists',
              name='nebr_bins',reducer=gob.sum_reduce)
    g.add_map_job( peek.ContactFit.contact_fit, (),
              requires=['strange_bins','nebr_bins'] )
    g.add_map_job(peek.lat_tile, saver='split_save')
    g.add_map_job(peek.StrangerDists.stranger_prob, 'lat_tile',
              requires=['contact_count','contact_fit'])
    g.add_cat('stranger_prob_cat','stranger_prob')
    g.add_map_job(peek.stranger_mat,'stranger_prob_cat',encoding='npz')

    def train_set(key, clump):
        return int(key)//20!=int(clump)
    def eval_set(key, clump):
        return int(key)//20==int(clump)

    # prep
    g.add_map_job(prep.NeighborsDict.nebrs_d,'pred_users',
              requires=['mloc_blur','lookup_leafs'])
    g.add_clump(train_set, 'nebrs_d', name='nebrs_train')
    g.add_clump(eval_set, 'nebrs_d', name='nebrs_eval')

    # mdist_curves and utc_offset
    g.add_map_job(prep.mdist_real,'nebrs_d')
    g.add_clump(train_set, 'mdist_real', name='mdist_train')
    g.add_map_job(prep.mdist_curves,'mdist_train')
    g.add_map_job(prep.UtcOffset.utc_offset, 'nebrs_train')

    # the predictor
    g.add_map_job(fl.nebr_vect,'nebrs_d')
    g.add_clump(train_set, 'nebr_vect', name='nvect_train')
    g.add_clump(eval_set, 'nebr_vect', name='nvect_eval')

    g.add_map_job(fl.nebr_clf,'nvect_train',encoding='pkl')
    # FIXME : why don't we need vectors for nvect_eval?
    g.add_map_job(peek.ContactFit.vect_fit, 'nvect_train',
              requires=['strange_bins','nebr_clf'] )
    g.add_map_job(fl.Predictors.predictions,'nebrs_eval',
              requires=['stranger_mat','mdist_curves','vect_fit','utc_offset','contact_fit'])
    g.add_map_job(graph.gr_preds,'predictions')
    g.add_map_job(fl.eval_preds,'predictions')

def make_gob(args):
    path = os.path.join(os.path.dirname(__file__),'data')
    if not args or args.single or args.input:
        env = gob.SimpleFileEnv(path,log_crashes=True)
    else:
        env = gob.MultiProcEnv(path,log_crashes=True,log_level=logging.INFO)
    my_gob = gob.Gob(env)
    create_jobs(my_gob)
    return my_gob

def setup(args):
    Model.database = utils.mongo(args.mongo or settings.region)
    logging.basicConfig(level=logging.INFO)
    if numpy:
        numpy.set_printoptions(precision=3, linewidth=160)

def inspect(job_name, source_set):
    """
    run a map reduce job on a single input and return its results.

    This is useful for poking around in ipython.
    """
    my_gob = make_gob(None)
    job = my_gob.jobs[job_name]
    funcs = job.runnable_funcs(my_gob.env)
    return my_gob.env.map_reduce_save(job, source_set, funcs)

def print_graph_dot(my_gob):
    srcs = ( (source.name,job)
             for job in my_gob.jobs
             for source in my_gob.jobs[job].sources
            )
    reqs = ( (req,job)
             for job in my_gob.jobs
             for req in my_gob.jobs[job].requires
            )
    dg = nx.DiGraph(itertools.chain(srcs,reqs))
    nx.write_dot(dg,'gob_deps.dot')

def run(my_gob,args):
    if args.dot:
        print_graph_dot(my_gob)
    if args.input:
        # just run the function for one input
        job = my_gob.jobs[args.job[0]]
        input_paths = [args.input.split(',')]
        my_gob.env.run(job,input_paths=input_paths)
    else:
        for cmd in args.job:
            if args.rm or args.force:
                my_gob.clear_job(cmd)
            if not args.rm:
                my_gob.run_job(cmd)

def main(*argv):
    # set globals so that we can %run this in ipython
    global my_gob, env
    args = parse_args(argv or None)
    my_gob = make_gob(args)
    setup(args)
    run(my_gob,args)
    env = my_gob.env

if __name__ == '__main__':
    main()
<commit_msg>updateing gb.py deps<commit_after>#!/usr/bin/env python

import logging
import argparse
import os.path
import itertools

from maroon import Model

from settings import settings
from explore import peek, sprawl, fixgis, graph
from predict import prep, fl
from base import gob
from base import utils

try:
    import networkx as nx
except ImportError:
    pass

try:
    import numpy
except ImportError:
    numpy = None

def parse_args(argv):
    parser = argparse.ArgumentParser(description='Run gob jobs.')
    parser.add_argument('job',nargs='*')
    parser.add_argument('-s','--single',action="store_true",
                        help='run in a single process')
    parser.add_argument('-m','--mongo')
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-i','--input')
    group.add_argument('--rm',action="store_true")
    group.add_argument('--force',action="store_true")
    group.add_argument('--dot',action="store_true")
    return parser.parse_args(argv)

def create_jobs(g):
    # the crawler
    g.add_source(utils.read_json, name='geotweets')
    g.add_map_job(sprawl.parse_geotweets,'geotweets',saver='split_save')
    g.add_map_job(sprawl.mloc_users,'parse_geotweets')
    g.add_map_job(sprawl.EdgeFinder.find_contacts,'mloc_users',
              reducer=gob.set_reduce)
    g.add_map_job(sprawl.contact_split,'find_contacts',saver='split_save')

    g.add_map_job(fixgis.gnp_gps,'mloc_users')
    g.add_cat('cat_gnp_gps','gnp_gps')
    g.add_map_job(fixgis.mdists,'cat_gnp_gps')
    g.add_map_job(sprawl.ContactLookup.lookup_contacts,'contact_split',procs=15)
    g.add_map_job(sprawl.mloc_uids,'mloc_users')
    g.add_map_job(sprawl.trash_extra_mloc,'mloc_uids')
    g.add_map_job(sprawl.MDistFixer.fix_mloc_mdists,'mloc_uids',requires=['mdists'])

    g.add_map_job(sprawl.pick_nebrs,'mloc_uids',
              requires=['lookup_contacts','mdists','trash_extra_mloc',
                        'fix_mloc_mdists'],
              reducer=gob.set_reduce,
              )
    g.add_map_job(sprawl.nebr_split, 'pick_nebrs', saver='split_save')
    g.add_map_job(sprawl.EdgeFinder.find_leafs,'nebr_split',reducer=gob.set_reduce)
    g.add_map_job(sprawl.contact_split,'find_leafs',
              name='leaf_split',saver='split_save')
    g.add_map_job(sprawl.saved_users,saver='split_save')
    g.add_map_job(sprawl.ContactLookup.lookup_contacts, 'leaf_split',
              name='lookup_leafs',requires=['saved_users'])

    g.add_map_job(peek.contact_blur,'nebr_split',
                  requires=['lookup_leafs'],reducer=gob.avg_reduce)

    # the graphs
    g.add_map_job(peek.geo_ats,saver='split_save',requires=['find_leafs'])
    g.add_cat('cat_geo_ats','geo_ats')
    g.add_map_job(peek.at_tuples,'cat_geo_ats',saver='split_save')
    g.add_map_job(peek.geo_ated,'at_tuples')

    g.add_map_job(prep.pred_users,'mloc_uids')
    g.add_map_job(peek.EdgesDict.edges_d,'pred_users',requires=['geo_ats'],procs=4)
    g.add_map_job(peek.edge_dists,'edges_d',reducer=gob.join_reduce)
    g.add_map_job(peek.rfrd_dists,'edges_d',requires=['contact_blur'])
    g.add_cat('cat_rfrd_dists','rfrd_dists')
    g.add_map_job(peek.edge_leaf_dists,'edges_d')
    g.add_map_job(graph.graph_edge_types_cuml,'edge_dists')
    g.add_map_job(graph.graph_edge_types_prot,'edge_dists')
    g.add_map_job(graph.graph_edge_types_norm,'edge_dists')
    g.add_map_job(graph.graph_com_types,'edge_dists')
    g.add_map_job(graph.graph_edge_count,'cat_rfrd_dists')
    g.add_map_job(graph.graph_locals,'cat_rfrd_dists')
    g.add_cat('cat_edges_d','edges_d')
    g.add_map_job(graph.graph_local_groups,'cat_edges_d')

    g.add_map_job(peek.rfr_triads,'pred_users')
    g.add_cat('cat_rfr_triads','rfr_triads')
    g.add_map_job(graph.near_triads,'cat_rfr_triads')

    # add noise to location field of geolocated users
    g.add_map_job(peek.contact_mdist,'contact_split')
    g.add_map_job(peek.contact_mdist,'mloc_uids',name='mloc_mdist')
    g.add_map_job(prep.MlocBlur.mloc_blur,
              requires=['mloc_mdist','contact_mdist'])

    # fb predictor
    g.add_map_job(peek.contact_count,'contact_split',reducer=gob.sum_reduce)
    g.add_map_job(peek.mloc_tile,'mloc_uids',reducer=gob.join_reduce)
    g.add_map_job(peek.tile_split, 'mloc_tile', saver='split_save')
    g.add_map_job(peek.nebr_dists, 'tile_split')
    g.add_map_job(peek.StrangerDists.stranger_dists, 'tile_split',
              requires=['contact_count'])
    g.add_map_job(peek.strange_nebr_bins, 'stranger_dists',
              name='strange_bins',reducer=gob.sum_reduce)
    g.add_map_job(peek.strange_nebr_bins, 'nebr_dists',
              name='nebr_bins',reducer=gob.sum_reduce)
    g.add_map_job( peek.ContactFit.contact_fit, (),
              requires=['strange_bins','nebr_bins'] )
    g.add_map_job(peek.lat_tile, saver='split_save')
    g.add_map_job(peek.StrangerDists.stranger_prob, 'lat_tile',
              requires=['contact_count','contact_fit'])
    g.add_cat('stranger_prob_cat','stranger_prob')
    g.add_map_job(peek.stranger_mat,'stranger_prob_cat',encoding='npz')

    def train_set(key, clump):
        return int(key)//20!=int(clump)
    def eval_set(key, clump):
        return int(key)//20==int(clump)

    # prep
    g.add_map_job(prep.NeighborsDict.nebrs_d,'pred_users',
              requires=['mloc_blur','lookup_leafs','contact_blur'])
    g.add_clump(train_set, 'nebrs_d', name='nebrs_train')
    g.add_clump(eval_set, 'nebrs_d', name='nebrs_eval')

    # mdist_curves and utc_offset
    g.add_map_job(prep.mdist_real,'nebrs_d')
    g.add_clump(train_set, 'mdist_real', name='mdist_train')
    g.add_map_job(prep.mdist_curves,'mdist_train')
    g.add_map_job(prep.UtcOffset.utc_offset, 'nebrs_train')

    # the predictor
    g.add_map_job(fl.nebr_vect,'nebrs_d')
    g.add_clump(train_set, 'nebr_vect', name='nvect_train')
    g.add_clump(eval_set, 'nebr_vect', name='nvect_eval')

    g.add_map_job(fl.nebr_clf,'nvect_train',encoding='pkl')
    # FIXME : why don't we need vectors for nvect_eval?
    g.add_map_job(peek.ContactFit.vect_fit, 'nvect_train',
              requires=['strange_bins','nebr_clf'] )
    g.add_map_job(fl.Predictors.predictions,'nebrs_eval',
              requires=['stranger_mat','mdist_curves','vect_fit','utc_offset','contact_fit'])
    g.add_map_job(graph.gr_preds,'predictions')
    g.add_map_job(fl.eval_preds,'predictions')

def make_gob(args):
    path = os.path.join(os.path.dirname(__file__),'data')
    if not args or args.single or args.input:
        env = gob.SimpleFileEnv(path,log_crashes=True)
    else:
        env = gob.MultiProcEnv(path,log_crashes=True,log_level=logging.INFO)
    my_gob = gob.Gob(env)
    create_jobs(my_gob)
    return my_gob

def setup(args):
    Model.database = utils.mongo(args.mongo or settings.region)
    logging.basicConfig(level=logging.INFO)
    if numpy:
        numpy.set_printoptions(precision=3, linewidth=160)

def inspect(job_name, source_set):
    """
    run a map reduce job on a single input and return its results.

    This is useful for poking around in ipython.
    """
    my_gob = make_gob(None)
    job = my_gob.jobs[job_name]
    funcs = job.runnable_funcs(my_gob.env)
    return my_gob.env.map_reduce_save(job, source_set, funcs)

def print_graph_dot(my_gob):
    srcs = ( (source.name,job)
             for job in my_gob.jobs
             for source in my_gob.jobs[job].sources
            )
    reqs = ( (req,job)
             for job in my_gob.jobs
             for req in my_gob.jobs[job].requires
            )
    dg = nx.DiGraph(itertools.chain(srcs,reqs))
    nx.write_dot(dg,'gob_deps.dot')

def run(my_gob,args):
    if args.dot:
        print_graph_dot(my_gob)
    if args.input:
        # just run the function for one input
        job = my_gob.jobs[args.job[0]]
        input_paths = [args.input.split(',')]
        my_gob.env.run(job,input_paths=input_paths)
    else:
        for cmd in args.job:
            if args.rm or args.force:
                my_gob.clear_job(cmd)
            if not args.rm:
                my_gob.run_job(cmd)

def main(*argv):
    # set globals so that we can %run this in ipython
    global my_gob, env
    args = parse_args(argv or None)
    my_gob = make_gob(args)
    setup(args)
    run(my_gob,args)
    env = my_gob.env

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>import re
from collections import namedtuple
import itertools

N = 9
W = N + 1
# Represent a board as a string, with '.' empty, 'X' to play, 'O' other player.
# Whitespace is used as a border (to avoid IndexError when computing neighbors)

# A Coordinate `c` is an int: an index into the board.
# A Move is a (Coordinate c | None).

# Human-readable move notation: columns go from A to T left to right
# rows go from 1 to 19 from bottom to top.
#  ABCD
#  ____\n
# 4....\n
# 3....\n
# 2....\n
# 1....

COLUMNS = 'ABCDEFGHJKLMNOPQRST'

EMPTY_BOARD = '\n'.join(
    [' ' * N] +
    ['.' * N for i in range(N)] +
    [' ' * W])

SWAP_COLORS = str.maketrans('XO', 'OX')

def load_board(string):
    string = re.sub(r'[^OX\.#]+', '', string)
    assert len(string) == N ** 2, "Board to load didn't have right dimensions"
    return '\n'.join([' ' * N] + [string[k*N:(k+1)*N] for k in range(N)] + [' ' * W])

def parse_coords(s):
    if s == 'pass':
        return None
    s = s.upper()
    col = COLUMNS.index(s[0])
    rows_from_top = N - int(s[1:])
    return W + (W * rows_from_top) + col

def neighbors(c):
    return [c+1, c-1, c+W, c-W]

def place_stone(board, color, c):
    return board[:c] + color + board[c+1:]

def capture_stones(board, stones):
    b = bytearray(board, encoding='ascii')
    for s in stones:
        b[s] = ord('.')
    return str(b, encoding='ascii')

def flood_fill(board, c):
    'From a starting coordinate c, flood-fill the board with a #'
    b = bytearray(board, encoding='ascii')
    color = b[c]
    entire_group = [c]
    frontier = [c]
    while frontier:
        current = frontier.pop()
        b[current] = ord('#')
        for n in neighbors(current):
            if b[n] == color:
                frontier.append(n)
                entire_group.append(n)
    return str(b, encoding='ascii'), set(entire_group)

def find_neighbors(color, board, stones):
    'Find all neighbors of a set of stones of a given color'
    potential_neighbors = set(itertools.chain(*(neighbors(s) for s in stones)))
    return {c for c in potential_neighbors if board[c] == color}

def find_liberties(board, stones):
    'Given a board and a set of stones, find liberties of those stones'
    return find_neighbors('.', board, stones)

class Group(namedtuple('Group', 'stones liberties')):
    '''
    stones: a set of Coordinates belonging to this group
    liberties: a set of Coordinates that are empty and adjacent to this group.
    '''
    pass

def deduce_groups(board):
    'Given a board, return a 2-tuple; a list of groups for each player'
    def find_groups(board, color):
        groups = []
        while color in board:
            c = board.index(color)
            board, stones = flood_fill(board, c)
            liberties = find_liberties(board, stones)
            groups.append(Group(stones=stones, liberties=liberties))
        return groups

    return find_groups(board, 'X'), find_groups(board, 'O')

def update_groups(board, existing_X_groups, existing_O_groups, c):
    '''
    When a move is played, update the list of groups and their liberties.
    This means possibly appending the new move to a group, creating a new 1-stone group, or merging existing groups.
    The new move should be of color X.
    The board should represent the state after the move has been played at `c`.
    '''
    assert board[c] == 'X'

    updated_X_groups, groups_to_merge = [], []
    for g in existing_X_groups:
        if c in g.liberties:
            groups_to_merge.append(g)
        else:
            updated_X_groups.append(g)

    new_stones = set([c])
    new_liberties = set(n for n in neighbors(c) if board[n] == '.')
    for g in groups_to_merge:
        new_stones = new_stones | g.stones
        new_liberties = new_liberties | g.liberties
    new_liberties = new_liberties - set([c])
    updated_X_groups.append(Group(stones=new_stones, liberties=new_liberties))

    updated_O_groups = []
    for g in existing_O_groups:
        if c in g.liberties:
            updated_O_groups.append(Group(stones=g.stones, liberties=g.liberties - {c}))
        else:
            updated_O_groups.append(g)

    return updated_X_groups, updated_O_groups

class Position(namedtuple('Position', 'board n caps groups ko')):
    '''
    board: a string representation of the board
    n: an int representing moves played so far
    caps: a (int, int) tuple of captures; caps[0] is the person to play.
    groups: a (list(Group), list(Group)) tuple of lists of Groups; groups[0] represents the groups of the person to play.
    ko: a Move
    '''
    @staticmethod
    def initial_state():
        return Position(EMPTY_BOARD, n=0, caps=(0, 0), groups=(set(), set()), ko=None)

    def pass_move(self):
        return Position(self.board.translate(SWAP_COLORS), self.n+1, (self.caps[1], self.caps[0]), (self.groups[1], self.groups[0]), None)

    def play_move(self, c):
        if c is None:
            return self.pass_move()
        if c == self.ko:
            return None

        working_board = place_stone(self.board, 'X', c)
        new_X_groups, new_O_groups = update_groups(working_board, self.groups[0], self.groups[1], c)

        # process opponent's captures first, then your own suicides.
        # As stones are removed, liberty counts become inaccurate.
        O_captures = set()
        X_suicides = set()
        surviving_O_groups = []
        surviving_X_groups = []
        for group in new_O_groups:
            if not group.liberties:
                O_captures |= group.stones
                working_board = capture_stones(working_board, group.stones)
            else:
                surviving_O_groups.append(group)

        if O_captures:
            # recalculate liberties for groups adjacent to a captured O group
            coords_with_updates = find_neighbors('X', working_board, O_captures)
            final_O_groups = surviving_O_groups
            final_X_groups = [g if not (g.stones & coords_with_updates)
                else Group(stones=g.stones, liberties=find_liberties(working_board, g.stones))
                for g in new_X_groups]
        else:
            # suicide can only happen if no O captures were made
            for group in new_X_groups:
                if not group.liberties:
                    X_suicides |= group.stones
                    working_board = capture_stones(working_board, group.stones)
                else:
                    surviving_X_groups.append(group)
            # recalculate liberties for groups adjacent to a suicidal X group
            coords_with_updates = find_neighbors('O', working_board, X_suicides)
            final_X_groups = surviving_X_groups
            final_O_groups = [g if not (g.stones & coords_with_updates)
                else Group(stones=g.stones, liberties=find_liberties(working_board, g.stones))
                for g in new_O_groups]

        return Position(
            board=working_board.translate(SWAP_COLORS),
            n=self.n + 1,
            caps=(self.caps[1] + len(X_suicides), self.caps[0] + len(O_captures)),
            groups=(final_O_groups, final_X_groups),
            ko=None
        )
<commit_msg>Don't allow suicides in play_move<commit_after>import re
from collections import namedtuple
import itertools

N = 9
W = N + 1
# Represent a board as a string, with '.' empty, 'X' to play, 'O' other player.
# Whitespace is used as a border (to avoid IndexError when computing neighbors)

# A Coordinate `c` is an int: an index into the board.
# A Move is a (Coordinate c | None).

# Human-readable move notation: columns go from A to T left to right
# rows go from 1 to 19 from bottom to top.
#  ABCD
#  ____\n
# 4....\n
# 3....\n
# 2....\n
# 1....

COLUMNS = 'ABCDEFGHJKLMNOPQRST'

EMPTY_BOARD = '\n'.join(
    [' ' * N] +
    ['.' * N for i in range(N)] +
    [' ' * W])

SWAP_COLORS = str.maketrans('XO', 'OX')

def load_board(string):
    string = re.sub(r'[^OX\.#]+', '', string)
    assert len(string) == N ** 2, "Board to load didn't have right dimensions"
    return '\n'.join([' ' * N] + [string[k*N:(k+1)*N] for k in range(N)] + [' ' * W])

def parse_coords(s):
    if s == 'pass':
        return None
    s = s.upper()
    col = COLUMNS.index(s[0])
    rows_from_top = N - int(s[1:])
    return W + (W * rows_from_top) + col

def neighbors(c):
    return [c+1, c-1, c+W, c-W]

def place_stone(board, color, c):
    return board[:c] + color + board[c+1:]

def capture_stones(board, stones):
    b = bytearray(board, encoding='ascii')
    for s in stones:
        b[s] = ord('.')
    return str(b, encoding='ascii')

def flood_fill(board, c):
    'From a starting coordinate c, flood-fill the board with a #'
    b = bytearray(board, encoding='ascii')
    color = b[c]
    entire_group = [c]
    frontier = [c]
    while frontier:
        current = frontier.pop()
        b[current] = ord('#')
        for n in neighbors(current):
            if b[n] == color:
                frontier.append(n)
                entire_group.append(n)
    return str(b, encoding='ascii'), set(entire_group)

def find_neighbors(color, board, stones):
    'Find all neighbors of a set of stones of a given color'
    potential_neighbors = set(itertools.chain(*(neighbors(s) for s in stones)))
    return {c for c in potential_neighbors if board[c] == color}

def find_liberties(board, stones):
    'Given a board and a set of stones, find liberties of those stones'
    return find_neighbors('.', board, stones)

class Group(namedtuple('Group', 'stones liberties')):
    '''
    stones: a set of Coordinates belonging to this group
    liberties: a set of Coordinates that are empty and adjacent to this group.
    '''
    pass

def deduce_groups(board):
    'Given a board, return a 2-tuple; a list of groups for each player'
    def find_groups(board, color):
        groups = []
        while color in board:
            c = board.index(color)
            board, stones = flood_fill(board, c)
            liberties = find_liberties(board, stones)
            groups.append(Group(stones=stones, liberties=liberties))
        return groups

    return find_groups(board, 'X'), find_groups(board, 'O')

def update_groups(board, existing_X_groups, existing_O_groups, c):
    '''
    When a move is played, update the list of groups and their liberties.
    This means possibly appending the new move to a group, creating a new 1-stone group, or merging existing groups.
    The new move should be of color X.
    The board should represent the state after the move has been played at `c`.
    '''
    assert board[c] == 'X'

    updated_X_groups, groups_to_merge = [], []
    for g in existing_X_groups:
        if c in g.liberties:
            groups_to_merge.append(g)
        else:
            updated_X_groups.append(g)

    new_stones = {c}
    new_liberties = set(n for n in neighbors(c) if board[n] == '.')
    for g in groups_to_merge:
        new_stones = new_stones | g.stones
        new_liberties = new_liberties | g.liberties
    new_liberties = new_liberties - {c}
    updated_X_groups.append(Group(stones=new_stones, liberties=new_liberties))

    updated_O_groups = []
    for g in existing_O_groups:
        if c in g.liberties:
            updated_O_groups.append(Group(stones=g.stones, liberties=g.liberties - {c}))
        else:
            updated_O_groups.append(g)

    return updated_X_groups, updated_O_groups

class Position(namedtuple('Position', 'board n caps groups ko')):
    '''
    board: a string representation of the board
    n: an int representing moves played so far
    caps: a (int, int) tuple of captures; caps[0] is the person to play.
    groups: a (list(Group), list(Group)) tuple of lists of Groups; groups[0] represents the groups of the person to play.
    ko: a Move
    '''
    @staticmethod
    def initial_state():
        return Position(EMPTY_BOARD, n=0, caps=(0, 0), groups=(set(), set()), ko=None)

    def pass_move(self):
        return Position(self.board.translate(SWAP_COLORS), self.n+1, (self.caps[1], self.caps[0]), (self.groups[1], self.groups[0]), None)

    def play_move(self, c):
        # Obeys CGOS Rules of Play. In short:
        # No suicides
        # Chinese/area scoring
        # Positional superko (this is very crudely approximate at the moment.)
        if c is None:
            return self.pass_move()
        if c == self.ko:
            return None

        working_board = place_stone(self.board, 'X', c)
        new_X_groups, new_O_groups = update_groups(working_board, self.groups[0], self.groups[1], c)

        # process opponent's captures first, then your own suicides.
        # As stones are removed, liberty counts become inaccurate.
        O_captures = set()
        surviving_O_groups = []
        for group in new_O_groups:
            if not group.liberties:
                O_captures |= group.stones
                working_board = capture_stones(working_board, group.stones)
            else:
                surviving_O_groups.append(group)

        if O_captures:
            # recalculate liberties for groups adjacent to a captured O group
            coords_with_updates = find_neighbors('X', working_board, O_captures)
            final_O_groups = surviving_O_groups
            final_X_groups = [g if not (g.stones & coords_with_updates)
                else Group(stones=g.stones, liberties=find_liberties(working_board, g.stones))
                for g in new_X_groups]
        else:
            # suicide can only happen if no O captures were made
            for group in new_X_groups:
                if not group.liberties:
                    # suicides are illegal!
                    return None

        return Position(
            board=working_board.translate(SWAP_COLORS),
            n=self.n + 1,
            caps=(self.caps[1], self.caps[0] + len(O_captures)),
            groups=(final_O_groups, final_X_groups),
            ko=None
        )
<|endoftext|>


<commit_before>#!/usr/bin/python
# Copyright 2015 UnitedStack, Inc.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import logging
import sys

from cliff.command import Command
from cliff.lister import Lister
from steth.stethclient.utils import Logger
from steth.stethclient.utils import setup_server

SETUP_LINK_IP_PRE = "192.168.100."

class TearDownLink(Command):
    "Delete a link"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(TearDownLink, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('interface', default='.')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        self.log.debug('Agent is %s' % parsed_args.agent)
        self.log.debug('Interface is %s' % parsed_args.interface)
        server = setup_server(parsed_args.agent)
        try:
            res = server.teardown_link(parsed_args.interface)
        except Exception as e:
            self.log.error('Error %s has occured because: %s' % (res, e))

class SetUpLink(Lister):
    "Setup a link"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(SetUpLink, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('interface', default='eth0')
        parser.add_argument('cidr', default='.')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        self.log.debug('Agent is %s' % parsed_args.agent)
        self.log.debug('Interface is %s' % parsed_args.interface)
        self.log.debug('Cidr is %s' % parsed_args.cidr)
        server = setup_server(parsed_args.agent)
        try:
            # Setup Link
            server.setup_link(parsed_args.interface,
                              parsed_args.cidr)
            # Get Link info
            res = server.get_interface(parsed_args.interface)
            self.log.debug('Response is %s' % res)
            if res['code'] == 1:
                Logger.log_fail(res['message'])
                sys.exit()
            if res['code'] == 0:
                return (('Field', 'Value'),
                        ((k, v) for k, v in res['data'].items()))
        except Exception as e:
            self.log.error('Agent %s return error: %s!' % parsed_args.agent, e)
            sys.exit()

class GetInterface(Lister):
    "A test function that show a message"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(GetInterface, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('interface', default='eth0')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        self.log.debug('Agent is %s' % parsed_args.agent)
        self.log.debug('Interface is %s' % parsed_args.interface)
        server = setup_server(parsed_args.agent)
        try:
            res = server.get_interface(parsed_args.interface)
            self.log.debug('Response is %s' % res)
            if res['code'] == 1:
                Logger.log_fail(res['message'])
                sys.exit()
            if res['code'] == 0:
                return (('Field', 'Value'),
                        ((k, v) for k, v in res['data'].items()))
        except:
            self.log.error('Agent %s return error!' % parsed_args.agent)
            sys.exit()

class AddVlanToInterface(Lister):
    "Setup a link"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(AddVlanToInterface, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('interface', default='eth0')
        parser.add_argument('vlan_id', default='1124')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        self.log.debug('Agent is %s' % parsed_args.agent)
        self.log.debug('Interface is %s' % parsed_args.interface)
        self.log.debug('Vlan_id is %s' % parsed_args.vlan_id)
        server = setup_server(parsed_args.agent)
        try:
            # Setup Link
            server.add_vlan_to_interface(parsed_args.interface,
                                         parsed_args.vlan_id)
            # Get Link info
            new_interface = parsed_args.interface + '.' + parsed_args.vlan_id
            res = server.get_interface(new_interface)
            self.log.debug('Response is %s' % res)
            if res['code'] == 1:
                Logger.log_fail(res['message'])
                sys.exit()
            if res['code'] == 0:
                return (('Field', 'Value'),
                        ((k, v) for k, v in res['data'].items()))
        except Exception as e:
            self.log.error('Agent %s return error: %s!' % parsed_args.agent, e)
            sys.exit()

class AgentPing(Lister):
    "Ping a destination from one agent"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(AgentPing, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('destination', default='1.2.4.8')
        parser.add_argument('--count', nargs='?', default='2')
        parser.add_argument('--timeout', nargs='?', default='2')
        parser.add_argument('--interface', nargs='?', default='eth0')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        server = setup_server(parsed_args.agent)
        try:
            dest = parsed_args.destination.split(',')
            res = server.ping(ips=dest,
                              count=parsed_args.count,
                              timeout=parsed_args.timeout,
                              interface=parsed_args.interface)
            self.log.debug('Response is %s' % res)
            if res['code'] == 1:
                Logger.log_fail(res['message'])
                sys.exit()
            if res['code'] == 0:
                return (('Destination', 'Packet Loss (%)'),
                        ((k, v) for k, v in res['data'].items()))
        except Exception as e:
            self.log.error('Agent %s return error: %s!' % parsed_args.agent, e)
            sys.exit()

class CheckPortsOnBr(Lister):
    "Check a port if exists on a ovs bridge"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(CheckPortsOnBr, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('bridge', default='br-int')
        parser.add_argument('port', default='br-int')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        server = setup_server(parsed_args.agent)
        try:
            res = server.check_ports_on_br(parsed_args.bridge,
                                           parsed_args.port)
            self.log.debug('Response is %s' % res)
            if res['code'] == 1:
                Logger.log_fail(res['message'])
                sys.exit()
            if res['code'] == 0:
                return (('Port', 'Exists'),
                        ((k, v) for k, v in res['data'].items()))
        except Exception as e:
            self.log.error('Agent %s return error: %s!' % parsed_args.agent, e)
            sys.exit()

class CheckVlanInterface(Lister):
    """Check vlan if exists in switch"""
    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(CheckVlanInterface, self).get_parser(prog_name)
        parser.add_argument('agentA', default='bad')
        parser.add_argument('agentB', default='bad')
        parser.add_argument('interface', default='eth0')
        parser.add_argument('vlan_id', default='1124')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        serverA = setup_server(parsed_args.agentA)
        serverB = setup_server(parsed_args.agentB)
        try:
            interface = parsed_args.interface + '.' + parsed_args.vlan_id
            # First of all, check the interface if exists
            resA = serverA.get_interface(interface)
            resB = serverB.get_interface(interface)
            if resA['code'] == 1:
                msg = "Agent: %s has no interface named %s!" % (
                    parsed_args.agentA, interface)
                Logger.log_fail(msg)
                sys.exit()
            if resB['code'] == 1:
                msg = "Agent: %s has no interface named %s!" % (
                    parsed_args.agentB, interface)
                Logger.log_fail(msg)
                sys.exit()
            # add vlan interface in each agent
            resA = serverA.add_vlan_to_interface(parsed_args.interface,
                                                 parsed_args.vlan_id)
            self.log.debug('Response is %s' % resA)
            resB = serverB.add_vlan_to_interface(parsed_args.interface,
                                                 parsed_args.vlan_id)
            self.log.debug('Response is %s' % resB)
            Logger.log_normal(('AgentA and agentB has already added the '
                               'interface %s ') % (interface))
            # setup link in each agent
            ipA = SETUP_LINK_IP_PRE + parsed_args.agentA.split('-')[1] + '/24'
            resA = serverA.setup_link(interface, ipA)
            self.log.debug('Response is %s' % resA)
            ipB = SETUP_LINK_IP_PRE + parsed_args.agentB.split('-')[1] + '/24'
            resB = serverB.setup_link(interface, ipB)
            self.log.debug('Response is %s' % resB)
            Logger.log_normal(('AgentA and agentB has already setup the '
                               'IP %s and IP %s') % (ipA, ipB))
            # ping a agent from exists IP to check connectivity
            res = serverA.ping(ips=[ipB])
            # teardown the interface in each agent to clean all resources
            resA = serverA.teardown_link(interface)
            self.log.debug('Response is %s' % resA)
            resB = serverB.teardown_link(interface)
            self.log.debug('Response is %s' % resB)
            Logger.log_normal(('AgentA and agentB has already deleted the'
                               'vlan %s in %s') % (parsed_args.vlan_id,
                                                   parsed_args.interface))
            if res['code'] == 0:
                return (('Destination', 'Packet Loss (%)'),
                        ((k, v) for k, v in res['data'].items()))
        except Exception as e:
            self.log.error('Agent %s return error: %s!' % parsed_args.agent, e)
            sys.exit()
<commit_msg>Catch error msg when setup link<commit_after>#!/usr/bin/python
# Copyright 2015 UnitedStack, Inc.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import logging
import sys

from cliff.command import Command
from cliff.lister import Lister
from steth.stethclient.utils import Logger
from steth.stethclient.utils import setup_server

SETUP_LINK_IP_PRE = "192.168.100."

class TearDownLink(Command):
    "Delete a link"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(TearDownLink, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('interface', default='.')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        self.log.debug('Agent is %s' % parsed_args.agent)
        self.log.debug('Interface is %s' % parsed_args.interface)
        server = setup_server(parsed_args.agent)
        try:
            res = server.teardown_link(parsed_args.interface)
        except Exception as e:
            self.log.error('Error %s has occured because: %s' % (res, e))

class SetUpLink(Lister):
    "Setup a link"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(SetUpLink, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('interface', default='eth0')
        parser.add_argument('cidr', default='.')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        self.log.debug('Agent is %s' % parsed_args.agent)
        self.log.debug('Interface is %s' % parsed_args.interface)
        self.log.debug('Cidr is %s' % parsed_args.cidr)
        server = setup_server(parsed_args.agent)
        # Setup Link
        server.setup_link(parsed_args.interface,
                          parsed_args.cidr)
        # Get Link info
        res = server.get_interface(parsed_args.interface)
        self.log.debug('Response is %s' % res)
        if res['code'] == 1:
            Logger.log_fail(res['message'])
            sys.exit()
        if res['code'] == 0:
            return (('Field', 'Value'),
                    ((k, v) for k, v in res['data'].items()))
        return (['Error Mssage', ' '], [('message', res['message'])])

class GetInterface(Lister):
    "A test function that show a message"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(GetInterface, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('interface', default='eth0')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        self.log.debug('Agent is %s' % parsed_args.agent)
        self.log.debug('Interface is %s' % parsed_args.interface)
        server = setup_server(parsed_args.agent)
        try:
            res = server.get_interface(parsed_args.interface)
            self.log.debug('Response is %s' % res)
            if res['code'] == 1:
                Logger.log_fail(res['message'])
                sys.exit()
            if res['code'] == 0:
                return (('Field', 'Value'),
                        ((k, v) for k, v in res['data'].items()))
        except:
            self.log.error('Agent %s return error!' % parsed_args.agent)
            sys.exit()

class AddVlanToInterface(Lister):
    "Setup a link"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(AddVlanToInterface, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('interface', default='eth0')
        parser.add_argument('vlan_id', default='1124')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        self.log.debug('Agent is %s' % parsed_args.agent)
        self.log.debug('Interface is %s' % parsed_args.interface)
        self.log.debug('Vlan_id is %s' % parsed_args.vlan_id)
        server = setup_server(parsed_args.agent)
        try:
            # Setup Link
            server.add_vlan_to_interface(parsed_args.interface,
                                         parsed_args.vlan_id)
            # Get Link info
            new_interface = parsed_args.interface + '.' + parsed_args.vlan_id
            res = server.get_interface(new_interface)
            self.log.debug('Response is %s' % res)
            if res['code'] == 1:
                Logger.log_fail(res['message'])
                sys.exit()
            if res['code'] == 0:
                return (('Field', 'Value'),
                        ((k, v) for k, v in res['data'].items()))
        except Exception as e:
            self.log.error('Agent %s return error: %s!' % parsed_args.agent, e)
            sys.exit()

class AgentPing(Lister):
    "Ping a destination from one agent"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(AgentPing, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('destination', default='1.2.4.8')
        parser.add_argument('--count', nargs='?', default='2')
        parser.add_argument('--timeout', nargs='?', default='2')
        parser.add_argument('--interface', nargs='?', default='eth0')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        server = setup_server(parsed_args.agent)
        try:
            dest = parsed_args.destination.split(',')
            res = server.ping(ips=dest,
                              count=parsed_args.count,
                              timeout=parsed_args.timeout,
                              interface=parsed_args.interface)
            self.log.debug('Response is %s' % res)
            if res['code'] == 1:
                Logger.log_fail(res['message'])
                sys.exit()
            if res['code'] == 0:
                return (('Destination', 'Packet Loss (%)'),
                        ((k, v) for k, v in res['data'].items()))
        except Exception as e:
            self.log.error('Agent %s return error: %s!' % parsed_args.agent, e)
            sys.exit()

class CheckPortsOnBr(Lister):
    "Check a port if exists on a ovs bridge"

    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(CheckPortsOnBr, self).get_parser(prog_name)
        parser.add_argument('agent', default='bad')
        parser.add_argument('bridge', default='br-int')
        parser.add_argument('port', default='br-int')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        server = setup_server(parsed_args.agent)
        try:
            res = server.check_ports_on_br(parsed_args.bridge,
                                           parsed_args.port)
            self.log.debug('Response is %s' % res)
            if res['code'] == 1:
                Logger.log_fail(res['message'])
                sys.exit()
            if res['code'] == 0:
                return (('Port', 'Exists'),
                        ((k, v) for k, v in res['data'].items()))
        except Exception as e:
            self.log.error('Agent %s return error: %s!' % parsed_args.agent, e)
            sys.exit()

class CheckVlanInterface(Lister):
    """Check vlan if exists in switch"""
    log = logging.getLogger(__name__)

    def get_parser(self, prog_name):
        parser = super(CheckVlanInterface, self).get_parser(prog_name)
        parser.add_argument('agentA', default='bad')
        parser.add_argument('agentB', default='bad')
        parser.add_argument('interface', default='eth0')
        parser.add_argument('vlan_id', default='1124')
        return parser

    def take_action(self, parsed_args):
        self.log.debug('Get parsed_args: %s' % parsed_args)
        serverA = setup_server(parsed_args.agentA)
        serverB = setup_server(parsed_args.agentB)
        try:
            interface = parsed_args.interface + '.' + parsed_args.vlan_id
            # First of all, check the interface if exists
            resA = serverA.get_interface(interface)
            resB = serverB.get_interface(interface)
            if resA['code'] == 1:
                msg = "Agent: %s has no interface named %s!" % (
                    parsed_args.agentA, interface)
                Logger.log_fail(msg)
                sys.exit()
            if resB['code'] == 1:
                msg = "Agent: %s has no interface named %s!" % (
                    parsed_args.agentB, interface)
                Logger.log_fail(msg)
                sys.exit()
            # add vlan interface in each agent
            resA = serverA.add_vlan_to_interface(parsed_args.interface,
                                                 parsed_args.vlan_id)
            self.log.debug('Response is %s' % resA)
            resB = serverB.add_vlan_to_interface(parsed_args.interface,
                                                 parsed_args.vlan_id)
            self.log.debug('Response is %s' % resB)
            Logger.log_normal(('AgentA and agentB has already added the '
                               'interface %s ') % (interface))
            # setup link in each agent
            ipA = SETUP_LINK_IP_PRE + parsed_args.agentA.split('-')[1] + '/24'
            resA = serverA.setup_link(interface, ipA)
            self.log.debug('Response is %s' % resA)
            ipB = SETUP_LINK_IP_PRE + parsed_args.agentB.split('-')[1] + '/24'
            resB = serverB.setup_link(interface, ipB)
            self.log.debug('Response is %s' % resB)
            Logger.log_normal(('AgentA and agentB has already setup the '
                               'IP %s and IP %s') % (ipA, ipB))
            # ping a agent from exists IP to check connectivity
            res = serverA.ping(ips=[ipB])
            # teardown the interface in each agent to clean all resources
            resA = serverA.teardown_link(interface)
            self.log.debug('Response is %s' % resA)
            resB = serverB.teardown_link(interface)
            self.log.debug('Response is %s' % resB)
            Logger.log_normal(('AgentA and agentB has already deleted the'
                               'vlan %s in %s') % (parsed_args.vlan_id,
                                                   parsed_args.interface))
            if res['code'] == 0:
                return (('Destination', 'Packet Loss (%)'),
                        ((k, v) for k, v in res['data'].items()))
        except Exception as e:
            self.log.error('Agent %s return error: %s!' % parsed_args.agent, e)
            sys.exit()
<|endoftext|>


<commit_before># Copyright 2019 Sergio Teruel <sergio.teruel@tecnativa.com>
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).
{
    "name": "Stock Barcodes",
    "summary": "It provides read barcode on stock operations.",
    "version": "13.0.1.0.1",
    "author": "Tecnativa, " "Odoo Community Association (OCA)",
    "website": "https://github.com/OCA/stock-logistics-barcode",
    "license": "AGPL-3",
    "category": "Extra Tools",
    "depends": ["barcodes", "stock"],
    "data": [
        "security/ir.model.access.csv",
        "views/assets.xml",
        "views/stock_inventory_views.xml",
        "views/stock_location_views.xml",
        "views/stock_picking_views.xml",
        "wizard/stock_production_lot_views.xml",
        "wizard/stock_barcodes_read_views.xml",
        "wizard/stock_barcodes_read_inventory_views.xml",
        "wizard/stock_barcodes_read_picking_views.xml",
    ],
    "installable": True,
}
<commit_msg>stock_barcodes 13.0.1.1.0<commit_after># Copyright 2019 Sergio Teruel <sergio.teruel@tecnativa.com>
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).
{
    "name": "Stock Barcodes",
    "summary": "It provides read barcode on stock operations.",
    "version": "13.0.1.1.0",
    "author": "Tecnativa, " "Odoo Community Association (OCA)",
    "website": "https://github.com/OCA/stock-logistics-barcode",
    "license": "AGPL-3",
    "category": "Extra Tools",
    "depends": ["barcodes", "stock"],
    "data": [
        "security/ir.model.access.csv",
        "views/assets.xml",
        "views/stock_inventory_views.xml",
        "views/stock_location_views.xml",
        "views/stock_picking_views.xml",
        "wizard/stock_production_lot_views.xml",
        "wizard/stock_barcodes_read_views.xml",
        "wizard/stock_barcodes_read_inventory_views.xml",
        "wizard/stock_barcodes_read_picking_views.xml",
    ],
    "installable": True,
}
<|endoftext|>