<commit_before>import numpy as np
from scipy.special import expit, logit
from scipy.stats import bernoulli as dist

class Bernoulli(object):
    n_params = 1

    def __init__(self, params):
        self.logit = params[0]
        self.prob = expit(self.logit)
        self.dist = dist(p=self.prob)

    def __getattr__(self, name):
        if name in dir(self.dist):
            return getattr(self.dist, name)
        return None

    def nll(self, Y):
        return -self.dist.logpmf(Y).mean()

    def D_nll(self, Y):
        D_1 = -np.exp(-self.logit) * self.prob
        D_0 = np.exp(-self.logit) * (self.prob ** 2) / (1 - self.prob)
        return (Y * D_1 + (1 - Y) * D_0)[:,np.newaxis]

    def fisher_info(self):
        FI = np.zeros((self.logit.shape[0], 1, 1))
        FI[:, 0, 0] = self.prob * np.exp(-self.logit) / (1 - self.prob)
        return FI

    def fit(Y):
        return np.array([logit(np.mean(Y))])
<commit_msg>Add Brier score, derivative, metric for Bernoulli.<commit_after>import numpy as np
from scipy.special import expit, logit
from scipy.stats import bernoulli as dist

class Bernoulli(object):
    n_params = 1

    def __init__(self, params):
        self.logit = params[0]
        self.prob = expit(self.logit)
        self.dist = dist(p=self.prob)

    def __getattr__(self, name):
        if name in dir(self.dist):
            return getattr(self.dist, name)
        return None

    def nll(self, Y):
        return -self.dist.logpmf(Y).mean()

    def D_nll(self, Y):
        D_1 = -np.exp(-self.logit) * self.prob
        D_0 = np.exp(-self.logit) * (self.prob ** 2) / (1 - self.prob)
        return (Y * D_1 + (1 - Y) * D_0)[:,np.newaxis]

    def fisher_info(self):
        FI = np.zeros((self.logit.shape[0], 1, 1))
        FI[:, 0, 0] = self.prob * np.exp(-self.logit) / (1 - self.prob)
        return FI

    def crps(self, Y):
        return ((self.prob - Y) ** 2).mean()

    def D_crps(self, Y):
        D = 2 * (self.prob - Y) * self.prob ** 2 * np.exp(-self.logit)
        return D[:,np.newaxis]

    def crps_metric(self):
        M = np.zeros((self.logit.shape[0], 1, 1))
        M[:, 0, 0] = 2 * self.prob ** 2 * np.exp(-2 * self.logit) * \
                     (1 + (self.prob / (1 - self.prob)) ** 2)
        return M

    def fit(Y):
        return np.array([logit(np.mean(Y))])
<|endoftext|>


<commit_before>#!/usr/bin/env python
## category Misc
## desc Strips sequences from reference FASTA files
'''
Strips alternative assemblies from reference FASTA files

This removes all sequences by looking for specific substrings in their name
(defaults to '_')

'''

import sys
import os
from eta import eta_open_iter

def strip_fasta(fname, substr='_'):
	good = False
	name = ''
	def callback():
		return name
	for line in eta_open_iter(fname, callback=callback):
		if line[0] == '>':
			name = line[1:].strip().split(' ')[0]
			if substr in name:
				good = False
			else:
				good = True
		if good:
			sys.stdout.write(line)

def usage(msg=None):
	if msg:
		sys.stdout.write('%s\n' % msg)
	sys.stdout.write(__doc__)
	sys.stdout.write('Usage: ngsutils strip_fasta filename.fa [string-to-search-for]\n\n')
	sys.exit(1)

if __name__ == '__main__':
	fname = None
	substr = '_'
	for arg in sys.argv[1:]:
		if not fname:
			if os.path.exists(arg):
				fname = arg
			else:
				usage('%s missing!' % arg)
		else:
			substr = arg

	if not fname:
		usage()

	strip_fasta(fname, substr)<commit_msg>more callback info<commit_after>#!/usr/bin/env python
## category Misc
## desc Strips sequences from reference FASTA files
'''
Strips alternative assemblies from reference FASTA files

This removes all sequences by looking for specific substrings in their name
(defaults to '_')

'''

import sys
import os
from eta import eta_open_iter

def strip_fasta(fname, substr='_'):
	good = False
	name = ''
	def callback():
		if good:
			return name
		else:
			return '%s (removed)' % name

	for line in eta_open_iter(fname, callback=callback):
		if line[0] == '>':
			name = line[1:].strip().split(' ')[0]
			if substr in name:
				good = False
			else:
				good = True
		if good:
			sys.stdout.write(line)

def usage(msg=None):
	if msg:
		sys.stdout.write('%s\n' % msg)
	sys.stdout.write(__doc__)
	sys.stdout.write('Usage: ngsutils strip_fasta filename.fa [string-to-search-for]\n\n')
	sys.exit(1)

if __name__ == '__main__':
	fname = None
	substr = '_'
	for arg in sys.argv[1:]:
		if not fname:
			if os.path.exists(arg):
				fname = arg
			else:
				usage('%s missing!' % arg)
		else:
			substr = arg

	if not fname:
		usage()

	strip_fasta(fname, substr)<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import, unicode_literals
import random
from datetime import datetime

from wechatpy.pay.utils import get_external_ip
from wechatpy.pay.base import BaseWeChatPayAPI

class WeChatRedpack(BaseWeChatPayAPI):

    def send(self, user_id, total_amount, send_name, act_name,
             wishing, remark, total_num=1, client_ip=None,
             nick_name=None, min_value=None,
             max_value=None, out_trade_no=None, logo_imgurl=None):
        """
        发送现金红包

        :param user_id: 接收红包的用户在公众号下的 openid
        :param total_amount: 红包金额，单位分
        :param send_name: 商户名称
        :param nick_name: 可选，提供方名称，默认和商户名称相同
        :param act_name: 活动名称
        :param wishing: 红包祝福语
        :param remark: 备注
        :param client_ip: 可选，调用接口的机器 IP 地址
        :param total_num: 可选，红包发放总人数，默认为 1
        :param min_value: 可选，最小红包金额，单位分
        :param max_value: 可选，最大红包金额，单位分
        :param out_trade_no: 可选，商户订单号，默认会自动生成
        :param logo_imgurl: 可选，商户 Logo 的 URL
        :return: 返回的结果数据字典
        """
        if not out_trade_no:
            now = datetime.now()
            out_trade_no = '{0}{1}{2}'.format(
                self.mch_id,
                now.strftime('%Y%m%d%H%M%S'),
                random.randint(1000, 10000)
            )
        data = {
            'wxappid': self.appid,
            're_openid': user_id,
            'total_amount': total_amount,
            'nick_name': nick_name or send_name,
            'send_name': send_name,
            'act_name': act_name,
            'wishing': wishing,
            'remark': remark,
            'client_ip': client_ip or get_external_ip(),
            'total_num': total_num,
            'min_value': min_value or total_amount,
            'max_value': max_value or total_amount,
            'mch_billno': out_trade_no,
            'logo_imgurl': logo_imgurl,
        }
        return self._post('mmpaymkttransfers/sendredpack', data=data)

    def send_group(self, user_id, total_amount, send_name, act_name, wishing,
                   remark, total_num, client_ip=None, amt_type="ALL_RAND",
                   amt_list=None, out_trade_no=None,
                   logo_imgurl=None, watermark_imgurl=None,
                   banner_imgurl=None):
        """
        发送裂变红包

        :param user_id: 接收红包的用户在公众号下的 openid
        :param total_amount: 红包金额，单位分
        :param send_name: 商户名称
        :param act_name: 活动名称
        :param wishing: 红包祝福语
        :param remark: 备注
        :param total_num: 红包发放总人数
        :param client_ip: 可选，调用接口的机器 IP 地址
        :param amt_type: 可选，红包金额设置方式
                         ALL_RAND—全部随机,商户指定总金额和红包发放总人数，由微信支付随机计算出各红包金额
                         ALL_SPECIFIED—全部自定义
                         SEED_SPECIFIED—种子红包自定义，其他随机
        :param amt_list: 可选，各红包具体金额，自定义金额时必须设置，单位分
        :param out_trade_no: 可选，商户订单号，默认会自动生成
        :param logo_imgurl: 可选，商户 Logo 的 URL
        :param watermark_imgurl: 可选，背景水印图片 URL
        :param banner_imgurl: 红包详情页面的 banner 图片 URL
        :return: 返回的结果数据字典
        """
        if not out_trade_no:
            now = datetime.now()
            out_trade_no = '{0}{1}{2}'.format(
                self._client.mch_id,
                now.strftime('%Y%m%d%H%M%S'),
                random.randint(1000, 10000)
            )
        data = {
            'wxappid': self.appid,
            're_openid': user_id,
            'total_amount': total_amount,
            'send_name': send_name,
            'act_name': act_name,
            'wishing': wishing,
            'remark': remark,
            'total_num': total_num,
            'client_ip': client_ip or get_external_ip(),
            'amt_type': amt_type,
            'amt_list': amt_list,
            'mch_billno': out_trade_no,
            'logo_imgurl': logo_imgurl,
            'watermark_imgurl': watermark_imgurl,
            'banner_imgurl': banner_imgurl,
        }
        return self._post('mmpaymkttransfers/sendgroupredpack', data=data)

    def query(self, out_trade_no, bill_type='MCHT'):
        """
        查询红包发放记录

        :param out_trade_no: 商户订单号
        :param bill_type: 可选，订单类型，目前固定为 MCHT
        :return: 返回的红包发放记录信息
        """
        data = {
            'mch_billno': out_trade_no,
            'bill_type': bill_type,
            'appid': self.appid,
        }
        return self._post('mmpaymkttransfers/gethbinfo', data=data)
<commit_msg>根据最新的微信支付文档增加删除了一些参数 (#201)<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import, unicode_literals
import random
from datetime import datetime

from wechatpy.pay.utils import get_external_ip
from wechatpy.pay.base import BaseWeChatPayAPI

class WeChatRedpack(BaseWeChatPayAPI):

    def send(self, user_id, total_amount, send_name, act_name,
             wishing, remark, total_num=1, client_ip=None,
             out_trade_no=None, scene_id=None, consume_mch_id=None):
        """
        发送现金红包

        :param user_id: 接收红包的用户在公众号下的 openid
        :param total_amount: 红包金额，单位分
        :param send_name: 商户名称
        :param act_name: 活动名称
        :param wishing: 红包祝福语
        :param remark: 备注
        :param client_ip: 可选，调用接口的机器 IP 地址
        :param total_num: 可选，红包发放总人数，默认为 1
        :param out_trade_no: 可选，商户订单号，默认会自动生成
        :param scene_id: 可选，发放红包使用场景，红包金额大于200时必传
        :param consume_mch_id: 可选，资金授权商户号。服务商替特约商户发放时使用
        :return: 返回的结果数据字典
        """
        if not out_trade_no:
            now = datetime.now()
            out_trade_no = '{0}{1}{2}'.format(
                self.mch_id,
                now.strftime('%Y%m%d%H%M%S'),
                random.randint(1000, 10000)
            )
        data = {
            'wxappid': self.appid,
            're_openid': user_id,
            'total_amount': total_amount,
            'send_name': send_name,
            'act_name': act_name,
            'wishing': wishing,
            'remark': remark,
            'client_ip': client_ip or get_external_ip(),
            'total_num': total_num,
            'mch_billno': out_trade_no,
            'scene_id': scene_id,
            'risk_info': None,
            'consume_mch_id': consume_mch_id
        }
        return self._post('mmpaymkttransfers/sendredpack', data=data)

    def send_group(self, user_id, total_amount, send_name, act_name, wishing,
                   remark, total_num, client_ip=None, amt_type="ALL_RAND",
                   out_trade_no=None, scene_id=None, consume_mch_id=None):
        """
        发送裂变红包

        :param user_id: 接收红包的用户在公众号下的 openid
        :param total_amount: 红包金额，单位分
        :param send_name: 商户名称
        :param act_name: 活动名称
        :param wishing: 红包祝福语
        :param remark: 备注
        :param total_num: 红包发放总人数
        :param client_ip: 可选，调用接口的机器 IP 地址
        :param amt_type: 可选，红包金额设置方式
                         ALL_RAND—全部随机,商户指定总金额和红包发放总人数，由微信支付随机计算出各红包金额
        :param out_trade_no: 可选，商户订单号，默认会自动生成
        :param scene_id: 可选，发放红包使用场景，红包金额大于200时必传
        :param consume_mch_id: 可选，资金授权商户号。服务商替特约商户发放时使用
        :return: 返回的结果数据字典
        """
        if not out_trade_no:
            now = datetime.now()
            out_trade_no = '{0}{1}{2}'.format(
                self._client.mch_id,
                now.strftime('%Y%m%d%H%M%S'),
                random.randint(1000, 10000)
            )
        data = {
            'wxappid': self.appid,
            're_openid': user_id,
            'total_amount': total_amount,
            'send_name': send_name,
            'act_name': act_name,
            'wishing': wishing,
            'remark': remark,
            'total_num': total_num,
            'client_ip': client_ip or get_external_ip(),
            'amt_type': amt_type,
            'mch_billno': out_trade_no,
            'scene_id': scene_id,
            'risk_info': None,
            'consume_mch_id': consume_mch_id
        }
        return self._post('mmpaymkttransfers/sendgroupredpack', data=data)

    def query(self, out_trade_no, bill_type='MCHT'):
        """
        查询红包发放记录

        :param out_trade_no: 商户订单号
        :param bill_type: 可选，订单类型，目前固定为 MCHT
        :return: 返回的红包发放记录信息
        """
        data = {
            'mch_billno': out_trade_no,
            'bill_type': bill_type,
            'appid': self.appid,
        }
        return self._post('mmpaymkttransfers/gethbinfo', data=data)
<|endoftext|>


<commit_before>"""
setuptools plugins
"""

import sys
try:
	# Prefer the Python 2 version of configparser because Python 3
	#  backport has issues when non-unicode strings are passed.
	# See issue #10 for details
	import ConfigParser as configparser
except ImportError:
	import configparser

from . import managers

__all__ = ['file_finder', 'version_calc']

def file_finder(dirname="."):
	"""
	Find the files in ``dirname`` under Mercurial version control
	according to the setuptools spec (see
	http://peak.telecommunity.com/DevCenter/setuptools#adding-support-for-other-revision-control-systems
	).
	"""
	import distutils.log
	dirname = dirname or '.'
	try:
		valid_mgrs = managers.HGRepoManager.get_valid_managers(dirname)
		valid_mgrs = managers.HGRepoManager.existing_only(valid_mgrs)
		for mgr in valid_mgrs:
			try:
				return mgr.find_all_files()
			except Exception:
				e = sys.exc_info()[1]
				distutils.log.warn("hgtools.%s could not find files: %s",
					mgr, e)
	except Exception:
		e = sys.exc_info()[1]
		distutils.log.warn("Unexpected error finding valid managers in "
			"hgtools.file_finder_plugin: %s", e)
	return []

def patch_egg_info(force_hg_version=False):
	"""
	A hack to replace egg_info.tagged_version with a wrapped version
	that will use the mercurial version if indicated.

	`force_hg_version` is used for hgtools itself.
	"""
	from setuptools.command.egg_info import egg_info
	from pkg_resources import safe_version
	import functools
	orig_ver = egg_info.tagged_version

	@functools.wraps(orig_ver)
	def tagged_version(self):
		using_hg_version = (
			force_hg_version
			or getattr(self.distribution, 'use_hg_version', False)
		)
		if force_hg_version:
			# disable patched `tagged_version` to avoid affecting
			#  subsequent installs in the same interpreter instance.
			egg_info.tagged_version = orig_ver
		if using_hg_version:
			result = safe_version(self.distribution.get_version())
		else:
			result = orig_ver(self)
		self.tag_build = result
		return result
	egg_info.tagged_version = tagged_version

def _calculate_version(mgr, options):
	default_increment = options.get('increment')
	repo_exists = bool(mgr.find_root())
	return (mgr.get_current_version(default_increment)
		if repo_exists else default_increment)

def calculate_version(options={}):
	# The version is cached in the tag_build value in setup.cfg (so that
	#  sdist packages will have a copy of the version as determined at
	#  the build environment).
	parser = configparser.ConfigParser()
	parser.read('setup.cfg')
	has_tag_build = (parser.has_section('egg_info')
		and 'tag_build' in parser.options('egg_info'))
	if has_tag_build:
		# a cached version is available, so use it.
		version = parser.get('egg_info', 'tag_build')
	else:
		# We don't have a version stored in tag_build, so calculate
		#  the version using an HGRepoManager.
		mgr = managers.HGRepoManager.get_first_valid_manager()
		version_handler = options.get('version_handler', _calculate_version)
		version = version_handler(mgr, options)
	return version

def version_calc(dist, attr, value):
	"""
	Handler for parameter to setup(use_hg_version=value)
	attr should be 'use_hg_version'
	bool(value) should be true to invoke this plugin.
	value may optionally be a dict and supply options to the plugin.
	"""
	if not value or not attr == 'use_hg_version': return
	options = value if isinstance(value, dict) else {}
	dist.metadata.version = calculate_version(options)
	patch_egg_info()
<commit_msg>Catch BaseException in plugin<commit_after>"""
setuptools plugins
"""

import sys
try:
	# Prefer the Python 2 version of configparser because Python 3
	#  backport has issues when non-unicode strings are passed.
	# See issue #10 for details
	import ConfigParser as configparser
except ImportError:
	import configparser

from . import managers

__all__ = ['file_finder', 'version_calc']

def file_finder(dirname="."):
	"""
	Find the files in ``dirname`` under Mercurial version control
	according to the setuptools spec (see
	http://peak.telecommunity.com/DevCenter/setuptools#adding-support-for-other-revision-control-systems
	).
	"""
	import distutils.log
	dirname = dirname or '.'
	try:
		valid_mgrs = managers.HGRepoManager.get_valid_managers(dirname)
		valid_mgrs = managers.HGRepoManager.existing_only(valid_mgrs)
		for mgr in valid_mgrs:
			try:
				return mgr.find_all_files()
			except Exception:
				e = sys.exc_info()[1]
				distutils.log.warn("hgtools.%s could not find files: %s",
					mgr, e)
	except BaseException:
		e = sys.exc_info()[1]
		distutils.log.warn("Unexpected error finding valid managers in "
			"hgtools.file_finder_plugin: %s", e)
	return []

def patch_egg_info(force_hg_version=False):
	"""
	A hack to replace egg_info.tagged_version with a wrapped version
	that will use the mercurial version if indicated.

	`force_hg_version` is used for hgtools itself.
	"""
	from setuptools.command.egg_info import egg_info
	from pkg_resources import safe_version
	import functools
	orig_ver = egg_info.tagged_version

	@functools.wraps(orig_ver)
	def tagged_version(self):
		using_hg_version = (
			force_hg_version
			or getattr(self.distribution, 'use_hg_version', False)
		)
		if force_hg_version:
			# disable patched `tagged_version` to avoid affecting
			#  subsequent installs in the same interpreter instance.
			egg_info.tagged_version = orig_ver
		if using_hg_version:
			result = safe_version(self.distribution.get_version())
		else:
			result = orig_ver(self)
		self.tag_build = result
		return result
	egg_info.tagged_version = tagged_version

def _calculate_version(mgr, options):
	default_increment = options.get('increment')
	repo_exists = bool(mgr.find_root())
	return (mgr.get_current_version(default_increment)
		if repo_exists else default_increment)

def calculate_version(options={}):
	# The version is cached in the tag_build value in setup.cfg (so that
	#  sdist packages will have a copy of the version as determined at
	#  the build environment).
	parser = configparser.ConfigParser()
	parser.read('setup.cfg')
	has_tag_build = (parser.has_section('egg_info')
		and 'tag_build' in parser.options('egg_info'))
	if has_tag_build:
		# a cached version is available, so use it.
		version = parser.get('egg_info', 'tag_build')
	else:
		# We don't have a version stored in tag_build, so calculate
		#  the version using an HGRepoManager.
		mgr = managers.HGRepoManager.get_first_valid_manager()
		version_handler = options.get('version_handler', _calculate_version)
		version = version_handler(mgr, options)
	return version

def version_calc(dist, attr, value):
	"""
	Handler for parameter to setup(use_hg_version=value)
	attr should be 'use_hg_version'
	bool(value) should be true to invoke this plugin.
	value may optionally be a dict and supply options to the plugin.
	"""
	if not value or not attr == 'use_hg_version': return
	options = value if isinstance(value, dict) else {}
	dist.metadata.version = calculate_version(options)
	patch_egg_info()
<|endoftext|>


<commit_before># coding: utf-8
from __future__ import absolute_import, unicode_literals

import functools
import logging
import sys

from celery import shared_task
from django_fsm import TransitionNotAllowed
from django.db import transaction, DatabaseError
import six

from nodeconductor.iaas import models

logger = logging.getLogger(__name__)

class StateChangeError(RuntimeError):
    pass

def _mock_processing(instance_uuid, should_fail=False):
    if should_fail:
        raise Exception('It\'s not my day')

    import time
    time.sleep(10)

    # update some values
    with transaction.atomic():
        try:
            instance = models.Instance.objects.get(uuid=instance_uuid)
            instance.ips = '1.2.3.4, 10.10.10.10'
            instance.save()
        except models.Instance.DoesNotExist:
            raise Exception('Error updating VM instance')

#TODO: extract to core
# noinspection PyProtectedMember
def set_state(model_class, uuid, transition):
    """
    Atomically change state of a model_class instance.

    Handles edge cases:
    * model instance missing in the database
    * concurrent database update of a model instance
    * model instance being in a state that forbids transition

    Raises StateChangeError in case of any of the above.

    Example:

    .. code-block:: python
        # models.py
        from django.db import
        from nodeconductor.core.models import UuidMixin

        class Worker(UuidMixin, models.Model):
            state = FSMField(default='idle')

            @transition(field=state, source='idle', target='working')
            def start_working(self):
                pass

        # views.py
        from django.shortcuts import render_to_response
        from . import models

        def begin_work(worker_uuid):
            try:
                set_state(models.Worker, worker_uuid, 'start_working')
            except StateChangeError:
                return render_to_response('failed to start working')
            else:
                return render_to_response('started working')

    :param model_class: model class of an instance to change state
    :type model_class: django.db.models.Model
    :param uuid: identifier of the model_class instance
    :type uuid: str
    :param transition: name of model's method to trigger transition
    :type transition: str
    :raises: StateChangeError
    """
    logged_operation = transition.replace('_', ' ')
    entity_name = model_class._meta.model_name

    logger.info(
        'About to start %s %s with uuid %s',
        logged_operation, entity_name, uuid
    )

    try:
        with transaction.atomic():
            entity = model_class._default_manager.get(uuid=uuid)

            # TODO: Make sure that the transition method actually exists
            transition = getattr(entity, transition)
            transition()

            entity.save()
    except model_class.DoesNotExist:
        msg = 'Could not perform %s %s with uuid, %s has gone' %\
            (logged_operation, entity_name, uuid, entity_name)
        # There's nothing we can do here to save the state of an entity
        logger.error(msg)

        six.reraise(StateChangeError, StateChangeError(msg), sys.exc_info()[2])
    except DatabaseError:
        # Transaction failed to commit, most likely due to concurrent update
        msg = 'Could not perform %s %s with uuid %s due to concurrent update' %\
              (logged_operation, entity_name, uuid)
        logger.error(msg)

        six.reraise(StateChangeError, StateChangeError(msg), sys.exc_info()[2])
    except TransitionNotAllowed:
        msg = 'Could not perform %s %s with uuid %s, transition not allowed' %\
              (logged_operation, entity.get_state_display(), uuid)
        # Leave the entity intact
        logger.error(msg)
        six.reraise(StateChangeError, StateChangeError(msg), sys.exc_info()[2])

    # TODO: Emit high level event log entry
    logger.info(
        'Managed to finish %s %s with uuid %s',
        logged_operation, entity_name, uuid
    )

def tracked_processing(model_class, processing_state, desired_state, error_state='erred'):
    def decorator(processing_fn):
        @functools.wraps(processing_fn)
        def wrapped(*args, **kwargs):
            # XXX: This is very fragile :(
            try:
                uuid = kwargs['uuid']
            except KeyError:
                uuid = args[0]

            set_entity_state = functools.partial(set_state, model_class, uuid)

            try:
                set_entity_state(processing_state)

                # We should handle all exceptions here so that processing
                # can concentrate on positive flow

                # noinspection PyBroadException
                try:
                    processing_fn(*args, **kwargs)
                except Exception:
                    # noinspection PyProtectedMember
                    logger.exception(
                        'Failed to finish %s %s with uuid %s',
                        processing_state, model_class._meta.model_name, uuid
                    )

                    set_entity_state(error_state)
                else:
                    set_entity_state(desired_state)
            except StateChangeError:
                # No logging is needed since set_state already logged everything
                pass

        return wrapped
    return decorator

@shared_task
@tracked_processing(models.Instance, processing_state='provisioning', desired_state='online')
def schedule_provisioning(instance_uuid):
    _mock_processing(instance_uuid)

@shared_task
@tracked_processing(models.Instance, processing_state='stopping', desired_state='offline')
def schedule_stopping(instance_uuid):
    _mock_processing(instance_uuid)

@shared_task
@tracked_processing(models.Instance, processing_state='starting', desired_state='online')
def schedule_starting(instance_uuid):
    _mock_processing(instance_uuid)

@shared_task
@tracked_processing(models.Instance, processing_state='deleting', desired_state='deleted')
def schedule_deleting(instance_uuid):
    _mock_processing(instance_uuid)
<commit_msg>Fix logging arguments<commit_after># coding: utf-8
from __future__ import absolute_import, unicode_literals

import functools
import logging
import sys

from celery import shared_task
from django_fsm import TransitionNotAllowed
from django.db import transaction, DatabaseError
import six

from nodeconductor.iaas import models

logger = logging.getLogger(__name__)

class StateChangeError(RuntimeError):
    pass

def _mock_processing(instance_uuid, should_fail=False):
    if should_fail:
        raise Exception('It\'s not my day')

    import time
    time.sleep(10)

    # update some values
    with transaction.atomic():
        try:
            instance = models.Instance.objects.get(uuid=instance_uuid)
            instance.ips = '1.2.3.4, 10.10.10.10'
            instance.save()
        except models.Instance.DoesNotExist:
            raise Exception('Error updating VM instance')

#TODO: extract to core
# noinspection PyProtectedMember
def set_state(model_class, uuid, transition):
    """
    Atomically change state of a model_class instance.

    Handles edge cases:
    * model instance missing in the database
    * concurrent database update of a model instance
    * model instance being in a state that forbids transition

    Raises StateChangeError in case of any of the above.

    Example:

    .. code-block:: python
        # models.py
        from django.db import
        from nodeconductor.core.models import UuidMixin

        class Worker(UuidMixin, models.Model):
            state = FSMField(default='idle')

            @transition(field=state, source='idle', target='working')
            def start_working(self):
                pass

        # views.py
        from django.shortcuts import render_to_response
        from . import models

        def begin_work(worker_uuid):
            try:
                set_state(models.Worker, worker_uuid, 'start_working')
            except StateChangeError:
                return render_to_response('failed to start working')
            else:
                return render_to_response('started working')

    :param model_class: model class of an instance to change state
    :type model_class: django.db.models.Model
    :param uuid: identifier of the model_class instance
    :type uuid: str
    :param transition: name of model's method to trigger transition
    :type transition: str
    :raises: StateChangeError
    """
    logged_operation = transition.replace('_', ' ')
    entity_name = model_class._meta.model_name

    logger.info(
        'About to start %s %s with uuid %s',
        logged_operation, entity_name, uuid
    )

    try:
        with transaction.atomic():
            entity = model_class._default_manager.get(uuid=uuid)

            # TODO: Make sure that the transition method actually exists
            transition = getattr(entity, transition)
            transition()

            entity.save()
    except model_class.DoesNotExist:
        msg = 'Could not perform %s %s with uuid %s. Instance has gone' %\
            (logged_operation, entity_name, uuid)
        # There's nothing we can do here to save the state of an entity
        logger.error(msg)

        six.reraise(StateChangeError, StateChangeError(msg), sys.exc_info()[2])
    except DatabaseError:
        # Transaction failed to commit, most likely due to concurrent update
        msg = 'Could not perform %s %s with uuid %s due to concurrent update' %\
              (logged_operation, entity_name, uuid)
        logger.error(msg)

        six.reraise(StateChangeError, StateChangeError(msg), sys.exc_info()[2])
    except TransitionNotAllowed:
        msg = 'Could not perform %s %s with uuid %s, transition not allowed' %\
              (logged_operation, entity_name, uuid)
        # Leave the entity intact
        logger.error(msg)
        six.reraise(StateChangeError, StateChangeError(msg), sys.exc_info()[2])

    # TODO: Emit high level event log entry
    logger.info(
        'Managed to finish %s %s with uuid %s',
        logged_operation, entity_name, uuid
    )

def tracked_processing(model_class, processing_state, desired_state, error_state='erred'):
    def decorator(processing_fn):
        @functools.wraps(processing_fn)
        def wrapped(*args, **kwargs):
            # XXX: This is very fragile :(
            try:
                uuid = kwargs['uuid']
            except KeyError:
                uuid = args[0]

            set_entity_state = functools.partial(set_state, model_class, uuid)

            try:
                set_entity_state(processing_state)

                # We should handle all exceptions here so that processing
                # can concentrate on positive flow

                # noinspection PyBroadException
                try:
                    processing_fn(*args, **kwargs)
                except Exception:
                    # noinspection PyProtectedMember
                    logger.exception(
                        'Failed to finish %s %s with uuid %s',
                        processing_state, model_class._meta.model_name, uuid
                    )

                    set_entity_state(error_state)
                else:
                    set_entity_state(desired_state)
            except StateChangeError:
                # No logging is needed since set_state already logged everything
                pass

        return wrapped
    return decorator

@shared_task
@tracked_processing(models.Instance, processing_state='provisioning', desired_state='online')
def schedule_provisioning(instance_uuid):
    _mock_processing(instance_uuid)

@shared_task
@tracked_processing(models.Instance, processing_state='stopping', desired_state='offline')
def schedule_stopping(instance_uuid):
    _mock_processing(instance_uuid)

@shared_task
@tracked_processing(models.Instance, processing_state='starting', desired_state='online')
def schedule_starting(instance_uuid):
    _mock_processing(instance_uuid)

@shared_task
@tracked_processing(models.Instance, processing_state='deleting', desired_state='deleted')
def schedule_deleting(instance_uuid):
    _mock_processing(instance_uuid)
<|endoftext|>


<commit_before>from python_terraform import *
import pytest
import os
import logging
import re

logging.basicConfig(level=logging.DEBUG)
current_path = os.path.dirname(os.path.realpath(__file__))

STRING_CASES = [
     [
         lambda x: x.generate_cmd_string('apply', 'the_folder',
                                         no_color=IsFlagged),
         "terraform apply -no-color the_folder"
     ],
     [
         lambda x: x.generate_cmd_string('push', 'path', vcs=True,
                                         token='token',
                                         atlas_address='url'),
         "terraform push -vcs=true -token=token -atlas-address=url path"
     ],
 ]

CMD_CASES = [
    ['method', 'expected_output'],
    [
        [
            lambda x: x.cmd('plan', 'var_to_output', no_color=IsFlagged, var={'test_var': 'test'}) ,
            "doesn't need to do anything"
        ]
    ]
]

class TestTerraform(object):
    def teardown_method(self, method):
        """ teardown any state that was previously setup with a setup_method
        call.
        """

        def purge(dir, pattern):
            for f in os.listdir(dir):
                if re.search(pattern, f):
                    if os.path.isfile(f):
                        os.remove(os.path.join(dir, f))

        purge('.', '.tfstate')

    @pytest.mark.parametrize([
                 "method", "expected"
             ], STRING_CASES)
    def test_generate_cmd_string(self, method, expected):
        tf = Terraform()
        result = method(tf)

        strs = expected.split()
        for s in strs:
            assert s in result

    @pytest.mark.parametrize(*CMD_CASES)
    def test_cmd(self, method, expected_output):
        tf = Terraform(working_dir=current_path)
        ret, out, err = method(tf)
        assert expected_output in out
        assert ret == 0

    @pytest.mark.parametrize(
        ("folder", "variables", "var_files", "expected_output"),
        [
            ("var_to_output", {'test_var': 'test'}, None, "test_output=test"),
            ("var_to_output", {'test_list_var': ['c', 'd']}, None, "test_list_output=[c,d]"),
            ("var_to_output", {'test_map_var': {"c": "c", "d": "d"}}, None, "test_map_output={a=ab=bc=cd=d}"),
            ("var_to_output", {'test_map_var': {"c": "c", "d": "d"}}, 'var_to_output/test_map_var.json', "test_map_output={a=ab=bc=cd=de=ef=f}")
        ])
    def test_apply(self, folder, variables, var_files, expected_output):
        tf = Terraform(variables=variables, var_file=var_files)
        ret, out, err = tf.apply(folder)
        assert ret == 0
        assert expected_output in out.replace('\n', '').replace(' ', '')
        assert err == ''

    def test_state_data(self):
        cwd = os.path.join(current_path, 'test_tfstate_file')
        tf = Terraform(working_dir=cwd, state='tfstate.test')
        tf.read_state_file()
        assert tf.tfstate.modules[0]['path'] == ['root']

    def test_pre_load_state_data(self):
        cwd = os.path.join(current_path, 'test_tfstate_file')
        tf = Terraform(working_dir=cwd, state='tfstate.test')
        assert tf.tfstate.modules[0]['path'] == ['root']

    @pytest.mark.parametrize(
        ("folder", 'variables'),
        [
            ("var_to_output", {'test_var': 'test'})
        ]
    )
    def test_override_default(self, folder, variables):
        tf = Terraform(variables=variables)
        ret, out, err = tf.apply(folder, var={'test_var': 'test2'},
                                 no_color=IsNotFlagged)
        out = out.replace('\n', '')
        assert '\x1b[0m\x1b[1m\x1b[32mApply' in out

    def test_get_output(self):
        tf = Terraform(variables={'test_var': 'test'})
        tf.apply('var_to_output')
        assert tf.output('test_output') == 'test'

    def test_destroy(self):
        tf = Terraform(variables={'test_var': 'test'})
        ret, out, err = tf.destroy('var_to_output')
        assert ret == 0
        assert 'Destroy complete! Resources: 0 destroyed.' in out

    def test_fmt(self):
        tf = Terraform(variables={'test_var': 'test'})
        ret, out, err = tf.fmt(diff=True)
        assert ret == 0
<commit_msg>fix test working directory<commit_after>from python_terraform import *
import pytest
import os
import logging
import re

logging.basicConfig(level=logging.DEBUG)
current_path = os.path.dirname(os.path.realpath(__file__))

STRING_CASES = [
     [
         lambda x: x.generate_cmd_string('apply', 'the_folder',
                                         no_color=IsFlagged),
         "terraform apply -no-color the_folder"
     ],
     [
         lambda x: x.generate_cmd_string('push', 'path', vcs=True,
                                         token='token',
                                         atlas_address='url'),
         "terraform push -vcs=true -token=token -atlas-address=url path"
     ],
 ]

CMD_CASES = [
    ['method', 'expected_output'],
    [
        [
            lambda x: x.cmd('plan', 'var_to_output', no_color=IsFlagged, var={'test_var': 'test'}) ,
            "doesn't need to do anything"
        ]
    ]
]

class TestTerraform(object):
    def teardown_method(self, method):
        """ teardown any state that was previously setup with a setup_method
        call.
        """

        def purge(dir, pattern):
            for f in os.listdir(dir):
                if re.search(pattern, f):
                    if os.path.isfile(f):
                        os.remove(os.path.join(dir, f))

        purge('.', '.tfstate')

    @pytest.mark.parametrize([
                 "method", "expected"
             ], STRING_CASES)
    def test_generate_cmd_string(self, method, expected):
        tf = Terraform(working_dir=current_path)
        result = method(tf)

        strs = expected.split()
        for s in strs:
            assert s in result

    @pytest.mark.parametrize(*CMD_CASES)
    def test_cmd(self, method, expected_output):
        tf = Terraform(working_dir=current_path)
        ret, out, err = method(tf)
        assert expected_output in out
        assert ret == 0

    @pytest.mark.parametrize(
        ("folder", "variables", "var_files", "expected_output"),
        [
            ("var_to_output", {'test_var': 'test'}, None, "test_output=test"),
            ("var_to_output", {'test_list_var': ['c', 'd']}, None, "test_list_output=[c,d]"),
            ("var_to_output", {'test_map_var': {"c": "c", "d": "d"}}, None, "test_map_output={a=ab=bc=cd=d}"),
            ("var_to_output", {'test_map_var': {"c": "c", "d": "d"}}, 'var_to_output/test_map_var.json', "test_map_output={a=ab=bc=cd=de=ef=f}")
        ])
    def test_apply(self, folder, variables, var_files, expected_output):
        tf = Terraform(working_dir=current_path, variables=variables, var_file=var_files)
        ret, out, err = tf.apply(folder)
        assert ret == 0
        assert expected_output in out.replace('\n', '').replace(' ', '')
        assert err == ''

    def test_state_data(self):
        cwd = os.path.join(current_path, 'test_tfstate_file')
        tf = Terraform(working_dir=cwd, state='tfstate.test')
        tf.read_state_file()
        assert tf.tfstate.modules[0]['path'] == ['root']

    def test_pre_load_state_data(self):
        cwd = os.path.join(current_path, 'test_tfstate_file')
        tf = Terraform(working_dir=cwd, state='tfstate.test')
        assert tf.tfstate.modules[0]['path'] == ['root']

    @pytest.mark.parametrize(
        ("folder", 'variables'),
        [
            ("var_to_output", {'test_var': 'test'})
        ]
    )
    def test_override_default(self, folder, variables):
        tf = Terraform(working_dir=current_path, variables=variables)
        ret, out, err = tf.apply(folder, var={'test_var': 'test2'},
                                 no_color=IsNotFlagged)
        out = out.replace('\n', '')
        assert '\x1b[0m\x1b[1m\x1b[32mApply' in out

    def test_get_output(self):
        tf = Terraform(working_dir=current_path, variables={'test_var': 'test'})
        tf.apply('var_to_output')
        assert tf.output('test_output') == 'test'

    def test_destroy(self):
        tf = Terraform(working_dir=current_path, variables={'test_var': 'test'})
        ret, out, err = tf.destroy('var_to_output')
        assert ret == 0
        assert 'Destroy complete! Resources: 0 destroyed.' in out

    def test_fmt(self):
        tf = Terraform(working_dir=current_path, variables={'test_var': 'test'})
        ret, out, err = tf.fmt(diff=True)
        assert ret == 0
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

__version__ = '0.2.20'
__author__ = 'Paul Inkenbrandt'
__name__ = 'wellapplication'

from transport import transport
from usgsGis import usgs
from chem import WQP
from graphs import piper, fdc, gantt
import MannKendall
import avgMeths
<commit_msg>version 0.2.21<commit_after># -*- coding: utf-8 -*-

__version__ = '0.2.21'
__author__ = 'Paul Inkenbrandt'
__name__ = 'wellapplication'

from transport import transport
from usgsGis import usgs
from chem import WQP
from graphs import piper, fdc, gantt
import MannKendall
import avgMeths
<|endoftext|>


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

import argparse
import logging.config
import sys
import time

from nodepool import nodedb
from nodepool import nodepool
from prettytable import PrettyTable

class NodePoolCmd(object):
    def __init__(self):
        self.args = None

    def parse_arguments(self):
        parser = argparse.ArgumentParser(description='Node pool.')
        parser.add_argument('-c', dest='config',
                            default='/etc/nodepool/nodepool.yaml',
                            help='path to config file')
        parser.add_argument('--version', dest='version', action='store_true',
                            help='show version')
        parser.add_argument('--debug', dest='debug', action='store_true',
                            help='show DEBUG level logging')

        subparsers = parser.add_subparsers(title='commands',
                                           description='valid commands',
                                           help='additional help')

        cmd_list = subparsers.add_parser('list', help='list nodes')
        cmd_list.set_defaults(func=self.list)
        cmd_image_list = subparsers.add_parser(
            'image-list', help='list images from providers')
        cmd_image_list.set_defaults(func=self.image_list)

        cmd_dib_image_list = subparsers.add_parser(
            'dib-image-list',
            help='list images built with diskimage-builder')
        cmd_dib_image_list.set_defaults(func=self.dib_image_list)

        cmd_image_update = subparsers.add_parser(
            'image-update',
            help='rebuild the image and upload to provider')
        cmd_image_update.add_argument(
            'provider',
            help='provider name (`all` for uploading to all providers)')
        cmd_image_update.add_argument('image', help='image name')
        cmd_image_update.set_defaults(func=self.image_update)

        cmd_image_build = subparsers.add_parser(
            'image-build',
            help='build image using diskimage-builder')
        cmd_image_build.add_argument('image', help='image name')
        cmd_image_build.set_defaults(func=self.image_build)

        cmd_alien_list = subparsers.add_parser(
            'alien-list',
            help='list nodes not accounted for by nodepool')
        cmd_alien_list.set_defaults(func=self.alien_list)
        cmd_alien_list.add_argument('provider', help='provider name',
                                    nargs='?')

        cmd_alien_image_list = subparsers.add_parser(
            'alien-image-list',
            help='list images not accounted for by nodepool')
        cmd_alien_image_list.set_defaults(func=self.alien_image_list)
        cmd_alien_image_list.add_argument('provider', help='provider name',
                                          nargs='?')

        cmd_hold = subparsers.add_parser(
            'hold',
            help='place a node in the HOLD state')
        cmd_hold.set_defaults(func=self.hold)
        cmd_hold.add_argument('id', help='node id')

        cmd_delete = subparsers.add_parser(
            'delete',
            help='place a node in the DELETE state')
        cmd_delete.set_defaults(func=self.delete)
        cmd_delete.add_argument('id', help='node id')
        cmd_delete.add_argument('--now',
                                action='store_true',
                                help='delete the node in the foreground')

        cmd_image_delete = subparsers.add_parser(
            'image-delete',
            help='delete an image')
        cmd_image_delete.set_defaults(func=self.image_delete)
        cmd_image_delete.add_argument('id', help='image id')

        cmd_dib_image_delete = subparsers.add_parser(
            'dib-image-delete',
            help='delete image built with diskimage-builder')
        cmd_dib_image_delete.set_defaults(func=self.dib_image_delete)
        cmd_dib_image_delete.add_argument('id', help='dib image id')

        cmd_image_upload = subparsers.add_parser(
            'image-upload',
            help='upload an image to a provider ')
        cmd_image_upload.set_defaults(func=self.image_upload)
        cmd_image_upload.add_argument(
            'provider',
            help='provider name (`all` for uploading to all providers)',
            nargs='?', default='all')
        cmd_image_upload.add_argument('image', help='image name')

        self.args = parser.parse_args()

    def setup_logging(self):
        if self.args.debug:
            logging.basicConfig(level=logging.DEBUG,
                                format='%(asctime)s %(levelname)s %(name)s: '
                                       '%(message)s')
        else:
            logging.basicConfig(level=logging.INFO,
                                format='%(asctime)s %(levelname)s %(name)s: '
                                       '%(message)s')

    def list(self, node_id=None):
        t = PrettyTable(["ID", "Provider", "AZ", "Label", "Target", "Hostname",
                         "NodeName", "Server ID", "IP", "State",
                         "Age (hours)"])
        t.align = 'l'
        now = time.time()
        with self.pool.getDB().getSession() as session:
            for node in session.getNodes():
                if node_id and node.id != node_id:
                    continue
                t.add_row([node.id, node.provider_name, node.az,
                           node.label_name, node.target_name, node.hostname,
                           node.nodename, node.external_id, node.ip,
                           nodedb.STATE_NAMES[node.state],
                           '%.02f' % ((now - node.state_time) / 3600)])
            print t

    def dib_image_list(self):
        t = PrettyTable(["ID", "Image", "Filename", "Version",
                         "State", "Age (hours)"])
        t.align = 'l'
        now = time.time()
        with self.pool.getDB().getSession() as session:
            for image in session.getDibImages():
                t.add_row([image.id, image.image_name,
                           image.filename, image.version,
                           nodedb.STATE_NAMES[image.state],
                           '%.02f' % ((now - image.state_time) / 3600)])
            print t

    def image_list(self):
        t = PrettyTable(["ID", "Provider", "Image", "Hostname", "Version",
                         "Image ID", "Server ID", "State", "Age (hours)"])
        t.align = 'l'
        now = time.time()
        with self.pool.getDB().getSession() as session:
            for image in session.getSnapshotImages():
                t.add_row([image.id, image.provider_name, image.image_name,
                           image.hostname, image.version,
                           image.external_id, image.server_external_id,
                           nodedb.STATE_NAMES[image.state],
                           '%.02f' % ((now - image.state_time) / 3600)])
            print t

    def image_update(self):
        with self.pool.getDB().getSession() as session:
            self.pool.reconfigureManagers(self.pool.config)
            label = self.pool.config.labels.get(self.args.image)

            def provider_save_image(provider_name):
                """Method to save an image to the given provider.

                This works either by uploading a previously build dib
                image or by spinning up a node and taking a snapshot
                of that node.

                This is a closure around the DB session and label being
                updated.
                """
                if provider_name in label.diskimage_providers:
                    self.pool.uploadImage(session, provider_name,
                                          label.name)
                elif provider_name in label.snapshot_providers:
                    self.pool.updateImage(session, provider_name,
                                          label.name)
                elif provider_name in self.pool.config.providers:
                    raise Exception("Provider %s does not provide label %s"
                                    % (provider_name, label.name))
                else:
                    raise Exception("Provider %s does not exist"
                                    % provider_name)

            if label and self.args.provider == 'all':
                if label.diskimage_providers:
                    self.image_build()
                for provider_name in label.providers:
                    provider_save_image(provider_name)
            elif label:
                if label.diskimage_providers.get(self.args.provider):
                    self.image_build()
                provider_save_image(self.args.provider)
            else:
                raise Exception("Trying to update unknown label: %s"
                                % self.args.image)

    def image_build(self):
        if self.args.image not in self.pool.config.diskimages:
            # only can build disk images, not snapshots
            raise Exception("Trying to build a non disk-image-builder "
                            "image: %s" % self.args.image)

        self.pool.reconfigureImageBuilder()
        self.pool.buildImage(self.pool.config.diskimages[self.args.image])
        self.pool.waitForBuiltImages()

    def image_upload(self):
        self.pool.reconfigureManagers(self.pool.config)
        if not self.args.image in self.pool.config.diskimages:
            # only can build disk images, not snapshots
            raise Exception("Trying to upload a non disk-image-builder "
                            "image: %s" % self.args.image)

        with self.pool.getDB().getSession() as session:
            if self.args.provider == 'all':
                # iterate for all providers listed in label
                for provider in self.pool.config.providers.values():
                    for image in provider.images.values():
                        if self.args.image == image.name:
                            self.pool.uploadImage(session, provider.name,
                                                  self.args.image)
                            break
            else:
                self.pool.uploadImage(session, self.args.provider,
                                      self.args.image)

    def alien_list(self):
        self.pool.reconfigureManagers(self.pool.config)

        t = PrettyTable(["Provider", "Hostname", "Server ID", "IP"])
        t.align = 'l'
        with self.pool.getDB().getSession() as session:
            for provider in self.pool.config.providers.values():
                if (self.args.provider and
                        provider.name != self.args.provider):
                    continue
                manager = self.pool.getProviderManager(provider)

                for server in manager.listServers():
                    if not session.getNodeByExternalID(
                            provider.name, server['id']):
                        t.add_row([provider.name, server['name'], server['id'],
                                   server['public_v4']])
        print t

    def alien_image_list(self):
        self.pool.reconfigureManagers(self.pool.config)

        t = PrettyTable(["Provider", "Name", "Image ID"])
        t.align = 'l'
        with self.pool.getDB().getSession() as session:
            for provider in self.pool.config.providers.values():
                if (self.args.provider and
                        provider.name != self.args.provider):
                    continue
                manager = self.pool.getProviderManager(provider)

                for image in manager.listImages():
                    if image['metadata'].get('image_type') == 'snapshot':
                        if not session.getSnapshotImageByExternalID(
                                provider.name, image['id']):
                            t.add_row([provider.name, image['name'],
                                       image['id']])
        print t

    def hold(self):
        node_id = None
        with self.pool.getDB().getSession() as session:
            node = session.getNode(self.args.id)
            node.state = nodedb.HOLD
            node_id = node.id
        self.list(node_id=node_id)

    def delete(self):
        self.pool.reconfigureManagers(self.pool.config)
        with self.pool.getDB().getSession() as session:
            node = session.getNode(self.args.id)
            if not node:
                print "Node %s not found." % self.args.id
            elif self.args.now:
                self.pool._deleteNode(session, node)
            else:
                node.state = nodedb.DELETE
                self.list(node_id=node.id)

    def dib_image_delete(self):
        self.pool.reconfigureManagers(self.pool.config)
        with self.pool.getDB().getSession() as session:
            dib_image = session.getDibImage(self.args.id)
            self.pool.deleteDibImage(dib_image)

    def image_delete(self):
        self.pool.reconfigureManagers(self.pool.config)
        self.pool.deleteImage(self.args.id)

    def main(self):
        if self.args.version:
            from nodepool.version import version_info as npc_version_info
            print "Nodepool version: %s" % npc_version_info.version_string()
            return(0)

        self.pool = nodepool.NodePool(self.args.config)
        config = self.pool.loadConfig()
        self.pool.reconfigureDatabase(config)
        self.pool.setConfig(config)
        self.args.func()

def main():
    npc = NodePoolCmd()
    npc.parse_arguments()
    npc.setup_logging()
    return npc.main()

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Better image checking in update_image command<commit_after>#!/usr/bin/env python
#
# Copyright 2013 OpenStack Foundation
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

import argparse
import logging.config
import sys
import time

from nodepool import nodedb
from nodepool import nodepool
from prettytable import PrettyTable

class NodePoolCmd(object):
    def __init__(self):
        self.args = None

    def parse_arguments(self):
        parser = argparse.ArgumentParser(description='Node pool.')
        parser.add_argument('-c', dest='config',
                            default='/etc/nodepool/nodepool.yaml',
                            help='path to config file')
        parser.add_argument('--version', dest='version', action='store_true',
                            help='show version')
        parser.add_argument('--debug', dest='debug', action='store_true',
                            help='show DEBUG level logging')

        subparsers = parser.add_subparsers(title='commands',
                                           description='valid commands',
                                           help='additional help')

        cmd_list = subparsers.add_parser('list', help='list nodes')
        cmd_list.set_defaults(func=self.list)
        cmd_image_list = subparsers.add_parser(
            'image-list', help='list images from providers')
        cmd_image_list.set_defaults(func=self.image_list)

        cmd_dib_image_list = subparsers.add_parser(
            'dib-image-list',
            help='list images built with diskimage-builder')
        cmd_dib_image_list.set_defaults(func=self.dib_image_list)

        cmd_image_update = subparsers.add_parser(
            'image-update',
            help='rebuild the image and upload to provider')
        cmd_image_update.add_argument(
            'provider',
            help='provider name (`all` for uploading to all providers)')
        cmd_image_update.add_argument('image', help='image name')
        cmd_image_update.set_defaults(func=self.image_update)

        cmd_image_build = subparsers.add_parser(
            'image-build',
            help='build image using diskimage-builder')
        cmd_image_build.add_argument('image', help='image name')
        cmd_image_build.set_defaults(func=self.image_build)

        cmd_alien_list = subparsers.add_parser(
            'alien-list',
            help='list nodes not accounted for by nodepool')
        cmd_alien_list.set_defaults(func=self.alien_list)
        cmd_alien_list.add_argument('provider', help='provider name',
                                    nargs='?')

        cmd_alien_image_list = subparsers.add_parser(
            'alien-image-list',
            help='list images not accounted for by nodepool')
        cmd_alien_image_list.set_defaults(func=self.alien_image_list)
        cmd_alien_image_list.add_argument('provider', help='provider name',
                                          nargs='?')

        cmd_hold = subparsers.add_parser(
            'hold',
            help='place a node in the HOLD state')
        cmd_hold.set_defaults(func=self.hold)
        cmd_hold.add_argument('id', help='node id')

        cmd_delete = subparsers.add_parser(
            'delete',
            help='place a node in the DELETE state')
        cmd_delete.set_defaults(func=self.delete)
        cmd_delete.add_argument('id', help='node id')
        cmd_delete.add_argument('--now',
                                action='store_true',
                                help='delete the node in the foreground')

        cmd_image_delete = subparsers.add_parser(
            'image-delete',
            help='delete an image')
        cmd_image_delete.set_defaults(func=self.image_delete)
        cmd_image_delete.add_argument('id', help='image id')

        cmd_dib_image_delete = subparsers.add_parser(
            'dib-image-delete',
            help='delete image built with diskimage-builder')
        cmd_dib_image_delete.set_defaults(func=self.dib_image_delete)
        cmd_dib_image_delete.add_argument('id', help='dib image id')

        cmd_image_upload = subparsers.add_parser(
            'image-upload',
            help='upload an image to a provider ')
        cmd_image_upload.set_defaults(func=self.image_upload)
        cmd_image_upload.add_argument(
            'provider',
            help='provider name (`all` for uploading to all providers)',
            nargs='?', default='all')
        cmd_image_upload.add_argument('image', help='image name')

        self.args = parser.parse_args()

    def setup_logging(self):
        if self.args.debug:
            logging.basicConfig(level=logging.DEBUG,
                                format='%(asctime)s %(levelname)s %(name)s: '
                                       '%(message)s')
        else:
            logging.basicConfig(level=logging.INFO,
                                format='%(asctime)s %(levelname)s %(name)s: '
                                       '%(message)s')

    def list(self, node_id=None):
        t = PrettyTable(["ID", "Provider", "AZ", "Label", "Target", "Hostname",
                         "NodeName", "Server ID", "IP", "State",
                         "Age (hours)"])
        t.align = 'l'
        now = time.time()
        with self.pool.getDB().getSession() as session:
            for node in session.getNodes():
                if node_id and node.id != node_id:
                    continue
                t.add_row([node.id, node.provider_name, node.az,
                           node.label_name, node.target_name, node.hostname,
                           node.nodename, node.external_id, node.ip,
                           nodedb.STATE_NAMES[node.state],
                           '%.02f' % ((now - node.state_time) / 3600)])
            print t

    def dib_image_list(self):
        t = PrettyTable(["ID", "Image", "Filename", "Version",
                         "State", "Age (hours)"])
        t.align = 'l'
        now = time.time()
        with self.pool.getDB().getSession() as session:
            for image in session.getDibImages():
                t.add_row([image.id, image.image_name,
                           image.filename, image.version,
                           nodedb.STATE_NAMES[image.state],
                           '%.02f' % ((now - image.state_time) / 3600)])
            print t

    def image_list(self):
        t = PrettyTable(["ID", "Provider", "Image", "Hostname", "Version",
                         "Image ID", "Server ID", "State", "Age (hours)"])
        t.align = 'l'
        now = time.time()
        with self.pool.getDB().getSession() as session:
            for image in session.getSnapshotImages():
                t.add_row([image.id, image.provider_name, image.image_name,
                           image.hostname, image.version,
                           image.external_id, image.server_external_id,
                           nodedb.STATE_NAMES[image.state],
                           '%.02f' % ((now - image.state_time) / 3600)])
            print t

    def image_update(self):
        with self.pool.getDB().getSession() as session:
            self.pool.reconfigureManagers(self.pool.config)
            if self.args.image not in self.pool.config.images_in_use:
                raise Exception("Image specified, %s, is not in use."
                                % self.args.image)

            if self.args.provider == 'all':
                dib_image_built = False
                for provider in self.pool.config.providers.values():
                    image = provider.images.get(self.args.image)
                    if image and image.diskimage:
                        if not dib_image_built:
                            self.image_build()
                            dib_image_built = True
                        self.pool.uploadImage(session, provider.name,
                                              image.name)
                    elif image:
                        self.pool.updateImage(session, provider.name,
                                              image.name)
            else:
                provider = self.pool.config.providers.get(self.args.provider)
                if not provider:
                    raise Exception("Provider %s does not exist"
                                    % self.args.provider)
                image = provider.images.get(self.args.image)
                if image and image.diskimage:
                    self.image_build()
                    self.pool.uploadImage(session, provider.name, image.name)
                elif image:
                    self.pool.updateImage(session, provider.name, image.name)
                else:
                    raise Exception("Image %s not in use by provider %s"
                                    % (self.args.image, self.args.provider))

    def image_build(self):
        if self.args.image not in self.pool.config.diskimages:
            # only can build disk images, not snapshots
            raise Exception("Trying to build a non disk-image-builder "
                            "image: %s" % self.args.image)

        self.pool.reconfigureImageBuilder()
        self.pool.buildImage(self.pool.config.diskimages[self.args.image])
        self.pool.waitForBuiltImages()

    def image_upload(self):
        self.pool.reconfigureManagers(self.pool.config)
        if not self.args.image in self.pool.config.diskimages:
            # only can build disk images, not snapshots
            raise Exception("Trying to upload a non disk-image-builder "
                            "image: %s" % self.args.image)

        with self.pool.getDB().getSession() as session:
            if self.args.provider == 'all':
                # iterate for all providers listed in label
                for provider in self.pool.config.providers.values():
                    for image in provider.images.values():
                        if self.args.image == image.name:
                            self.pool.uploadImage(session, provider.name,
                                                  self.args.image)
                            break
            else:
                self.pool.uploadImage(session, self.args.provider,
                                      self.args.image)

    def alien_list(self):
        self.pool.reconfigureManagers(self.pool.config)

        t = PrettyTable(["Provider", "Hostname", "Server ID", "IP"])
        t.align = 'l'
        with self.pool.getDB().getSession() as session:
            for provider in self.pool.config.providers.values():
                if (self.args.provider and
                        provider.name != self.args.provider):
                    continue
                manager = self.pool.getProviderManager(provider)

                for server in manager.listServers():
                    if not session.getNodeByExternalID(
                            provider.name, server['id']):
                        t.add_row([provider.name, server['name'], server['id'],
                                   server['public_v4']])
        print t

    def alien_image_list(self):
        self.pool.reconfigureManagers(self.pool.config)

        t = PrettyTable(["Provider", "Name", "Image ID"])
        t.align = 'l'
        with self.pool.getDB().getSession() as session:
            for provider in self.pool.config.providers.values():
                if (self.args.provider and
                        provider.name != self.args.provider):
                    continue
                manager = self.pool.getProviderManager(provider)

                for image in manager.listImages():
                    if image['metadata'].get('image_type') == 'snapshot':
                        if not session.getSnapshotImageByExternalID(
                                provider.name, image['id']):
                            t.add_row([provider.name, image['name'],
                                       image['id']])
        print t

    def hold(self):
        node_id = None
        with self.pool.getDB().getSession() as session:
            node = session.getNode(self.args.id)
            node.state = nodedb.HOLD
            node_id = node.id
        self.list(node_id=node_id)

    def delete(self):
        self.pool.reconfigureManagers(self.pool.config)
        with self.pool.getDB().getSession() as session:
            node = session.getNode(self.args.id)
            if not node:
                print "Node %s not found." % self.args.id
            elif self.args.now:
                self.pool._deleteNode(session, node)
            else:
                node.state = nodedb.DELETE
                self.list(node_id=node.id)

    def dib_image_delete(self):
        self.pool.reconfigureManagers(self.pool.config)
        with self.pool.getDB().getSession() as session:
            dib_image = session.getDibImage(self.args.id)
            self.pool.deleteDibImage(dib_image)

    def image_delete(self):
        self.pool.reconfigureManagers(self.pool.config)
        self.pool.deleteImage(self.args.id)

    def main(self):
        if self.args.version:
            from nodepool.version import version_info as npc_version_info
            print "Nodepool version: %s" % npc_version_info.version_string()
            return(0)

        self.pool = nodepool.NodePool(self.args.config)
        config = self.pool.loadConfig()
        self.pool.reconfigureDatabase(config)
        self.pool.setConfig(config)
        self.args.func()

def main():
    npc = NodePoolCmd()
    npc.parse_arguments()
    npc.setup_logging()
    return npc.main()

if __name__ == "__main__":
    sys.exit(main())
<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import division
import roslib
import rospy
import rosparam
import copy
import cv
import cv2
import numpy as np
import threading
import dynamic_reconfigure.server
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image
from std_msgs.msg import Float32, Header, String

from multi_tracker.msg import Contourinfo, Contourlist, DeltaVid
from multi_tracker.msg import Trackedobject, Trackedobjectlist
from multi_tracker.srv import resetBackgroundService
import os

import image_processing

import matplotlib.pyplot as plt

# for basler ace cameras, use camera_aravis
# https://github.com/ssafarik/camera_aravis
# rosrun camera_aravis camnode
# default image: /camera/image_raw

# for firefley cameras, camera1394 does not provide timestamps but otherwise works
# use point grey drivers
# http://wiki.ros.org/pointgrey_camera_driver
# rosrun pointgrey_camera_driver camera_node
# default image: /camera/image_mono

# The main tracking class, a ROS node
class DeCompressor:
    def __init__(self):
        '''
        Default image_topic for:
            Basler ace cameras with camera_aravis driver: camera/image_raw
            Pt Grey Firefly cameras with pt grey driver : camera/image_mono
        '''
        # default parameters (parameter server overides them)

        # initialize the node
        rospy.init_node('delta_decompressor')
        self.nodename = rospy.get_name().rstrip('/')

        # Publishers - publish contours
        self.pubDeltaVid = rospy.Publisher('/camera/image_decompressed', Image, queue_size=30)
        self.subDeltaVid = rospy.Subscriber('/multi_tracker/delta_video', DeltaVid, self.delta_image_callback, queue_size=30)

        self.directory = rospy.get_param('/multi_tracker/delta_video/directory', default='')

        self.cvbridge = CvBridge()

        self.backgroundImage = None
        self.background_img_filename = 'none'

    def delta_image_callback(self, delta_vid):
        if self.background_img_filename != delta_vid.background_image:
            self.background_img_filename = delta_vid.background_image
            basename = os.path.basename(self.background_img_filename)
            directory_with_basename = os.path.join(self.directory, basename)
            self.backgroundImage = cv2.imread(directory_with_basename, cv2.CV_8UC1)
            try:
                self.backgroundImage = self.backgroundImage.reshape([self.backgroundImage.shape[0], self.backgroundImage[1], 1]) # for hydro
            except:
                pass # for indigo

        new_image = copy.copy(self.backgroundImage)

        if delta_vid.values is not None:
            if len(delta_vid.values) > 0:
                try:
                    new_image[delta_vid.xpixels, delta_vid.ypixels, 0] = delta_vid.values # for hydro
                except:
                    new_image[delta_vid.xpixels, delta_vid.ypixels] = delta_vid.values # for indigo

        image_message = self.cvbridge.cv2_to_imgmsg(new_image, encoding="mono8")
        image_message.header.stamp = delta_vid.header.stamp
        self.pubDeltaVid.publish(image_message)

    def Main(self):
        rospy.spin()

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

if __name__ == '__main__':

    decompressor = DeCompressor()
    decompressor.Main()
<commit_msg>corrected framenumber<commit_after>#!/usr/bin/env python
from __future__ import division
import roslib
import rospy
import rosparam
import copy
import cv
import cv2
import numpy as np
import threading
import dynamic_reconfigure.server
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image
from std_msgs.msg import Float32, Header, String

from multi_tracker.msg import Contourinfo, Contourlist, DeltaVid
from multi_tracker.msg import Trackedobject, Trackedobjectlist
from multi_tracker.srv import resetBackgroundService
import os

import image_processing

import matplotlib.pyplot as plt

# for basler ace cameras, use camera_aravis
# https://github.com/ssafarik/camera_aravis
# rosrun camera_aravis camnode
# default image: /camera/image_raw

# for firefley cameras, camera1394 does not provide timestamps but otherwise works
# use point grey drivers
# http://wiki.ros.org/pointgrey_camera_driver
# rosrun pointgrey_camera_driver camera_node
# default image: /camera/image_mono

# The main tracking class, a ROS node
class DeCompressor:
    def __init__(self):
        '''
        Default image_topic for:
            Basler ace cameras with camera_aravis driver: camera/image_raw
            Pt Grey Firefly cameras with pt grey driver : camera/image_mono
        '''
        # default parameters (parameter server overides them)

        # initialize the node
        rospy.init_node('delta_decompressor')
        self.nodename = rospy.get_name().rstrip('/')

        # Publishers - publish contours
        self.pubDeltaVid = rospy.Publisher('/camera/image_decompressed', Image, queue_size=30)
        self.subDeltaVid = rospy.Subscriber('/multi_tracker/delta_video', DeltaVid, self.delta_image_callback, queue_size=30)

        self.directory = rospy.get_param('/multi_tracker/delta_video/directory', default='')

        self.cvbridge = CvBridge()

        self.backgroundImage = None
        self.background_img_filename = 'none'

    def delta_image_callback(self, delta_vid):
        if self.background_img_filename != delta_vid.background_image:
            self.background_img_filename = delta_vid.background_image
            basename = os.path.basename(self.background_img_filename)
            directory_with_basename = os.path.join(self.directory, basename)
            self.backgroundImage = cv2.imread(directory_with_basename, cv2.CV_8UC1)
            try:
                self.backgroundImage = self.backgroundImage.reshape([self.backgroundImage.shape[0], self.backgroundImage[1], 1]) # for hydro
            except:
                pass # for indigo

        new_image = copy.copy(self.backgroundImage)

        if delta_vid.values is not None:
            if len(delta_vid.values) > 0:
                try:
                    new_image[delta_vid.xpixels, delta_vid.ypixels, 0] = delta_vid.values # for hydro
                except:
                    new_image[delta_vid.xpixels, delta_vid.ypixels] = delta_vid.values # for indigo

        image_message = self.cvbridge.cv2_to_imgmsg(new_image, encoding="mono8")
        image_message.header = delta_vid.header
        self.pubDeltaVid.publish(image_message)

    def Main(self):
        rospy.spin()

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

if __name__ == '__main__':

    decompressor = DeCompressor()
    decompressor.Main()
<|endoftext|>


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

import ConfigParser
import os
import socket
import re
import tarfile
import zipfile
import time
try:
   from subprocess import *
   use_popen2 = False
except:
   from popen2 import *
   use_popen2 = True
from cStringIO import StringIO

SUCCESS = 0
FAILURE = 1

# Generic enumerated type
class wf_type:
   def __init__(self,list=[]):
      count = 0
      for item in list.split():
         vars(self)[item] = count
         count = count + 1

   def __setattr__(self, name, value):
      raise ValueError('Modifying values is not allowed')

   def __delattr__(self, name):
      raise ValueError('Deleting entries is not allowed')

# The type of requests from condor to act upon
condor_wf_types = wf_type('get_work prepare_job reply_claim_accept reply_claim_reject update_job_status exit_exit exit_remove exit_hold exit_evict')

# The information about the job
class condor_wf(object):
    def set_type(self, thetype):
           self.__type__ = thetype

    def get_type(self):
        return self.__type__

    type = property(get_type, set_type)

    def set_data(self, string):
           self.__data__ = string

    def get_data(self):
        return self.__data__

    data = property(get_data, set_data)

# General failure exception
class GeneralError(Exception):
   def __init__(self, *msgs):
      self.msgs = msgs

# Configuration exception
class config_err(Exception):
   def __init__(self, msg):
      self.msg = msg

def run_cmd(cmd, args, environ=None):
   """Runs the command specified in 'cmd' with the given 'args' using
      either the subprocess module or the open2 module, depending on which
      is supported.  The subprocess module is favored.  Returns
      (return_code, stdout, stderr)"""
   retcode = -1
   std_out = ''
   std_err = ''
   if environ != None:
      try:
         for var in environ.keys():
            os.environ[var] = environ[var]
      except:
         return ([-1, None, None])

   if use_popen2 == False:
      cmd_list = [cmd]
      cmd_list += args.split(' ')
      obj = Popen(cmd_list, stdout=PIPE, stderr=PIPE)
      (std_out, std_err) = obj.communicate()
      retcode = obj.returncode
   else:
      obj = Popen3('%s %s' % (cmd, args), True)
      retcode = obj.wait()
      try:
         std_out = obj.fromchild.readlines()[0]
         obj.fromchild.close()
      except:
         pass
      try:
         obj.tochild.close()
      except:
         pass
      try:
         std_err = obj.childerr.readlines()[0]
         obj.childerr.close()
      except:
         pass

   if environ != None:
      for var in environ.keys():
         try:
            del os.environ[var]
         except:
            pass
   return (retcode, std_out, std_err)

def read_condor_config(subsys, attr_list):
   """ Uses condor_config_val to look up values in condor's configuration.
       First looks for subsys_param, then for the newer subsys.param.
       Returns map(param, value)"""
   config = {}
   if attr_list == []:
      (rcode, value, stderr) = run_cmd('condor_config_val', '%s' % (subsys))
      if rcode == 0:
         config[subsys.lower()] = value.strip()
      else:
         # Config value not found.  Raise an exception
         raise config_err('"%s" is not defined' % subsys)
   else:
      for attr in attr_list:
         (rcode, value, stderr) = run_cmd('condor_config_val', '%s_%s' % (subsys, attr))
         if rcode == 0:
            config[attr.lower()] = value.strip()
         else:
            # Try the newer <subsys>.param form
            (rcode, value, stderr) = run_cmd('condor_config_val', '%s.%s' % (subsys, attr))
            if rcode == 0:
               config[attr.lower()] = value.strip()
            else:
               # Config value not found.  Raise an exception
               raise config_err('"%s_%s" is not defined' % (subsys, attr))
   return config

def read_config_file(config_file, section):
   """Given configuration file and section names, returns a list of all value
     	 pairs in the config file as a dictionary"""
   parser = ConfigParser.ConfigParser()

   # Parse the config file
   if os.path.exists(config_file) == False:
      raise config_err('%s configuration file does not exist.' % config_file)
   else:
      try:
         parser.readfp(open(config_file))
         items = parser.items(section)
      except Exception, error:
         raise config_err('Problems reading %s.' % config_file, str(error))

      # Take the list of lists and convert into a dictionary
      dict = {}
      for list in items:
         dict[list[0]] = list[1].strip()
      return dict

def socket_read_all(sock):
   """Read all data waiting to be read on the given socket.  The first attempt
      to read will be a blocking call, and all subsequent reads will be
      non-blocking."""
   msg = ''
   old_timeout = sock.gettimeout()
   data = sock.recv(1024)
   sock.settimeout(0.1)
   try:
      while len(data) != 0:
         msg += data
         data = sock.recv(1024)
   except socket.timeout, timeout:
      pass
   except Exception, error:
      close_socket(sock)
      if error != None:
         raise GeneralError('socket error %d: %s' % (error[0], error[1]))
   sock.settimeout(old_timeout)
   return msg

def close_socket(connection):
   """Performs a shutdown and a close on the socket.  Any errors
      are logged to the system logging service.  Raises a GeneralError
      if any errors are encountered"""
   try:
      try:
         # python 2.3 doesn't have SHUT_WR defined, so use it's value (1)
         connection.shutdown(1)
      except Exception, error:
         if error != None:
            raise GeneralError('socket error %d: %s' % (error[0], error[1]))
      try:
         data = connection.recv(4096)
         while len(data) != 0:
            data = connection.recv(4096)
      except:
         pass
   finally:
      connection.close()
      connection = None

def log_messages(level, exception, logger_name=''):
   """Logs messages in the passed GeneralError exception to the
      system logger"""
   logger = logging.getLogger(logger_name)
   for msg in exception.msgs:
      if (msg != ''):
         logger.log(level, msg)

def write_file(filename, data):
   """Writes the given data into the given filename in pieces to account for
      large files"""
   file_ptr = open(filename, 'wb')
   file_str = StringIO(data)
   buffer_length = 2 ** 20
   contents = file_str.read(buffer_length)
   while contents:
      file_ptr.write(contents)
      contents = file_str.read(buffer_length)
   file_ptr.flush()
   file_ptr.close()

def grep(pattern, data):
   """Returns the first instance found of pattern in data.  If pattern
      doesn't exist in data, None is returned.  If data contains groups,
      returns the matching subgroups instead"""
   for line in data.split('\n'):
      match = re.match(pattern, line)
      if match != None:
         break
   if match != None and match.groups() != None:
      found = match.groups()
   else:
      found = match
   return(found)

def zip_extract(filename):
   """Unzips the contents of the filename given, preserving permissions and
      ownership"""
   zip = zipfile.ZipFile(filename, 'r')
   contents = zip.namelist()

   # Loop through the archive names and extract the contents
   # Make sure to create higher level directories before creating a lower
   # level file or directory (if applicable)
   contents.sort()
   for item in contents:
      if item.endswith('/'):
         # This item is all directories, so recursively create them if
         # they don't already exist
         if not os.path.exists(item):
            os.makedirs(item)
      else:
         # This is a directory + file or a file by itself, so create the
         # directory hierarchy first (if applicable) then extract the file.
         dirs = item.split('/')
         dir_structure = ''
         for count in range(0,len(dirs)-1):
            full_path = os.path.join(dir_structure, dirs[count])
            if dirs[count] != '' and not os.path.exists(full_path):
               os.mkdir(full_path)
            dir_structure = full_path

         write_file(item, zip.read(item))

   # Set the perserved permissions and timestamp for the extracted
   # files/directories
   for name in contents:
      info = zip.getinfo(name)
      file_time = time.mktime(info.date_time + (0, 0, -1))
      os.chmod(name, info.external_attr >> 16L)
      os.utime(name, (file_time, file_time))

def tarball_extract(filename):
   """Extracts the tarball (.tar.gz) given"""
   tarball = tarfile.open(filename, 'r:gz')
   for file in tarball.getnames():
      tarball.extract(file)
   tarball.close()
<commit_msg>Added importing of logging into functions.py<commit_after>#   Copyright 2008 Red Hat, Inc.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import ConfigParser
import os
import socket
import re
import tarfile
import zipfile
import logging
import time
try:
   from subprocess import *
   use_popen2 = False
except:
   from popen2 import *
   use_popen2 = True
from cStringIO import StringIO

SUCCESS = 0
FAILURE = 1

# Generic enumerated type
class wf_type:
   def __init__(self,list=[]):
      count = 0
      for item in list.split():
         vars(self)[item] = count
         count = count + 1

   def __setattr__(self, name, value):
      raise ValueError('Modifying values is not allowed')

   def __delattr__(self, name):
      raise ValueError('Deleting entries is not allowed')

# The type of requests from condor to act upon
condor_wf_types = wf_type('get_work prepare_job reply_claim_accept reply_claim_reject update_job_status exit_exit exit_remove exit_hold exit_evict')

# The information about the job
class condor_wf(object):
    def set_type(self, thetype):
           self.__type__ = thetype

    def get_type(self):
        return self.__type__

    type = property(get_type, set_type)

    def set_data(self, string):
           self.__data__ = string

    def get_data(self):
        return self.__data__

    data = property(get_data, set_data)

# General failure exception
class GeneralError(Exception):
   def __init__(self, *msgs):
      self.msgs = msgs

# Configuration exception
class config_err(Exception):
   def __init__(self, msg):
      self.msg = msg

def run_cmd(cmd, args, environ=None):
   """Runs the command specified in 'cmd' with the given 'args' using
      either the subprocess module or the open2 module, depending on which
      is supported.  The subprocess module is favored.  Returns
      (return_code, stdout, stderr)"""
   retcode = -1
   std_out = ''
   std_err = ''
   if environ != None:
      try:
         for var in environ.keys():
            os.environ[var] = environ[var]
      except:
         return ([-1, None, None])

   if use_popen2 == False:
      cmd_list = [cmd]
      cmd_list += args.split(' ')
      obj = Popen(cmd_list, stdout=PIPE, stderr=PIPE)
      (std_out, std_err) = obj.communicate()
      retcode = obj.returncode
   else:
      obj = Popen3('%s %s' % (cmd, args), True)
      retcode = obj.wait()
      try:
         std_out = obj.fromchild.readlines()[0]
         obj.fromchild.close()
      except:
         pass
      try:
         obj.tochild.close()
      except:
         pass
      try:
         std_err = obj.childerr.readlines()[0]
         obj.childerr.close()
      except:
         pass

   if environ != None:
      for var in environ.keys():
         try:
            del os.environ[var]
         except:
            pass
   return (retcode, std_out, std_err)

def read_condor_config(subsys, attr_list):
   """ Uses condor_config_val to look up values in condor's configuration.
       First looks for subsys_param, then for the newer subsys.param.
       Returns map(param, value)"""
   config = {}
   if attr_list == []:
      (rcode, value, stderr) = run_cmd('condor_config_val', '%s' % (subsys))
      if rcode == 0:
         config[subsys.lower()] = value.strip()
      else:
         # Config value not found.  Raise an exception
         raise config_err('"%s" is not defined' % subsys)
   else:
      for attr in attr_list:
         (rcode, value, stderr) = run_cmd('condor_config_val', '%s_%s' % (subsys, attr))
         if rcode == 0:
            config[attr.lower()] = value.strip()
         else:
            # Try the newer <subsys>.param form
            (rcode, value, stderr) = run_cmd('condor_config_val', '%s.%s' % (subsys, attr))
            if rcode == 0:
               config[attr.lower()] = value.strip()
            else:
               # Config value not found.  Raise an exception
               raise config_err('"%s_%s" is not defined' % (subsys, attr))
   return config

def read_config_file(config_file, section):
   """Given configuration file and section names, returns a list of all value
     	 pairs in the config file as a dictionary"""
   parser = ConfigParser.ConfigParser()

   # Parse the config file
   if os.path.exists(config_file) == False:
      raise config_err('%s configuration file does not exist.' % config_file)
   else:
      try:
         parser.readfp(open(config_file))
         items = parser.items(section)
      except Exception, error:
         raise config_err('Problems reading %s.' % config_file, str(error))

      # Take the list of lists and convert into a dictionary
      dict = {}
      for list in items:
         dict[list[0]] = list[1].strip()
      return dict

def socket_read_all(sock):
   """Read all data waiting to be read on the given socket.  The first attempt
      to read will be a blocking call, and all subsequent reads will be
      non-blocking."""
   msg = ''
   old_timeout = sock.gettimeout()
   data = sock.recv(1024)
   sock.settimeout(0.1)
   try:
      while len(data) != 0:
         msg += data
         data = sock.recv(1024)
   except socket.timeout, timeout:
      pass
   except Exception, error:
      close_socket(sock)
      if error != None:
         raise GeneralError('socket error %d: %s' % (error[0], error[1]))
   sock.settimeout(old_timeout)
   return msg

def close_socket(connection):
   """Performs a shutdown and a close on the socket.  Any errors
      are logged to the system logging service.  Raises a GeneralError
      if any errors are encountered"""
   try:
      try:
         # python 2.3 doesn't have SHUT_WR defined, so use it's value (1)
         connection.shutdown(1)
      except Exception, error:
         if error != None:
            raise GeneralError('socket error %d: %s' % (error[0], error[1]))
      try:
         data = connection.recv(4096)
         while len(data) != 0:
            data = connection.recv(4096)
      except:
         pass
   finally:
      connection.close()
      connection = None

def log_messages(level, exception, logger_name=''):
   """Logs messages in the passed GeneralError exception to the
      system logger"""
   logger = logging.getLogger(logger_name)
   for msg in exception.msgs:
      if (msg != ''):
         logger.log(level, msg)

def write_file(filename, data):
   """Writes the given data into the given filename in pieces to account for
      large files"""
   file_ptr = open(filename, 'wb')
   file_str = StringIO(data)
   buffer_length = 2 ** 20
   contents = file_str.read(buffer_length)
   while contents:
      file_ptr.write(contents)
      contents = file_str.read(buffer_length)
   file_ptr.flush()
   file_ptr.close()

def grep(pattern, data):
   """Returns the first instance found of pattern in data.  If pattern
      doesn't exist in data, None is returned.  If data contains groups,
      returns the matching subgroups instead"""
   for line in data.split('\n'):
      match = re.match(pattern, line)
      if match != None:
         break
   if match != None and match.groups() != None:
      found = match.groups()
   else:
      found = match
   return(found)

def zip_extract(filename):
   """Unzips the contents of the filename given, preserving permissions and
      ownership"""
   zip = zipfile.ZipFile(filename, 'r')
   contents = zip.namelist()

   # Loop through the archive names and extract the contents
   # Make sure to create higher level directories before creating a lower
   # level file or directory (if applicable)
   contents.sort()
   for item in contents:
      if item.endswith('/'):
         # This item is all directories, so recursively create them if
         # they don't already exist
         if not os.path.exists(item):
            os.makedirs(item)
      else:
         # This is a directory + file or a file by itself, so create the
         # directory hierarchy first (if applicable) then extract the file.
         dirs = item.split('/')
         dir_structure = ''
         for count in range(0,len(dirs)-1):
            full_path = os.path.join(dir_structure, dirs[count])
            if dirs[count] != '' and not os.path.exists(full_path):
               os.mkdir(full_path)
            dir_structure = full_path

         write_file(item, zip.read(item))

   # Set the perserved permissions and timestamp for the extracted
   # files/directories
   for name in contents:
      info = zip.getinfo(name)
      file_time = time.mktime(info.date_time + (0, 0, -1))
      os.chmod(name, info.external_attr >> 16L)
      os.utime(name, (file_time, file_time))

def tarball_extract(filename):
   """Extracts the tarball (.tar.gz) given"""
   tarball = tarfile.open(filename, 'r:gz')
   for file in tarball.getnames():
      tarball.extract(file)
   tarball.close()
<|endoftext|>