<commit_before>from __future__ import print_function, unicode_literals, absolute_import
from builtins import dict, str
import logging
import numpy as np
import networkx
import itertools
from copy import deepcopy
from pysb import kappa
from pysb import Observable
from pysb.core import as_complex_pattern
from indra.statements import *
from indra.assemblers import pysb_assembler as pa

logger = logging.getLogger('model_checker')

class ModelChecker(object):
    """Check a PySB model against a set of INDRA statements."""

    def __init__(self, model, statements):
        self.model = model
        self.statements = statements
        self._im = None

    def get_im(self):
        if self._im:
            return self._im
        if not self.model:
            raise Exception("Cannot get influence map if there is no model.")
        else:
            self._im = kappa.influence_map(self.model)
            self._im.is_multigraph = lambda: False
            return self._im

    def check_model(self):
        results = []
        for stmt in self.statements:
            result = self.check_statement(stmt)
            results.append((stmt, result))
        return results

    def check_statement(self, stmt):
        if isinstance(stmt, Modification):
            return self.check_modification(stmt)
        else:
            return False

    def check_modification(self, stmt):
        # Identify the observable we're looking for in the model, which
        # may not exist!
        # The observable is the modified form of the substrate
        logger.info('Checking stmt: %s' % stmt)
        enz_mp = pa.get_monomer_pattern(self.model, stmt.enz)
        if type(stmt) in (Dephosphorylation,):
            unmodify_stmt = True
        else:
            unmodify_stmt = False
        modified_sub = _add_modification_to_agent(stmt.sub, 'phosphorylation',
                                                  stmt.residue, stmt.position)
        # Check if the site and corresponding state exists in the
        # model--if not, then it won't be observed. If we try to get a monomer
        # pattern for a monomer that doesn't have the specified site, PySB
        # throws a (generic) Exception; we catch this as an indication that
        # the site doesn't exist, and hence the modification is not observed,
        # and return False
        try:
            sub_mp = pa.get_monomer_pattern(self.model, modified_sub)
        except Exception as e:
            logger.info("Couldn't create monomer pattern: %s" % e)
            return False
        # Generate the influence map
        # Find rules in the model corresponding to the inputs and outputs
        input_rules = match_lhs(enz_mp, self.model.rules)
        logger.info('Found %s input rules matching %s' %
                    (len(input_rules), str(enz_mp)))
        # Get the production and consumption rules for the target
        target_prod_rules = find_production_rules(sub_mp, self.model.rules)
        target_cons_rules = find_consumption_rules(sub_mp, self.model.rules)
        # The final list contains a list of rules along with their "polarities"
        # (production = 1, consumption = -1)
        target_rules = (list(zip(target_prod_rules, [1]*len(target_prod_rules)))
                    + list(zip(target_cons_rules, [-1]*len(target_cons_rules))))
        logger.info('Found %s target rules matching %s' %
                    (len(target_rules), str(sub_mp)))
        if input_rules and target_rules:
            # Generate the influence map
            paths = []
            for input_rule, target_rule_info in \
                            itertools.product(input_rules, target_rules):
                (target_rule, rule_polarity) = target_rule_info
                try:
                    sp_gen = networkx.shortest_simple_paths(self.get_im(),
                                                       input_rule.name,
                                                       target_rule.name)
                    # Iterate over paths until we find one with positive
                    # polarity
                    for path_ix, sp in enumerate(sp_gen):
                        need_positive_path = False if unmodify_stmt else True
                        if need_positive_path == \
                                positive_path(self.get_im(), sp, rule_polarity):
                            logger.info('Found non-repeating path %s: %s' %
                                        (path_ix, str(sp)))
                            logger.info('Rule polarity of target: %s' %
                                        rule_polarity)
                            paths.append(sp)
                            break
                except networkx.NetworkXNoPath as nopath:
                    pass
            if paths:
                return True
            else:
                return False
        else:
            return False

def _add_modification_to_agent(agent, mod_type, residue, position):
    new_mod = ModCondition(mod_type, residue, position)
    # Check if this modification already exists
    for old_mod in agent.mods:
        if old_mod.equals(new_mod):
            return agent
    new_agent = deepcopy(agent)
    new_agent.mods.append(new_mod)
    return new_agent

def match_lhs(cp, rules):
    rule_matches = []
    for rule in rules:
        reactant_pattern = rule.rule_expression.reactant_pattern
        for rule_cp in reactant_pattern.complex_patterns:
            if cp_embeds_into(rule_cp, cp):
                rule_matches.append(rule)
                break
    return rule_matches

def match_rhs(cp, rules):
    rule_matches = []
    for rule in rules:
        product_pattern = rule.rule_expression.product_pattern
        for rule_cp in product_pattern.complex_patterns:
            if cp_embeds_into(rule_cp, cp):
                rule_matches.append(rule)
                break
    return rule_matches

def find_production_rules(cp, rules):
    # Find rules where the CP matches the left hand side
    lhs_rule_set = set(match_lhs(cp, rules))
    # Now find rules where the CP matches the right hand side
    rhs_rule_set = set(match_rhs(cp, rules))
    # Production rules are rules where there is a match on the right hand
    # side but not on the left hand side
    prod_rules = list(rhs_rule_set.difference(lhs_rule_set))
    return prod_rules

def find_consumption_rules(cp, rules):
    # Find rules where the CP matches the left hand side
    lhs_rule_set = set(match_lhs(cp, rules))
    # Now find rules where the CP matches the right hand side
    rhs_rule_set = set(match_rhs(cp, rules))
    # Consumption rules are rules where there is a match on the left hand
    # side but not on the right hand side
    cons_rules = list(lhs_rule_set.difference(rhs_rule_set))
    return cons_rules

def cp_embeds_into(cp1, cp2):
    # Check that any state in cp2 is matched in cp2
    # If the thing we're matching to is just a monomer pattern, that makes
    # things easier--we just need to find the corresponding monomer pattern
    # in cp1
    cp1 = as_complex_pattern(cp1)
    cp2 = as_complex_pattern(cp2)
    if len(cp2.monomer_patterns) == 1:
        mp2 = cp2.monomer_patterns[0]
        # Iterate over the monomer patterns in cp1 and see if there is one
        # that has the same name
        for mp1 in cp1.monomer_patterns:
            if mp_embeds_into(mp1, mp2):
                return True
    return False

def mp_embeds_into(mp1, mp2):
    sc_matches = []
    if mp1.monomer.name != mp2.monomer.name:
        return False
    # Check that all conditions in mp2 are met in mp1
    for site_name, site_state in mp2.site_conditions.items():
        if site_name not in mp1.site_conditions or \
           site_state != mp1.site_conditions[site_name]:
            return False
    return True

def positive_path(im, path, rule_polarity):
    # This doesn't address the effect of the rules themselves on the
    # observables of interest--just the effects of the rules on each other
    if rule_polarity not in (1, -1):
        raise ValueError('rule_polarity must be 1 or -1.')
    if len(path) == 1:
        return True if rule_polarity == 1 else False
    edge_polarities = []
    for rule_ix in range(len(path) - 1):
        from_rule = path[rule_ix]
        to_rule = path[rule_ix + 1]
        edge = im.get_edge(from_rule, to_rule)
        if edge.attr.get('color') is None:
            raise Exception('No color attribute for edge.')
        elif edge.attr['color'] == 'green':
            edge_polarities.append(1)
        elif edge.attr['color'] == 'red':
            edge_polarities.append(-1)
        else:
            raise Exception('Unexpected edge color: %s' % edge.attr['color'])
    # Compute and return the overall path polarity
    path_polarity = np.prod(edge_polarities) * rule_polarity
    assert path_polarity == 1 or path_polarity == -1
    return True if path_polarity == 1 else False

<commit_msg>Break after finding one path<commit_after>from __future__ import print_function, unicode_literals, absolute_import
from builtins import dict, str
import logging
import numpy as np
import networkx
import itertools
from copy import deepcopy
from pysb import kappa
from pysb import Observable
from pysb.core import as_complex_pattern
from indra.statements import *
from indra.assemblers import pysb_assembler as pa

logger = logging.getLogger('model_checker')

class ModelChecker(object):
    """Check a PySB model against a set of INDRA statements."""

    def __init__(self, model, stmts_to_check=None):
        self.model = model
        if stmts_to_check:
            self.statements = stmts_to_check
        else:
            self.statements = []
        self._im = None

    def get_im(self):
        if self._im:
            return self._im
        if not self.model:
            raise Exception("Cannot get influence map if there is no model.")
        else:
            self._im = kappa.influence_map(self.model)
            self._im.is_multigraph = lambda: False
            return self._im

    def check_model(self):
        results = []
        for stmt in self.statements:
            result = self.check_statement(stmt)
            results.append((stmt, result))
        return results

    def check_statement(self, stmt):
        if isinstance(stmt, Modification):
            return self.check_modification(stmt)
        else:
            return False

    def check_modification(self, stmt):
        # Identify the observable we're looking for in the model, which
        # may not exist!
        # The observable is the modified form of the substrate
        logger.info('Checking stmt: %s' % stmt)
        enz_mp = pa.get_monomer_pattern(self.model, stmt.enz)
        if type(stmt) in (Dephosphorylation,):
            unmodify_stmt = True
        else:
            unmodify_stmt = False
        modified_sub = _add_modification_to_agent(stmt.sub, 'phosphorylation',
                                                  stmt.residue, stmt.position)
        # Check if the site and corresponding state exists in the
        # model--if not, then it won't be observed. If we try to get a monomer
        # pattern for a monomer that doesn't have the specified site, PySB
        # throws a (generic) Exception; we catch this as an indication that
        # the site doesn't exist, and hence the modification is not observed,
        # and return False
        try:
            sub_mp = pa.get_monomer_pattern(self.model, modified_sub)
        except Exception as e:
            logger.info("Couldn't create monomer pattern: %s" % e)
            return False
        # Generate the influence map
        # Find rules in the model corresponding to the inputs and outputs
        input_rules = match_lhs(enz_mp, self.model.rules)
        logger.info('Found %s input rules matching %s' %
                    (len(input_rules), str(enz_mp)))
        # Get the production and consumption rules for the target
        target_prod_rules = find_production_rules(sub_mp, self.model.rules)
        target_cons_rules = find_consumption_rules(sub_mp, self.model.rules)
        # The final list contains a list of rules along with their "polarities"
        # (production = 1, consumption = -1)
        target_rules = (list(zip(target_prod_rules, [1]*len(target_prod_rules)))
                    + list(zip(target_cons_rules, [-1]*len(target_cons_rules))))
        logger.info('Found %s target rules matching %s' %
                    (len(target_rules), str(sub_mp)))
        if input_rules and target_rules:
            # Generate the influence map
            paths = []
            # Look for any rule that is both in the input and target set--
            # indicates that constraint is satisfied by a single rule
            for input_rule, target_rule_info in \
                            itertools.product(input_rules, target_rules):
                if paths:
                    break
                (target_rule, rule_polarity) = target_rule_info
                try:
                    sp_gen = networkx.shortest_simple_paths(self.get_im(),
                                                       input_rule.name,
                                                       target_rule.name)
                    # Iterate over paths until we find one with positive
                    # polarity
                    for path_ix, sp in enumerate(sp_gen):
                        need_positive_path = False if unmodify_stmt else True
                        if need_positive_path == \
                                positive_path(self.get_im(), sp, rule_polarity):
                            logger.info('Found non-repeating path %s: %s' %
                                        (path_ix, str(sp)))
                            logger.info('Rule polarity of target: %s' %
                                        rule_polarity)
                            paths.append(sp)
                            break
                except networkx.NetworkXNoPath as nopath:
                    pass
            if paths:
                return True
            else:
                return False
        else:
            return False

def _add_modification_to_agent(agent, mod_type, residue, position):
    new_mod = ModCondition(mod_type, residue, position)
    # Check if this modification already exists
    for old_mod in agent.mods:
        if old_mod.equals(new_mod):
            return agent
    new_agent = deepcopy(agent)
    new_agent.mods.append(new_mod)
    return new_agent

def match_lhs(cp, rules):
    rule_matches = []
    for rule in rules:
        reactant_pattern = rule.rule_expression.reactant_pattern
        for rule_cp in reactant_pattern.complex_patterns:
            if cp_embeds_into(rule_cp, cp):
                rule_matches.append(rule)
                break
    return rule_matches

def match_rhs(cp, rules):
    rule_matches = []
    for rule in rules:
        product_pattern = rule.rule_expression.product_pattern
        for rule_cp in product_pattern.complex_patterns:
            if cp_embeds_into(rule_cp, cp):
                rule_matches.append(rule)
                break
    return rule_matches

def find_production_rules(cp, rules):
    # Find rules where the CP matches the left hand side
    lhs_rule_set = set(match_lhs(cp, rules))
    # Now find rules where the CP matches the right hand side
    rhs_rule_set = set(match_rhs(cp, rules))
    # Production rules are rules where there is a match on the right hand
    # side but not on the left hand side
    prod_rules = list(rhs_rule_set.difference(lhs_rule_set))
    return prod_rules

def find_consumption_rules(cp, rules):
    # Find rules where the CP matches the left hand side
    lhs_rule_set = set(match_lhs(cp, rules))
    # Now find rules where the CP matches the right hand side
    rhs_rule_set = set(match_rhs(cp, rules))
    # Consumption rules are rules where there is a match on the left hand
    # side but not on the right hand side
    cons_rules = list(lhs_rule_set.difference(rhs_rule_set))
    return cons_rules

def cp_embeds_into(cp1, cp2):
    # Check that any state in cp2 is matched in cp2
    # If the thing we're matching to is just a monomer pattern, that makes
    # things easier--we just need to find the corresponding monomer pattern
    # in cp1
    cp1 = as_complex_pattern(cp1)
    cp2 = as_complex_pattern(cp2)
    if len(cp2.monomer_patterns) == 1:
        mp2 = cp2.monomer_patterns[0]
        # Iterate over the monomer patterns in cp1 and see if there is one
        # that has the same name
        for mp1 in cp1.monomer_patterns:
            if mp_embeds_into(mp1, mp2):
                return True
    return False

def mp_embeds_into(mp1, mp2):
    sc_matches = []
    if mp1.monomer.name != mp2.monomer.name:
        return False
    # Check that all conditions in mp2 are met in mp1
    for site_name, site_state in mp2.site_conditions.items():
        if site_name not in mp1.site_conditions or \
           site_state != mp1.site_conditions[site_name]:
            return False
    return True

def positive_path(im, path, rule_polarity):
    # This doesn't address the effect of the rules themselves on the
    # observables of interest--just the effects of the rules on each other
    if rule_polarity not in (1, -1):
        raise ValueError('rule_polarity must be 1 or -1.')
    if len(path) == 1:
        return True if rule_polarity == 1 else False
    edge_polarities = []
    for rule_ix in range(len(path) - 1):
        from_rule = path[rule_ix]
        to_rule = path[rule_ix + 1]
        edge = im.get_edge(from_rule, to_rule)
        if edge.attr.get('color') is None:
            raise Exception('No color attribute for edge.')
        elif edge.attr['color'] == 'green':
            edge_polarities.append(1)
        elif edge.attr['color'] == 'red':
            edge_polarities.append(-1)
        else:
            raise Exception('Unexpected edge color: %s' % edge.attr['color'])
    # Compute and return the overall path polarity
    path_polarity = np.prod(edge_polarities) * rule_polarity
    assert path_polarity == 1 or path_polarity == -1
    return True if path_polarity == 1 else False

<|endoftext|>


<commit_before>"""
QemuXML file ``/etc/libvirt/qemu/*.xml``
----------------------------------------
"""
from .. import XMLParser, parser
from insights.specs import Specs

@parser(Specs.qemu_xml)
class QemuXML(XMLParser):
    """This class parses xml files under ``/etc/libvirt/qemu/`` using
    ``XMLParser`` base parser.

    Sample file::

        <!--
        WARNING: THIS IS AN AUTO-GENERATED FILE. CHANGES TO IT ARE LIKELY TO BE
        OVERWRITTEN AND LOST. Changes to this xml configuration should be made using:
          virsh edit 05-s00c06h0
        or other application using the libvirt API.
        -->

        <domain type='kvm'>
          <name>05-s00c06h0</name>
          <uuid>02cf0bba-2bd6-11e7-8337-e4115b9a50d0</uuid>
          <memory unit='KiB'>12582912</memory>
          <currentMemory unit='KiB'>12582912</currentMemory>
          <vcpu placement='static'>4</vcpu>
          <cputune>
            <vcpupin vcpu='0' cpuset='1'/>
            <vcpupin vcpu='1' cpuset='2'/>
            <vcpupin vcpu='2' cpuset='3'/>
            <vcpupin vcpu='3' cpuset='4'/>
            <emulatorpin cpuset='1-4'/>
          </cputune>
          <numatune>
            <memory mode='strict' nodeset='0-1'/>
            <memnode cellid='0' mode='strict' nodeset='0'/>
            <memnode cellid='1' mode='strict' nodeset='1'/>
          </numatune>
          <os>
            <type arch='x86_64' machine='pc-i440fx-rhel7.0.0'>hvm</type>
            <boot dev='hd'/>
            <boot dev='network'/>
            <bootmenu enable='yes' timeout='1000'/>
            <bios useserial='yes' rebootTimeout='0'/>
          </os>
          <features>
            <acpi/>
            <apic/>
            <pae/>
          </features>
          <cpu>
            <numa>
              <cell id='0' cpus='0-1' memory='6291456' unit='KiB'/>
              <cell id='1' cpus='2-3' memory='6291456' unit='KiB'/>
            </numa>
          </cpu>
          <clock offset='utc'/>
          <on_poweroff>destroy</on_poweroff>
          <on_reboot>restart</on_reboot>
          <on_crash>restart</on_crash>
          <devices>
            <emulator>/usr/libexec/qemu-kvm</emulator>
            <disk type='file' device='disk'>
              <driver name='qemu' type='raw' cache='none' io='threads'/>
              <source file='/var/lib/libvirt/images/05-s00c06h0_1.img'/>
              <target dev='vda' bus='virtio'/>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x05' function='0x0'/>
            </disk>
            <controller type='usb' index='0'>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x01' function='0x2'/>
            </controller>
            <controller type='pci' index='0' model='pci-root'/>
            <controller type='virtio-serial' index='0'>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x03' function='0x0'/>
            </controller>
            <interface type='hostdev' managed='yes'>
              <mac address='b2:59:73:15:00:00'/>
              <source>
                <address type='pci' domain='0x0000' bus='0x04' slot='0x10' function='0x0'/>
              </source>
              <rom bar='on' file='/opt/vcp/share/ipxe/808610ed.rom'/>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x04' function='0x0'/>
            </interface>
            <interface type='hostdev' managed='yes'>
              <mac address='b2:59:73:15:00:01'/>
              <source>
                <address type='pci' domain='0x0000' bus='0x04' slot='0x10' function='0x1'/>
              </source>
              <rom bar='on' file='/opt/vcp/share/ipxe/808610ed.rom'/>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x08' function='0x0'/>
            </interface>
            <serial type='pty'>
              <target port='0'/>
            </serial>
            <console type='pty'>
              <target type='serial' port='0'/>
            </console>
            <channel type='pipe'>
              <source path='/var/lib/libvirt/qemu/channels/FROM-05-s00c06h0'/>
              <target type='virtio' name='virtio2host'/>
              <address type='virtio-serial' controller='0' bus='0' port='1'/>
            </channel>
            <channel type='pipe'>
              <source path='/var/lib/libvirt/qemu/channels/HGC-05-s00c06h0'/>
              <target type='virtio' name='virtio_host_guest_check'/>
              <address type='virtio-serial' controller='0' bus='0' port='2'/>
            </channel>
            <input type='mouse' bus='ps2'/>
            <input type='keyboard' bus='ps2'/>
            <graphics type='vnc' port='-1' autoport='yes'>
              <listen type='address'/>
            </graphics>
            <video>
              <model type='cirrus' vram='16384' heads='1' primary='yes'/>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x02' function='0x0'/>
            </video>
            <watchdog model='i6300esb' action='reset'>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x07' function='0x0'/>
            </watchdog>
            <memballoon model='virtio'>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x06' function='0x0'/>
            </memballoon>
          </devices>
        </domain>

    Examples:
        >>> xml_numa.file_name == 'vm.xml'
        True
        >>> xml_numa.vm_name == '05-s00c06h0'
        True
        >>> memnode = xml_numa.get_elements('./numatune/memnode', None)
        >>> len(memnode[0].items()) == 3
        True
        >>> len(memnode[1].items()) == 3
        True
        >>> memnode[0].get('cellid') == '0'
        True
        >>> memnode[1].get('mode') == 'strict'
        True
    """
    def parse_dom(self):
        if self.dom is None:
            return
        else:
            domain = {}
            for child in self.dom:
                if not child.getchildren():
                    domain[child.tag] = child.text
                else:
                    domain[child.tag] = [c.items() for c in child.getchildren()]

            return domain

    @property
    def vm_name(self):
        return self.data.get('name', None)
<commit_msg>Avoid to use `getchildren` method (#1593)<commit_after>"""
QemuXML file ``/etc/libvirt/qemu/*.xml``
----------------------------------------
"""
from .. import XMLParser, parser
from insights.specs import Specs

@parser(Specs.qemu_xml)
class QemuXML(XMLParser):
    """This class parses xml files under ``/etc/libvirt/qemu/`` using
    ``XMLParser`` base parser.

    Sample file::

        <!--
        WARNING: THIS IS AN AUTO-GENERATED FILE. CHANGES TO IT ARE LIKELY TO BE
        OVERWRITTEN AND LOST. Changes to this xml configuration should be made using:
          virsh edit 05-s00c06h0
        or other application using the libvirt API.
        -->

        <domain type='kvm'>
          <name>05-s00c06h0</name>
          <uuid>02cf0bba-2bd6-11e7-8337-e4115b9a50d0</uuid>
          <memory unit='KiB'>12582912</memory>
          <currentMemory unit='KiB'>12582912</currentMemory>
          <vcpu placement='static'>4</vcpu>
          <cputune>
            <vcpupin vcpu='0' cpuset='1'/>
            <vcpupin vcpu='1' cpuset='2'/>
            <vcpupin vcpu='2' cpuset='3'/>
            <vcpupin vcpu='3' cpuset='4'/>
            <emulatorpin cpuset='1-4'/>
          </cputune>
          <numatune>
            <memory mode='strict' nodeset='0-1'/>
            <memnode cellid='0' mode='strict' nodeset='0'/>
            <memnode cellid='1' mode='strict' nodeset='1'/>
          </numatune>
          <os>
            <type arch='x86_64' machine='pc-i440fx-rhel7.0.0'>hvm</type>
            <boot dev='hd'/>
            <boot dev='network'/>
            <bootmenu enable='yes' timeout='1000'/>
            <bios useserial='yes' rebootTimeout='0'/>
          </os>
          <features>
            <acpi/>
            <apic/>
            <pae/>
          </features>
          <cpu>
            <numa>
              <cell id='0' cpus='0-1' memory='6291456' unit='KiB'/>
              <cell id='1' cpus='2-3' memory='6291456' unit='KiB'/>
            </numa>
          </cpu>
          <clock offset='utc'/>
          <on_poweroff>destroy</on_poweroff>
          <on_reboot>restart</on_reboot>
          <on_crash>restart</on_crash>
          <devices>
            <emulator>/usr/libexec/qemu-kvm</emulator>
            <disk type='file' device='disk'>
              <driver name='qemu' type='raw' cache='none' io='threads'/>
              <source file='/var/lib/libvirt/images/05-s00c06h0_1.img'/>
              <target dev='vda' bus='virtio'/>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x05' function='0x0'/>
            </disk>
            <controller type='usb' index='0'>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x01' function='0x2'/>
            </controller>
            <controller type='pci' index='0' model='pci-root'/>
            <controller type='virtio-serial' index='0'>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x03' function='0x0'/>
            </controller>
            <interface type='hostdev' managed='yes'>
              <mac address='b2:59:73:15:00:00'/>
              <source>
                <address type='pci' domain='0x0000' bus='0x04' slot='0x10' function='0x0'/>
              </source>
              <rom bar='on' file='/opt/vcp/share/ipxe/808610ed.rom'/>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x04' function='0x0'/>
            </interface>
            <interface type='hostdev' managed='yes'>
              <mac address='b2:59:73:15:00:01'/>
              <source>
                <address type='pci' domain='0x0000' bus='0x04' slot='0x10' function='0x1'/>
              </source>
              <rom bar='on' file='/opt/vcp/share/ipxe/808610ed.rom'/>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x08' function='0x0'/>
            </interface>
            <serial type='pty'>
              <target port='0'/>
            </serial>
            <console type='pty'>
              <target type='serial' port='0'/>
            </console>
            <channel type='pipe'>
              <source path='/var/lib/libvirt/qemu/channels/FROM-05-s00c06h0'/>
              <target type='virtio' name='virtio2host'/>
              <address type='virtio-serial' controller='0' bus='0' port='1'/>
            </channel>
            <channel type='pipe'>
              <source path='/var/lib/libvirt/qemu/channels/HGC-05-s00c06h0'/>
              <target type='virtio' name='virtio_host_guest_check'/>
              <address type='virtio-serial' controller='0' bus='0' port='2'/>
            </channel>
            <input type='mouse' bus='ps2'/>
            <input type='keyboard' bus='ps2'/>
            <graphics type='vnc' port='-1' autoport='yes'>
              <listen type='address'/>
            </graphics>
            <video>
              <model type='cirrus' vram='16384' heads='1' primary='yes'/>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x02' function='0x0'/>
            </video>
            <watchdog model='i6300esb' action='reset'>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x07' function='0x0'/>
            </watchdog>
            <memballoon model='virtio'>
              <address type='pci' domain='0x0000' bus='0x00' slot='0x06' function='0x0'/>
            </memballoon>
          </devices>
        </domain>

    Examples:
        >>> xml_numa.file_name == 'vm.xml'
        True
        >>> xml_numa.vm_name == '05-s00c06h0'
        True
        >>> memnode = xml_numa.get_elements('./numatune/memnode', None)
        >>> len(memnode[0].items()) == 3
        True
        >>> len(memnode[1].items()) == 3
        True
        >>> memnode[0].get('cellid') == '0'
        True
        >>> memnode[1].get('mode') == 'strict'
        True
    """
    def parse_dom(self):
        if self.dom is None:
            return
        else:
            domain = {}
            for child in self.dom:
                if len(child) == 0:
                    domain[child.tag] = child.text
                else:
                    domain[child.tag] = [c.items() for c in child]

            return domain

    @property
    def vm_name(self):
        return self.data.get('name', None)
<|endoftext|>


<commit_before>import colour
import csv
import json
import os
import pprint

# write to file as json, csv, markdown, plaintext or print table

def write_data(data, user, format=None):
    if format is not None:

        directory = './data/'

        if not os.path.exists(directory):
            os.makedirs(directory)

        f = open(directory + user + '.' + format, 'w')

        if format == 'json':
            f.write(json.dumps(data, indent=4))
        elif format == 'csv':
            keys = data[0].keys()
            dw = csv.DictWriter(f, fieldnames=keys)
            dw.writeheader()
            dw.writerows(data)
        elif format == 'md':
            f.write('## %s - GitHub repositories\n' % user)
            for row in data:
                f.write(
                    '#### {}\n\n{}  \n_{}_, {} star(s)\n\n'.format(row['name'],
                                                                   row['desc'],
                                                                   row['lang'],
                                                                   row['stars']))
        elif format == 'txt':
            f.write('%s - GitHub repositories\n\n' % user)
            for row in data:
                f.write('{}\n{}\n{}, {} star(s)\n\n'.format(row['name'],
                                                            row['desc'],
                                                            row['lang'],
                                                            row['stars']))

        f.close()
<commit_msg>Print table if no file format provided<commit_after>import csv
import json
import os
from tabulate import tabulate

def write_data(d, u, f=None):
    if f is not None:

        directory = './data/'

        if not os.path.exists(directory):
            os.makedirs(directory)

        file = open(directory + u + '.' + f, 'w')

        if f == 'json':
            file.write(json.dumps(d, indent=4))
        elif f == 'csv':
            keys = d[0].keys()
            dw = csv.DictWriter(file, fieldnames=keys)
            dw.writeheader()
            dw.writerows(d)
        elif f == 'md':
            file.write('## %s - GitHub repositories\n' % u)
            for row in d:
                file.write(
                    '#### {}\n\n{}  \n_{}_, {} star(s)\n\n'.format(row['name'],
                                                                   row['desc'],
                                                                   row['lang'],
                                                                   row['stars']))
        elif f == 'txt':
            file.write('%s - GitHub repositories\n\n' % u)
            for row in d:
                file.write('{}\n{}\n{}, {} star(s)\n\n'.format(row['name'],
                                                               row['desc'],
                                                               row['lang'],
                                                               row['stars']))
        file.close()
    else:
        print(tabulate(d, headers="keys"))
<|endoftext|>


<commit_before>import os
import shutil
import sys
import traceback

# This file needs to remain Python 2.7 compatible
# due to the underlying cctbx installer logic

installer_path = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
libtbx_path = os.path.join(installer_path, "lib")
if libtbx_path not in sys.path:
    sys.path.append(libtbx_path)
from libtbx.auto_build import install_distribution

class installer(install_distribution.installer):
    organization = "dials"
    product_name = "DIALS"
    dest_dir_prefix = "dials"
    make_apps = []
    configure_modules = ["dials", "xia2", "iota", "prime"]
    include_gui_packages = True
    base_package_options = ["--dials"]
    installer_dir = installer_path
    modules = [
        # hot
        "annlib",
        "ccp4io",
        # base
        "cbflib",
        "cctbx_project",
        "gui_resources",
        "ccp4io_adaptbx",
        "annlib_adaptbx",
        # dials
        "dxtbx",
        "dials",
        "xia2",
        "iota",
        "prime",
    ]
    flags = list(install_distribution.installer.flags)
    try:
        flags.remove("create_versioned_dispatchers")
    except ValueError:
        pass

    def product_specific_preinstallation_hook(self):
        prefix = os.path.abspath(self.options.prefix)
        if prefix.startswith(installer_path):
            sys.exit(
                "Invalid installation option: --prefix={givenprefix}\n\n"
                "Please install DIALS to a location outside of the installer directory.\n"
                "Suggested alternative: --prefix={suggestedprefix}".format(
                    givenprefix=self.options.prefix,
                    suggestedprefix=os.path.dirname(prefix),
                )
            )

    def reconfigure(self, log=None, *args, **kwargs):
        """Intercept any errors and print log excerpt"""
        try:
            return super().reconfigure(log=log, *args, **kwargs)
        except Exception:
            if not self.options.verbose:
                print("\n" + " -=-" * 20)
                print("\nAn error occured during installation\n")
                print("Excerpt from installation log:")
                with open(log.name) as fh:
                    for line in fh.readlines()[-30:]:
                        print(" :", line, end="")
                print("\nThis led to ", end="")
                sys.stdout.flush()
            traceback.print_exc()
            print("\n")
            sys.exit(
                "Please report this installation error to dials-support@lists.sourceforge.net"
            )

    def product_specific_prepackage_hook(self, directory):
        """
        Remove irrelevant files from installer.
        """
        self.print_header("Deflating installer")

        suffixes = ["B", "KB", "MB", "GB", "TB", "PB"]

        def humansize(nbytes):
            if nbytes == 0:
                return "0 B"
            i = 0
            while nbytes >= 1024 and i < len(suffixes) - 1:
                nbytes /= 1024.0
                i += 1
            f = f"{nbytes:.2f}".rstrip("0").rstrip(".")
            return f"{f} {suffixes[i]}"

        self._cleaned_size, self._cleaned_files = 0, 0

        def rmdir(subdir):
            fullpath = os.path.join(directory, subdir)
            if not os.path.exists(fullpath):
                print("Skipping", " " * 26, subdir)
                return
            num_files, total_size = 0, 0
            for dirpath, dirnames, filenames in os.walk(fullpath):
                for f in filenames:
                    fp = os.path.join(dirpath, f)
                    total_size += os.path.getsize(fp)
                    num_files += 1
            print(
                "Removing %9s, %4d files from %s"
                % (humansize(total_size), num_files, subdir)
            )
            shutil.rmtree(fullpath)
            self._cleaned_size = self._cleaned_size + total_size
            self._cleaned_files = self._cleaned_files + num_files

        def rmext(subdir, extension):
            fullpath = os.path.join(directory, subdir)
            if not os.path.exists(fullpath):
                print("Skipping *%s" % extension, " " * 26, subdir)
                return
            filelist, total_size = [], 0
            for dirpath, dirnames, filenames in os.walk(fullpath):
                for f in filenames:
                    if f.endswith(extension):
                        fp = os.path.join(dirpath, f)
                        filelist.append(fp)
                        total_size += os.path.getsize(fp)
            print(
                "Removing %9s, %4d %s files from %s"
                % (humansize(total_size), len(filelist), extension, subdir)
            )
            for f in filelist:
                os.remove(f)
            self._cleaned_size = self._cleaned_size + total_size
            self._cleaned_files = self._cleaned_files + len(filelist)

        def rmfile(filename):
            fullpath = os.path.join(directory, filename)
            if not os.path.exists(fullpath):
                print("Skipping", " " * 26, filename)
                return
            filesize = os.path.getsize(fullpath)
            print(f"Removing {humansize(filesize):>9}, file {filename}")
            os.remove(fullpath)
            self._cleaned_size = self._cleaned_size + filesize
            self._cleaned_files = self._cleaned_files + 1

        # Deduce matplotlib path
        # (base/lib/python2.??/site-packages/matplotlib-????/matplotlib)
        # (base/Python.framework/Versions/?.?/lib/python?.?/site-packages/matplotlib-(...) on MacOS)
        try:
            import inspect

            import matplotlib

            matplotpath = os.path.dirname(
                os.path.dirname(inspect.getsourcefile(matplotlib))
            )
            relpath = []
            matplotpath, d = os.path.split(matplotpath)
            relpath.append(d)
            while d and (d != "base"):
                matplotpath, d = os.path.split(matplotpath)
                relpath.append(d)
            if d == "base":
                relpath.reverse()
                # delete matplotlib tests
                matplotpath = os.path.join(*relpath)
                rmdir(os.path.join(matplotpath, "matplotlib", "tests"))
                rmdir(os.path.join(matplotpath, "mpl_toolkits", "tests"))

                # ...while we're here
                sitepath = os.path.dirname(matplotpath)
                rmdir(os.path.join(sitepath, "numpy/core/tests"))
                rmdir(os.path.join(sitepath, "numpy/doc"))
                rmdir(os.path.join(sitepath, "numpy/distutils/tests"))
                rmdir(os.path.join(sitepath, "numpy/f2py/docs"))

                pythonpath = os.path.dirname(sitepath)
                rmdir(os.path.join(pythonpath, "test"))
        except Exception:
            print("Could not deduce python package paths")

        rmdir("base/man")
        rmdir("base/share/doc")
        rmdir("base/share/gtk-doc")
        rmdir("base/share/hdf5_examples")
        rmdir("base/share/man")
        rmdir("build/dials_data")
        rmdir("build/precommitbx")
        rmdir("build/regression_data")
        rmdir("build/xia2_regression")
        rmext("build", ".o")
        for f in ("setpaths", "setpaths_debug", "setpaths_all", "unsetpaths"):
            for ext in (".sh", ".csh"):
                rmfile(os.path.join("build", f + ext))
        for p in (
            "chrono",
            "date_time",
            "detail",
            "filesystem",
            "program_options",
            "python",
            "system",
            "thread",
            "timer",
        ):
            rmdir(os.path.join("modules/boost/libs", p, "example"))
            rmdir(os.path.join("modules/boost/libs", p, "doc"))
            rmdir(os.path.join("modules/boost/libs", p, "test"))
            rmdir(os.path.join("modules/boost/libs", p, "tutorial"))
        rmdir("modules/boost/libs/date_time/xmldoc")
        rmdir("modules/cbflib/doc")
        rmdir("modules/cbflib/examples")
        rmdir("modules/cbflib/ply-3.2/doc")
        rmdir("modules/cbflib/ply-3.2/example")
        rmdir("modules/cbflib/ply-3.2/test")
        rmext("modules/cbflib", ".cbf")
        rmdir("modules/clipper/examples")
        print("-" * 60)
        print(
            "Deleted %d files, decrufting installation by %s\n"
            % (self._cleaned_files, humansize(self._cleaned_size))
        )

if __name__ == "__main__":
    installer(sys.argv[1:]).install()
<commit_msg>Restore python 2.7-compatibility for dials_installer.py<commit_after>from __future__ import absolute_import, division, print_function

import os
import shutil
import sys
import traceback

# This file needs to remain Python 2.7 compatible
# due to the underlying cctbx installer logic

installer_path = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
libtbx_path = os.path.join(installer_path, "lib")
if libtbx_path not in sys.path:
    sys.path.append(libtbx_path)
from libtbx.auto_build import install_distribution

class installer(install_distribution.installer):
    organization = "dials"
    product_name = "DIALS"
    dest_dir_prefix = "dials"
    make_apps = []
    configure_modules = ["dials", "xia2", "iota", "prime"]
    include_gui_packages = True
    base_package_options = ["--dials"]
    installer_dir = installer_path
    modules = [
        # hot
        "annlib",
        "ccp4io",
        # base
        "cbflib",
        "cctbx_project",
        "gui_resources",
        "ccp4io_adaptbx",
        "annlib_adaptbx",
        # dials
        "dxtbx",
        "dials",
        "xia2",
        "iota",
        "prime",
    ]
    flags = list(install_distribution.installer.flags)
    try:
        flags.remove("create_versioned_dispatchers")
    except ValueError:
        pass

    def product_specific_preinstallation_hook(self):
        prefix = os.path.abspath(self.options.prefix)
        if prefix.startswith(installer_path):
            sys.exit(
                "Invalid installation option: --prefix={givenprefix}\n\n"
                "Please install DIALS to a location outside of the installer directory.\n"
                "Suggested alternative: --prefix={suggestedprefix}".format(
                    givenprefix=self.options.prefix,
                    suggestedprefix=os.path.dirname(prefix),
                )
            )

    def reconfigure(self, log=None, *args, **kwargs):
        """Intercept any errors and print log excerpt"""
        try:
            return super(installer, self).reconfigure(log=log, *args, **kwargs)
        except Exception:
            if not self.options.verbose:
                print("\n" + " -=-" * 20)
                print("\nAn error occured during installation\n")
                print("Excerpt from installation log:")
                with open(log.name, "r") as fh:
                    for line in fh.readlines()[-30:]:
                        print(" :", line, end="")
                print("\nThis led to ", end="")
                sys.stdout.flush()
            traceback.print_exc()
            print("\n")
            sys.exit(
                "Please report this installation error to dials-support@lists.sourceforge.net"
            )

    def product_specific_prepackage_hook(self, directory):
        """
        Remove irrelevant files from installer.
        """
        self.print_header("Deflating installer")

        suffixes = ["B", "KB", "MB", "GB", "TB", "PB"]

        def humansize(nbytes):
            if nbytes == 0:
                return "0 B"
            i = 0
            while nbytes >= 1024 and i < len(suffixes) - 1:
                nbytes /= 1024.0
                i += 1
            f = ("%.2f" % nbytes).rstrip("0").rstrip(".")
            return "%s %s" % (f, suffixes[i])

        self._cleaned_size, self._cleaned_files = 0, 0

        def rmdir(subdir):
            fullpath = os.path.join(directory, subdir)
            if not os.path.exists(fullpath):
                print("Skipping", " " * 26, subdir)
                return
            num_files, total_size = 0, 0
            for dirpath, dirnames, filenames in os.walk(fullpath):
                for f in filenames:
                    fp = os.path.join(dirpath, f)
                    total_size += os.path.getsize(fp)
                    num_files += 1
            print(
                "Removing %9s, %4d files from %s"
                % (humansize(total_size), num_files, subdir)
            )
            shutil.rmtree(fullpath)
            self._cleaned_size = self._cleaned_size + total_size
            self._cleaned_files = self._cleaned_files + num_files

        def rmext(subdir, extension):
            fullpath = os.path.join(directory, subdir)
            if not os.path.exists(fullpath):
                print("Skipping *%s" % extension, " " * 26, subdir)
                return
            filelist, total_size = [], 0
            for dirpath, dirnames, filenames in os.walk(fullpath):
                for f in filenames:
                    if f.endswith(extension):
                        fp = os.path.join(dirpath, f)
                        filelist.append(fp)
                        total_size += os.path.getsize(fp)
            print(
                "Removing %9s, %4d %s files from %s"
                % (humansize(total_size), len(filelist), extension, subdir)
            )
            for f in filelist:
                os.remove(f)
            self._cleaned_size = self._cleaned_size + total_size
            self._cleaned_files = self._cleaned_files + len(filelist)

        def rmfile(filename):
            fullpath = os.path.join(directory, filename)
            if not os.path.exists(fullpath):
                print("Skipping", " " * 26, filename)
                return
            filesize = os.path.getsize(fullpath)
            print("Removing %9s, file %s" % (humansize(filesize), filename))
            os.remove(fullpath)
            self._cleaned_size = self._cleaned_size + filesize
            self._cleaned_files = self._cleaned_files + 1

        # Deduce matplotlib path
        # (base/lib/python2.??/site-packages/matplotlib-????/matplotlib)
        # (base/Python.framework/Versions/?.?/lib/python?.?/site-packages/matplotlib-(...) on MacOS)
        try:
            import inspect

            import matplotlib

            matplotpath = os.path.dirname(
                os.path.dirname(inspect.getsourcefile(matplotlib))
            )
            relpath = []
            matplotpath, d = os.path.split(matplotpath)
            relpath.append(d)
            while d and (d != "base"):
                matplotpath, d = os.path.split(matplotpath)
                relpath.append(d)
            if d == "base":
                relpath.reverse()
                # delete matplotlib tests
                matplotpath = os.path.join(*relpath)
                rmdir(os.path.join(matplotpath, "matplotlib", "tests"))
                rmdir(os.path.join(matplotpath, "mpl_toolkits", "tests"))

                # ...while we're here
                sitepath = os.path.dirname(matplotpath)
                rmdir(os.path.join(sitepath, "numpy/core/tests"))
                rmdir(os.path.join(sitepath, "numpy/doc"))
                rmdir(os.path.join(sitepath, "numpy/distutils/tests"))
                rmdir(os.path.join(sitepath, "numpy/f2py/docs"))

                pythonpath = os.path.dirname(sitepath)
                rmdir(os.path.join(pythonpath, "test"))
        except Exception:
            print("Could not deduce python package paths")

        rmdir("base/man")
        rmdir("base/share/doc")
        rmdir("base/share/gtk-doc")
        rmdir("base/share/hdf5_examples")
        rmdir("base/share/man")
        rmdir("build/dials_data")
        rmdir("build/precommitbx")
        rmdir("build/regression_data")
        rmdir("build/xia2_regression")
        rmext("build", ".o")
        for f in ("setpaths", "setpaths_debug", "setpaths_all", "unsetpaths"):
            for ext in (".sh", ".csh"):
                rmfile(os.path.join("build", f + ext))
        for p in (
            "chrono",
            "date_time",
            "detail",
            "filesystem",
            "program_options",
            "python",
            "system",
            "thread",
            "timer",
        ):
            rmdir(os.path.join("modules/boost/libs", p, "example"))
            rmdir(os.path.join("modules/boost/libs", p, "doc"))
            rmdir(os.path.join("modules/boost/libs", p, "test"))
            rmdir(os.path.join("modules/boost/libs", p, "tutorial"))
        rmdir("modules/boost/libs/date_time/xmldoc")
        rmdir("modules/cbflib/doc")
        rmdir("modules/cbflib/examples")
        rmdir("modules/cbflib/ply-3.2/doc")
        rmdir("modules/cbflib/ply-3.2/example")
        rmdir("modules/cbflib/ply-3.2/test")
        rmext("modules/cbflib", ".cbf")
        rmdir("modules/clipper/examples")
        print("-" * 60)
        print(
            "Deleted %d files, decrufting installation by %s\n"
            % (self._cleaned_files, humansize(self._cleaned_size))
        )

if __name__ == "__main__":
    installer(sys.argv[1:]).install()
<|endoftext|>


<commit_before># encoding: utf-8
"""
request.py

Created by Thomas Mangin on 2012-02-27.
Copyright (c) 2012 Exa Networks. All rights reserved.
"""

class Request (object):
	def __init__ (self,request):
		self.raw = request
		method, self.uri, version = request.split()
		self.method = method.upper()

		version = version.split('/')[-1]
		if version not in ('1.0', '1.1') and '.' in version:
			major, minor = version.split('.', 1)
			if major.isdigit() and minor.isdigit():
				version = str(int(major)) + '.' + str(int(minor))

		self.version = version

	def parse (self):
		# protocol
		if '://'  in self.uri:
			protocol, uri = self.uri.split('://', 1)
			# we have :// in the path
			if '/' in protocol:
				self.protocol = 'http'
				uri = self.uri
			else:
				self.protocol = protocol
		else:
			self.protocol = 'http'
			uri = self.uri

		# ipv6 host
		if uri.startswith('[') and ']' in uri:
			self.host, remaining = uri[1:].split(']', 1)
			if '/' in remaining:
				port,path = remaining.split('/')
				self.path = '/' + path
				if port:
					if not port.startswith(':'):
						raise ValueError('Malformed headers, ipv6 address was followed by an invalid port')

					self.port = self._checkport(port[1:])
				else:
					self.port = '80'

				return self
			else:
				self.port = self._checkport(remaining)
				self.path = ''
				return self

		# split on path
		if '/' in uri:
			# we have a path
			host, path = uri.split('/', 1)
			self.path = '/' + path
			if ':' in host:
				self.host,port = host.split(':',1)
				self.port = self._checkport(port)
			else:
				self.host = host
				self.port = '80'
		else:
			self.path = '/'
			if ':' in uri:
				self.host,port = uri.split(':',1)
				self.port = self._checkport(port)
			else:
				self.host = uri
				self.port = '80'
		return self

	def _checkport (self,port):
		if port and not port.isdigit():
			raise ValueError('Malformed headers')
		return port

	def __str__ (self):
		if self.protocol == 'http':
			return self.method + ' ' + self.path + ' HTTP/' + self.version

		return self.raw

<commit_msg>Change order of string comparisons to perform the most likely first<commit_after># encoding: utf-8
"""
request.py

Created by Thomas Mangin on 2012-02-27.
Copyright (c) 2012 Exa Networks. All rights reserved.
"""

class Request (object):
	def __init__ (self,request):
		self.raw = request
		method, self.uri, version = request.split()
		self.method = method.upper()

		version = version.split('/')[-1]
		if version not in ('1.1', '1.0') and '.' in version:
			major, minor = version.split('.', 1)
			if major.isdigit() and minor.isdigit():
				version = str(int(major)) + '.' + str(int(minor))

		self.version = version

	def parse (self):
		# protocol
		if '://'  in self.uri:
			protocol, uri = self.uri.split('://', 1)
			# we have :// in the path
			if '/' in protocol:
				self.protocol = 'http'
				uri = self.uri
			else:
				self.protocol = protocol
		else:
			self.protocol = 'http'
			uri = self.uri

		# ipv6 host
		if uri.startswith('[') and ']' in uri:
			self.host, remaining = uri[1:].split(']', 1)
			if '/' in remaining:
				port,path = remaining.split('/')
				self.path = '/' + path
				if port:
					if not port.startswith(':'):
						raise ValueError('Malformed headers, ipv6 address was followed by an invalid port')

					self.port = self._checkport(port[1:])
				else:
					self.port = '80'

				return self
			else:
				self.port = self._checkport(remaining)
				self.path = ''
				return self

		# split on path
		if '/' in uri:
			# we have a path
			host, path = uri.split('/', 1)
			self.path = '/' + path
			if ':' in host:
				self.host,port = host.split(':',1)
				self.port = self._checkport(port)
			else:
				self.host = host
				self.port = '80'
		else:
			self.path = '/'
			if ':' in uri:
				self.host,port = uri.split(':',1)
				self.port = self._checkport(port)
			else:
				self.host = uri
				self.port = '80'
		return self

	def _checkport (self,port):
		if port and not port.isdigit():
			raise ValueError('Malformed headers')
		return port

	def __str__ (self):
		if self.protocol == 'http':
			return self.method + ' ' + self.path + ' HTTP/' + self.version

		return self.raw

<|endoftext|>


<commit_before>from ..core import Condition, Machine, Transition

import logging
import asyncio

_LOGGER = logging.getLogger(__name__)
_LOGGER.addHandler(logging.NullHandler())

class AsyncCondition(Condition):

    async def check(self, event_data):
        predicate = event_data.machine.resolve_callable(self.func, event_data)
        if asyncio.iscoroutinefunction(predicate):
            if event_data.machine.send_event:
                return await predicate(event_data) == self.target
            else:
                return await predicate(*event_data.args, **event_data.kwargs) == self.target
        else:
            print("SIMPLE RETURN")
            return super(AsyncCondition, self).check(event_data)

class AsyncTransition(Transition):

    condition_cls = AsyncCondition

    def _eval_conditions(self, event_data):
        loop = asyncio.get_event_loop()
        res = loop.run_until_complete(
            asyncio.gather(*[
                cond.check(event_data) for cond in self.conditions
            ])
        )
        if not all(res):
            _LOGGER.debug("%sTransition condition failed: Transition halted.", event_data.machine.name)
            return False
        return True

class AsyncMachine(Machine):

    transition_cls = AsyncTransition

    def callbacks(self, funcs, event_data):
        loop = asyncio.get_event_loop()
        loop.run_until_complete(
            asyncio.gather(*[
                event_data.machine.callback(func, event_data) for func in funcs
            ])
        )

    async def callback(self, func, event_data):
        func = self.resolve_callable(func, event_data)

        if self.send_event:
            if asyncio.iscoroutinefunction(func):
                await func(event_data)
            else:
                func(event_data)
        else:
            if asyncio.iscoroutinefunction(func):
                await func(*event_data.args, **event_data.kwargs)
            else:
                func(*event_data.args, **event_data.kwargs)
<commit_msg>remove debug print call<commit_after>from ..core import Condition, Machine, Transition

import logging
import asyncio

_LOGGER = logging.getLogger(__name__)
_LOGGER.addHandler(logging.NullHandler())

class AsyncCondition(Condition):

    async def check(self, event_data):
        predicate = event_data.machine.resolve_callable(self.func, event_data)
        if asyncio.iscoroutinefunction(predicate):
            if event_data.machine.send_event:
                return await predicate(event_data) == self.target
            else:
                return await predicate(*event_data.args, **event_data.kwargs) == self.target
        else:
            return super(AsyncCondition, self).check(event_data)

class AsyncTransition(Transition):

    condition_cls = AsyncCondition

    def _eval_conditions(self, event_data):
        loop = asyncio.get_event_loop()
        res = loop.run_until_complete(
            asyncio.gather(*[
                cond.check(event_data) for cond in self.conditions
            ])
        )
        if not all(res):
            _LOGGER.debug("%sTransition condition failed: Transition halted.", event_data.machine.name)
            return False
        return True

class AsyncMachine(Machine):

    transition_cls = AsyncTransition

    def callbacks(self, funcs, event_data):
        loop = asyncio.get_event_loop()
        loop.run_until_complete(
            asyncio.gather(*[
                event_data.machine.callback(func, event_data) for func in funcs
            ])
        )

    async def callback(self, func, event_data):
        func = self.resolve_callable(func, event_data)

        if self.send_event:
            if asyncio.iscoroutinefunction(func):
                await func(event_data)
            else:
                func(event_data)
        else:
            if asyncio.iscoroutinefunction(func):
                await func(*event_data.args, **event_data.kwargs)
            else:
                func(*event_data.args, **event_data.kwargs)
<|endoftext|>


<commit_before>"""
:mod:`disco.worker` -- Python Worker Interface
==============================================

        for record in task.input(open=...):
                pass

    :type  input: list of inputs or list of list of inputs
    :param input: Each input must be specified in one of the following ways:

                   * ``http://www.example.com/data`` - any HTTP address
                   * ``disco://cnode03/bigtxt/file_name`` - Disco address. Refers to ``cnode03:/var/disco/bigtxt/file_name``. Currently this is an alias for ``http://cnode03:[DISCO_PORT]/bigtxt/file_name``.
                   * ``dir://cnode03/jobname/`` - Result directory. This format is used by Disco internally.
                   * ``/home/bob/bigfile.txt`` - a local file. Note that the file must either exist on all the nodes or you must make sure that the job is run only on the nodes where the file exists. Due to these restrictions, this form has only limited use.
                   * ``raw://some_string`` - pseudo-address; instead of fetching data from a remote source, use ``some_string`` in the address as data. Useful for specifying dummy inputs for generator maps.
                   * ``tag://tagname`` - a tag stored in :ref:`DDFS` (*Added in version 0.3*)

                  (*Added in version 0.3.2*)
                  Tags can be token protected.
                  For the data in tags to be used as job inputs,
                  the tags should be resolved into the constituent urls or replica sets,
                  and provided as the value of the input parameter.

                  (*Added in version 0.2.2*):
                  An input entry can be a list of inputs:
                  This lets you specify redundant versions of an input file.
                  If a list of redundant inputs is specified,
                  the scheduler chooses the input that is located on the node
                  with the lowest load at the time of scheduling.
                  Redundant inputs are tried one by one until the task succeeds.
                  Redundant inputs require that the *map* function is specified.
"""
import cPickle, os, sys, traceback

class Worker(dict):
    """
    :type  required_files: list of paths or dict
    :param required_files: additional files that are required by the worker.
                           Either a list of paths to files to include,
                           or a dictionary which contains items of the form
                           ``(filename, filecontents)``.
                           (*Added in version 0.2.3*)

    :type  required_modules: list of modules or module names
    :param required_modules: required modules to send to the worker.
                             Can also be a list of module objects.
                             (*Changed in version 0.4*):

    :type  scheduler: dict
    :param scheduler: options for the job scheduler.
                      The following keys are supported:

                       * *max_cores* - use this many cores at most
                                       (applies to both map and reduce).

                                       Default is ``2**31``.

                       * *force_local* - always run task on the node where
                                         input data is located;
                                         never use HTTP to access data remotely.

                       * *force_remote* - never run task on the node where input
                                          data is located;
                                          always use HTTP to access data remotely.

                      (*Added in version 0.2.4*)
    """
    def __init__(self, **kwargs):
        super(Worker, self).__init__(self.defaults())
        self.update(kwargs)

    @property
    def bin(self):
        return os.path.join('lib', '%s.py' % self.__module__.replace('.', '/'))

    def defaults(self):
        return {'map': None,
                'merge_partitions': False, # XXX: maybe deprecated
                'reduce': None,
                'required_files': {},
                'required_modules': None,
                'save': False,
                'scheduler': {},
                'partitions': 1,  # move to classic once partitions are dynamic
                'profile': False}

    def getitem(self, key, job, **jobargs):
        if key in jobargs:
            return jobargs[key]
        elif hasattr(job, key):
            return getattr(job, key)
        return self.get(key)

    def jobdict(self, job, **jobargs):
        from disco.util import inputlist, ispartitioned, read_index
        def get(key):
            return self.getitem(key, job, **jobargs)
        has_map = bool(get('map'))
        has_reduce = bool(get('reduce'))
        input = inputlist(get('input') or (),
                          partition=None if has_map else False,
                          settings=job.settings)

        # -- nr_reduces --
        # ignored if there is not actually a reduce specified
        # XXX: master should always handle this
        if has_map:
            # partitioned map has N reduces; non-partitioned map has 1 reduce
            nr_reduces = get('partitions') or 1
        elif ispartitioned(input):
            # no map, with partitions: len(dir://) specifies nr_reduces
            nr_reduces = 1 + max(int(id)
                                 for dir in input
                                 for id, url in read_index(dir))
        else:
            # no map, without partitions can only have 1 reduce
            nr_reduces = 1

        if get('merge_partitions'):
            nr_reduces = 1

        return {'input': input,
                'worker': self.bin,
                'map?': has_map,
                'reduce?': has_reduce,
                'profile?': get('profile'),
                'nr_reduces': nr_reduces,
                'prefix': job.name,
                'scheduler': get('scheduler'),
                'owner': job.settings['DISCO_JOB_OWNER']}

    def jobenvs(self, job, **jobargs):
        settings = job.settings
        settings['LC_ALL'] = 'C'
        settings['LD_LIBRARY_PATH'] = 'lib'
        settings['PYTHONPATH'] = ':'.join((settings.get('PYTHONPATH', ''), 'lib'))
        return settings

    def jobhome(self, job, **jobargs):
        jobzip = self.jobzip(job, **jobargs)
        jobzip.close()
        return jobzip.dumps()

    def jobzip(self, job, **jobargs):
        from disco import __file__ as discopath
        from disco.fileutils import DiscoZipFile
        from disco.util import iskv
        def get(key):
            return self.getitem(key, job, **jobargs)
        jobzip = DiscoZipFile()
        jobzip.writepy(os.path.dirname(discopath), 'lib')
        jobzip.writemodule(job.__module__)
        jobzip.writemodule(self.__module__)
        if isinstance(get('required_files'), dict):
            for path, bytes in get('required_files').iteritems():
                    jobzip.writebytes(path, bytes)
        else:
            for path in get('required_files'):
                jobzip.writepath(path)
        for mod in get('required_modules') or ():
            jobzip.writemodule((mod[0] if iskv(mod) else mod), 'lib')
        return jobzip

    def jobdata(self, job, **jobargs):
        return cPickle.dumps((self, job, jobargs), -1)

    def start(self, job, task):
        from disco.sysutil import set_mem_limit
        set_mem_limit(job.settings['DISCO_WORKER_MAX_MEM'])
        if self['profile']:
            from cProfile import runctx
            name = 'profile-%s' % task.uid
            path = task.path(name)
            runctx('self.run(task)', globals(), locals(), path)
            task.put(name, open(path).read())
        else:
            self.run(task)
        self.end(task)

    def run(self, task):
        self[task.mode](task)

    def end(self, task):
        from disco.events import Status
        if not self['save'] or (task.mode == 'map' and self['reduce']):
            task.send()
            Status("Results sent to master").send()
        else:
            task.save()
            Status("Results saved to DDFS").send()

    @classmethod
    def unpack(cls, jobpack):
        try:
            from imp import find_module, load_module
            __disco__ = load_module('__disco__', *find_module('__main__', ['lib']))
            sys.modules['__main__'].__dict__.update(__disco__.__dict__)
        except ImportError:
            pass
        worker, job, jobargs = cPickle.loads(jobpack.jobdata)
        for key in worker:
            worker[key] = worker.getitem(key, job, **jobargs)
        return worker, job

    @classmethod
    def main(cls):
        from disco.error import DataError
        from disco.events import AnnouncePID, WorkerDone, DataUnavailable, TaskFailed
        from disco.job import JobPack
        from disco.task import Task
        from disco.fileutils import NonBlockingInput
        from disco.util import MessageWriter
        try:
            sys.stdin = NonBlockingInput(sys.stdin, timeout=600)
            sys.stdout = MessageWriter()
            AnnouncePID(str(os.getpid())).send()
            worker, job = cls.unpack(JobPack.request())
            worker.start(job, Task.request())
            WorkerDone().send()
        except (DataError, EnvironmentError, MemoryError), e:
            # check the number of open file descriptors (under proc), warn if close to max
            # http://stackoverflow.com/questions/899038/getting-the-highest-allocated-file-descriptor
            # also check for other known reasons for error, such as if disk is full
            DataUnavailable(traceback.format_exc()).send()
            raise
        except Exception, e:
            TaskFailed(MessageWriter.force_utf8(traceback.format_exc())).send()
            raise
<commit_msg>bugfix: force jobenvs to strings and only include those which have defaults<commit_after>"""
:mod:`disco.worker` -- Python Worker Interface
==============================================

        for record in task.input(open=...):
                pass

    :type  input: list of inputs or list of list of inputs
    :param input: Each input must be specified in one of the following ways:

                   * ``http://www.example.com/data`` - any HTTP address
                   * ``disco://cnode03/bigtxt/file_name`` - Disco address. Refers to ``cnode03:/var/disco/bigtxt/file_name``. Currently this is an alias for ``http://cnode03:[DISCO_PORT]/bigtxt/file_name``.
                   * ``dir://cnode03/jobname/`` - Result directory. This format is used by Disco internally.
                   * ``/home/bob/bigfile.txt`` - a local file. Note that the file must either exist on all the nodes or you must make sure that the job is run only on the nodes where the file exists. Due to these restrictions, this form has only limited use.
                   * ``raw://some_string`` - pseudo-address; instead of fetching data from a remote source, use ``some_string`` in the address as data. Useful for specifying dummy inputs for generator maps.
                   * ``tag://tagname`` - a tag stored in :ref:`DDFS` (*Added in version 0.3*)

                  (*Added in version 0.3.2*)
                  Tags can be token protected.
                  For the data in tags to be used as job inputs,
                  the tags should be resolved into the constituent urls or replica sets,
                  and provided as the value of the input parameter.

                  (*Added in version 0.2.2*):
                  An input entry can be a list of inputs:
                  This lets you specify redundant versions of an input file.
                  If a list of redundant inputs is specified,
                  the scheduler chooses the input that is located on the node
                  with the lowest load at the time of scheduling.
                  Redundant inputs are tried one by one until the task succeeds.
                  Redundant inputs require that the *map* function is specified.
"""
import cPickle, os, sys, traceback

class Worker(dict):
    """
    :type  required_files: list of paths or dict
    :param required_files: additional files that are required by the worker.
                           Either a list of paths to files to include,
                           or a dictionary which contains items of the form
                           ``(filename, filecontents)``.
                           (*Added in version 0.2.3*)

    :type  required_modules: list of modules or module names
    :param required_modules: required modules to send to the worker.
                             Can also be a list of module objects.
                             (*Changed in version 0.4*):

    :type  scheduler: dict
    :param scheduler: options for the job scheduler.
                      The following keys are supported:

                       * *max_cores* - use this many cores at most
                                       (applies to both map and reduce).

                                       Default is ``2**31``.

                       * *force_local* - always run task on the node where
                                         input data is located;
                                         never use HTTP to access data remotely.

                       * *force_remote* - never run task on the node where input
                                          data is located;
                                          always use HTTP to access data remotely.

                      (*Added in version 0.2.4*)
    """
    def __init__(self, **kwargs):
        super(Worker, self).__init__(self.defaults())
        self.update(kwargs)

    @property
    def bin(self):
        return os.path.join('lib', '%s.py' % self.__module__.replace('.', '/'))

    def defaults(self):
        return {'map': None,
                'merge_partitions': False, # XXX: maybe deprecated
                'reduce': None,
                'required_files': {},
                'required_modules': None,
                'save': False,
                'scheduler': {},
                'partitions': 1,  # move to classic once partitions are dynamic
                'profile': False}

    def getitem(self, key, job, **jobargs):
        if key in jobargs:
            return jobargs[key]
        elif hasattr(job, key):
            return getattr(job, key)
        return self.get(key)

    def jobdict(self, job, **jobargs):
        from disco.util import inputlist, ispartitioned, read_index
        def get(key):
            return self.getitem(key, job, **jobargs)
        has_map = bool(get('map'))
        has_reduce = bool(get('reduce'))
        input = inputlist(get('input') or (),
                          partition=None if has_map else False,
                          settings=job.settings)

        # -- nr_reduces --
        # ignored if there is not actually a reduce specified
        # XXX: master should always handle this
        if has_map:
            # partitioned map has N reduces; non-partitioned map has 1 reduce
            nr_reduces = get('partitions') or 1
        elif ispartitioned(input):
            # no map, with partitions: len(dir://) specifies nr_reduces
            nr_reduces = 1 + max(int(id)
                                 for dir in input
                                 for id, url in read_index(dir))
        else:
            # no map, without partitions can only have 1 reduce
            nr_reduces = 1

        if get('merge_partitions'):
            nr_reduces = 1

        return {'input': input,
                'worker': self.bin,
                'map?': has_map,
                'reduce?': has_reduce,
                'profile?': get('profile'),
                'nr_reduces': nr_reduces,
                'prefix': job.name,
                'scheduler': get('scheduler'),
                'owner': job.settings['DISCO_JOB_OWNER']}

    def jobenvs(self, job, **jobargs):
        settings = job.settings
        settings['LC_ALL'] = 'C'
        settings['LD_LIBRARY_PATH'] = 'lib'
        settings['PYTHONPATH'] = ':'.join((settings.get('PYTHONPATH', ''), 'lib'))
        return dict((k, str(settings[k])) for k in settings.defaults)

    def jobhome(self, job, **jobargs):
        jobzip = self.jobzip(job, **jobargs)
        jobzip.close()
        return jobzip.dumps()

    def jobzip(self, job, **jobargs):
        from disco import __file__ as discopath
        from disco.fileutils import DiscoZipFile
        from disco.util import iskv
        def get(key):
            return self.getitem(key, job, **jobargs)
        jobzip = DiscoZipFile()
        jobzip.writepy(os.path.dirname(discopath), 'lib')
        jobzip.writemodule(job.__module__)
        jobzip.writemodule(self.__module__)
        if isinstance(get('required_files'), dict):
            for path, bytes in get('required_files').iteritems():
                    jobzip.writebytes(path, bytes)
        else:
            for path in get('required_files'):
                jobzip.writepath(path)
        for mod in get('required_modules') or ():
            jobzip.writemodule((mod[0] if iskv(mod) else mod), 'lib')
        return jobzip

    def jobdata(self, job, **jobargs):
        return cPickle.dumps((self, job, jobargs), -1)

    def start(self, job, task):
        from disco.sysutil import set_mem_limit
        set_mem_limit(job.settings['DISCO_WORKER_MAX_MEM'])
        if self['profile']:
            from cProfile import runctx
            name = 'profile-%s' % task.uid
            path = task.path(name)
            runctx('self.run(task)', globals(), locals(), path)
            task.put(name, open(path).read())
        else:
            self.run(task)
        self.end(task)

    def run(self, task):
        self[task.mode](task)

    def end(self, task):
        from disco.events import Status
        if not self['save'] or (task.mode == 'map' and self['reduce']):
            task.send()
            Status("Results sent to master").send()
        else:
            task.save()
            Status("Results saved to DDFS").send()

    @classmethod
    def unpack(cls, jobpack):
        try:
            from imp import find_module, load_module
            __disco__ = load_module('__disco__', *find_module('__main__', ['lib']))
            sys.modules['__main__'].__dict__.update(__disco__.__dict__)
        except ImportError:
            pass
        worker, job, jobargs = cPickle.loads(jobpack.jobdata)
        for key in worker:
            worker[key] = worker.getitem(key, job, **jobargs)
        return worker, job

    @classmethod
    def main(cls):
        from disco.error import DataError
        from disco.events import AnnouncePID, WorkerDone, DataUnavailable, TaskFailed
        from disco.job import JobPack
        from disco.task import Task
        from disco.fileutils import NonBlockingInput
        from disco.util import MessageWriter
        try:
            sys.stdin = NonBlockingInput(sys.stdin, timeout=600)
            sys.stdout = MessageWriter()
            AnnouncePID(str(os.getpid())).send()
            worker, job = cls.unpack(JobPack.request())
            worker.start(job, Task.request())
            WorkerDone().send()
        except (DataError, EnvironmentError, MemoryError), e:
            # check the number of open file descriptors (under proc), warn if close to max
            # http://stackoverflow.com/questions/899038/getting-the-highest-allocated-file-descriptor
            # also check for other known reasons for error, such as if disk is full
            DataUnavailable(traceback.format_exc()).send()
            raise
        except Exception, e:
            TaskFailed(MessageWriter.force_utf8(traceback.format_exc())).send()
            raise
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# This software may be modified and distributed under the terms
# of the MIT license.  See the LICENSE file for details.

import socket
import ssl
import sys

import pylogbeat

from logstash_async.utils import ichunked

class TimeoutNotSet(object):
    pass

class UdpTransport(object):

    _keep_connection = False

    # ----------------------------------------------------------------------
    # pylint: disable=unused-argument
    def __init__(self, host, port, timeout=TimeoutNotSet, **kwargs):
        self._host = host
        self._port = port
        self._timeout = timeout
        self._sock = None

    # ----------------------------------------------------------------------
    def send(self, events, use_logging=False):  # pylint: disable=unused-argument
        # Ideally we would keep the socket open but this is risky because we might not notice
        # a broken TCP connection and send events into the dark.
        # On UDP we push into the dark by design :)
        self._create_socket()
        try:
            self._send(events)
        finally:
            self._close()

    # ----------------------------------------------------------------------
    def _create_socket(self):
        if self._sock is not None:
            return

        # from logging.handlers.DatagramHandler
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        if self._timeout is not TimeoutNotSet:
            self._sock.settimeout(self._timeout)

    # ----------------------------------------------------------------------
    def _send(self, events):
        for event in events:
            self._send_via_socket(event)

    # ----------------------------------------------------------------------
    def _send_via_socket(self, data):
        data_to_send = self._convert_data_to_send(data)
        self._sock.sendto(data_to_send, (self._host, self._port))

    # ----------------------------------------------------------------------
    def _convert_data_to_send(self, data):
        if sys.version_info < (3, 0):
            return data
        # Python3
        elif not isinstance(data, bytes):
            return bytes(data, 'utf-8')

        return data

    # ----------------------------------------------------------------------
    def _close(self, force=False):
        if not self._keep_connection or force:
            if self._sock:
                self._sock.close()
                self._sock = None

    # ----------------------------------------------------------------------
    def close(self):
        self._close(force=True)

class TcpTransport(UdpTransport):

    # ----------------------------------------------------------------------
    def __init__(  # pylint: disable=too-many-arguments
            self,
            host,
            port,
            ssl_enable,
            ssl_verify,
            keyfile,
            certfile,
            ca_certs,
            timeout=TimeoutNotSet):
        super(TcpTransport, self).__init__(host, port)
        self._ssl_enable = ssl_enable
        self._ssl_verify = ssl_verify
        self._keyfile = keyfile
        self._certfile = certfile
        self._ca_certs = ca_certs
        self._timeout = timeout

    # ----------------------------------------------------------------------
    def _create_socket(self):
        if self._sock is not None:
            return

        # from logging.handlers.SocketHandler
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if self._timeout is not TimeoutNotSet:
            self._sock.settimeout(self._timeout)

        try:
            self._sock.connect((self._host, self._port))
            # non-SSL
            if not self._ssl_enable:
                return
            # SSL
            cert_reqs = ssl.CERT_REQUIRED
            if not self._ssl_verify:
                if self._ca_certs:
                    cert_reqs = ssl.CERT_OPTIONAL
                else:
                    cert_reqs = ssl.CERT_NONE
            self._sock = ssl.wrap_socket(
                self._sock,
                keyfile=self._keyfile,
                certfile=self._certfile,
                ca_certs=self._ca_certs,
                cert_reqs=cert_reqs)
        except socket.error:
            self._close()
            raise

    # ----------------------------------------------------------------------
    def _send_via_socket(self, data):
        data_to_send = self._convert_data_to_send(data)
        self._sock.sendall(data_to_send)

class BeatsTransport(object):

    _batch_size = 10

    # ----------------------------------------------------------------------
    def __init__(  # pylint: disable=too-many-arguments
            self,
            host,
            port,
            ssl_enable,
            ssl_verify,
            keyfile,
            certfile,
            ca_certs,
            timeout=TimeoutNotSet):
        timeout_ = None if timeout is not TimeoutNotSet else timeout
        self._client_arguments = dict(
            host=host,
            port=port,
            timeout=timeout_,
            ssl_enable=ssl_enable,
            ssl_verify=ssl_verify,
            keyfile=keyfile,
            certfile=certfile,
            ca_certs=ca_certs)

    # ----------------------------------------------------------------------
    def close(self):
        pass  # nothing to do

    # ----------------------------------------------------------------------
    def send(self, events, use_logging=False):
        client = pylogbeat.PyLogBeatClient(use_logging=use_logging, **self._client_arguments)
        with client:
            for events_subset in ichunked(events, self._batch_size):
                client.send(events_subset)
<commit_msg>Fix socket timeout setting ignored for filebeat (#50)<commit_after># -*- coding: utf-8 -*-
#
# This software may be modified and distributed under the terms
# of the MIT license.  See the LICENSE file for details.

import socket
import ssl
import sys

import pylogbeat

from logstash_async.utils import ichunked

class TimeoutNotSet(object):
    pass

class UdpTransport(object):

    _keep_connection = False

    # ----------------------------------------------------------------------
    # pylint: disable=unused-argument
    def __init__(self, host, port, timeout=TimeoutNotSet, **kwargs):
        self._host = host
        self._port = port
        self._timeout = timeout
        self._sock = None

    # ----------------------------------------------------------------------
    def send(self, events, use_logging=False):  # pylint: disable=unused-argument
        # Ideally we would keep the socket open but this is risky because we might not notice
        # a broken TCP connection and send events into the dark.
        # On UDP we push into the dark by design :)
        self._create_socket()
        try:
            self._send(events)
        finally:
            self._close()

    # ----------------------------------------------------------------------
    def _create_socket(self):
        if self._sock is not None:
            return

        # from logging.handlers.DatagramHandler
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        if self._timeout is not TimeoutNotSet:
            self._sock.settimeout(self._timeout)

    # ----------------------------------------------------------------------
    def _send(self, events):
        for event in events:
            self._send_via_socket(event)

    # ----------------------------------------------------------------------
    def _send_via_socket(self, data):
        data_to_send = self._convert_data_to_send(data)
        self._sock.sendto(data_to_send, (self._host, self._port))

    # ----------------------------------------------------------------------
    def _convert_data_to_send(self, data):
        if sys.version_info < (3, 0):
            return data
        # Python3
        elif not isinstance(data, bytes):
            return bytes(data, 'utf-8')

        return data

    # ----------------------------------------------------------------------
    def _close(self, force=False):
        if not self._keep_connection or force:
            if self._sock:
                self._sock.close()
                self._sock = None

    # ----------------------------------------------------------------------
    def close(self):
        self._close(force=True)

class TcpTransport(UdpTransport):

    # ----------------------------------------------------------------------
    def __init__(  # pylint: disable=too-many-arguments
            self,
            host,
            port,
            ssl_enable,
            ssl_verify,
            keyfile,
            certfile,
            ca_certs,
            timeout=TimeoutNotSet):
        super(TcpTransport, self).__init__(host, port)
        self._ssl_enable = ssl_enable
        self._ssl_verify = ssl_verify
        self._keyfile = keyfile
        self._certfile = certfile
        self._ca_certs = ca_certs
        self._timeout = timeout

    # ----------------------------------------------------------------------
    def _create_socket(self):
        if self._sock is not None:
            return

        # from logging.handlers.SocketHandler
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if self._timeout is not TimeoutNotSet:
            self._sock.settimeout(self._timeout)

        try:
            self._sock.connect((self._host, self._port))
            # non-SSL
            if not self._ssl_enable:
                return
            # SSL
            cert_reqs = ssl.CERT_REQUIRED
            if not self._ssl_verify:
                if self._ca_certs:
                    cert_reqs = ssl.CERT_OPTIONAL
                else:
                    cert_reqs = ssl.CERT_NONE
            self._sock = ssl.wrap_socket(
                self._sock,
                keyfile=self._keyfile,
                certfile=self._certfile,
                ca_certs=self._ca_certs,
                cert_reqs=cert_reqs)
        except socket.error:
            self._close()
            raise

    # ----------------------------------------------------------------------
    def _send_via_socket(self, data):
        data_to_send = self._convert_data_to_send(data)
        self._sock.sendall(data_to_send)

class BeatsTransport(object):

    _batch_size = 10

    # ----------------------------------------------------------------------
    def __init__(  # pylint: disable=too-many-arguments
            self,
            host,
            port,
            ssl_enable,
            ssl_verify,
            keyfile,
            certfile,
            ca_certs,
            timeout=TimeoutNotSet):
        timeout_ = None if timeout is TimeoutNotSet else timeout
        self._client_arguments = dict(
            host=host,
            port=port,
            timeout=timeout_,
            ssl_enable=ssl_enable,
            ssl_verify=ssl_verify,
            keyfile=keyfile,
            certfile=certfile,
            ca_certs=ca_certs)

    # ----------------------------------------------------------------------
    def close(self):
        pass  # nothing to do

    # ----------------------------------------------------------------------
    def send(self, events, use_logging=False):
        client = pylogbeat.PyLogBeatClient(use_logging=use_logging, **self._client_arguments)
        with client:
            for events_subset in ichunked(events, self._batch_size):
                client.send(events_subset)
<|endoftext|>