Python Utils.NodeSet类代码示例

    def update(self):
        """Update runtime progress info"""
        wrbwinfo = ''
        if self.bytes_written > 0:
            bandwidth = self.bytes_written/(time.time() - self.start_time)
            wrbwinfo = " write: %s/s" % human_bi_bytes_unit(bandwidth)

        gws = self.task.gateways.keys()
        if gws:
            # tree mode
            act_targets = NodeSet()
            for gw, (chan, metaworkers) in self.task.gateways.iteritems():
                act_targets.updaten(mw.gwtargets[gw] for mw in metaworkers)
            cnt = len(act_targets) + len(self.task._engine.clients()) - len(gws)
            gwinfo = ' gw %d' % len(gws)
        else:
            cnt = len(self.task._engine.clients())
            gwinfo = ''
        if self.bytes_written > 0 or cnt != self.cnt_last:
            self.cnt_last = cnt
            # display completed/total clients
            towrite = 'clush: %*d/%*d%s%s\r' % (self.tslen, self.total - cnt,
                                                self.tslen, self.total, gwinfo,
                                                wrbwinfo)
            self.wholelen = len(towrite)
            sys.stderr.write(towrite)
            self.started = True

    def ev_close(self, worker):
        # Worker is closing -- it's time to gather results...
        self._runtimer_finalize(worker)
        # Display command output, try to order buffers by rc
        nodesetify = lambda v: (v[0], NodeSet._fromlist1(v[1]))
        cleaned = False
        for _rc, nodelist in sorted(worker.iter_retcodes()):
            ns_remain = NodeSet._fromlist1(nodelist)
            # Then order by node/nodeset (see bufnodeset_cmp)
            for buf, nodeset in sorted(map(nodesetify,
                                           worker.iter_buffers(nodelist)),
                                       cmp=bufnodeset_cmp):
                if not cleaned:
                    # clean runtimer line before printing first result
                    self._runtimer_clean()
                    cleaned = True
                self._display.print_gather(nodeset, buf)
                ns_remain.difference_update(nodeset)
            if ns_remain:
                self._display.print_gather_finalize(ns_remain)
        self._display.flush()

        self._close_common(worker)

        # Notify main thread to update its prompt
        self.update_prompt(worker)

def display(tree, disp, gather, trace_mode, enable_nodeset_key):
    """nicely display MsgTree instance `tree' content according to
    `disp' Display object and `gather' boolean flag"""
    out = sys.stdout
    try:
        if trace_mode:
            display_tree(tree, disp, out)
        else:
            if gather:
                if enable_nodeset_key:
                    # lambda to create a NodeSet from keys returned by walk()
                    ns_getter = lambda x: NodeSet.fromlist(x[1])
                    for nodeset in sorted(imap(ns_getter, tree.walk()),
                                          cmp=nodeset_cmp):
                        disp.print_gather(nodeset, tree[nodeset[0]])
                else:
                    for msg, key in tree.walk():
                        disp.print_gather_keys(key, msg)
            else:
                if enable_nodeset_key:
                    # nodes are automagically sorted by NodeSet
                    for node in NodeSet.fromlist(tree.keys()).nsiter():
                        disp.print_gather(node, tree[str(node)])
                else:
                    for key in tree.keys():
                        disp.print_gather_keys([ key ], tree[key])
    finally:
        out.flush()

class LiveGatherOutputHandler(GatherOutputHandler):
    """Live line-gathered output event handler class."""

    def __init__(self, display, nodes):
        assert nodes is not None, "cannot gather local command"
        GatherOutputHandler.__init__(self, display)
        self._nodes = NodeSet(nodes)
        self._nodecnt = dict.fromkeys(self._nodes, 0)
        self._mtreeq = []
        self._offload = 0

    def ev_read(self, worker):
        # Read new line from node
        node = worker.current_node
        self._nodecnt[node] += 1
        cnt = self._nodecnt[node]
        if len(self._mtreeq) < cnt:
            self._mtreeq.append(MsgTree())
        self._mtreeq[cnt - self._offload - 1].add(node, worker.current_msg)
        self._live_line(worker)

    def ev_hup(self, worker):
        if self._mtreeq and worker.current_node not in self._mtreeq[0]:
            # forget a node that doesn't answer to continue live line
            # gathering anyway
            self._nodes.remove(worker.current_node)
            self._live_line(worker)

    def _live_line(self, worker):
        # if all nodes have replied, display gathered line
        while self._mtreeq and len(self._mtreeq[0]) == len(self._nodes):
            mtree = self._mtreeq.pop(0)
            self._offload += 1
            self._runtimer_clean()
            nodesetify = lambda v: (v[0], NodeSet.fromlist(v[1]))
            for buf, nodeset in sorted(map(nodesetify, mtree.walk()),
                                       cmp=bufnodeset_cmp):
                self._display.print_gather(nodeset, buf)
            self._runtimer_set_dirty()

    def ev_close(self, worker):
        # Worker is closing -- it's time to gather results...
        self._runtimer_finalize(worker)

        for mtree in self._mtreeq:
            nodesetify = lambda v: (v[0], NodeSet.fromlist(v[1]))
            for buf, nodeset in sorted(map(nodesetify, mtree.walk()),
                                       cmp=bufnodeset_cmp):
                self._display.print_gather(nodeset, buf)

        self._close_common(worker)

        # Notify main thread to update its prompt
        self.update_prompt(worker)

    def _close_common(self, worker):
        verbexit = VERB_QUIET
        if self._display.maxrc:
            verbexit = VERB_STD
        # Display return code if not ok ( != 0)
        for rc, nodelist in worker.iter_retcodes():
            if rc != 0:
                ns = NodeSet._fromlist1(nodelist)
                self._display.vprint_err(verbexit, \
                    "clush: %s: exited with exit code %d" % (ns, rc))

        # Display nodes that didn't answer within command timeout delay
        if worker.num_timeout() > 0:
            self._display.vprint_err(verbexit, "clush: %s: command timeout" % \
                NodeSet._fromlist1(worker.iter_keys_timeout()))

 def __init__(self, display, nodes):
     assert nodes is not None, "cannot gather local command"
     GatherOutputHandler.__init__(self, display)
     self._nodes = NodeSet(nodes)
     self._nodecnt = dict.fromkeys(self._nodes, 0)
     self._mtreeq = []
     self._offload = 0

 def update(self):
     gws = self.task.gateways.keys()
     if gws:
         # tree mode
         act_targets = NodeSet()
         for gw, (chan, metaworkers) in self.task.gateways.iteritems():
             act_targets.updaten(mw.gwtargets[gw] for mw in metaworkers)
         cnt = len(act_targets) + len(self.task._engine.clients()) - len(gws)
         gwinfo = " gw %d" % len(gws)
     else:
         cnt = len(self.task._engine.clients())
         gwinfo = ""
     if cnt != self.cnt_last:
         self.cnt_last = cnt
         # display completed/total clients
         towrite = "clush: %*d/%*d%s\r" % (self.tslen, self.total - cnt, self.tslen, self.total, gwinfo)
         self.wholelen = len(towrite)
         sys.stderr.write(towrite)
         self.started = True

 def _live_line(self, worker):
     # if all nodes have replied, display gathered line
     while self._mtreeq and len(self._mtreeq[0]) == len(self._nodes):
         mtree = self._mtreeq.pop(0)
         self._offload += 1
         self._runtimer_clean()
         nodesetify = lambda v: (v[0], NodeSet.fromlist(v[1]))
         for buf, nodeset in sorted(map(nodesetify, mtree.walk()), cmp=bufnodeset_cmp):
             self._display.print_gather(nodeset, buf)
         self._runtimer_set_dirty()

def display_tree(tree, disp, out):
    """display sub-routine for clubak -T (msgtree trace mode)"""
    togh = True
    offset = 2
    reldepth = -offset
    reldepths = {}
    line_mode = disp.line_mode
    for msgline, keys, depth, nchildren in tree.walk_trace():
        if togh:
            if depth in reldepths:
                reldepth = reldepths[depth]
            else:
                reldepth = reldepths[depth] = reldepth + offset
            if line_mode:
                out.write("%s:\n" % NodeSet.fromlist(keys))
            else:
                out.write("%s\n" % \
                    (disp.format_header(NodeSet.fromlist(keys), reldepth)))
        out.write("%s%s\n" % (" " * reldepth, msgline))
        togh = nchildren != 1

    def ev_close(self, worker):
        # Worker is closing -- it's time to gather results...
        self._runtimer_finalize(worker)

        for mtree in self._mtreeq:
            nodesetify = lambda v: (v[0], NodeSet.fromlist(v[1]))
            for buf, nodeset in sorted(map(nodesetify, mtree.walk()), cmp=bufnodeset_cmp):
                self._display.print_gather(nodeset, buf)

        self._close_common(worker)

        # Notify main thread to update its prompt
        self.update_prompt(worker)

    def ev_close(self, worker):
        # Worker is closing -- it's time to gather results...
        self._runtimer_finalize(worker)
        assert worker.current_node is not None, "cannot gather local command"
        # Display command output, try to order buffers by rc
        nodesetify = lambda v: (v[0], NodeSet.fromlist(v[1]))
        cleaned = False
        for rc, nodelist in worker.iter_retcodes():
            # Then order by node/nodeset (see bufnodeset_cmp)
            for buf, nodeset in sorted(map(nodesetify,
                                           worker.iter_buffers(nodelist)),
                                       cmp=bufnodeset_cmp):
                if not cleaned:
                    # clean runtimer line before printing first result
                    self._runtimer_clean()
                    cleaned = True
                self._display.print_gather(nodeset, buf)

        self._close_common(worker)

        # Notify main thread to update its prompt
        self.update_prompt(worker)

def nodeset():
    """script subroutine"""
    class_set = NodeSet
    usage = "%prog [COMMAND] [OPTIONS] [ns1 [-ixX] ns2|...]"

    parser = OptionParser(usage)
    parser.install_nodeset_commands()
    parser.install_nodeset_operations()
    parser.install_nodeset_options()
    (options, args) = parser.parse_args()

    if options.debug:
        STD_GROUP_RESOLVER.set_verbosity(1)

    # Check for command presence
    cmdcount = int(options.count) + int(options.expand) + \
               int(options.fold) + int(bool(options.list)) + \
               int(options.regroup) + int(options.groupsources)
    if not cmdcount:
        parser.error("No command specified.")
    elif cmdcount > 1:
        parser.error("Multiple commands not allowed.")

    if options.rangeset:
        class_set = RangeSet

    if options.all or options.regroup:
        assert class_set == NodeSet, "-a/-r only supported in NodeSet mode"

    if options.groupsource and not options.quiet and \
       (class_set == RangeSet or options.groupsources):
        print >> sys.stderr, "WARNING: option group source \"%s\" ignored" \
                                % options.groupsource

    # The list command doesn't need any NodeSet, check for it first.
    if options.list > 0:
        list_level = options.list
        for group in grouplist(options.groupsource):
            if options.groupsource and not options.groupbase:
                nsgroup = "@%s:%s" % (options.groupsource, group)
            else:
                nsgroup = "@%s" % group
            if list_level == 1:
                print nsgroup
            else:
                nodes = NodeSet(nsgroup)
                if list_level == 2:     # -ll ?
                    print "%s %s" % (nsgroup, nodes)
                else:                   # -lll ?
                    print "%s %s %d" % (nsgroup, nodes, len(nodes))
        return
    # Also, the groupsources command simply lists group sources.
    elif options.groupsources:
        if options.quiet:
            dispdefault = ""    # don't show (default) if quiet is set
        else:
            dispdefault = " (default)"
        for src in STD_GROUP_RESOLVER.sources():
            print "%s%s" % (src, dispdefault)
            dispdefault = ""
        return

    # We want -s <groupsource> to act as a substition of default groupsource
    # (ie. it's not necessary to prefix group names by this group source).
    if options.groupsource:
        STD_GROUP_RESOLVER.default_sourcename = options.groupsource

    # Instantiate RangeSet or NodeSet object
    xset = class_set(autostep=options.autostep)

    if options.all:
        # Include all nodes from external node groups support.
        xset.update(NodeSet.fromall()) # uses default_sourcename
    elif not args:
        # No need to specify '-' to read stdin if no argument at all.
        process_stdin(xset.update, xset.__class__, options.autostep)

    # Apply first operations (before first non-option)
    for nodes in options.and_nodes:
        if nodes == '-':
            process_stdin(xset.intersection_update, xset.__class__,
                          options.autostep)
        else:
            xset.intersection_update(class_set(nodes,
                                               autostep=options.autostep))
    for nodes in options.sub_nodes:
        if nodes == '-':
            process_stdin(xset.difference_update, xset.__class__,
                          options.autostep)
        else:
            xset.difference_update(class_set(nodes, autostep=options.autostep))
    for nodes in options.xor_nodes:
        if nodes == '-':
            process_stdin(xset.symmetric_difference_update, xset.__class__,
                          options.autostep)
        else:
            xset.symmetric_difference_update(class_set(nodes, \
                                             autostep=options.autostep))

    # Finish xset computing from args
#.........这里部分代码省略.........

def nodeset():
    """script subroutine"""
    class_set = NodeSet
    usage = "%prog [COMMAND] [OPTIONS] [ns1 [-ixX] ns2|...]"

    parser = OptionParser(usage)
    parser.install_nodeset_commands()
    parser.install_nodeset_operations()
    parser.install_nodeset_options()
    (options, args) = parser.parse_args()

    group_resolver = std_group_resolver()

    if options.debug:
        logging.basicConfig(level=logging.DEBUG)

    # Check for command presence
    cmdcount = int(options.count) + int(options.expand) + \
               int(options.fold) + int(bool(options.list)) + \
               int(bool(options.listall)) + int(options.regroup) + \
               int(options.groupsources)
    if not cmdcount:
        parser.error("No command specified.")
    elif cmdcount > 1:
        parser.error("Multiple commands not allowed.")

    if options.rangeset:
        class_set = RangeSet

    if options.all or options.regroup:
        if class_set != NodeSet:
            parser.error("-a/-r only supported in NodeSet mode")

    if options.maxsplit is not None and options.contiguous:
        parser.error("incompatible splitting options (split, contiguous)")

    if options.maxsplit is None:
        options.maxsplit = 1

    if options.axis and (not options.fold or options.rangeset):
        parser.error("--axis option is only supported when folding nodeset")

    if options.groupsource and not options.quiet and class_set == RangeSet:
        print >> sys.stderr, "WARNING: option group source \"%s\" ignored" \
                                % options.groupsource

    # We want -s <groupsource> to act as a substition of default groupsource
    # (ie. it's not necessary to prefix group names by this group source).
    if options.groupsource:
        group_resolver.default_source_name = options.groupsource

    # The groupsources command simply lists group sources.
    if options.groupsources:
        if options.quiet:
            dispdefault = ""    # don't show (default) if quiet is set
        else:
            dispdefault = " (default)"
        for src in group_resolver.sources():
            print "%s%s" % (src, dispdefault)
            dispdefault = ""
        return

    autostep = options.autostep

    # Do not use autostep for computation when a percentage or the special
    # value 'auto' is specified. Real autostep value is set post-process.
    if type(autostep) is float or autostep == 'auto':
        autostep = None

    # Instantiate RangeSet or NodeSet object
    xset = class_set(autostep=autostep)

    if options.all:
        # Include all nodes from external node groups support.
        xset.update(NodeSet.fromall()) # uses default_source when set

    if not args and not options.all and not (options.list or options.listall):
        # No need to specify '-' to read stdin in these cases
        process_stdin(xset.update, xset.__class__, autostep)

    # Apply first operations (before first non-option)
    for nodes in options.and_nodes:
        if nodes == '-':
            process_stdin(xset.intersection_update, xset.__class__, autostep)
        else:
            xset.intersection_update(class_set(nodes, autostep=autostep))
    for nodes in options.sub_nodes:
        if nodes == '-':
            process_stdin(xset.difference_update, xset.__class__, autostep)
        else:
            xset.difference_update(class_set(nodes, autostep=autostep))
    for nodes in options.xor_nodes:
        if nodes == '-':
            process_stdin(xset.symmetric_difference_update, xset.__class__,
                          autostep)
        else:
            xset.symmetric_difference_update(class_set(nodes,
                                                       autostep=autostep))

    # Finish xset computing from args
#.........这里部分代码省略.........

def ttyloop(task, nodeset, timeout, display, remote):
    """Manage the interactive prompt to run command"""
    readline_avail = False
    interactive = task.default("USER_interactive")
    if interactive:
        try:
            import readline
            readline_setup()
            readline_avail = True
        except ImportError:
            pass
        display.vprint(VERB_STD, \
            "Enter 'quit' to leave this interactive mode")

    rc = 0
    ns = NodeSet(nodeset)
    ns_info = True
    cmd = ""
    while task.default("USER_running") or \
            (interactive and cmd.lower() != 'quit'):
        try:
            # Set SIGUSR1 handler if needed
            if task.default("USER_handle_SIGUSR1"):
                signal.signal(signal.SIGUSR1, signal_handler)

            if task.default("USER_interactive") and \
                    not task.default("USER_running"):
                if ns_info:
                    display.vprint(VERB_QUIET, \
                                   "Working with nodes: %s" % ns)
                    ns_info = False
                prompt = "clush> "
            else:
                prompt = ""
            try:
                cmd = raw_input(prompt)
                assert cmd is not None, "Result of raw_input() is None!"
            finally:
                signal.signal(signal.SIGUSR1, signal.SIG_IGN)
        except EOFError:
            print
            return
        except UpdatePromptException:
            if task.default("USER_interactive"):
                continue
            return
        except KeyboardInterrupt, kbe:
            # Caught SIGINT here (main thread) but the signal will also reach
            # subprocesses (that will most likely kill them)
            if display.gather:
                # Suspend task, so we can safely access its data from here
                task.suspend()

                # If USER_running is not set, the task had time to finish,
                # that could mean all subprocesses have been killed and all
                # handlers have been processed.
                if not task.default("USER_running"):
                    # let's clush_excepthook handle the rest
                    raise kbe

                # If USER_running is set, the task didn't have time to finish
                # its work, so we must print something for the user...
                print_warn = False

                # Display command output, but cannot order buffers by rc
                nodesetify = lambda v: (v[0], NodeSet._fromlist1(v[1]))
                for buf, nodeset in sorted(map(nodesetify, task.iter_buffers()),
                                           cmp=bufnodeset_cmp):
                    if not print_warn:
                        print_warn = True
                        display.vprint_err(VERB_STD, \
                            "Warning: Caught keyboard interrupt!")
                    display.print_gather(nodeset, buf)

                # Return code handling
                verbexit = VERB_QUIET
                if display.maxrc:
                    verbexit = VERB_STD
                ns_ok = NodeSet()
                for rc, nodelist in task.iter_retcodes():
                    ns_ok.add(NodeSet._fromlist1(nodelist))
                    if rc != 0:
                        # Display return code if not ok ( != 0)
                        ns = NodeSet._fromlist1(nodelist)
                        display.vprint_err(verbexit, \
                            "clush: %s: exited with exit code %s" % (ns, rc))
                # Add uncompleted nodeset to exception object
                kbe.uncompleted_nodes = ns - ns_ok

                # Display nodes that didn't answer within command timeout delay
                if task.num_timeout() > 0:
                    display.vprint_err(verbexit, \
                        "clush: %s: command timeout" % \
                            NodeSet._fromlist1(task.iter_keys_timeout()))
            raise kbe

        if task.default("USER_running"):
            ns_reg, ns_unreg = NodeSet(), NodeSet()
            for client in task._engine.clients():
                if client.registered:
#.........这里部分代码省略.........

def main(args=sys.argv):
    """clush script entry point"""
    sys.excepthook = clush_excepthook

    # Default values
    nodeset_base, nodeset_exclude = NodeSet(), NodeSet()

    #
    # Argument management
    #
    usage = "%prog [options] command"

    parser = OptionParser(usage)

    parser.add_option("--nostdin", action="store_true", dest="nostdin",
                      help="don't watch for possible input from stdin")

    parser.install_nodes_options()
    parser.install_display_options(verbose_options=True)
    parser.install_filecopy_options()
    parser.install_ssh_options()

    (options, args) = parser.parse_args(args[1:])

    #
    # Load config file and apply overrides
    #
    config = ClushConfig(options)

    # Should we use ANSI colors for nodes?
    if config.color == "auto":
        color = sys.stdout.isatty() and (options.gatherall or \
                                         sys.stderr.isatty())
    else:
        color = config.color == "always"

    # Create and configure display object.
    display = Display(options, config, color)

    #
    # Compute the nodeset
    #
    if options.nodes:
        nodeset_base = NodeSet.fromlist(options.nodes)
    if options.exclude:
        nodeset_exclude = NodeSet.fromlist(options.exclude)

    if options.groupsource:
        # Be sure -a/g -s source work as espected.
        STD_GROUP_RESOLVER.default_sourcename = options.groupsource

    # FIXME: add public API to enforce engine
    Task._std_default['engine'] = options.engine

    # Do we have nodes group?
    task = task_self()
    task.set_info("debug", config.verbosity > 1)
    if config.verbosity == VERB_DEBUG:
        STD_GROUP_RESOLVER.set_verbosity(1)
    if options.nodes_all:
        all_nodeset = NodeSet.fromall()
        display.vprint(VERB_DEBUG, "Adding nodes from option -a: %s" % \
                                   all_nodeset)
        nodeset_base.add(all_nodeset)

    if options.group:
        grp_nodeset = NodeSet.fromlist(options.group,
                                       resolver=NOGROUP_RESOLVER)
        for grp in grp_nodeset:
            addingrp = NodeSet("@" + grp)
            display.vprint(VERB_DEBUG, \
                "Adding nodes from option -g %s: %s" % (grp, addingrp))
            nodeset_base.update(addingrp)

    if options.exgroup:
        grp_nodeset = NodeSet.fromlist(options.exgroup,
                                       resolver=NOGROUP_RESOLVER)
        for grp in grp_nodeset:
            removingrp = NodeSet("@" + grp)
            display.vprint(VERB_DEBUG, \
                "Excluding nodes from option -X %s: %s" % (grp, removingrp))
            nodeset_exclude.update(removingrp)

    # Do we have an exclude list? (-x ...)
    nodeset_base.difference_update(nodeset_exclude)
    if len(nodeset_base) < 1:
        parser.error('No node to run on.')

    # Set open files limit.
    set_fdlimit(config.fd_max, display)

    #
    # Task management
    #
    # check for clush interactive mode
    interactive = not len(args) and \
                  not (options.copy or options.rcopy)
    # check for foreground ttys presence (input)
    stdin_isafgtty = sys.stdin.isatty() and \
        os.tcgetpgrp(sys.stdin.fileno()) == os.getpgrp()
#.........这里部分代码省略.........