<commit_before># Copyright (c) 2006-2011 Mitch Garnaat http://garnaat.org/
# Copyright (c) 2011, Eucalyptus Systems, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish, dis-
# tribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the fol-
# lowing conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.

"""
Represents an EC2 Security Group
"""
from boto.ec2.ec2object import TaggedEC2Object
from boto.exception import BotoClientError

class SecurityGroup(TaggedEC2Object):

    def __init__(self, connection=None, owner_id=None,
                 name=None, description=None, id=None):
        TaggedEC2Object.__init__(self, connection)
        self.id = id
        self.owner_id = owner_id
        self.name = name
        self.description = description
        self.vpc_id = None
        self.rules = IPPermissionsList()
        self.rules_egress = IPPermissionsList()

    def __repr__(self):
        return 'SecurityGroup:%s' % self.name

    def startElement(self, name, attrs, connection):
        retval = TaggedEC2Object.startElement(self, name, attrs, connection)
        if retval is not None:
            return retval
        if name == 'ipPermissions':
            return self.rules
        elif name == 'ipPermissionsEgress':
            return self.rules_egress
        else:
            return None

    def endElement(self, name, value, connection):
        if name == 'ownerId':
            self.owner_id = value
        elif name == 'groupId':
            self.id = value
        elif name == 'groupName':
            self.name = value
        elif name == 'vpcId':
            self.vpc_id = value
        elif name == 'groupDescription':
            self.description = value
        elif name == 'ipRanges':
            pass
        elif name == 'return':
            if value == 'false':
                self.status = False
            elif value == 'true':
                self.status = True
            else:
                raise Exception(
                    'Unexpected value of status %s for group %s'%(
                        value,
                        self.name
                    )
                )
        else:
            setattr(self, name, value)

    def delete(self):
        return self.connection.delete_security_group(group_id=self.id)

    def add_rule(self, ip_protocol, from_port, to_port,
                 src_group_name, src_group_owner_id, cidr_ip, src_group_group_id):
        """
        Add a rule to the SecurityGroup object.  Note that this method
        only changes the local version of the object.  No information
        is sent to EC2.
        """
        rule = IPPermissions(self)
        rule.ip_protocol = ip_protocol
        rule.from_port = from_port
        rule.to_port = to_port
        self.rules.append(rule)
        rule.add_grant(src_group_name, src_group_owner_id, cidr_ip, src_group_group_id)

    def remove_rule(self, ip_protocol, from_port, to_port,
                    src_group_name, src_group_owner_id, cidr_ip, src_group_group_id):
        """
        Remove a rule to the SecurityGroup object.  Note that this method
        only changes the local version of the object.  No information
        is sent to EC2.
        """
        target_rule = None
        for rule in self.rules:
            if rule.ip_protocol == ip_protocol:
                if rule.from_port == from_port:
                    if rule.to_port == to_port:
                        target_rule = rule
                        target_grant = None
                        for grant in rule.grants:
                            if grant.name == src_group_name or grant.group_id == src_group_group_id:
                                if grant.owner_id == src_group_owner_id:
                                    if grant.cidr_ip == cidr_ip:
                                        target_grant = grant
                        if target_grant:
                            rule.grants.remove(target_grant)
        if len(rule.grants) == 0:
            self.rules.remove(target_rule)

    def authorize(self, ip_protocol=None, from_port=None, to_port=None,
                  cidr_ip=None, src_group=None):
        """
        Add a new rule to this security group.
        You need to pass in either src_group_name
        OR ip_protocol, from_port, to_port,
        and cidr_ip.  In other words, either you are authorizing another
        group or you are authorizing some ip-based rule.

        :type ip_protocol: string
        :param ip_protocol: Either tcp | udp | icmp

        :type from_port: int
        :param from_port: The beginning port number you are enabling

        :type to_port: int
        :param to_port: The ending port number you are enabling

        :type cidr_ip: string or list of strings
        :param cidr_ip: The CIDR block you are providing access to.
                        See http://en.wikipedia.org/wiki/Classless_Inter-Domain_Routing

        :type src_group: :class:`boto.ec2.securitygroup.SecurityGroup` or
                         :class:`boto.ec2.securitygroup.GroupOrCIDR`
        :param src_group: The Security Group you are granting access to.

        :rtype: bool
        :return: True if successful.
        """
        if src_group:
            cidr_ip = None
            src_group_name = src_group.name
            src_group_owner_id = src_group.owner_id
            if hasattr(src_group, 'group_id'):
                src_group_group_id = src_group.group_id
            else:
                src_group_group_id = src_group.id
        else:
            src_group_name = None
            src_group_owner_id = None
            src_group_group_id = None
        status = self.connection.authorize_security_group(self.name,
                                                          src_group_name,
                                                          src_group_owner_id,
                                                          ip_protocol,
                                                          from_port,
                                                          to_port,
                                                          cidr_ip,
                                                          None,
                                                          src_group_group_id)
        if status:
            if type(cidr_ip) != list:
                cidr_ip = [cidr_ip]
            for single_cidr_ip in cidr_ip:
                self.add_rule(ip_protocol, from_port, to_port, src_group_name,
                              src_group_owner_id, single_cidr_ip, src_group_group_id)

        return status

    def revoke(self, ip_protocol=None, from_port=None, to_port=None,
               cidr_ip=None, src_group=None):
        if src_group:
            cidr_ip=None
            src_group_name = src_group.name
            src_group_owner_id = src_group.owner_id
            if hasattr(src_group, 'group_id'):
                src_group_group_id = src_group.group_id
            else:
                src_group_group_id = src_group.id
        else:
            src_group_name = None
            src_group_owner_id = None
            src_group_group_id = None
        status = self.connection.revoke_security_group(self.name,
                                                       src_group_name,
                                                       src_group_owner_id,
                                                       ip_protocol,
                                                       from_port,
                                                       to_port,
                                                       cidr_ip,
                                                       None,
                                                       src_group_group_id)
        if status:
            self.remove_rule(ip_protocol, from_port, to_port, src_group_name,
                             src_group_owner_id, cidr_ip, src_group_group_id)
        return status

    def copy_to_region(self, region, name=None):
        """
        Create a copy of this security group in another region.
        Note that the new security group will be a separate entity
        and will not stay in sync automatically after the copy
        operation.

        :type region: :class:`boto.ec2.regioninfo.RegionInfo`
        :param region: The region to which this security group will be copied.

        :type name: string
        :param name: The name of the copy.  If not supplied, the copy
                     will have the same name as this security group.

        :rtype: :class:`boto.ec2.securitygroup.SecurityGroup`
        :return: The new security group.
        """
        if region.name == self.region:
            raise BotoClientError('Unable to copy to the same Region')
        conn_params = self.connection.get_params()
        rconn = region.connect(**conn_params)
        sg = rconn.create_security_group(name or self.name, self.description)
        source_groups = []
        for rule in self.rules:
            for grant in rule.grants:
                grant_nom = grant.name or grant.group_id
                if grant_nom:
                    if grant_nom not in source_groups:
                        source_groups.append(grant_nom)
                        sg.authorize(None, None, None, None, grant)
                else:
                    sg.authorize(rule.ip_protocol, rule.from_port, rule.to_port,
                                 grant.cidr_ip)
        return sg

    def instances(self):
        """
        Find all of the current instances that are running within this
        security group.

        :rtype: list of :class:`boto.ec2.instance.Instance`
        :return: A list of Instance objects
        """
        # It would be more efficient to do this with filters now
        # but not all services that implement EC2 API support filters.
        instances = []
        rs = self.connection.get_all_instances()
        for reservation in rs:
            uses_group = [g.name for g in reservation.groups if g.name == self.name]
            if uses_group:
                instances.extend(reservation.instances)
        return instances

class IPPermissionsList(list):

    def startElement(self, name, attrs, connection):
        if name == 'item':
            self.append(IPPermissions(self))
            return self[-1]
        return None

    def endElement(self, name, value, connection):
        pass

class IPPermissions(object):

    def __init__(self, parent=None):
        self.parent = parent
        self.ip_protocol = None
        self.from_port = None
        self.to_port = None
        self.grants = []

    def __repr__(self):
        return 'IPPermissions:%s(%s-%s)' % (self.ip_protocol,
                                            self.from_port, self.to_port)

    def startElement(self, name, attrs, connection):
        if name == 'item':
            self.grants.append(GroupOrCIDR(self))
            return self.grants[-1]
        return None

    def endElement(self, name, value, connection):
        if name == 'ipProtocol':
            self.ip_protocol = value
        elif name == 'fromPort':
            self.from_port = value
        elif name == 'toPort':
            self.to_port = value
        else:
            setattr(self, name, value)

    def add_grant(self, name=None, owner_id=None, cidr_ip=None, group_id=None):
        grant = GroupOrCIDR(self)
        grant.owner_id = owner_id
        grant.group_id = group_id
        grant.name = name
        grant.cidr_ip = cidr_ip
        self.grants.append(grant)
        return grant

class GroupOrCIDR(object):

    def __init__(self, parent=None):
        self.owner_id = None
        self.group_id = None
        self.name = None
        self.cidr_ip = None

    def __repr__(self):
        if self.cidr_ip:
            return '%s' % self.cidr_ip
        else:
            return '%s-%s' % (self.name or self.group_id, self.owner_id)

    def startElement(self, name, attrs, connection):
        return None

    def endElement(self, name, value, connection):
        if name == 'userId':
            self.owner_id = value
        elif name == 'groupId':
            self.group_id = value
        elif name == 'groupName':
            self.name = value
        if name == 'cidrIp':
            self.cidr_ip = value
        else:
            setattr(self, name, value)
<commit_msg>use id instead of name in authorize() and revoke()<commit_after># Copyright (c) 2006-2011 Mitch Garnaat http://garnaat.org/
# Copyright (c) 2011, Eucalyptus Systems, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish, dis-
# tribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the fol-
# lowing conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.

"""
Represents an EC2 Security Group
"""
from boto.ec2.ec2object import TaggedEC2Object
from boto.exception import BotoClientError

class SecurityGroup(TaggedEC2Object):

    def __init__(self, connection=None, owner_id=None,
                 name=None, description=None, id=None):
        TaggedEC2Object.__init__(self, connection)
        self.id = id
        self.owner_id = owner_id
        self.name = name
        self.description = description
        self.vpc_id = None
        self.rules = IPPermissionsList()
        self.rules_egress = IPPermissionsList()

    def __repr__(self):
        return 'SecurityGroup:%s' % self.name

    def startElement(self, name, attrs, connection):
        retval = TaggedEC2Object.startElement(self, name, attrs, connection)
        if retval is not None:
            return retval
        if name == 'ipPermissions':
            return self.rules
        elif name == 'ipPermissionsEgress':
            return self.rules_egress
        else:
            return None

    def endElement(self, name, value, connection):
        if name == 'ownerId':
            self.owner_id = value
        elif name == 'groupId':
            self.id = value
        elif name == 'groupName':
            self.name = value
        elif name == 'vpcId':
            self.vpc_id = value
        elif name == 'groupDescription':
            self.description = value
        elif name == 'ipRanges':
            pass
        elif name == 'return':
            if value == 'false':
                self.status = False
            elif value == 'true':
                self.status = True
            else:
                raise Exception(
                    'Unexpected value of status %s for group %s'%(
                        value,
                        self.name
                    )
                )
        else:
            setattr(self, name, value)

    def delete(self):
        return self.connection.delete_security_group(group_id=self.id)

    def add_rule(self, ip_protocol, from_port, to_port,
                 src_group_name, src_group_owner_id, cidr_ip, src_group_group_id):
        """
        Add a rule to the SecurityGroup object.  Note that this method
        only changes the local version of the object.  No information
        is sent to EC2.
        """
        rule = IPPermissions(self)
        rule.ip_protocol = ip_protocol
        rule.from_port = from_port
        rule.to_port = to_port
        self.rules.append(rule)
        rule.add_grant(src_group_name, src_group_owner_id, cidr_ip, src_group_group_id)

    def remove_rule(self, ip_protocol, from_port, to_port,
                    src_group_name, src_group_owner_id, cidr_ip, src_group_group_id):
        """
        Remove a rule to the SecurityGroup object.  Note that this method
        only changes the local version of the object.  No information
        is sent to EC2.
        """
        target_rule = None
        for rule in self.rules:
            if rule.ip_protocol == ip_protocol:
                if rule.from_port == from_port:
                    if rule.to_port == to_port:
                        target_rule = rule
                        target_grant = None
                        for grant in rule.grants:
                            if grant.name == src_group_name or grant.group_id == src_group_group_id:
                                if grant.owner_id == src_group_owner_id:
                                    if grant.cidr_ip == cidr_ip:
                                        target_grant = grant
                        if target_grant:
                            rule.grants.remove(target_grant)
        if len(rule.grants) == 0:
            self.rules.remove(target_rule)

    def authorize(self, ip_protocol=None, from_port=None, to_port=None,
                  cidr_ip=None, src_group=None):
        """
        Add a new rule to this security group.
        You need to pass in either src_group_name
        OR ip_protocol, from_port, to_port,
        and cidr_ip.  In other words, either you are authorizing another
        group or you are authorizing some ip-based rule.

        :type ip_protocol: string
        :param ip_protocol: Either tcp | udp | icmp

        :type from_port: int
        :param from_port: The beginning port number you are enabling

        :type to_port: int
        :param to_port: The ending port number you are enabling

        :type cidr_ip: string or list of strings
        :param cidr_ip: The CIDR block you are providing access to.
                        See http://en.wikipedia.org/wiki/Classless_Inter-Domain_Routing

        :type src_group: :class:`boto.ec2.securitygroup.SecurityGroup` or
                         :class:`boto.ec2.securitygroup.GroupOrCIDR`
        :param src_group: The Security Group you are granting access to.

        :rtype: bool
        :return: True if successful.
        """
        src_group_name = None
        src_group_owner_id = None
        src_group_group_id = None
        if src_group:
            cidr_ip = None
            if not self.vpc_id:
                src_group_name = src_group.name
            src_group_owner_id = src_group.owner_id
            if hasattr(src_group, 'group_id'):
                src_group_group_id = src_group.group_id
            else:
                src_group_group_id = src_group.id
        status = self.connection.authorize_security_group(None,
                                                          None,
                                                          src_group_owner_id,
                                                          ip_protocol,
                                                          from_port,
                                                          to_port,
                                                          cidr_ip,
                                                          self.id,
                                                          src_group_group_id)
        if status:
            if type(cidr_ip) != list:
                cidr_ip = [cidr_ip]
            for single_cidr_ip in cidr_ip:
                self.add_rule(ip_protocol, from_port, to_port, src_group_name,
                              src_group_owner_id, single_cidr_ip, src_group_group_id)
        return status

    def revoke(self, ip_protocol=None, from_port=None, to_port=None,
               cidr_ip=None, src_group=None):
        src_group_name = None
        src_group_owner_id = None
        src_group_group_id = None
        if src_group:
            cidr_ip=None
            if not self.vpc_id:
                src_group_name = src_group.name
            src_group_owner_id = src_group.owner_id
            if hasattr(src_group, 'group_id'):
                src_group_group_id = src_group.group_id
            else:
                src_group_group_id = src_group.id
        status = self.connection.revoke_security_group(None,
                                                       None,
                                                       src_group_owner_id,
                                                       ip_protocol,
                                                       from_port,
                                                       to_port,
                                                       cidr_ip,
                                                       self.id,
                                                       src_group_group_id)
        if status:
            self.remove_rule(ip_protocol, from_port, to_port, src_group_name,
                             src_group_owner_id, cidr_ip, src_group_group_id)
        return status

    def copy_to_region(self, region, name=None):
        """
        Create a copy of this security group in another region.
        Note that the new security group will be a separate entity
        and will not stay in sync automatically after the copy
        operation.

        :type region: :class:`boto.ec2.regioninfo.RegionInfo`
        :param region: The region to which this security group will be copied.

        :type name: string
        :param name: The name of the copy.  If not supplied, the copy
                     will have the same name as this security group.

        :rtype: :class:`boto.ec2.securitygroup.SecurityGroup`
        :return: The new security group.
        """
        if region.name == self.region:
            raise BotoClientError('Unable to copy to the same Region')
        conn_params = self.connection.get_params()
        rconn = region.connect(**conn_params)
        sg = rconn.create_security_group(name or self.name, self.description)
        source_groups = []
        for rule in self.rules:
            for grant in rule.grants:
                grant_nom = grant.name or grant.group_id
                if grant_nom:
                    if grant_nom not in source_groups:
                        source_groups.append(grant_nom)
                        sg.authorize(None, None, None, None, grant)
                else:
                    sg.authorize(rule.ip_protocol, rule.from_port, rule.to_port,
                                 grant.cidr_ip)
        return sg

    def instances(self):
        """
        Find all of the current instances that are running within this
        security group.

        :rtype: list of :class:`boto.ec2.instance.Instance`
        :return: A list of Instance objects
        """
        # It would be more efficient to do this with filters now
        # but not all services that implement EC2 API support filters.
        instances = []
        rs = self.connection.get_all_instances()
        for reservation in rs:
            uses_group = [g.name for g in reservation.groups if g.name == self.name]
            if uses_group:
                instances.extend(reservation.instances)
        return instances

class IPPermissionsList(list):

    def startElement(self, name, attrs, connection):
        if name == 'item':
            self.append(IPPermissions(self))
            return self[-1]
        return None

    def endElement(self, name, value, connection):
        pass

class IPPermissions(object):

    def __init__(self, parent=None):
        self.parent = parent
        self.ip_protocol = None
        self.from_port = None
        self.to_port = None
        self.grants = []

    def __repr__(self):
        return 'IPPermissions:%s(%s-%s)' % (self.ip_protocol,
                                            self.from_port, self.to_port)

    def startElement(self, name, attrs, connection):
        if name == 'item':
            self.grants.append(GroupOrCIDR(self))
            return self.grants[-1]
        return None

    def endElement(self, name, value, connection):
        if name == 'ipProtocol':
            self.ip_protocol = value
        elif name == 'fromPort':
            self.from_port = value
        elif name == 'toPort':
            self.to_port = value
        else:
            setattr(self, name, value)

    def add_grant(self, name=None, owner_id=None, cidr_ip=None, group_id=None):
        grant = GroupOrCIDR(self)
        grant.owner_id = owner_id
        grant.group_id = group_id
        grant.name = name
        grant.cidr_ip = cidr_ip
        self.grants.append(grant)
        return grant

class GroupOrCIDR(object):

    def __init__(self, parent=None):
        self.owner_id = None
        self.group_id = None
        self.name = None
        self.cidr_ip = None

    def __repr__(self):
        if self.cidr_ip:
            return '%s' % self.cidr_ip
        else:
            return '%s-%s' % (self.name or self.group_id, self.owner_id)

    def startElement(self, name, attrs, connection):
        return None

    def endElement(self, name, value, connection):
        if name == 'userId':
            self.owner_id = value
        elif name == 'groupId':
            self.group_id = value
        elif name == 'groupName':
            self.name = value
        if name == 'cidrIp':
            self.cidr_ip = value
        else:
            setattr(self, name, value)
<|endoftext|>


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

import sys, os, time, traceback, shlex, re

from pyomo.environ import *
from pyomo.opt import SolverFactory, SolverStatus, TerminationCondition

from utilities import create_model, _ArgumentParser

def main(args=sys.argv[1:]):

    # build a module list based on configuration options passed in,
    # and add the current module (to register define_arguments callback)
    modules = get_module_list(args)

    # Define the model
    model = create_model(modules, args=args)

    # Add any suffixes specified on the command line (usually only iis)
    add_extra_suffixes(model)

    # get a list of modules to iterate through
    iterate_modules = get_iteration_list(model)

    print "\n\n======================================================================="
    print "arguments:"
    print " ".join(k+"="+repr(v) for k, v in model.options.__dict__.items() if v)
    print "modules:", modules
    if iterate_modules:
        print "iterate_modules", iterate_modules
    print "======================================================================="

    # create an instance
    instance = model.load_inputs()

    # make sure the outputs_dir exists (used by some modules during iterate)
    if not os.path.exists(instance.options.outputs_dir):
        os.makedirs(instance.options.outputs_dir)

    # solve the model
    if iterate_modules:
        print "iterating model..."
        iterate(instance, iterate_modules)
    else:
        print "solving model..."
        solve(instance)
        print "finished solving"

    # report/save results
    instance.post_solve()

def iterate(m, iterate_modules, depth=0):
    """Iterate through all modules listed in the iterate_list (usually iterate.txt),
    if any. If there is no iterate_list, then this will just solve the model once.

    If it exists, the iterate_list contains one row per level of iteration,
    and each row contains a list of modules to test for iteration at that level
    (these can be separated with commas, spaces or tabs).
    The model will run through the levels like nested loops, running the lowest level
    till it converges, then advancing the next higher level by one step, then running the
    lowest level to convergence/completion again, repeating until all levels are complete.
    During each iteration, the pre_iterate() and post_iterate() functions of each specified
    module (if they exist) will be called before and after solving. When a module is
    converged or completed, its post_iterate() function should return True.
    All modules specified in the iterate_list should also be loaded via the module_list
    or include_module(s) arguments.
    """

    # create or truncate the iteration tree
    if depth == 0:
        m.iteration_node = []

    if depth == len(iterate_modules):
        # asked to converge at the deepest level
        # just preprocess to reflect all changes and then solve
        m.preprocess()
        solve(m)
    else:
        # iterate until converged at the current level

        # note: the modules in iterate_modules were also specified in the model's
        # module list, and have already been loaded, so they are accessible via sys.modules
        current_modules = [sys.modules[module_name] for module_name in iterate_modules[depth]]
        # truncate the iteration tree at the current level
        m.iteration_node = m.iteration_node[:depth] + [0]

        j = 0
        converged = False
        while not converged:
            # take one step at the current level
            if m.options.max_iter is not None and j >= m.options.max_iter:
                break

            # record the current iteration number and node for use by modules
            # (e.g., to name files or reset/index params)
            m.iteration_number = j
            m.iteration_node[-1] = j

            converged = True
            # pre-iterate modules at this level
            for module in current_modules:
                if hasattr(module, 'pre_iterate'):
                    converged = module.pre_iterate(m) and converged

            # converge the deeper-level modules, if any (inner loop)
            iterate(m, iterate_modules, depth=depth+1)

            # post-iterate modules at this level
            for module in current_modules:
                if hasattr(module, 'post_iterate'):
                    converged = module.post_iterate(m) and converged

            j += 1
        if converged:
            print "Iteration of {ms} was completed after {j} rounds.".format(ms=iterate_modules[depth], j=j)
        else:
            print "Iteration of {ms} was stopped after {j} iterations without convergence.".format(ms=iterate_modules[depth], j=j)
    return

def define_arguments(argparser):
    # callback function to define model configuration arguments while the model is built

    # add standard module arguments (not used later, but this adds them to the help)
    add_module_args(argparser)

    # iteration options
    argparser.add_argument(
        "--iterate-list", default=None,
        help="Text file with a list of modules to iterate until converged (default is iterate.txt); "
             "each row is one level of iteration, and there can be multiple modules on each row"
    )
    argparser.add_argument(
        "--max-iter", type=int, default=None,
        help="Maximum number of iterations to complete at each level for iterated models"
    )

    # scenario information
    argparser.add_argument(
        "--scenario-name", default="", help="Name of research scenario represented by this model"
    )

    # note: pyomo has a --solver-suffix option but it is not clear
    # whether that does the same thing as --suffix defined here,
    # so we don't reuse the same name.
    argparser.add_argument("--suffixes", "--suffix", nargs="+", default=[],
        help="Extra suffixes to add to the model and exchange with the solver (e.g., iis, rc, dual, or slack)")

    # Define solver-related arguments
    # These are a subset of the arguments offered by "pyomo solve --solver=cplex --help"
    argparser.add_argument("--solver", default="glpk",
        help='Name of Pyomo solver to use for the model (default is "glpk")')
    argparser.add_argument("--solver-io", default=None, help="Method for Pyomo to use to communicate with solver")
    # note: pyomo has a --solver-options option but it is not clear
    # whether that does the same thing as --solver-options-string so we don't reuse the same name.
    argparser.add_argument("--solver-options-string", default=None,
        help='A quoted string of options to pass to the model solver. Each option must be of the form option=value. '
            '(e.g., --solver-options-string "mipgap=0.001 primalopt advance=2 threads=1")')
    argparser.add_argument("--keepfiles", action='store_true', default=None,
        help="Keep temporary files produced by the solver (may be useful with --symbolic-solver-labels)")
    argparser.add_argument(
        "--stream-output", "--stream-solver", action='store_true', dest="tee", default=None,
        help="Display information from the solver about its progress (usually combined with a suitable --solver-options string)")
    argparser.add_argument(
        "--symbolic-solver-labels", action='store_true', default=None,
        help='Use symbol names derived from the model when interfacing with the solver. '
            'See "pyomo solve --solver=x --help" for more details.')
    argparser.add_argument("--tempdir", default=None,
        help='The name of a directory to hold temporary files produced by the solver. '
             'This is usually paired with --keepfiles and --symbolic-solver-labels.')

    # NOTE: the following could potentially be made into standard arguments for all models,
    # e.g. by defining them in a define_standard_arguments() function in switch.utilities.py

    # Define input/output options
    argparser.add_argument("--inputs-dir", default="inputs",
        help='Directory containing input files (default is "inputs")')
    argparser.add_argument("--outputs-dir", default="outputs",
        help='Directory to write output files (default is "outputs")')

    # General purpose arguments
    argparser.add_argument(
        '--verbose', '-v', default=False, action='store_true',
        help='Show information about model preparation and solution')

def add_module_args(parser):
    parser.add_argument(
        "--module-list", default=None,
        help='Text file with a list of modules to include in the model (default is "modules.txt")'
    )
    parser.add_argument(
        "--include-modules", "--include-module", dest="include_modules", nargs='+', default=[],
        help="Module(s) to add to the model in addition to any specified with --module-list"
    )
    parser.add_argument(
        "--exclude-modules", "--exclude-module", dest="exclude_modules", nargs='+', default=[],
        help="Module(s) to remove from the model after processing --module-list and --include-modules"
    )

def get_module_list(args):
    # parse module options
    parser = _ArgumentParser(allow_abbrev=False, add_help=False)
    add_module_args(parser)
    module_options = parser.parse_known_args(args=args)[0]

    # identify modules to load
    module_list_file = module_options.module_list
    if module_list_file is None and os.path.exists("modules.txt"):
        module_list_file = "modules.txt"
    if module_list_file is None:
        modules = []
    else:
        # if it exists, the module list contains one module name per row (no .py extension)
        # we strip whitespace from either end (because those errors can be annoyingly hard to debug),
        # but otherwise take the module names as given.
        with open(module_list_file) as f:
            modules = [r.strip() for r in f.read().splitlines()]

    modules.extend(module_options.include_modules)
    for module_name in module_options.exclude_modules:
        modules.remove(module_name)

    # add the current module, since it has callbacks, e.g. define_arguments for iteration and suffixes
    modules.append(__name__)

    print "module list:", modules

    return modules

def get_iteration_list(m):
    # Identify modules to iterate until convergence (if any)
    iterate_list_file = m.options.iterate_list
    if iterate_list_file is None and os.path.exists("iterate.txt"):
        iterate_list_file = "iterate.txt"
    if iterate_list_file is None:
        iterate_modules = []
    else:
        with open(iterate_list_file) as f:
            iterate_rows = f.read().splitlines()
        iterate_modules = [re.sub("[ \t,]+", " ", r).split(" ") for r in iterate_rows]
    return iterate_modules

def get_option_file_args():
    args = []
    # retrieve base arguments from options.txt (if present)
    # note: these can be on multiple lines to ease editing,
    # and lines can be commented out with #
    if os.path.exists("options.txt"):
        with open("options.txt") as f:
            base_options = f.read().splitlines()
        for r in base_options:
            if not r.lstrip().startswith("#"):
                args.extend(shlex.split(r))
    return args

# Generic argument-related code; could potentially be moved to utilities.py
# if we want to make these standard parts of Switch.

def add_extra_suffixes(model):
    """
    Add any suffix objects requested in the configuration options.
    We assume they will be used for import or export of floating-point values
    note: modules that need suffixes should normally just create them (possibly
    checking whether they already exist first). Then solve() will automatically
    pass them to the solver.
    """
    for suffix in model.options.suffixes:
        setattr(model, suffix, Suffix(direction=Suffix.IMPORT_EXPORT))

def solve(model):
    if not hasattr(model, "solver"):
        # Create a solver object the first time in. We don't do this until a solve is
        # requested, because sometimes a different solve function may be used,
        # with its own solver object (e.g., with runph or a parallel solver server).
        # In those cases, we don't want to go through the expense of creating an
        # unused solver object, or get errors if the solver options are invalid.
        model.solver = SolverFactory(model.options.solver, solver_io=model.options.solver_io)

        # patch for Pyomo < 4.2
        # note: Pyomo added an options_string argument to solver.solve() in Pyomo 4.2 rev 10587.
        # (See https://software.sandia.gov/trac/pyomo/browser/pyomo/trunk/pyomo/opt/base/solvers.py?rev=10587 )
        # This is misreported in the documentation as options=, but options= actually accepts a dictionary.
        if model.options.solver_options_string and not hasattr(model.solver, "_options_string_to_dict"):
            for k, v in _options_string_to_dict(model.options.solver_options_string).items():
                model.solver.options[k] = v

    # get solver arguments (if any)
    if hasattr(model, "options"):
        solver_args = dict(
            options_string=model.options.solver_options_string,
            keepfiles=model.options.keepfiles,
            tee=model.options.tee,
            symbolic_solver_labels=model.options.symbolic_solver_labels
        )
        # drop all the unspecified options
        solver_args = {k: v for (k, v) in solver_args.items() if v is not None}
    else:
        solver_args={}

    # Automatically send all defined suffixes to the solver
    solver_args["suffixes"] = [c.cname() for c in model.component_objects() if isinstance(c, Suffix)]
    # note: the next few lines are faster than the line above, but seem risky:
    # i = m._ctypes.get(Suffix, [None])[0]
    # solver_args["suffixes"] = []
    # while i is not None:
    #     c, i = m._decl_order[i]
    #     solver_args[suffixes].append(c.cname())

    # patch for Pyomo < 4.2
    if not hasattr(model.solver, "_options_string_to_dict"):
        solver_args.pop("options_string", "")

    # solve the model
    if model.options.verbose:
        print "solving model..."
    start = time.time()

    if model.options.tempdir is not None:
        # from https://software.sandia.gov/downloads/pub/pyomo/PyomoOnlineDocs.html#_changing_the_temporary_directory
        from pyutilib.services import TempFileManager
        TempfileManager.tempdir = model.options.tempdir

    results = model.solver.solve(model, **solver_args)

    if model.options.verbose:
        print "Total time in solver: {t}s".format(t=time.time()-start)

    # check for errors
    model.solutions.load_from(results)
    if results.solver.termination_condition == pyomo.opt.TerminationCondition.infeasible:
        if hasattr(model, "iis"):
            print "Model was infeasible; irreducible infeasible set (IIS) returned by solver:"
            print "\n".join(c.cname() for c in m.iis)
        else:
            print "Model was infeasible; if the solver can generate an irreducible infeasible set,"
            print "more information may be available by calling this script with --suffixes iis ..."
        raise RuntimeError("Infeasible model")

    return results

# taken from https://software.sandia.gov/trac/pyomo/browser/pyomo/trunk/pyomo/opt/base/solvers.py?rev=10784
# This can be removed when all users are on Pyomo 4.2
import pyutilib
def _options_string_to_dict(istr):
    ans = {}
    istr = istr.strip()
    if not istr:
        return ans
    if istr[0] == "'" or istr[0] == '"':
        istr = eval(istr)
    tokens = pyutilib.misc.quote_split('[ ]+',istr)
    for token in tokens:
        index = token.find('=')
        if index is -1:
            raise ValueError(
                "Solver options must have the form option=value: '%s'" % istr)
        try:
            val = eval(token[(index+1):])
        except:
            val = token[(index+1):]
        ans[token[:index]] = val
    return ans

###############

if __name__ == "__main__":
    # combine default arguments read from options.txt file with
    # additional arguments specified on the command line
    args = get_option_file_args()
    # add any command-line arguments
    args.extend(sys.argv[1:])
    main(args)

<commit_msg>allow comments in modules.txt; properly assign tempdir with --tempdir<commit_after>#!/usr/bin/env python

import sys, os, time, traceback, shlex, re

from pyomo.environ import *
from pyomo.opt import SolverFactory, SolverStatus, TerminationCondition

from utilities import create_model, _ArgumentParser

def main(args=sys.argv[1:]):

    # build a module list based on configuration options passed in,
    # and add the current module (to register define_arguments callback)
    modules = get_module_list(args)

    # Define the model
    model = create_model(modules, args=args)

    # Add any suffixes specified on the command line (usually only iis)
    add_extra_suffixes(model)

    # get a list of modules to iterate through
    iterate_modules = get_iteration_list(model)

    print "\n\n======================================================================="
    print "arguments:"
    print " ".join(k+"="+repr(v) for k, v in model.options.__dict__.items() if v)
    print "modules:", modules
    if iterate_modules:
        print "iterate_modules", iterate_modules
    print "======================================================================="

    # create an instance
    instance = model.load_inputs()

    # make sure the outputs_dir exists (used by some modules during iterate)
    if not os.path.exists(instance.options.outputs_dir):
        os.makedirs(instance.options.outputs_dir)

    # solve the model
    if iterate_modules:
        print "iterating model..."
        iterate(instance, iterate_modules)
    else:
        print "solving model..."
        solve(instance)
        print "finished solving"

    # report/save results
    instance.post_solve()

def iterate(m, iterate_modules, depth=0):
    """Iterate through all modules listed in the iterate_list (usually iterate.txt),
    if any. If there is no iterate_list, then this will just solve the model once.

    If it exists, the iterate_list contains one row per level of iteration,
    and each row contains a list of modules to test for iteration at that level
    (these can be separated with commas, spaces or tabs).
    The model will run through the levels like nested loops, running the lowest level
    till it converges, then advancing the next higher level by one step, then running the
    lowest level to convergence/completion again, repeating until all levels are complete.
    During each iteration, the pre_iterate() and post_iterate() functions of each specified
    module (if they exist) will be called before and after solving. When a module is
    converged or completed, its post_iterate() function should return True.
    All modules specified in the iterate_list should also be loaded via the module_list
    or include_module(s) arguments.
    """

    # create or truncate the iteration tree
    if depth == 0:
        m.iteration_node = []

    if depth == len(iterate_modules):
        # asked to converge at the deepest level
        # just preprocess to reflect all changes and then solve
        m.preprocess()
        solve(m)
    else:
        # iterate until converged at the current level

        # note: the modules in iterate_modules were also specified in the model's
        # module list, and have already been loaded, so they are accessible via sys.modules
        current_modules = [sys.modules[module_name] for module_name in iterate_modules[depth]]
        # truncate the iteration tree at the current level
        m.iteration_node = m.iteration_node[:depth] + [0]

        j = 0
        converged = False
        while not converged:
            # take one step at the current level
            if m.options.max_iter is not None and j >= m.options.max_iter:
                break

            # record the current iteration number and node for use by modules
            # (e.g., to name files or reset/index params)
            m.iteration_number = j
            m.iteration_node[-1] = j

            converged = True
            # pre-iterate modules at this level
            for module in current_modules:
                if hasattr(module, 'pre_iterate'):
                    converged = module.pre_iterate(m) and converged

            # converge the deeper-level modules, if any (inner loop)
            iterate(m, iterate_modules, depth=depth+1)

            # post-iterate modules at this level
            for module in current_modules:
                if hasattr(module, 'post_iterate'):
                    converged = module.post_iterate(m) and converged

            j += 1
        if converged:
            print "Iteration of {ms} was completed after {j} rounds.".format(ms=iterate_modules[depth], j=j)
        else:
            print "Iteration of {ms} was stopped after {j} iterations without convergence.".format(ms=iterate_modules[depth], j=j)
    return

def define_arguments(argparser):
    # callback function to define model configuration arguments while the model is built

    # add standard module arguments (not used later, but this adds them to the help)
    add_module_args(argparser)

    # iteration options
    argparser.add_argument(
        "--iterate-list", default=None,
        help="Text file with a list of modules to iterate until converged (default is iterate.txt); "
             "each row is one level of iteration, and there can be multiple modules on each row"
    )
    argparser.add_argument(
        "--max-iter", type=int, default=None,
        help="Maximum number of iterations to complete at each level for iterated models"
    )

    # scenario information
    argparser.add_argument(
        "--scenario-name", default="", help="Name of research scenario represented by this model"
    )

    # note: pyomo has a --solver-suffix option but it is not clear
    # whether that does the same thing as --suffix defined here,
    # so we don't reuse the same name.
    argparser.add_argument("--suffixes", "--suffix", nargs="+", default=[],
        help="Extra suffixes to add to the model and exchange with the solver (e.g., iis, rc, dual, or slack)")

    # Define solver-related arguments
    # These are a subset of the arguments offered by "pyomo solve --solver=cplex --help"
    argparser.add_argument("--solver", default="glpk",
        help='Name of Pyomo solver to use for the model (default is "glpk")')
    argparser.add_argument("--solver-io", default=None, help="Method for Pyomo to use to communicate with solver")
    # note: pyomo has a --solver-options option but it is not clear
    # whether that does the same thing as --solver-options-string so we don't reuse the same name.
    argparser.add_argument("--solver-options-string", default=None,
        help='A quoted string of options to pass to the model solver. Each option must be of the form option=value. '
            '(e.g., --solver-options-string "mipgap=0.001 primalopt advance=2 threads=1")')
    argparser.add_argument("--keepfiles", action='store_true', default=None,
        help="Keep temporary files produced by the solver (may be useful with --symbolic-solver-labels)")
    argparser.add_argument(
        "--stream-output", "--stream-solver", action='store_true', dest="tee", default=None,
        help="Display information from the solver about its progress (usually combined with a suitable --solver-options string)")
    argparser.add_argument(
        "--symbolic-solver-labels", action='store_true', default=None,
        help='Use symbol names derived from the model when interfacing with the solver. '
            'See "pyomo solve --solver=x --help" for more details.')
    argparser.add_argument("--tempdir", default=None,
        help='The name of a directory to hold temporary files produced by the solver. '
             'This is usually paired with --keepfiles and --symbolic-solver-labels.')

    # NOTE: the following could potentially be made into standard arguments for all models,
    # e.g. by defining them in a define_standard_arguments() function in switch.utilities.py

    # Define input/output options
    argparser.add_argument("--inputs-dir", default="inputs",
        help='Directory containing input files (default is "inputs")')
    argparser.add_argument("--outputs-dir", default="outputs",
        help='Directory to write output files (default is "outputs")')

    # General purpose arguments
    argparser.add_argument(
        '--verbose', '-v', default=False, action='store_true',
        help='Show information about model preparation and solution')

def add_module_args(parser):
    parser.add_argument(
        "--module-list", default=None,
        help='Text file with a list of modules to include in the model (default is "modules.txt")'
    )
    parser.add_argument(
        "--include-modules", "--include-module", dest="include_modules", nargs='+', default=[],
        help="Module(s) to add to the model in addition to any specified with --module-list"
    )
    parser.add_argument(
        "--exclude-modules", "--exclude-module", dest="exclude_modules", nargs='+', default=[],
        help="Module(s) to remove from the model after processing --module-list and --include-modules"
    )

def get_module_list(args):
    # parse module options
    parser = _ArgumentParser(allow_abbrev=False, add_help=False)
    add_module_args(parser)
    module_options = parser.parse_known_args(args=args)[0]

    # identify modules to load
    module_list_file = module_options.module_list
    if module_list_file is None and os.path.exists("modules.txt"):
        module_list_file = "modules.txt"
    if module_list_file is None:
        modules = []
    else:
        # if it exists, the module list contains one module name per row (no .py extension)
        # we strip whitespace from either end (because those errors can be annoyingly hard to debug).
        # We also omit blank lines and lines that start with "#"
        # Otherwise take the module names as given.
        with open(module_list_file) as f:
            modules = [r.strip() for r in f.read().splitlines()]
        modules = [m for m in modules if m and not m.startswith("#")]

    modules.extend(module_options.include_modules)
    for module_name in module_options.exclude_modules:
        modules.remove(module_name)

    # add the current module, since it has callbacks, e.g. define_arguments for iteration and suffixes
    modules.append(__name__)

    print "module list:", modules

    return modules

def get_iteration_list(m):
    # Identify modules to iterate until convergence (if any)
    iterate_list_file = m.options.iterate_list
    if iterate_list_file is None and os.path.exists("iterate.txt"):
        iterate_list_file = "iterate.txt"
    if iterate_list_file is None:
        iterate_modules = []
    else:
        with open(iterate_list_file) as f:
            iterate_rows = f.read().splitlines()
        iterate_modules = [re.sub("[ \t,]+", " ", r).split(" ") for r in iterate_rows]
    return iterate_modules

def get_option_file_args():
    args = []
    # retrieve base arguments from options.txt (if present)
    # note: these can be on multiple lines to ease editing,
    # and lines can be commented out with #
    if os.path.exists("options.txt"):
        with open("options.txt") as f:
            base_options = f.read().splitlines()
        for r in base_options:
            if not r.lstrip().startswith("#"):
                args.extend(shlex.split(r))
    return args

# Generic argument-related code; could potentially be moved to utilities.py
# if we want to make these standard parts of Switch.

def add_extra_suffixes(model):
    """
    Add any suffix objects requested in the configuration options.
    We assume they will be used for import or export of floating-point values
    note: modules that need suffixes should normally just create them (possibly
    checking whether they already exist first). Then solve() will automatically
    pass them to the solver.
    """
    for suffix in model.options.suffixes:
        setattr(model, suffix, Suffix(direction=Suffix.IMPORT_EXPORT))

def solve(model):
    if not hasattr(model, "solver"):
        # Create a solver object the first time in. We don't do this until a solve is
        # requested, because sometimes a different solve function may be used,
        # with its own solver object (e.g., with runph or a parallel solver server).
        # In those cases, we don't want to go through the expense of creating an
        # unused solver object, or get errors if the solver options are invalid.
        model.solver = SolverFactory(model.options.solver, solver_io=model.options.solver_io)

        # patch for Pyomo < 4.2
        # note: Pyomo added an options_string argument to solver.solve() in Pyomo 4.2 rev 10587.
        # (See https://software.sandia.gov/trac/pyomo/browser/pyomo/trunk/pyomo/opt/base/solvers.py?rev=10587 )
        # This is misreported in the documentation as options=, but options= actually accepts a dictionary.
        if model.options.solver_options_string and not hasattr(model.solver, "_options_string_to_dict"):
            for k, v in _options_string_to_dict(model.options.solver_options_string).items():
                model.solver.options[k] = v

    # get solver arguments (if any)
    if hasattr(model, "options"):
        solver_args = dict(
            options_string=model.options.solver_options_string,
            keepfiles=model.options.keepfiles,
            tee=model.options.tee,
            symbolic_solver_labels=model.options.symbolic_solver_labels
        )
        # drop all the unspecified options
        solver_args = {k: v for (k, v) in solver_args.items() if v is not None}
    else:
        solver_args={}

    # Automatically send all defined suffixes to the solver
    solver_args["suffixes"] = [c.cname() for c in model.component_objects() if isinstance(c, Suffix)]
    # note: the next few lines are faster than the line above, but seem risky:
    # i = m._ctypes.get(Suffix, [None])[0]
    # solver_args["suffixes"] = []
    # while i is not None:
    #     c, i = m._decl_order[i]
    #     solver_args[suffixes].append(c.cname())

    # patch for Pyomo < 4.2
    if not hasattr(model.solver, "_options_string_to_dict"):
        solver_args.pop("options_string", "")

    # solve the model
    if model.options.verbose:
        print "solving model..."
    start = time.time()

    if model.options.tempdir is not None:
        # from https://software.sandia.gov/downloads/pub/pyomo/PyomoOnlineDocs.html#_changing_the_temporary_directory
        from pyutilib.services import TempfileManager
        TempfileManager.tempdir = model.options.tempdir

    results = model.solver.solve(model, **solver_args)

    if model.options.verbose:
        print "Total time in solver: {t}s".format(t=time.time()-start)

    # check for errors
    model.solutions.load_from(results)
    if results.solver.termination_condition == pyomo.opt.TerminationCondition.infeasible:
        if hasattr(model, "iis"):
            print "Model was infeasible; irreducible infeasible set (IIS) returned by solver:"
            print "\n".join(c.cname() for c in m.iis)
        else:
            print "Model was infeasible; if the solver can generate an irreducible infeasible set,"
            print "more information may be available by calling this script with --suffixes iis ..."
        raise RuntimeError("Infeasible model")

    return results

# taken from https://software.sandia.gov/trac/pyomo/browser/pyomo/trunk/pyomo/opt/base/solvers.py?rev=10784
# This can be removed when all users are on Pyomo 4.2
import pyutilib
def _options_string_to_dict(istr):
    ans = {}
    istr = istr.strip()
    if not istr:
        return ans
    if istr[0] == "'" or istr[0] == '"':
        istr = eval(istr)
    tokens = pyutilib.misc.quote_split('[ ]+',istr)
    for token in tokens:
        index = token.find('=')
        if index is -1:
            raise ValueError(
                "Solver options must have the form option=value: '%s'" % istr)
        try:
            val = eval(token[(index+1):])
        except:
            val = token[(index+1):]
        ans[token[:index]] = val
    return ans

###############

if __name__ == "__main__":
    # combine default arguments read from options.txt file with
    # additional arguments specified on the command line
    args = get_option_file_args()
    # add any command-line arguments
    args.extend(sys.argv[1:])
    main(args)

<|endoftext|>


<commit_before>#  Copyright (c) 2015-2016 Cisco Systems, Inc.
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to
#  deal in the Software without restriction, including without limitation the
#  rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
#  sell copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included in
#  all copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#  DEALINGS IN THE SOFTWARE.

import os
import os.path
import shutil

import pytest

from molecule import config
from molecule import core
from molecule import state

@pytest.helpers.register
def os_split(s):
    rest, tail = os.path.split(s)
    if rest in ('', os.path.sep):
        return tail,
    return os_split(rest) + (tail, )

@pytest.fixture()
def temp_files(tmpdir, request):
    def wrapper(fixtures=[]):
        d = tmpdir.mkdir(pytest.helpers.random_string())
        confs = []
        for index, fixture in enumerate(fixtures):
            c = d.join(os.extsep.join((fixture, 'yml')))
            c.write(request.getfuncargvalue(fixtures[index]))
            confs.append(c.strpath)

        def cleanup():
            shutil.rmtree(d.strpath)

        request.addfinalizer(cleanup)

        return confs

    return wrapper

@pytest.fixture()
def molecule_instance(temp_dir, temp_files, state_path_without_data):
    c = temp_files(fixtures=['molecule_vagrant_v1_config'])
    m = core.Molecule(config.ConfigV1(configs=c), {})
    m.state = state.State(state_file=state_path_without_data)
    m.main()

    return m

@pytest.fixture()
def molecule_instance_with_env_expansion(temp_dir, temp_files,
                                         state_path_without_data):
    os.environ['MOLECULE_TEST_01'] = 'vagrant'
    os.environ['MOLECULE_TEST_02'] = 'ubuntu'
    os.environ['MOLECULE_TEST_03'] = 'trusty64'
    os.environ['MOLECULE_TEST_04'] = 'example1'
    os.environ['MOLECULE_TEST_05'] = 'append_platform_to_hostname'
    c = temp_files(fixtures=['molecule_vagrant_v1_config_using_env'])
    m = core.Molecule(config.ConfigV1(configs=c), {})
    m.state = state.State(state_file=state_path_without_data)
    m.main()

    return m

@pytest.fixture()
def docker_molecule_instance(temp_dir, temp_files, state_path_without_data):
    c = temp_files(fixtures=['molecule_docker_v1_config'])
    m = core.Molecule(config.ConfigV1(configs=c), {})
    m.state = state.State(state_file=state_path_without_data)
    m.main()

    return m

@pytest.fixture()
def openstack_molecule_instance(temp_dir, temp_files, state_path_without_data):
    c = temp_files(fixtures=['molecule_openstack_v1_config'])
    m = core.Molecule(config.ConfigV1(configs=c), {})
    m.state = state.State(state_file=state_path_without_data)
    m.main()

    return m

@pytest.fixture()
def vagrant_molecule_instance(temp_dir, temp_files, state_path_without_data):
    c = temp_files(fixtures=['molecule_vagrant_v1_config'])
    m = core.Molecule(config.ConfigV1(configs=c), {})
    m.state = state.State(state_file=state_path_without_data)
    m.main()

    return m

@pytest.fixture()
def molecule_vagrant_v1_config(molecule_v1_section_data,
                               vagrant_v1_section_data,
                               ansible_v1_section_data):
    return reduce(lambda x, y: config.merge_dicts(x, y), [
        molecule_v1_section_data, vagrant_v1_section_data,
        ansible_v1_section_data
    ])

@pytest.fixture()
def molecule_vagrant_v1_config_using_env(molecule_v1_section_data,
                                         vagrant_v1_section_data_using_env,
                                         ansible_v1_section_data):
    return reduce(lambda x, y: config.merge_dicts(x, y), [
        molecule_v1_section_data, vagrant_v1_section_data_using_env,
        ansible_v1_section_data
    ])

@pytest.fixture()
def molecule_docker_v1_config(molecule_v1_section_data, docker_v1_section_data,
                              ansible_v1_section_data):
    return reduce(lambda x, y: config.merge_dicts(x, y), [
        molecule_v1_section_data, docker_v1_section_data,
        ansible_v1_section_data
    ])

@pytest.fixture()
def molecule_openstack_v1_config(molecule_v1_section_data,
                                 openstack_v1_section_data,
                                 ansible_v1_section_data):
    return reduce(lambda x, y: config.merge_dicts(x, y), [
        molecule_v1_section_data, openstack_v1_section_data,
        ansible_v1_section_data
    ])

@pytest.fixture()
def molecule_v1_section_data(state_path_without_data):
    return {
        'molecule': {
            'ignore_paths': ['.git', '.vagrant', '.molecule'],
            'serverspec_dir': 'spec',
            'testinfra_dir': 'tests',
            'goss_dir': 'tests',
            'goss_playbook': 'test_default.yml',
            'molecule_dir': 'test',
            'state_file': state_path_without_data,
            'vagrantfile_file': 'vagrantfile_file',
            'rakefile_file': 'rakefile_file',
            'vagrantfile_template': 'vagrantfile.j2',
            'raw_ssh_args': [
                '-o StrictHostKeyChecking=no',
                '-o UserKnownHostsFile=/dev/null'
            ],
            'test': {
                'sequence': [
                    'destroy', 'dependency', 'syntax', 'create', 'converge',
                    'idempotence', 'verify'
                ]
            }
        }
    }

@pytest.fixture()
def vagrant_v1_section_data():
    return {
        'vagrant': {
            'platforms': [{
                'name': 'ubuntu',
                'box': 'ubuntu/trusty64'
            }, {
                'name': 'centos7',
                'box': 'centos/7'
            }],
            'providers': [{
                'name': 'virtualbox',
                'type': 'virtualbox'
            }],
            'instances': [{
                'name': 'aio-01',
                'ansible_groups': ['example', 'example1'],
                'options': {
                    'append_platform_to_hostname': True
                }
            }, {
                'name': 'aio-02',
                'ansible_groups': ['example', 'example1'],
                'options': {
                    'append_platform_to_hostname': True
                }
            }]
        }
    }

@pytest.fixture()
def vagrant_v1_section_data_using_env():
    return {
        '${MOLECULE_TEST_01}': {
            'platforms': [{
                'name': 'ubuntu',
                'box': '${MOLECULE_TEST_02}/${MOLECULE_TEST_03}'
            }, {
                'name': 'centos7',
                'box': 'centos/7'
            }],
            'providers': [{
                'name': 'virtualbox',
                'type': 'virtualbox'
            }],
            'instances': [{
                'name': 'aio-01',
                'ansible_groups': ['example', '${MOLECULE_TEST_04}'],
                'options': {
                    '${MOLECULE_TEST_05}': True
                }
            }, {
                'name': 'aio-02',
                'ansible_groups': ['example', 'example1'],
                'options': {
                    'append_platform_to_hostname': True
                }
            }]
        }
    }

@pytest.fixture()
def docker_v1_section_data():
    return {
        'docker': {
            'containers': [{
                'name': 'test1',
                'image': 'ubuntu',
                'image_version': 'latest',
                'port_bindings': {
                    80: 80,
                    443: 443
                },
                'environment': {
                    'FOO': 'BAR',
                    'BAZ': 'QUX'
                },
                'options': {
                    'append_platform_to_hostname': True
                },
                'volume_mounts': ['/tmp/test1:/inside:rw'],
                'cap_add': ['SYS_ADMIN', 'SETPCAP'],
                'cap_drop': ['MKNOD'],
                'ansible_groups': ['group1'],
                'network_mode': 'bridge'
            }, {
                'name': 'test2',
                'image': 'ubuntu',
                'links': {
                    'test1': '80',
                },
                'image_version': 'latest',
                'ansible_groups': ['group2'],
                'network_mode': 'bridge',
                'command': '/bin/sh',
                'options': {
                    'append_platform_to_hostname': True
                },
            }]
        }
    }

@pytest.fixture()
def openstack_v1_section_data():
    return {
        'openstack': {
            'instances': [{
                'name': 'aio-01',
                'sshuser': 'root',
                'ansible_groups': ['example', 'example1'],
                'options': {
                    'append_platform_to_hostname': True
                }
            }]
        }
    }

@pytest.fixture()
def ansible_v1_section_data(playbook):
    return {
        'ansible': {
            'timeout': 30,
            'become': True,
            'become_user': False,
            'ask_become_pass': False,
            'ask_vault_pass': False,
            'vault_password_file': False,
            'limit': 'all',
            'verbose': True,
            'diff': True,
            'tags': False,
            'host_key_checking': False,
            'raw_ssh_args': [
                '-o UserKnownHostsFile=/dev/null', '-o IdentitiesOnly=yes',
                '-o ControlMaster=auto', '-o ControlPersist=60s'
            ],
            'config_file': 'config_file',
            'inventory_file': 'inventory_file',
            'playbook': playbook,
            'raw_env_vars': {
                'FOO': 'bar',
                'RETRY_FILES_ENABLED': 0,
            }
        },
        'verifier': {
            'name': 'testinfra',
            'options': {}
        },
        'dependency': {
            'name': 'galaxy',
            'options': {}
        },
        '_disabled': []
    }

@pytest.fixture()
def state_data():
    return {
        'converged': None,
        'created': True,
        'default_platform': None,
        'default_provider': None,
        'driver': None,
        'hosts': {},
        'multiple_platforms': None
    }

@pytest.fixture()
def state_path_with_data(temp_files):
    return temp_files(fixtures=['state_data'])[0]

@pytest.fixture()
def state_path_without_data(tmpdir, request):
    d = tmpdir.mkdir(pytest.helpers.random_string())
    c = d.join(os.extsep.join(('state', 'yml')))

    def cleanup():
        shutil.rmtree(d.strpath)

    request.addfinalizer(cleanup)

    return c.strpath

@pytest.fixture()
def playbook_data():
    return [{'hosts': 'all', 'tasks': [{'command': 'echo'}]}]

@pytest.fixture()
def playbook(temp_files):
    return temp_files(fixtures=['playbook_data'])[0]

@pytest.fixture()
def patched_ansible_playbook(mocker):
    m = mocker.patch('molecule.ansible_playbook.AnsiblePlaybook.execute')
    m.return_value = (None, None)

    return m

@pytest.fixture()
def patched_ansible_galaxy(patched_run_command):
    return patched_run_command

@pytest.fixture
def patched_print_info(mocker):
    return mocker.patch('molecule.util.print_info')

@pytest.fixture
def patched_print_debug(mocker):
    return mocker.patch('molecule.util.print_debug')

@pytest.fixture
def patched_print_error(mocker):
    return mocker.patch('molecule.util.print_error')

@pytest.fixture()
def patched_run_command(mocker):
    return mocker.patch('molecule.util.run_command')
<commit_msg>add network_mode bridge to test1/2 so as not to break link testing, test3 to test none and test4 to test host<commit_after>#  Copyright (c) 2015-2016 Cisco Systems, Inc.
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to
#  deal in the Software without restriction, including without limitation the
#  rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
#  sell copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included in
#  all copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#  DEALINGS IN THE SOFTWARE.

import os
import os.path
import shutil

import pytest

from molecule import config
from molecule import core
from molecule import state

@pytest.helpers.register
def os_split(s):
    rest, tail = os.path.split(s)
    if rest in ('', os.path.sep):
        return tail,
    return os_split(rest) + (tail, )

@pytest.fixture()
def temp_files(tmpdir, request):
    def wrapper(fixtures=[]):
        d = tmpdir.mkdir(pytest.helpers.random_string())
        confs = []
        for index, fixture in enumerate(fixtures):
            c = d.join(os.extsep.join((fixture, 'yml')))
            c.write(request.getfuncargvalue(fixtures[index]))
            confs.append(c.strpath)

        def cleanup():
            shutil.rmtree(d.strpath)

        request.addfinalizer(cleanup)

        return confs

    return wrapper

@pytest.fixture()
def molecule_instance(temp_dir, temp_files, state_path_without_data):
    c = temp_files(fixtures=['molecule_vagrant_v1_config'])
    m = core.Molecule(config.ConfigV1(configs=c), {})
    m.state = state.State(state_file=state_path_without_data)
    m.main()

    return m

@pytest.fixture()
def molecule_instance_with_env_expansion(temp_dir, temp_files,
                                         state_path_without_data):
    os.environ['MOLECULE_TEST_01'] = 'vagrant'
    os.environ['MOLECULE_TEST_02'] = 'ubuntu'
    os.environ['MOLECULE_TEST_03'] = 'trusty64'
    os.environ['MOLECULE_TEST_04'] = 'example1'
    os.environ['MOLECULE_TEST_05'] = 'append_platform_to_hostname'
    c = temp_files(fixtures=['molecule_vagrant_v1_config_using_env'])
    m = core.Molecule(config.ConfigV1(configs=c), {})
    m.state = state.State(state_file=state_path_without_data)
    m.main()

    return m

@pytest.fixture()
def docker_molecule_instance(temp_dir, temp_files, state_path_without_data):
    c = temp_files(fixtures=['molecule_docker_v1_config'])
    m = core.Molecule(config.ConfigV1(configs=c), {})
    m.state = state.State(state_file=state_path_without_data)
    m.main()

    return m

@pytest.fixture()
def openstack_molecule_instance(temp_dir, temp_files, state_path_without_data):
    c = temp_files(fixtures=['molecule_openstack_v1_config'])
    m = core.Molecule(config.ConfigV1(configs=c), {})
    m.state = state.State(state_file=state_path_without_data)
    m.main()

    return m

@pytest.fixture()
def vagrant_molecule_instance(temp_dir, temp_files, state_path_without_data):
    c = temp_files(fixtures=['molecule_vagrant_v1_config'])
    m = core.Molecule(config.ConfigV1(configs=c), {})
    m.state = state.State(state_file=state_path_without_data)
    m.main()

    return m

@pytest.fixture()
def molecule_vagrant_v1_config(molecule_v1_section_data,
                               vagrant_v1_section_data,
                               ansible_v1_section_data):
    return reduce(lambda x, y: config.merge_dicts(x, y), [
        molecule_v1_section_data, vagrant_v1_section_data,
        ansible_v1_section_data
    ])

@pytest.fixture()
def molecule_vagrant_v1_config_using_env(molecule_v1_section_data,
                                         vagrant_v1_section_data_using_env,
                                         ansible_v1_section_data):
    return reduce(lambda x, y: config.merge_dicts(x, y), [
        molecule_v1_section_data, vagrant_v1_section_data_using_env,
        ansible_v1_section_data
    ])

@pytest.fixture()
def molecule_docker_v1_config(molecule_v1_section_data, docker_v1_section_data,
                              ansible_v1_section_data):
    return reduce(lambda x, y: config.merge_dicts(x, y), [
        molecule_v1_section_data, docker_v1_section_data,
        ansible_v1_section_data
    ])

@pytest.fixture()
def molecule_openstack_v1_config(molecule_v1_section_data,
                                 openstack_v1_section_data,
                                 ansible_v1_section_data):
    return reduce(lambda x, y: config.merge_dicts(x, y), [
        molecule_v1_section_data, openstack_v1_section_data,
        ansible_v1_section_data
    ])

@pytest.fixture()
def molecule_v1_section_data(state_path_without_data):
    return {
        'molecule': {
            'ignore_paths': ['.git', '.vagrant', '.molecule'],
            'serverspec_dir': 'spec',
            'testinfra_dir': 'tests',
            'goss_dir': 'tests',
            'goss_playbook': 'test_default.yml',
            'molecule_dir': 'test',
            'state_file': state_path_without_data,
            'vagrantfile_file': 'vagrantfile_file',
            'rakefile_file': 'rakefile_file',
            'vagrantfile_template': 'vagrantfile.j2',
            'raw_ssh_args': [
                '-o StrictHostKeyChecking=no',
                '-o UserKnownHostsFile=/dev/null'
            ],
            'test': {
                'sequence': [
                    'destroy', 'dependency', 'syntax', 'create', 'converge',
                    'idempotence', 'verify'
                ]
            }
        }
    }

@pytest.fixture()
def vagrant_v1_section_data():
    return {
        'vagrant': {
            'platforms': [{
                'name': 'ubuntu',
                'box': 'ubuntu/trusty64'
            }, {
                'name': 'centos7',
                'box': 'centos/7'
            }],
            'providers': [{
                'name': 'virtualbox',
                'type': 'virtualbox'
            }],
            'instances': [{
                'name': 'aio-01',
                'ansible_groups': ['example', 'example1'],
                'options': {
                    'append_platform_to_hostname': True
                }
            }, {
                'name': 'aio-02',
                'ansible_groups': ['example', 'example1'],
                'options': {
                    'append_platform_to_hostname': True
                }
            }]
        }
    }

@pytest.fixture()
def vagrant_v1_section_data_using_env():
    return {
        '${MOLECULE_TEST_01}': {
            'platforms': [{
                'name': 'ubuntu',
                'box': '${MOLECULE_TEST_02}/${MOLECULE_TEST_03}'
            }, {
                'name': 'centos7',
                'box': 'centos/7'
            }],
            'providers': [{
                'name': 'virtualbox',
                'type': 'virtualbox'
            }],
            'instances': [{
                'name': 'aio-01',
                'ansible_groups': ['example', '${MOLECULE_TEST_04}'],
                'options': {
                    '${MOLECULE_TEST_05}': True
                }
            }, {
                'name': 'aio-02',
                'ansible_groups': ['example', 'example1'],
                'options': {
                    'append_platform_to_hostname': True
                }
            }]
        }
    }

@pytest.fixture()
def docker_v1_section_data():
    return {
        'docker': {
            'containers': [{
                'name': 'test1',
                'image': 'ubuntu',
                'image_version': 'latest',
                'port_bindings': {
                    80: 80,
                    443: 443
                },
                'environment': {
                    'FOO': 'BAR',
                    'BAZ': 'QUX'
                },
                'volume_mounts': ['/tmp/test1:/inside:rw'],
                'cap_add': ['SYS_ADMIN', 'SETPCAP'],
                'cap_drop': ['MKNOD'],
                'ansible_groups': ['group1'],
                'network_mode': 'bridge',
                'options': {
                    'append_platform_to_hostname': True
                },
            }, {
                'name': 'test2',
                'image': 'ubuntu',
                'links': {
                    'test1': '80',
                },
                'image_version': 'latest',
                'ansible_groups': ['group2'],
                'network_mode': 'bridge',
                'command': '/bin/sh',
                'options': {
                    'append_platform_to_hostname': True
                },
            }, {
                'name': 'test3',
                'image': 'ubuntu',
                'image_version': 'latest',
                'ansible_groups': ['group1'],
                'network_mode': 'none',
                'command': '/bin/sh',
                'options': {
                    'append_platform_to_hostname': True
                },
            }, {
                'name': 'test4',
                'image': 'ubuntu',
                'image_version': 'latest',
                'ansible_groups': ['group2'],
                'network_mode': 'host',
                'command': '/bin/sh',
                'options': {
                    'append_platform_to_hostname': True
                },
            }]
        }
    }

@pytest.fixture()
def openstack_v1_section_data():
    return {
        'openstack': {
            'instances': [{
                'name': 'aio-01',
                'sshuser': 'root',
                'ansible_groups': ['example', 'example1'],
                'options': {
                    'append_platform_to_hostname': True
                }
            }]
        }
    }

@pytest.fixture()
def ansible_v1_section_data(playbook):
    return {
        'ansible': {
            'timeout': 30,
            'become': True,
            'become_user': False,
            'ask_become_pass': False,
            'ask_vault_pass': False,
            'vault_password_file': False,
            'limit': 'all',
            'verbose': True,
            'diff': True,
            'tags': False,
            'host_key_checking': False,
            'raw_ssh_args': [
                '-o UserKnownHostsFile=/dev/null', '-o IdentitiesOnly=yes',
                '-o ControlMaster=auto', '-o ControlPersist=60s'
            ],
            'config_file': 'config_file',
            'inventory_file': 'inventory_file',
            'playbook': playbook,
            'raw_env_vars': {
                'FOO': 'bar',
                'RETRY_FILES_ENABLED': 0,
            }
        },
        'verifier': {
            'name': 'testinfra',
            'options': {}
        },
        'dependency': {
            'name': 'galaxy',
            'options': {}
        },
        '_disabled': []
    }

@pytest.fixture()
def state_data():
    return {
        'converged': None,
        'created': True,
        'default_platform': None,
        'default_provider': None,
        'driver': None,
        'hosts': {},
        'multiple_platforms': None
    }

@pytest.fixture()
def state_path_with_data(temp_files):
    return temp_files(fixtures=['state_data'])[0]

@pytest.fixture()
def state_path_without_data(tmpdir, request):
    d = tmpdir.mkdir(pytest.helpers.random_string())
    c = d.join(os.extsep.join(('state', 'yml')))

    def cleanup():
        shutil.rmtree(d.strpath)

    request.addfinalizer(cleanup)

    return c.strpath

@pytest.fixture()
def playbook_data():
    return [{'hosts': 'all', 'tasks': [{'command': 'echo'}]}]

@pytest.fixture()
def playbook(temp_files):
    return temp_files(fixtures=['playbook_data'])[0]

@pytest.fixture()
def patched_ansible_playbook(mocker):
    m = mocker.patch('molecule.ansible_playbook.AnsiblePlaybook.execute')
    m.return_value = (None, None)

    return m

@pytest.fixture()
def patched_ansible_galaxy(patched_run_command):
    return patched_run_command

@pytest.fixture
def patched_print_info(mocker):
    return mocker.patch('molecule.util.print_info')

@pytest.fixture
def patched_print_debug(mocker):
    return mocker.patch('molecule.util.print_debug')

@pytest.fixture
def patched_print_error(mocker):
    return mocker.patch('molecule.util.print_error')

@pytest.fixture()
def patched_run_command(mocker):
    return mocker.patch('molecule.util.run_command')
<|endoftext|>