<commit_before># -*- coding: utf-8 -*-
###############################################################################
#
#    Copyright (C) 2001-2014 Micronaet SRL (<http://www.micronaet.it>).
#
#    Original module for stock.move from:
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2010 Tiny SPRL (<http://tiny.be>).
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
###############################################################################

import os
import pdb
import sys
import logging
import openerp
import openerp.netsvc as netsvc
import openerp.addons.decimal_precision as dp
from openerp.osv import fields, osv, expression, orm
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from openerp import SUPERUSER_ID
from openerp import tools
from openerp.tools.translate import _
from openerp.tools.float_utils import float_round as round
from openerp.tools import (DEFAULT_SERVER_DATE_FORMAT,
    DEFAULT_SERVER_DATETIME_FORMAT,
    DATETIME_FORMATS_MAP,
    float_compare)

_logger = logging.getLogger(__name__)

class ResPartnerPricelistProduct(orm.Model):
    """ Class for manage particularity for partner - product
    """
    _name = 'res.partner.pricelist.product'
    _description = 'Partner product'
    _rec_name = 'product_id'
    _order = 'product_id'

    _columns = {
        'product_id': fields.many2one('product.product', 'Product'),
        'alias_id': fields.many2one('product.product', 'Alias'),
        'alias_name': fields.char(
            'Alias testo', size=50,
            help='Forzatura testuale del prodotto alias qualora non esistesse'
                 'il prodotto fisico',
        ),
        'date': fields.date('Data'),
        'deadline': fields.date('Scadenza'),
        # 'sell': fields.boolean('Sell rule'),
        # 'cost': fields.float('Cost', digits=(16, 2)), # for purchase?
        'price': fields.float('Prezzo', digits=(16, 2)),
        # TODO Currency
        'partner_id': fields.many2one('res.partner', 'Partner'),
        'note': fields.text('Note'),

        'pallet_weight': fields.integer('Peso pallet'),
        'packaging_id': fields.many2one(
            'product.packaging', 'Imballo',
            ondelete='set null'),
        # ---------------------------------------------------------------------
        # todo used for calculate packaging? XXX not used for now
        'product_ul': fields.many2one(
            'product.ul', 'Pack',
            ondelete='set null'),
        'load_qty': fields.float('Load q.ty', digits=(16, 2)),
        # ---------------------------------------------------------------------
        }

    _defaults = {
        'date': lambda *a: datetime.now().strftime(
            DEFAULT_SERVER_DATE_FORMAT),
        }

class ResPartner(orm.Model):
    """ Class partner
    """
    _inherit = 'res.partner'

    _columns = {
        'pricelist_product_ids': fields.one2many(
            'res.partner.pricelist.product', 'partner_id',
            'Pricelist products'),
        }

class SaleOrderLine(orm.Model):
    """ Add event for onchange in sale.order.line
    """
    _inherit = 'sale.order.line'

    def product_id_change(
            self, cr, uid, ids, pricelist, product, qty=0,
            uom=False, qty_uos=0, uos=False, name='', partner_id=False,
            lang=False, update_tax=True, date_order=False, packaging=False,
            fiscal_position=False, flag=False, context=None):
        """ Override function for set up extra fields as partner customization
        """
        context = context or {}
        res = super(SaleOrderLine, self).product_id_change(
            cr, uid, ids, pricelist=pricelist, product=product, qty=qty,
            uom=uom, qty_uos=qty_uos, uos=uos, name=name,
            partner_id=partner_id, lang=lang, update_tax=update_tax,
            date_order=date_order, packaging=packaging,
            fiscal_position=fiscal_position, flag=flag, context=context)

        if 'value' not in res:
            res['value'] = {}
        # Reset if partner or product not present:
        if not partner_id or not product:
            res['value'].update({
                'name': '',
                'alias_id': False,
                'price_unit': False,
                'product_packaging': False,
                'pallet_weight': False,
                'load_qty': False,  # todo remove?
                # todo
                })
            return res

        # Used pool:
        partner_pool = self.pool.get('res.partner')

        # Update field instead:
        partner_proxy = partner_pool.browse(
            cr, uid, partner_id, context=context)

        for item in partner_proxy.pricelist_product_ids:
            if item.product_id.id == product:
                name = item.alias_name or item.alias_id.name or \
                       item.product_id.name
                data = {
                    'name': name,
                    'alias_id': item.alias_id.id,
                    'price_unit': item.price,
                    # todo use first if not present in customization?
                    'product_packaging': item.packaging_id.id,
                    'pallet_weight':
                        item.pallet_weight or partner_proxy.pallet_weight,
                    # todo also pallet_weight for company if not present?

                    'load_qty': item.load_qty,  # todo remove?
                }
                res['value'].update(data)

                if 'warning' in res:
                    _logger.error(
                        'Remove warning message: \n%s' %
                        res['warning'].get('message'))
                    del(res['warning'])
                break
        return res

    def set_sale_line_as_default_for_partner(self, cr, uid, ids, context=None):
        """ Default for this partner
        """
        if context is None:
            context = {}
        setup_pool = self.pool.get('res.partner.pricelist.product')

        line = self.browse(cr, uid, ids, context=context)[0]
        partner_id = line.order_id.partner_id.id
        product_id = line.product_id.id

        setup_ids = setup_pool.search(cr, uid, [
            ('partner_id', '=', partner_id),
            ('product_id', '=', product_id),
        ], context=context)
        name = line.name.split(']')[-1].split('\n')[0].strip()

        data = {
            'partner_id': partner_id,
            'product_id': product_id,
            'alias_id': line.alias_id.id,
            'alias_name': name if line.product_id.name != name
            else '',
            'price': line.price_unit,
            'pallet_weight': line.pallet_weight,
            'packaging_id': line.product_packaging.id,
        }
        if context.get('force_only_mrp'):  # MRP not update price!
            _logger.warning('Updating only MRP data!')
            del(data['price'])

        if setup_ids:  # Update setup:
            setup_pool.write(cr, uid, setup_ids[0], data, context=context)
        else:  # Create new record:
            setup_pool.create(cr, uid, data, context=context)
        return True

    _columns = {
        'pallet_weight': fields.integer(
            'Peso pallet',
            help='Caricare per questo prodotto il pallet con questi Kg'),

        # todo remove?
        'load_qty': fields.float('Load q.ty', digits=(16, 2)),
        }

<commit_msg>generic bugfix (fast message)<commit_after># -*- coding: utf-8 -*-
###############################################################################
#
#    Copyright (C) 2001-2014 Micronaet SRL (<http://www.micronaet.it>).
#
#    Original module for stock.move from:
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2010 Tiny SPRL (<http://tiny.be>).
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
###############################################################################

import os
import pdb
import sys
import logging
import openerp
import openerp.netsvc as netsvc
import openerp.addons.decimal_precision as dp
from openerp.osv import fields, osv, expression, orm
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from openerp import SUPERUSER_ID
from openerp import tools
from openerp.tools.translate import _
from openerp.tools.float_utils import float_round as round
from openerp.tools import (DEFAULT_SERVER_DATE_FORMAT,
    DEFAULT_SERVER_DATETIME_FORMAT,
    DATETIME_FORMATS_MAP,
    float_compare)

_logger = logging.getLogger(__name__)

class ResPartnerPricelistProduct(orm.Model):
    """ Class for manage particularity for partner - product
    """
    _name = 'res.partner.pricelist.product'
    _description = 'Partner product'
    _rec_name = 'product_id'
    _order = 'product_id'

    _columns = {
        'product_id': fields.many2one('product.product', 'Product'),
        'alias_id': fields.many2one('product.product', 'Alias'),
        'alias_name': fields.char(
            'Alias testo', size=50,
            help='Forzatura testuale del prodotto alias qualora non esistesse'
                 'il prodotto fisico',
        ),
        'date': fields.date('Data'),
        'deadline': fields.date('Scadenza'),
        # 'sell': fields.boolean('Sell rule'),
        # 'cost': fields.float('Cost', digits=(16, 2)), # for purchase?
        'price': fields.float('Prezzo', digits=(16, 2)),
        # TODO Currency
        'partner_id': fields.many2one('res.partner', 'Partner'),
        'note': fields.text('Note'),

        'pallet_weight': fields.integer('Peso pallet'),
        'packaging_id': fields.many2one(
            'product.packaging', 'Imballo',
            ondelete='set null'),
        # ---------------------------------------------------------------------
        # todo used for calculate packaging? XXX not used for now
        'product_ul': fields.many2one(
            'product.ul', 'Pack',
            ondelete='set null'),
        'load_qty': fields.float('Load q.ty', digits=(16, 2)),
        # ---------------------------------------------------------------------
        }

    _defaults = {
        'date': lambda *a: datetime.now().strftime(
            DEFAULT_SERVER_DATE_FORMAT),
        }

class ResPartner(orm.Model):
    """ Class partner
    """
    _inherit = 'res.partner'

    _columns = {
        'pricelist_product_ids': fields.one2many(
            'res.partner.pricelist.product', 'partner_id',
            'Pricelist products'),
        }

class SaleOrderLine(orm.Model):
    """ Add event for onchange in sale.order.line
    """
    _inherit = 'sale.order.line'

    def product_id_change(
            self, cr, uid, ids, pricelist, product, qty=0,
            uom=False, qty_uos=0, uos=False, name='', partner_id=False,
            lang=False, update_tax=True, date_order=False, packaging=False,
            fiscal_position=False, flag=False, context=None):
        """ Override function for set up extra fields as partner customization
        """
        context = context or {}
        res = super(SaleOrderLine, self).product_id_change(
            cr, uid, ids, pricelist=pricelist, product=product, qty=qty,
            uom=uom, qty_uos=qty_uos, uos=uos, name=name,
            partner_id=partner_id, lang=lang, update_tax=update_tax,
            date_order=date_order, packaging=packaging,
            fiscal_position=fiscal_position, flag=flag, context=context)

        if 'value' not in res:
            res['value'] = {}
        # Reset if partner or product not present:
        if not partner_id or not product:
            res['value'].update({
                'name': '',
                'alias_id': False,
                'price_unit': False,
                'product_packaging': False,
                'pallet_weight': False,
                'load_qty': False,  # todo remove?
                # todo
                })
            return res

        # Used pool:
        partner_pool = self.pool.get('res.partner')

        # Update field instead:
        partner_proxy = partner_pool.browse(
            cr, uid, partner_id, context=context)

        for item in partner_proxy.pricelist_product_ids:
            if item.product_id.id == product:
                name = item.alias_name or item.alias_id.name or \
                       item.product_id.name
                data = {
                    'name': name,
                    'alias_id': item.alias_id.id,
                    'price_unit': item.price,
                    # todo use first if not present in customization?
                    'product_packaging': item.packaging_id.id,
                    'pallet_weight':
                        item.pallet_weight or partner_proxy.pallet_weight,
                    # todo also pallet_weight for company if not present?

                    'load_qty': item.load_qty,  # todo remove?
                }
                res['value'].update(data)

                if 'warning' in res:
                    _logger.error(
                        'Remove warning message: \n%s' %
                        res['warning'].get('message'))
                    del(res['warning'])
                break
        return res

    def set_sale_line_as_default_for_partner(self, cr, uid, ids, context=None):
        """ Default for this partner
        """
        if context is None:
            context = {}
        setup_pool = self.pool.get('res.partner.pricelist.product')

        line = self.browse(cr, uid, ids, context=context)[0]
        partner_id = line.order_id.partner_id.id
        product_id = line.product_id.id

        setup_ids = setup_pool.search(cr, uid, [
            ('partner_id', '=', partner_id),
            ('product_id', '=', product_id),
        ], context=context)
        name = line.name.split(']')[-1].split('\n')[0].strip()

        data = {
            'partner_id': partner_id,
            'product_id': product_id,
            'alias_id': line.alias_id.id,
            'alias_name': name if line.product_id.name != name
            else '',
            'price': line.price_unit,
            'pallet_weight': line.pallet_weight,
            'packaging_id': line.product_packaging.id,
        }

        if setup_ids:  # Update setup:
            if context.get('force_only_mrp'):  # MRP not update price!
                _logger.warning('Updating only MRP data!')
                del (data['price'])

            setup_pool.write(cr, uid, setup_ids[0], data, context=context)
        else:  # Create new record:
            setup_pool.create(cr, uid, data, context=context)
        return True

    _columns = {
        'pallet_weight': fields.integer(
            'Peso pallet',
            help='Caricare per questo prodotto il pallet con questi Kg'),

        # todo remove?
        'load_qty': fields.float('Load q.ty', digits=(16, 2)),
        }

<|endoftext|>


<commit_before>from rest_framework.views import APIView
from models import HaProxyConfigModel
from serializers import HaProxyConfigModelSerializer
from rest_framework.response import Response
from rest_framework.status import HTTP_201_CREATED
from api_core.exceptions import InvalidRequestException, DoesNotExistException, DuplicateEntryException
from django.db import IntegrityError
import json

class HaProxyConfigBuildView(APIView):
    """
    API view handling build process of sections in HaProxy configuration.
    """

    def get(self, request, checksum=None):
        """
        Method responding to GET request lists specific configuration section stored in a database whenever
        unique checksum string is provided. Otherwise list of all configuration sections is returned in response.
        :param request: request data
        :param checksum: unique identifier of configuration sections
        :return: rest_framework.response.Response containing serialized data
        """
        if checksum is not None:
            try:
                config = HaProxyConfigModel.objects.get(checksum=checksum)
                serializer = HaProxyConfigModelSerializer(config)
            except HaProxyConfigModel.DoesNotExist:
                raise DoesNotExistException()
        else:
            config = HaProxyConfigModel.objects.all()
            serializer = HaProxyConfigModelSerializer(config, many=True)
        return Response(serializer.data)

    def post(self, request):
        """
        Method responding to POST request creates new configuration section, validates input and
        stores it in a database.
        :param request: request data
        :return: rest_framework.response.Response containing serialized data
        """
        section = request.DATA.get('section', None)
        section_name = request.DATA.get('section_name', None)
        configuration = request.DATA.get('configuration', None)

        named_sections = ['frontend', 'backend', 'listen']
        if section in named_sections and not all([x is not None for x in [section, section_name, configuration]]):
            raise InvalidRequestException()
        else:
            try:
                json.loads(configuration)  # Should raise ValueError if configuration data are invalid
                config = HaProxyConfigModel(section=section, section_name=section_name, configuration=configuration)
                config.save()
            except IntegrityError:
                raise DuplicateEntryException()
            except ValueError:
                raise InvalidRequestException()

        return Response({'checksum': config.checksum}, status=HTTP_201_CREATED)

class HaProxyConfigGenerateView(APIView):

    def get(self, request):
        result = HaProxyConfigModel.objects.all()
        result.query.group_by = ['section', 'section_name']
        if result:
            for res in result:
                print res.checksum, res.section, res.section_name, res.configuration, res.meta

        return Response()<commit_msg>added sorting mechanism before config file generation<commit_after>from rest_framework.views import APIView
from models import HaProxyConfigModel
from serializers import HaProxyConfigModelSerializer
from rest_framework.response import Response
from rest_framework.status import HTTP_201_CREATED
from api_core.exceptions import InvalidRequestException, DoesNotExistException, DuplicateEntryException
from django.db import IntegrityError
from operator import methodcaller
import json

class HaProxyConfigBuildView(APIView):
    """
    API view handling build process of sections in HaProxy configuration.
    """

    def get(self, request, checksum=None):
        """
        Method responding to GET request lists specific configuration section stored in a database whenever
        unique checksum string is provided. Otherwise list of all configuration sections is returned in response.
        :param request: request data
        :param checksum: unique identifier of configuration sections
        :return: rest_framework.response.Response containing serialized data
        """
        if checksum is not None:
            try:
                config = HaProxyConfigModel.objects.get(checksum=checksum)
                serializer = HaProxyConfigModelSerializer(config)
            except HaProxyConfigModel.DoesNotExist:
                raise DoesNotExistException()
        else:
            config = HaProxyConfigModel.objects.all()
            serializer = HaProxyConfigModelSerializer(config, many=True)
        return Response(serializer.data)

    def post(self, request):
        """
        Method responding to POST request creates new configuration section, validates input and
        stores it in a database.
        :param request: request data
        :return: rest_framework.response.Response containing serialized data
        """
        section = request.DATA.get('section', None)
        section_name = request.DATA.get('section_name', None)
        configuration = request.DATA.get('configuration', None)

        named_sections = ['frontend', 'backend', 'listen']
        if section in named_sections and not all([x is not None for x in [section, section_name, configuration]]):
            raise InvalidRequestException()
        else:
            try:
                json.loads(configuration)  # Should raise ValueError if configuration data are invalid
                config = HaProxyConfigModel(section=section, section_name=section_name, configuration=configuration)
                config.save()
            except IntegrityError:
                raise DuplicateEntryException()
            except ValueError:
                raise InvalidRequestException()

        return Response({'checksum': config.checksum}, status=HTTP_201_CREATED)

class HaProxyConfigGenerateView(APIView):
    """
    API view handling generation process of HaProxy configuration file.
    """

    def get(self, request):
        result = HaProxyConfigModel.objects.all()
        result.query.group_by = ['section', 'section_name']

        if result:
            result = sorted(result, key=methodcaller('get_section_weight'))
            for res in result:
                print res.section, res.section_name

        return Response()<|endoftext|>


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

import yaml
import sys
import os
import subprocess
import datetime
import atexit
import syslog
import argparse

parser = argparse.ArgumentParser(description="Build middleman sites based "
                                             "on changes in git")
parser.add_argument("-f", "--force", help="force the rebuild",
                    action="store_true")
parser.add_argument("-n", "--dry-run", help="do not sync remotely",
                    action="store_true")
parser.add_argument("-d", "--debug", help="show debug output",
                    action="store_true")
parser.add_argument("config_file", help="yaml file for the builder config")

args = parser.parse_args()

def debug_print(message):
    if args.debug:
        print message

def refresh_checkout(checkout_dir):
    os.chdir(checkout_dir)
    subprocess.call(['git', 'fetch', '-q'])

def get_last_commit(checkout_dir):
    os.chdir(checkout_dir)
    r = subprocess.check_output(['git', 'ls-remote', '-q', '.',
                                 'refs/remotes/origin/HEAD'])
    return r.split()[0]

def get_last_commit_submodule(checkout_dir, submodule):
    os.chdir("{}/{}/".format(checkout_dir, submodule))
    r = subprocess.check_output(['git', 'ls-remote', '-q', '.',
                                 'refs/remotes/origin/HEAD'])
    return r.split()[0]

def get_submodules_checkout(checkout_dir):
    os.chdir(checkout_dir)
    result = []
    submodule_status = subprocess.check_output(['git', 'submodule', 'status'])
    for s in submodule_status.split('\n'):
        # there is a empty line at the end...
        if s:
            result.append(s.split()[1])
    return result

def load_config(config_file):
    if not os.path.exists(config_file):
        print "Error %s, do not exist" % config_file
        sys.exit(1)
    if not os.path.isfile(config_file):
        print "Error %s is not a file" % config_file
        sys.exit(1)

    with open(config_file) as f:
        config = yaml.safe_load(f)

    return config

def has_submodules(checkout_dir):
    os.chdir(checkout_dir)
    r = subprocess.check_output(['git', 'submodule', 'status'])
    return len(r) > 0

# TODO complete that
def notify_error(stage, error):
    print error
    sys.exit(3)

if not args.config_file:
    print "This script take only 1 single argument, the config file"
    sys.exit(1)

config = load_config(args.config_file)

if 'name' not in config:
    print "Incorrect config file: {}".format(args.config_file)
    sys.exit(1)

name = config['name']

lock_file = os.path.expanduser('{}/lock_{}'.format(
    os.environ.get('XDG_RUNTIME_DIR', '~'), name))
if os.path.exists(lock_file):
    # TODO verify if the PID in the file still exist
    print "Builder already running, exiting"
    sys.exit(2)

# TODO try/except, show a better error message
fd = os.open(lock_file, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
os.write(fd, str(os.getpid()))
os.close(fd)
atexit.register(os.unlink, lock_file)

status_file = os.path.expanduser('~/status_%s.yml' % name)
status = {}
if os.path.exists(status_file):
    with open(status_file) as f:
        status = yaml.safe_load(f)

checkout_dir = os.path.expanduser("~/%s" % name)
if not os.path.isdir(checkout_dir):
    os.unlink(lock_file)
    print "Checkout not existing, exiting"
    sys.exit(2)

refresh_checkout(checkout_dir)

start_build = False

last_build = datetime.datetime.fromtimestamp(
    int(status.get("last_build", "0")))

if ('regular_rebuild_interval' in config
        and datetime.datetime.now() - last_build
        > datetime.timedelta(hours=config['regular_rebuild_interval'])):
    start_build = True

current_commit = get_last_commit(checkout_dir)
if current_commit != status.get('last_build_commit', ''):
    start_build = True

submodule_commits = status.get('submodule_commits', {})
current_submodule_commits = {}
for submodule in get_submodules_checkout(checkout_dir):
    debug_print('Looking for %s' % submodule)
    current_submodule_commits[submodule] = \
        get_last_commit_submodule(checkout_dir, submodule)

    if current_submodule_commits[submodule] != submodule_commits.get(submodule, ''):
        start_build = True

if args.force:
    start_build = True

if not start_build:
    debug_print("Nothing to build")
    sys.exit(0)

syslog.syslog("Start the build of {}".format(name))

os.chdir(checkout_dir)
if config.get('force_checkout', False):
    subprocess.call(['git', 'checkout', '-f'])
else:
    subprocess.call(['git', 'pull', '--rebase'])

if has_submodules(checkout_dir):
    subprocess.call(['git', 'submodule', 'init'])
    subprocess.call(['git', 'submodule', 'sync'])

if config.get('update_submodule_head', False):
    subprocess.call(['git', 'submodule', 'foreach',
                     '"git pull -qf origin master"'])

os.environ['PATH'] = "/usr/local/bin:/srv/builder/bin:" + os.environ['PATH']
try:
    syslog.syslog("Build of {}: bundle install".format(name))
    result = subprocess.check_output(['bundle', 'install'])
except subprocess.CalledProcessError, C:
    notify_error('install', C.output)

try:
    syslog.syslog("Build of {}: bundle exec middleman build".format(name))
    result = subprocess.check_output(['bundle', 'exec', 'middleman',
                                      'build', '--verbose'])
except subprocess.CalledProcessError, C:
    notify_error('build', C.output)

if not args.dry_run:
    syslog.syslog("Build of {}: start sync".format(name))
    # TODO log the message
    if config['remote']:
        subprocess.call(['rsync', '-e', 'ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i ' + os.path.expanduser('~/.ssh/{}_id.rsa'.format(name)), '--delete-after', '-rqavz', '%s/build/' % checkout_dir, config['remote']])
    else:
        subprocess.call(['bundle', 'exec', 'middleman', 'deploy'])
    syslog.syslog("Build of {}: finish sync".format(name))
else:
    syslog.syslog("Build of {}: not syncing, dry-run".format(name))

status = {}
status['last_build_commit'] = current_commit
status['last_build'] = datetime.datetime.now().strftime("%s")
status['submodule_commits'] = current_submodule_commits

with open(status_file, 'w+') as f:
    f.write(yaml.dump(status, default_flow_style=False))
<commit_msg>Add a 'last_build_human' status, since I can't convert easily the date<commit_after>#!/usr/bin/python
#
# Copyright (c) 2015 Michael Scherer <mscherer@redhat.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# This script trigger a middleman rebuild if upstream git or any submodule
# did change

import yaml
import sys
import os
import subprocess
import datetime
import atexit
import syslog
import argparse

parser = argparse.ArgumentParser(description="Build middleman sites based "
                                             "on changes in git")
parser.add_argument("-f", "--force", help="force the rebuild",
                    action="store_true")
parser.add_argument("-n", "--dry-run", help="do not sync remotely",
                    action="store_true")
parser.add_argument("-d", "--debug", help="show debug output",
                    action="store_true")
parser.add_argument("config_file", help="yaml file for the builder config")

args = parser.parse_args()

def debug_print(message):
    if args.debug:
        print message

def refresh_checkout(checkout_dir):
    os.chdir(checkout_dir)
    subprocess.call(['git', 'fetch', '-q'])

def get_last_commit(checkout_dir):
    os.chdir(checkout_dir)
    r = subprocess.check_output(['git', 'ls-remote', '-q', '.',
                                 'refs/remotes/origin/HEAD'])
    return r.split()[0]

def get_last_commit_submodule(checkout_dir, submodule):
    os.chdir("{}/{}/".format(checkout_dir, submodule))
    r = subprocess.check_output(['git', 'ls-remote', '-q', '.',
                                 'refs/remotes/origin/HEAD'])
    return r.split()[0]

def get_submodules_checkout(checkout_dir):
    os.chdir(checkout_dir)
    result = []
    submodule_status = subprocess.check_output(['git', 'submodule', 'status'])
    for s in submodule_status.split('\n'):
        # there is a empty line at the end...
        if s:
            result.append(s.split()[1])
    return result

def load_config(config_file):
    if not os.path.exists(config_file):
        print "Error %s, do not exist" % config_file
        sys.exit(1)
    if not os.path.isfile(config_file):
        print "Error %s is not a file" % config_file
        sys.exit(1)

    with open(config_file) as f:
        config = yaml.safe_load(f)

    return config

def has_submodules(checkout_dir):
    os.chdir(checkout_dir)
    r = subprocess.check_output(['git', 'submodule', 'status'])
    return len(r) > 0

# TODO complete that
def notify_error(stage, error):
    print error
    sys.exit(3)

if not args.config_file:
    print "This script take only 1 single argument, the config file"
    sys.exit(1)

config = load_config(args.config_file)

if 'name' not in config:
    print "Incorrect config file: {}".format(args.config_file)
    sys.exit(1)

name = config['name']

lock_file = os.path.expanduser('{}/lock_{}'.format(
    os.environ.get('XDG_RUNTIME_DIR', '~'), name))
if os.path.exists(lock_file):
    # TODO verify if the PID in the file still exist
    print "Builder already running, exiting"
    sys.exit(2)

# TODO try/except, show a better error message
fd = os.open(lock_file, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
os.write(fd, str(os.getpid()))
os.close(fd)
atexit.register(os.unlink, lock_file)

status_file = os.path.expanduser('~/status_%s.yml' % name)
status = {}
if os.path.exists(status_file):
    with open(status_file) as f:
        status = yaml.safe_load(f)

checkout_dir = os.path.expanduser("~/%s" % name)
if not os.path.isdir(checkout_dir):
    os.unlink(lock_file)
    print "Checkout not existing, exiting"
    sys.exit(2)

refresh_checkout(checkout_dir)

start_build = False

last_build = datetime.datetime.fromtimestamp(
    int(status.get("last_build", "0")))

if ('regular_rebuild_interval' in config
        and datetime.datetime.now() - last_build
        > datetime.timedelta(hours=config['regular_rebuild_interval'])):
    start_build = True

current_commit = get_last_commit(checkout_dir)
if current_commit != status.get('last_build_commit', ''):
    start_build = True

submodule_commits = status.get('submodule_commits', {})
current_submodule_commits = {}
for submodule in get_submodules_checkout(checkout_dir):
    debug_print('Looking for %s' % submodule)
    current_submodule_commits[submodule] = \
        get_last_commit_submodule(checkout_dir, submodule)

    if current_submodule_commits[submodule] != submodule_commits.get(submodule, ''):
        start_build = True

if args.force:
    start_build = True

if not start_build:
    debug_print("Nothing to build")
    sys.exit(0)

syslog.syslog("Start the build of {}".format(name))

os.chdir(checkout_dir)
if config.get('force_checkout', False):
    subprocess.call(['git', 'checkout', '-f'])
else:
    subprocess.call(['git', 'pull', '--rebase'])

if has_submodules(checkout_dir):
    subprocess.call(['git', 'submodule', 'init'])
    subprocess.call(['git', 'submodule', 'sync'])

if config.get('update_submodule_head', False):
    subprocess.call(['git', 'submodule', 'foreach',
                     '"git pull -qf origin master"'])

os.environ['PATH'] = "/usr/local/bin:/srv/builder/bin:" + os.environ['PATH']
try:
    syslog.syslog("Build of {}: bundle install".format(name))
    result = subprocess.check_output(['bundle', 'install'])
except subprocess.CalledProcessError, C:
    notify_error('install', C.output)

try:
    syslog.syslog("Build of {}: bundle exec middleman build".format(name))
    result = subprocess.check_output(['bundle', 'exec', 'middleman',
                                      'build', '--verbose'])
except subprocess.CalledProcessError, C:
    notify_error('build', C.output)

if not args.dry_run:
    syslog.syslog("Build of {}: start sync".format(name))
    # TODO log the message
    if config['remote']:
        subprocess.call(['rsync', '-e', 'ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i ' + os.path.expanduser('~/.ssh/{}_id.rsa'.format(name)), '--delete-after', '-rqavz', '%s/build/' % checkout_dir, config['remote']])
    else:
        subprocess.call(['bundle', 'exec', 'middleman', 'deploy'])
    syslog.syslog("Build of {}: finish sync".format(name))
else:
    syslog.syslog("Build of {}: not syncing, dry-run".format(name))

status = {}
status['last_build_commit'] = current_commit
status['last_build'] = datetime.datetime.now().strftime("%s")
status['last_build_human'] = datetime.datetime.now().strftime("%c")
status['submodule_commits'] = current_submodule_commits

with open(status_file, 'w+') as f:
    f.write(yaml.dump(status, default_flow_style=False))
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
    extract FileInfo object for local files
'''

import os

from sync_app.sync_utils import get_md5

from sync_app.file_list import FileInfo, FileList

class FileInfoLocal(FileInfo):
    def __init__(self, fn='', md5='', fs=None):
        absfn = os.path.abspath(fn)
        if not os.path.exists(absfn):
            print 'ERROR'
            return False
        _url = 'file://%s' % absfn
        FileInfo.__init__(self, fn=absfn, url=_url, md5=md5, fs=fs)
        pass

    def get_md5(self):
        if os.path.exists(self.filename):
            return get_md5(self.filename)
        else:
            return None

    def get_stat(self):
        if os.path.exists(self.filename):
            return os.stat(self.filename)
        else:
            return None

class FileListLocal(FileList):
    def __init__(self, basepath='', subdirs=None, filelist=None):
        FileList.__init__(self, filelist=filelist)

    def fill_cache_file_list_local(self, directory):
        def parse_dir(arg, path, filelist):
            for fn in filelist:
                fullfn = os.path.abspath('%s/%s' % (path, fn)).replace('//', '/')
                if os.path.isfile(fullfn):
                    finfo = FileInfoLocal(fn=fullfn)
                    arg.append(finfo)

        if type(directory) == str:
            if os.path.isdir(directory):
                os.path.walk(directory, parse_dir, self)
            elif os.path.isfile(directory):
                add_file(directory)
        if type(directory) == list:
            for d in directory:
                if os.path.isdir(d):
                    os.path.walk(d, parse_dir, self)
                elif os.path.isfile(d):
                    add_file(d)
<commit_msg>function name follows convention<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
    extract FileInfo object for local files
'''

import os

from sync_app.sync_utils import get_md5

from sync_app.file_list import FileInfo, FileList

class FileInfoLocal(FileInfo):
    def __init__(self, fn='', md5='', fs=None):
        absfn = os.path.abspath(fn)
        if not os.path.exists(absfn):
            print 'ERROR'
            return False
        _url = 'file://%s' % absfn
        FileInfo.__init__(self, fn=absfn, url=_url, md5=md5, fs=fs)

    def get_md5(self):
        if os.path.exists(self.filename):
            return get_md5(self.filename)
        else:
            return None

    def get_stat(self):
        if os.path.exists(self.filename):
            return os.stat(self.filename)
        else:
            return None

class FileListLocal(FileList):
    def __init__(self, basepath='', subdirs=None, filelist=None):
        FileList.__init__(self, filelist=filelist)

    def fill_file_list_local(self, directory):
        def parse_dir(arg, path, filelist):
            for fn in filelist:
                fullfn = os.path.abspath('%s/%s' % (path, fn)).replace('//', '/')
                if os.path.isfile(fullfn):
                    finfo = FileInfoLocal(fn=fullfn)
                    arg.append(finfo)

        if type(directory) == str:
            if os.path.isdir(directory):
                os.path.walk(directory, parse_dir, self)
            elif os.path.isfile(directory):
                add_file(directory)
        if type(directory) == list:
            for d in directory:
                if os.path.isdir(d):
                    os.path.walk(d, parse_dir, self)
                elif os.path.isfile(d):
                    add_file(d)
<|endoftext|>


<commit_before>"""Interface to a series of frequently used mirth shell commands.

Project setup.py defines entry points for the functions below.

"""
import argparse
import getpass
import os
import tempfile

from pheme.util.config import Config
from pheme.warehouse.mirth_channel_transform import TransformManager

CHANNELS = ('PHEME_hl7_obx_insert',
            'PHEME_http_receiver',
            'PHEME_hl7_visit_insert',
            'PHEME_batchfile_consumer',
            'PHEME_hl7_dx_insert',
            'PHEME_hl7_obr_insert',
            'dump_to_disk',)

class MirthShell(object):
    """Sets up and executes common tasks via the mirth shell

    We have a few common patterns, such as deploying test and
    production versions of the mirth channels.  This manages pulling
    in configuration details and setting up the necessary paths to
    make it scriptable.

    """
    def __init__(self):
        self.config = Config()
        self.mirth_home = self.config.get('mirth', 'mirth_home')
        self.mirth_system_user = self.config.get('mirth', 'mirth_system_user')

    def write_script(self, script_file, imports=[], exports=[],):
        """Generate script for mirth shell to run

        :param script_file: An open file handle to a writable file
        :param imports: The mirth channel(s) XML export to import
        :param exports: The list of (channel name, output path) to
          export.

        Writes the necessary mirth shell instructions to the
        script_file so that the file can be executed via the mirth
        shell.  NB - this takes the liberty of relaxing the 0600 mode
        settings hardcoded in NamedTemporaryFile so mirth shell will
        have read access if running as a different user.

        Codetemplates are also imported or exported (if imports or
        exports are defined) into the same directory as the first named
        import / export in a file named "codetemplates.xml".

        """
        for channel in imports:
            assert(os.path.exists(channel))
            script_file.write("import %s force\n" % channel)
        for channel, output in exports:
            script_file.write("export %s %s\n" % (channel, output))
        if imports:
            codetemplates = os.path.join(os.path.dirname(imports[0]),
                                         "codetemplates.xml")
            script_file.write("importcodetemplates %s\n" % codetemplates)
            script_file.write("deploy\n")
            script_file.write("status\n")
        if exports:
            codetemplates = os.path.join(os.path.dirname(exports[0][1]),
                                          "codetemplates.xml")
            script_file.write("exportcodetemplates %s\n" % codetemplates)
        script_file.flush()
        script_file.seek(0)
        os.chmod(script_file.name, 0644)

    def execute_script(self, script):
        "Execute the given script via the mirth shell"
        try:
            orig_dir = os.getcwd()
            os.chdir(self.mirth_home)
            sudo = ''
            if getpass.getuser() != self.mirth_system_user:
                sudo = 'sudo -H -u %s' % self.mirth_system_user
            command = "%s ./mccommand -s %s" %\
                (sudo, script.name)
            print command
            os.system(command)
        except:
            raise
        finally:
            os.chdir(orig_dir)

def transform_channels():
    """Apply default transform to PHEME channels"""

    doc = """
    Mirth channels can be easily exported in XML format.  This utility
    provides a mechanims to alter an export for subsequent import.
    Useful for altering details such as database name and user
    authentication.

    NB - values defined in the project configuration file will be used
    unless provided as optional arguments.  See
    `pheme.util.config.Config`
    """
    config = Config()
    ap = argparse.ArgumentParser(description=doc)
    ap.add_argument("-d", "--database", dest="db",
                    default=config.get('warehouse', 'database'),
                    help="name of database (overrides "
                    "[warehouse]database)")
    ap.add_argument("-u", "--user", dest="user",
                    default=config.get('warehouse', 'database_user'),
                    help="database user (overrides "
                    "[warehouse]database_user)")
    ap.add_argument("-p", "--password", dest="password",
                    default=config.get('warehouse', 'database_password'),
                    help="database password (overrides [warehouse]"
                    "database_password)")
    ap.add_argument("--input_dir", dest="input_dir",
                    default=config.get('warehouse', 'input_dir'),
                    help="filesystem directory for channel to poll "
                    "(overrides [warehouse]input_dir)")
    ap.add_argument("--output_dir", dest="output_dir",
                    default=config.get('warehouse', 'output_dir'),
                    help="filesystem directory for channel output "
                    "(overrides [warehouse]output_dir)")
    ap.add_argument("--error_dir", dest="error_dir",
                    default=config.get('warehouse', 'error_dir'),
                    help="filesystem directory for channel errors "
                    "(overrides [warehouse]error_dir)")
    ap.add_argument("source_directory",
                    help="directory containing source channel "
                    "definition files")
    ap.add_argument("target_directory",
                    help="directory to write transformed channel "
                    "definition files")
    args = ap.parse_args()
    source_dir = os.path.realpath(args.source_directory)
    target_dir = os.path.realpath(args.target_directory)

    transformer = TransformManager(src=None,
                                   target_dir=target_dir,
                                   options=args)
    for c in CHANNELS:
        transformer.src = os.path.join(source_dir, '%s.xml' % c)
        transformer()

def deploy_channels():
    """Entry point to deploy the channels to mirth on localhost"""
    ap = argparse.ArgumentParser(description="deploy known PHEME channels "
                                 "and code templates to Mirth Connect")
    ap.add_argument("deploy_directory",
                    help="directory containing channel definition files")
    args = ap.parse_args()
    path = os.path.realpath(args.deploy_directory)

    imports = [os.path.join(path, '%s.xml' % c) for c in CHANNELS]
    ms = MirthShell()
    with tempfile.NamedTemporaryFile('w') as script_file:
        ms.write_script(script_file, imports=imports)
        ms.execute_script(script_file)

def export_channels():
    """Entry point to export the PHEME channels to named directory"""
    ap = argparse.ArgumentParser(description="export known PHEME channels "
                                 "and code templates from Mirth Connect")
    ap.add_argument("export_directory",
                    help="directory for exported files")
    args = ap.parse_args()
    path = os.path.realpath(args.export_directory)

    if not os.path.isdir(path):
        raise argparse.ArgumentTypeError("can't access %s" % path)
    exports = []
    for c in CHANNELS:
        exports.append((c, os.path.join(path, '%s.xml' % c)))

    ms = MirthShell()
    with tempfile.NamedTemporaryFile('w') as script_file:
        ms.write_script(script_file, exports=exports)
        ms.execute_script(script_file)
<commit_msg>Copy codetemplates during transformation for easy pickup during install of channels.<commit_after>"""Interface to a series of frequently used mirth shell commands.

Project setup.py defines entry points for the functions below.

"""
import argparse
import getpass
import os
import shutil
import tempfile

from pheme.util.config import Config
from pheme.warehouse.mirth_channel_transform import TransformManager

CHANNELS = ('PHEME_hl7_obx_insert',
            'PHEME_http_receiver',
            'PHEME_hl7_visit_insert',
            'PHEME_batchfile_consumer',
            'PHEME_hl7_dx_insert',
            'PHEME_hl7_obr_insert',
            'dump_to_disk',)

class MirthShell(object):
    """Sets up and executes common tasks via the mirth shell

    We have a few common patterns, such as deploying test and
    production versions of the mirth channels.  This manages pulling
    in configuration details and setting up the necessary paths to
    make it scriptable.

    """
    def __init__(self):
        self.config = Config()
        self.mirth_home = self.config.get('mirth', 'mirth_home')
        self.mirth_system_user = self.config.get('mirth', 'mirth_system_user')

    def write_script(self, script_file, imports=[], exports=[],):
        """Generate script for mirth shell to run

        :param script_file: An open file handle to a writable file
        :param imports: The mirth channel(s) XML export to import
        :param exports: The list of (channel name, output path) to
          export.

        Writes the necessary mirth shell instructions to the
        script_file so that the file can be executed via the mirth
        shell.  NB - this takes the liberty of relaxing the 0600 mode
        settings hardcoded in NamedTemporaryFile so mirth shell will
        have read access if running as a different user.

        Codetemplates are also imported or exported (if imports or
        exports are defined) into the same directory as the first named
        import / export in a file named "codetemplates.xml".

        """
        for channel in imports:
            assert(os.path.exists(channel))
            script_file.write("import %s force\n" % channel)
        for channel, output in exports:
            script_file.write("export %s %s\n" % (channel, output))
        if imports:
            codetemplates = os.path.join(os.path.dirname(imports[0]),
                                         "codetemplates.xml")
            script_file.write("importcodetemplates %s\n" % codetemplates)
            script_file.write("deploy\n")
            script_file.write("status\n")
        if exports:
            codetemplates = os.path.join(os.path.dirname(exports[0][1]),
                                         "codetemplates.xml")
            script_file.write("exportcodetemplates %s\n" % codetemplates)
        script_file.flush()
        script_file.seek(0)
        os.chmod(script_file.name, 0644)

    def execute_script(self, script):
        "Execute the given script via the mirth shell"
        try:
            orig_dir = os.getcwd()
            os.chdir(self.mirth_home)
            sudo = ''
            if getpass.getuser() != self.mirth_system_user:
                sudo = 'sudo -H -u %s' % self.mirth_system_user
            command = "%s ./mccommand -s %s" %\
                (sudo, script.name)
            print command
            os.system(command)
        except:
            raise
        finally:
            os.chdir(orig_dir)

def transform_channels():
    """Apply default transform to PHEME channels"""

    doc = """
    Mirth channels can be easily exported in XML format.  This utility
    provides a mechanims to alter an export for subsequent import.
    Useful for altering details such as database name and user
    authentication.

    NB - values defined in the project configuration file will be used
    unless provided as optional arguments.  See
    `pheme.util.config.Config`
    """
    config = Config()
    ap = argparse.ArgumentParser(description=doc)
    ap.add_argument("-d", "--database", dest="db",
                    default=config.get('warehouse', 'database'),
                    help="name of database (overrides "
                    "[warehouse]database)")
    ap.add_argument("-u", "--user", dest="user",
                    default=config.get('warehouse', 'database_user'),
                    help="database user (overrides "
                    "[warehouse]database_user)")
    ap.add_argument("-p", "--password", dest="password",
                    default=config.get('warehouse', 'database_password'),
                    help="database password (overrides [warehouse]"
                    "database_password)")
    ap.add_argument("--input_dir", dest="input_dir",
                    default=config.get('warehouse', 'input_dir'),
                    help="filesystem directory for channel to poll "
                    "(overrides [warehouse]input_dir)")
    ap.add_argument("--output_dir", dest="output_dir",
                    default=config.get('warehouse', 'output_dir'),
                    help="filesystem directory for channel output "
                    "(overrides [warehouse]output_dir)")
    ap.add_argument("--error_dir", dest="error_dir",
                    default=config.get('warehouse', 'error_dir'),
                    help="filesystem directory for channel errors "
                    "(overrides [warehouse]error_dir)")
    ap.add_argument("source_directory",
                    help="directory containing source channel "
                    "definition files")
    ap.add_argument("target_directory",
                    help="directory to write transformed channel "
                    "definition files")
    args = ap.parse_args()
    source_dir = os.path.realpath(args.source_directory)
    target_dir = os.path.realpath(args.target_directory)

    transformer = TransformManager(src=None,
                                   target_dir=target_dir,
                                   options=args)
    for c in CHANNELS:
        transformer.src = os.path.join(source_dir, '%s.xml' % c)
        transformer()
    # no transformation on codetemplates at this time - but the
    # importer expects the codetemplates.xml file to be in the same
    # directory, so copy it over.
    shutil.copy(os.path.join(source_dir, 'codetemplates.xml'), target_dir)

def deploy_channels():
    """Entry point to deploy the channels to mirth on localhost"""
    ap = argparse.ArgumentParser(description="deploy known PHEME channels "
                                 "and code templates to Mirth Connect")
    ap.add_argument("deploy_directory",
                    help="directory containing channel definition files")
    args = ap.parse_args()
    path = os.path.realpath(args.deploy_directory)

    imports = [os.path.join(path, '%s.xml' % c) for c in CHANNELS]
    ms = MirthShell()
    with tempfile.NamedTemporaryFile('w') as script_file:
        ms.write_script(script_file, imports=imports)
        ms.execute_script(script_file)

def export_channels():
    """Entry point to export the PHEME channels to named directory"""
    ap = argparse.ArgumentParser(description="export known PHEME channels "
                                 "and code templates from Mirth Connect")
    ap.add_argument("export_directory",
                    help="directory for exported files")
    args = ap.parse_args()
    path = os.path.realpath(args.export_directory)

    if not os.path.isdir(path):
        raise argparse.ArgumentTypeError("can't access %s" % path)
    exports = []
    for c in CHANNELS:
        exports.append((c, os.path.join(path, '%s.xml' % c)))

    ms = MirthShell()
    with tempfile.NamedTemporaryFile('w') as script_file:
        ms.write_script(script_file, exports=exports)
        ms.execute_script(script_file)
<|endoftext|>


<commit_before>from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
from .retrieve import HandlerBase
from .readers.spe import PrincetonSPEFile
from collections import deque
# import six
# from six.moves import range
import logging
import h5py
import numpy as np
import os.path

logger = logging.getLogger(__name__)

class AreaDetectorSPEHandler(HandlerBase):
    specs = {'AD_SPE'} | HandlerBase.specs

    def __init__(self, fpath, template, filename,
                 frame_per_point=1):
        self._path = fpath
        self._fpp = frame_per_point
        self._template = template
        self._filename = filename
        self._f_cache = dict()

    def __call__(self, point_number):
        # The file number is just the point_number

        out_stack = deque()
        if point_number not in self._f_cache:
            fname = self._template % (self._path,
                                      self._filename,
                                      point_number)
            spe_obj = PrincetonSPEFile(fname)
            self._f_cache[point_number] = spe_obj

        spe = self._f_cache[point_number]

        if len(spe) > 1:
            out_stack.append(spe.getData())
        else:
            out_stack.append(spe[0])

        # return stacked and squeezed results
        return np.dstack(out_stack).squeeze()

class DummyAreaDetectorHandler(HandlerBase):
    def __init__(self, fpath, frame_per_point=1, **kwargs):

        self._fpp = frame_per_point

    def __call__(self, **kwargs):
        out_stack = np.ones((self._fpp, 10, 10)) * np.nan
        # return stacked and squeezed results
        return out_stack.squeeze()

class _HDF5HandlerBase(HandlerBase):

    def open(self):
        if self._file:
            return
        self._file = h5py.File(self._filename)

    def close(self):
        super(_HDF5HandlerBase, self).close()
        self._file.close()
        self._file = None

class HDF5DatasetSliceHandler(_HDF5HandlerBase):
    """
    Handler for data stored in one Dataset of an HDF5 file.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    key : string
        key of the single HDF5 Dataset used by this Handler
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    def __init__(self, filename, key, frame_per_point=1):
        self._fpp = frame_per_point
        self._filename = filename
        self._key = key
        self._file = None
        self._dataset = None
        self.open()

    def __call__(self, point_number):
        # Don't read out the dataset until it is requested for the first time.
        if not self._dataset:
            self._dataset = self._file[self._key]
        start, stop = point_number * self._fpp, (point_number + 1) * self._fpp
        return self._dataset[start:stop].squeeze()

class AreaDetectorHDF5Handler(HDF5DatasetSliceHandler):
    """
    Handler for the 'AD_HDF5' spec used by Area Detectors.

    In this spec, the key (i.e., HDF5 dataset path) is always
    '/entry/data/data'.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    specs = {'AD_HDF5'} | HDF5DatasetSliceHandler.specs

    def __init__(self, filename, frame_per_point=1):
        hardcoded_key = '/entry/data/data'
        super(AreaDetectorHDF5Handler, self).__init__(
            filename=filename, key=hardcoded_key,
            frame_per_point=frame_per_point)

class _HdfMapsHandlerBase(_HDF5HandlerBase):
    """
    Reader for XRF data stored in hdf5 files.

    The data set is assumed to be in a group called MAPS and stored
    as a 3D array ordered [energy, x, y].

    Parameters
    ----------
    filename : str
        Path to physical location of file
    dset_path : str
        The path to the dataset inside of 'MAPS'
    """
    def __init__(self, filename, dset_path):
        self._filename = filename
        self._dset_path = dset_path
        self._file = None
        self._dset = None
        self.open()

    def open(self):
        """
        Open the file for reading.

        Provided as a stand alone function to allow re-opening of the handler
        """
        if self._file:
            return
        self._file = h5py.File(self._filename, mode='r')
        self._dset = self._file['/'.join(['MAPS', self._dset_path])]

    def __call__(self):

        if not self._file:
            raise RuntimeError("File is not open")

class HDFMapsSpectrumHandler(_HdfMapsHandlerBase):
    """
    Handler which selects energy spectrum from
    a MAPS XRF data product.
    """
    specs = {'MAPS_SPECTRUM'} | _HdfMapsHandlerBase.specs

    def __call__(self, x, y):
        """
        Return the spectrum at the x, y position

        Parameters
        ----------
        x : int
            raster index in the x direction

        y : int
            raster index in the y direction

        Returns
        -------
        spectrum : ndarray
            The MCA channels
        """
        super(HDFMapsSpectrumHandler, self).__call__()
        return self._dset[:, x, y]

class HDFMapsEnergyHandler(_HdfMapsHandlerBase):
    """
    Handler which select fixed-energy slices from
    a MAPS XRF data file.
    """
    specs = {'MAPS_PLANE'} | _HdfMapsHandlerBase.specs

    def __call__(self, e_index):
        """
        Return the raster plane at a fixed energy

        Parameters
        ----------
        e_index : int
            The index of the engery

        Returns
        -------
        plane : ndarray
            The raster image at a fixed energy.
        """
        super(HDFMapsEnergyHandler, self).__call__()
        return self._dset[e_index, :, :]

class NpyHandler(HandlerBase):
    """
    Class to deal with reading npy files

    Parameters
    ----------
    fpath : str
        Path to file

    mmap_mode : {'r', 'r+', c}, optional
        memmap mode to use to open file
    """
    specs = {'npy'} | HandlerBase.specs

    def __init__(self, filename, mmap_mode=None):
        self._mmap_mode = mmap_mode
        if not os.path.exists(filename):
            raise IOError("the requested file {fpath} does not exst")
        self._fpath = filename

    def __call__(self):
        return np.load(self._fpath, self._mmap_mode)
<commit_msg>ENH : Added multiple files<commit_after>from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
from .retrieve import HandlerBase
from .readers.spe import PrincetonSPEFile
from collections import deque
# import six
# from six.moves import range
import logging
import h5py
import numpy as np
import os.path

logger = logging.getLogger(__name__)

class AreaDetectorSPEHandler(HandlerBase):
    specs = {'AD_SPE'} | HandlerBase.specs

    def __init__(self, fpath, template, filename,
                 frame_per_point=1):
        self._path = fpath
        self._fpp = frame_per_point
        self._template = template
        self._filename = filename
        self._f_cache = dict()

    def __call__(self, point_number):
        # The file number is just the point_number
        start, stop = point_number * self._fpp, (point_number + 1) * self._fpp
        out_stack = deque()
        for file_number in range(start, stop):
            if file_number not in self._f_cache:
                fname = self._template % (self._path,
                                          self._filename,
                                          file_number)
                spe_obj = PrincetonSPEFile(fname)
                self._f_cache[point_number] = spe_obj

            spe = self._f_cache[file_number]

            if len(spe) > 1:
                out_stack.append(spe.getData())
            else:
                out_stack.append(spe[0])

        # return stacked and squeezed results
        return np.dstack(out_stack).squeeze()

class DummyAreaDetectorHandler(HandlerBase):
    def __init__(self, fpath, frame_per_point=1, **kwargs):

        self._fpp = frame_per_point

    def __call__(self, **kwargs):
        out_stack = np.ones((self._fpp, 10, 10)) * np.nan
        # return stacked and squeezed results
        return out_stack.squeeze()

class _HDF5HandlerBase(HandlerBase):

    def open(self):
        if self._file:
            return
        self._file = h5py.File(self._filename)

    def close(self):
        super(_HDF5HandlerBase, self).close()
        self._file.close()
        self._file = None

class HDF5DatasetSliceHandler(_HDF5HandlerBase):
    """
    Handler for data stored in one Dataset of an HDF5 file.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    key : string
        key of the single HDF5 Dataset used by this Handler
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    def __init__(self, filename, key, frame_per_point=1):
        self._fpp = frame_per_point
        self._filename = filename
        self._key = key
        self._file = None
        self._dataset = None
        self.open()

    def __call__(self, point_number):
        # Don't read out the dataset until it is requested for the first time.
        if not self._dataset:
            self._dataset = self._file[self._key]
        start, stop = point_number * self._fpp, (point_number + 1) * self._fpp
        return self._dataset[start:stop].squeeze()

class AreaDetectorHDF5Handler(HDF5DatasetSliceHandler):
    """
    Handler for the 'AD_HDF5' spec used by Area Detectors.

    In this spec, the key (i.e., HDF5 dataset path) is always
    '/entry/data/data'.

    Parameters
    ----------
    filename : string
        path to HDF5 file
    frame_per_point : integer, optional
        number of frames to return as one datum, default 1
    """
    specs = {'AD_HDF5'} | HDF5DatasetSliceHandler.specs

    def __init__(self, filename, frame_per_point=1):
        hardcoded_key = '/entry/data/data'
        super(AreaDetectorHDF5Handler, self).__init__(
            filename=filename, key=hardcoded_key,
            frame_per_point=frame_per_point)

class _HdfMapsHandlerBase(_HDF5HandlerBase):
    """
    Reader for XRF data stored in hdf5 files.

    The data set is assumed to be in a group called MAPS and stored
    as a 3D array ordered [energy, x, y].

    Parameters
    ----------
    filename : str
        Path to physical location of file
    dset_path : str
        The path to the dataset inside of 'MAPS'
    """
    def __init__(self, filename, dset_path):
        self._filename = filename
        self._dset_path = dset_path
        self._file = None
        self._dset = None
        self.open()

    def open(self):
        """
        Open the file for reading.

        Provided as a stand alone function to allow re-opening of the handler
        """
        if self._file:
            return
        self._file = h5py.File(self._filename, mode='r')
        self._dset = self._file['/'.join(['MAPS', self._dset_path])]

    def __call__(self):

        if not self._file:
            raise RuntimeError("File is not open")

class HDFMapsSpectrumHandler(_HdfMapsHandlerBase):
    """
    Handler which selects energy spectrum from
    a MAPS XRF data product.
    """
    specs = {'MAPS_SPECTRUM'} | _HdfMapsHandlerBase.specs

    def __call__(self, x, y):
        """
        Return the spectrum at the x, y position

        Parameters
        ----------
        x : int
            raster index in the x direction

        y : int
            raster index in the y direction

        Returns
        -------
        spectrum : ndarray
            The MCA channels
        """
        super(HDFMapsSpectrumHandler, self).__call__()
        return self._dset[:, x, y]

class HDFMapsEnergyHandler(_HdfMapsHandlerBase):
    """
    Handler which select fixed-energy slices from
    a MAPS XRF data file.
    """
    specs = {'MAPS_PLANE'} | _HdfMapsHandlerBase.specs

    def __call__(self, e_index):
        """
        Return the raster plane at a fixed energy

        Parameters
        ----------
        e_index : int
            The index of the engery

        Returns
        -------
        plane : ndarray
            The raster image at a fixed energy.
        """
        super(HDFMapsEnergyHandler, self).__call__()
        return self._dset[e_index, :, :]

class NpyHandler(HandlerBase):
    """
    Class to deal with reading npy files

    Parameters
    ----------
    fpath : str
        Path to file

    mmap_mode : {'r', 'r+', c}, optional
        memmap mode to use to open file
    """
    specs = {'npy'} | HandlerBase.specs

    def __init__(self, filename, mmap_mode=None):
        self._mmap_mode = mmap_mode
        if not os.path.exists(filename):
            raise IOError("the requested file {fpath} does not exst")
        self._fpath = filename

    def __call__(self):
        return np.load(self._fpath, self._mmap_mode)
<|endoftext|>


<commit_before>## protocol.py
# IRC low-level protocol constants and helpers.
import re
import collections

## Own definitions.

# While this *technically* is supposed to be 143/994, I've yet to see a server that actually uses those.
DEFAULT_PORT = 6667
DEFAULT_TLS_PORT = 6697

FALLBACK_ENCODING = 'iso-8859-1'

## Limits.
CHANNEL_LIMITS_GROUPS = { '#': frozenset('#&'), '&': frozenset('#&') }
CHANNEL_LIMITS = { frozenset('#&'): 10 }
LIST_LIMIT_GROUPS = { 'b': frozenset('b') }
LIST_LIMITS = { frozenset('b'): None }
PARAMETER_LIMIT = 15
MESSAGE_LENGTH_LIMIT = 512
CHANNEL_LENGTH_LIMIT = 200
NICKNAME_LENGTH_LIMIT = 8
TOPIC_LENGTH_LIMIT = 450

## Defaults.

BEHAVIOUR_NO_PARAMETER = 'noparam'
BEHAVIOUR_PARAMETER = 'param'
BEHAVIOUR_PARAMETER_ON_SET = 'param_set'
BEHAVIOUR_LIST = 'list'

CHANNEL_MODES = { 'o', 'p', 's', 'i', 't', 'n', 'b', 'v', 'm', 'r', 'k', 'l' }
CHANNEL_MODES_BEHAVIOUR = { BEHAVIOUR_LIST: { 'b' }, BEHAVIOUR_PARAMETER: { 'o', 'v' }, BEHAVIOUR_PARAMETER_ON_SET: { 'k', 'l' }, BEHAVIOUR_NO_PARAMETER: { 'p', 's', 'i', 't', 'n', 'm', 'r' } }
CHANNEL_PREFIXES = { '#', '&' }
CASE_MAPPINGS = { 'ascii', 'rfc1459', 'strict-rfc1459' }
CASE_MAPPING = 'rfc1459'
NICKNAME_PREFIXES = collections.OrderedDict([('@', 'o'), ('+', 'v')])
USER_MODES = { 'i', 'w', 's', 'o' }
# Maybe one day, user modes will have parameters...
USER_MODES_BEHAVIOUR = { BEHAVIOUR_NO_PARAMETER: { 'i', 'w', 's', 'o' } }

## Message parsing.

FORBIDDEN_CHARACTERS = { '\r', '\n', '\0' }

LINE_SEPARATOR = '\r\n'
MINIMAL_LINE_SEPARATOR = '\n'
USER_SEPARATOR = '!'
HOST_SEPARATOR = '@'

PRIVATE_CHANNEL_SIGIL = '@'
SECRET_CHANNEL_SIGIL = '*'
PUBLIC_CHANNEL_SIGIL = '='

ARGUMENT_SEPARATOR = re.compile(' +', re.UNICODE)
COMMAND_PATTERN = re.compile('^([a-zA-Z]+|[0-9]+)$', re.UNICODE)
TRAILING_PREFIX = ':'

## Errors.

class ProtocolViolation(Exception):
    """ An error that occurred while parsing or constructing an IRC message that violates the IRC protocol. """
    def __init__(self, msg, message):
        super().__init__(msg)
        self.irc_message = message

## Message parsing and construction.

class Message:
    def __init__(self, command, params, source=None, _valid=True, **kw):
        self.command = command
        self.params = params
        self.source = source
        self.kw = kw
        self._valid = _valid

    @classmethod
    def parse(cls, line, encoding='utf-8'):
        """
        Parse given line into IRC message structure.
        Returns a Message.
        """
        valid = True

        # Decode message.
        try:
            message = line.decode(encoding)
        except UnicodeDecodeError:
            # Try our fallback encoding.
            message = line.decode(FALLBACK_ENCODING)

        # Sanity check for message length.
        if len(message) > MESSAGE_LENGTH_LIMIT:
            valid = False

        # Strip message separator.
        if message.endswith(LINE_SEPARATOR):
            message = message[:-len(LINE_SEPARATOR)]
        elif message.endswith(MINIMAL_LINE_SEPARATOR):
            message = message[:-len(MINIMAL_LINE_SEPARATOR)]

        # Sanity check for forbidden characters.
        if any(ch in message for ch in FORBIDDEN_CHARACTERS):
            valid = False

        # Extract message sections.
        # Format: (:source)? command parameter*
        if message.startswith(':'):
            parts = ARGUMENT_SEPARATOR.split(message, 2)
        else:
            parts = [ None ] + ARGUMENT_SEPARATOR.split(message, 1)
        if len(parts) == 3:
            source, command, raw_params = parts
        elif len(parts) == 2:
            source, command, raw_params = parts + ['']
        else:
            raise ProtocolViolation('Improper IRC message format: not enough elements.', message=message)

        # Sanity check for command.
        if not COMMAND_PATTERN.match(command):
            valid = False

        # Extract parameters properly.
        # Format: (word|:sentence)*

        # Only parameter is a 'trailing' sentence.
        if raw_params.startswith(TRAILING_PREFIX):
            params = [ raw_params[len(TRAILING_PREFIX):] ]
        # We have a sentence in our parameters.
        elif ' ' + TRAILING_PREFIX in raw_params:
            index = raw_params.find(' ' + TRAILING_PREFIX)

             # Get all single-word parameters.
            params = ARGUMENT_SEPARATOR.split(raw_params[:index].rstrip(' '))
            # Extract last parameter as sentence
            params.append(raw_params[index + len(TRAILING_PREFIX) + 1:])
        # We have some parameters, but no sentences.
        elif raw_params:
            params = ARGUMENT_SEPARATOR.split(raw_params)
        # No parameters.
        else:
            params = []

        # Commands can be either [a-zA-Z]+ or [0-9]+.
        # In the former case, force it to uppercase.
        # In the latter case (a numeric command), try to represent it as such.
        try:
            command = int(command)
        except ValueError:
            command = command.upper()

        # Return parsed message.
        return Message(command, params, source=source, _valid=valid)

    def __str__(self):
        return self.construct()

    def construct(self):
        """ Construct a raw IRC message. """
        # Sanity check for command.
        command = str(self.command)
        if not COMMAND_PATTERN.match(command):
            raise ProtocolViolation('The constructed command does not follow the command pattern ({pat})'.format(pat=COMMAND_PATTERN.pattern), message=command)
        message = command.upper()

        # Add parameters.
        if not self.params:
            message += ' '
        for idx, param in enumerate(self.params):
            # Trailing parameter?
            if ' ' in param:
                if idx + 1 < len(self.params):
                    raise ProtocolViolation('Only the final parameter of an IRC message can be trailing and thus contain spaces.', message=param)
                message += ' ' + TRAILING_PREFIX + param
            # Regular paramter.
            else:
                message += ' ' + param

        # Prepend source.
        if self.source:
            message = ':' + self.source + ' ' + message

        # Sanity check for characters.
        if any(ch in message for ch in FORBIDDEN_CHARACTERS):
            raise ProtocolViolation('The constructed message contains forbidden characters ({chs}).'.format(chs=', '.join(FORBIDDEN_CHARACTERS)), message=message)

        # Sanity check for length.
        message += LINE_SEPARATOR
        if len(message) > MESSAGE_LENGTH_LIMIT:
            raise ProtocolViolation('The constructed message is too long. ({len} > {maxlen})'.format(len=len(message), maxlen=MESSAGE_LENGTH_LIMIT), message=message)

        return message

def normalize(input, case_mapping='rfc1459'):
    """ Normalize input according to case mapping. """
    if case_mapping not in CASE_MAPPINGS:
        raise ProtocolViolation('Unknown case mapping ({})'.format(case_mapping))

    input = input.lower()

    if case_mapping in ('rfc1459', 'rfc1459-strict'):
        input = input.replace('{', '[').replace('}', ']').replace('|', '\\')
    if case_mapping == 'rfc1459':
        input = input.replace('~', '^')

    return input

def equals(left, right, case_mapping='rfc1459'):
    """ Determine whether or not the given nicknames are equal in the given case mapping. """
    if case_mapping not in CASE_MAPPINGS:
        raise ProtocolViolation('Unknown case mapping ({})'.format(case_mapping))

    return normalize(left, case_mapping) == normalize(right, case_mapping)

# Parsing.

def parse_user(raw):
    """ Parse nick(!user(@host)?)? structure. """
    nick = raw
    user = None
    host = None

    # Attempt to extract host.
    if HOST_SEPARATOR in raw:
        raw, host = raw.split(HOST_SEPARATOR)
    # Attempt to extract user.
    if USER_SEPARATOR in raw:
        nick, user = raw.split(USER_SEPARATOR)

    return nick, user, host

def parse_user_modes(modes, current):
    """ Parse user mode changes and return updated mode dictionary. """
    return parse_modes(modes, current, behaviour=USER_MODES_BEHAVIOUR)

def parse_channel_modes(modes, current, behaviour=CHANNEL_MODES_BEHAVIOUR):
    """ Parse channel mode changes and return updated mode dictionary. """
    return parse_modes(modes, current, behaviour=behaviour)

def parse_modes(modes, current, behaviour):
    """ Parse mode change string(s) and return updated dictionary. """
    current = current.copy()

    # Iterate in a somewhat odd way over the list because we want to modify it during iteration.
    i = 0
    while i < len(modes):
        piece = modes[i]
        add = True
        sigiled = False

        for mode in piece:
            # Set mode to addition or deletion of modes.
            if mode == '+':
                add = True
                sigiled = True
                continue
            if mode == '-':
                add = False
                sigiled = True
                continue

            # Find mode behaviour.
            for type, affected in behaviour.items():
                if mode in affected:
                    break
            else:
                # If we don't have a behaviour for this mode, assume it has no parameters...
                type = BEHAVIOUR_NO_PARAMETER

            # Don't parse modes that are meant for list retrieval.
            if type == BEHAVIOUR_LIST and not sigiled:
                continue

            # Do we require a parameter?
            if type in (BEHAVIOUR_PARAMETER, BEHAVIOUR_LIST) or (type == BEHAVIOUR_PARAMETER_ON_SET and add):
                # Do we _have_ a parameter?
                if i + 1 == len(modes):
                    raise ProtocolViolation(
                        'Attempted to parse mode with parameter ({s}{mode}) but no parameters left in mode list.'.format(
                        mode=mode, s='+' if add else '-'), ' '.join(modes))
                param = modes.pop(i + 1)

            # Now update the actual mode dict with our new values.
            if type in (BEHAVIOUR_PARAMETER, BEHAVIOUR_LIST):
                # Add/remove parameter from list.
                if add:
                    if mode not in current:
                        current[mode] = []
                    current[mode].append(param)
                else:
                    if mode in current and param in current[mode]:
                        current[mode].remove(param)
            elif type == BEHAVIOUR_PARAMETER_ON_SET and add:
                # Simply set parameter.
                current[mode] = param
            else:
                # Simply add/remove option.
                if add:
                    current[mode] = True
                else:
                    if mode in current:
                        del current[mode]

        # Onto the next mode.
        i += 1

    return current

<commit_msg>Neatify parts handling in message parsing.<commit_after>## protocol.py
# IRC low-level protocol constants and helpers.
import re
import collections

## Own definitions.

# While this *technically* is supposed to be 143/994, I've yet to see a server that actually uses those.
DEFAULT_PORT = 6667
DEFAULT_TLS_PORT = 6697

FALLBACK_ENCODING = 'iso-8859-1'

## Limits.
CHANNEL_LIMITS_GROUPS = { '#': frozenset('#&'), '&': frozenset('#&') }
CHANNEL_LIMITS = { frozenset('#&'): 10 }
LIST_LIMIT_GROUPS = { 'b': frozenset('b') }
LIST_LIMITS = { frozenset('b'): None }
PARAMETER_LIMIT = 15
MESSAGE_LENGTH_LIMIT = 512
CHANNEL_LENGTH_LIMIT = 200
NICKNAME_LENGTH_LIMIT = 8
TOPIC_LENGTH_LIMIT = 450

## Defaults.

BEHAVIOUR_NO_PARAMETER = 'noparam'
BEHAVIOUR_PARAMETER = 'param'
BEHAVIOUR_PARAMETER_ON_SET = 'param_set'
BEHAVIOUR_LIST = 'list'

CHANNEL_MODES = { 'o', 'p', 's', 'i', 't', 'n', 'b', 'v', 'm', 'r', 'k', 'l' }
CHANNEL_MODES_BEHAVIOUR = { BEHAVIOUR_LIST: { 'b' }, BEHAVIOUR_PARAMETER: { 'o', 'v' }, BEHAVIOUR_PARAMETER_ON_SET: { 'k', 'l' }, BEHAVIOUR_NO_PARAMETER: { 'p', 's', 'i', 't', 'n', 'm', 'r' } }
CHANNEL_PREFIXES = { '#', '&' }
CASE_MAPPINGS = { 'ascii', 'rfc1459', 'strict-rfc1459' }
CASE_MAPPING = 'rfc1459'
NICKNAME_PREFIXES = collections.OrderedDict([('@', 'o'), ('+', 'v')])
USER_MODES = { 'i', 'w', 's', 'o' }
# Maybe one day, user modes will have parameters...
USER_MODES_BEHAVIOUR = { BEHAVIOUR_NO_PARAMETER: { 'i', 'w', 's', 'o' } }

## Message parsing.

FORBIDDEN_CHARACTERS = { '\r', '\n', '\0' }

LINE_SEPARATOR = '\r\n'
MINIMAL_LINE_SEPARATOR = '\n'
USER_SEPARATOR = '!'
HOST_SEPARATOR = '@'

PRIVATE_CHANNEL_SIGIL = '@'
SECRET_CHANNEL_SIGIL = '*'
PUBLIC_CHANNEL_SIGIL = '='

ARGUMENT_SEPARATOR = re.compile(' +', re.UNICODE)
COMMAND_PATTERN = re.compile('^([a-zA-Z]+|[0-9]+)$', re.UNICODE)
TRAILING_PREFIX = ':'

## Errors.

class ProtocolViolation(Exception):
    """ An error that occurred while parsing or constructing an IRC message that violates the IRC protocol. """
    def __init__(self, msg, message):
        super().__init__(msg)
        self.irc_message = message

## Message parsing and construction.

class Message:
    def __init__(self, command, params, source=None, _valid=True, **kw):
        self.command = command
        self.params = params
        self.source = source
        self.kw = kw
        self._valid = _valid

    @classmethod
    def parse(cls, line, encoding='utf-8'):
        """
        Parse given line into IRC message structure.
        Returns a Message.
        """
        valid = True

        # Decode message.
        try:
            message = line.decode(encoding)
        except UnicodeDecodeError:
            # Try our fallback encoding.
            message = line.decode(FALLBACK_ENCODING)

        # Sanity check for message length.
        if len(message) > MESSAGE_LENGTH_LIMIT:
            valid = False

        # Strip message separator.
        if message.endswith(LINE_SEPARATOR):
            message = message[:-len(LINE_SEPARATOR)]
        elif message.endswith(MINIMAL_LINE_SEPARATOR):
            message = message[:-len(MINIMAL_LINE_SEPARATOR)]

        # Sanity check for forbidden characters.
        if any(ch in message for ch in FORBIDDEN_CHARACTERS):
            valid = False

        # Extract message sections.
        # Format: (:source)? command parameter*
        if message.startswith(':'):
            parts = ARGUMENT_SEPARATOR.split(message, 2)
        else:
            parts = [ None ] + ARGUMENT_SEPARATOR.split(message, 1)

        if len(parts) == 3:
            source, command, raw_params = parts
        elif len(parts) == 2:
            source, command = parts
            raw_params = ''
        else:
            raise ProtocolViolation('Improper IRC message format: not enough elements.', message=message)

        # Sanity check for command.
        if not COMMAND_PATTERN.match(command):
            valid = False

        # Extract parameters properly.
        # Format: (word|:sentence)*

        # Only parameter is a 'trailing' sentence.
        if raw_params.startswith(TRAILING_PREFIX):
            params = [ raw_params[len(TRAILING_PREFIX):] ]
        # We have a sentence in our parameters.
        elif ' ' + TRAILING_PREFIX in raw_params:
            index = raw_params.find(' ' + TRAILING_PREFIX)

             # Get all single-word parameters.
            params = ARGUMENT_SEPARATOR.split(raw_params[:index].rstrip(' '))
            # Extract last parameter as sentence
            params.append(raw_params[index + len(TRAILING_PREFIX) + 1:])
        # We have some parameters, but no sentences.
        elif raw_params:
            params = ARGUMENT_SEPARATOR.split(raw_params)
        # No parameters.
        else:
            params = []

        # Commands can be either [a-zA-Z]+ or [0-9]+.
        # In the former case, force it to uppercase.
        # In the latter case (a numeric command), try to represent it as such.
        try:
            command = int(command)
        except ValueError:
            command = command.upper()

        # Return parsed message.
        return Message(command, params, source=source, _valid=valid)

    def __str__(self):
        return self.construct()

    def construct(self):
        """ Construct a raw IRC message. """
        # Sanity check for command.
        command = str(self.command)
        if not COMMAND_PATTERN.match(command):
            raise ProtocolViolation('The constructed command does not follow the command pattern ({pat})'.format(pat=COMMAND_PATTERN.pattern), message=command)
        message = command.upper()

        # Add parameters.
        if not self.params:
            message += ' '
        for idx, param in enumerate(self.params):
            # Trailing parameter?
            if ' ' in param:
                if idx + 1 < len(self.params):
                    raise ProtocolViolation('Only the final parameter of an IRC message can be trailing and thus contain spaces.', message=param)
                message += ' ' + TRAILING_PREFIX + param
            # Regular paramter.
            else:
                message += ' ' + param

        # Prepend source.
        if self.source:
            message = ':' + self.source + ' ' + message

        # Sanity check for characters.
        if any(ch in message for ch in FORBIDDEN_CHARACTERS):
            raise ProtocolViolation('The constructed message contains forbidden characters ({chs}).'.format(chs=', '.join(FORBIDDEN_CHARACTERS)), message=message)

        # Sanity check for length.
        message += LINE_SEPARATOR
        if len(message) > MESSAGE_LENGTH_LIMIT:
            raise ProtocolViolation('The constructed message is too long. ({len} > {maxlen})'.format(len=len(message), maxlen=MESSAGE_LENGTH_LIMIT), message=message)

        return message

def normalize(input, case_mapping='rfc1459'):
    """ Normalize input according to case mapping. """
    if case_mapping not in CASE_MAPPINGS:
        raise ProtocolViolation('Unknown case mapping ({})'.format(case_mapping))

    input = input.lower()

    if case_mapping in ('rfc1459', 'rfc1459-strict'):
        input = input.replace('{', '[').replace('}', ']').replace('|', '\\')
    if case_mapping == 'rfc1459':
        input = input.replace('~', '^')

    return input

def equals(left, right, case_mapping='rfc1459'):
    """ Determine whether or not the given nicknames are equal in the given case mapping. """
    if case_mapping not in CASE_MAPPINGS:
        raise ProtocolViolation('Unknown case mapping ({})'.format(case_mapping))

    return normalize(left, case_mapping) == normalize(right, case_mapping)

# Parsing.

def parse_user(raw):
    """ Parse nick(!user(@host)?)? structure. """
    nick = raw
    user = None
    host = None

    # Attempt to extract host.
    if HOST_SEPARATOR in raw:
        raw, host = raw.split(HOST_SEPARATOR)
    # Attempt to extract user.
    if USER_SEPARATOR in raw:
        nick, user = raw.split(USER_SEPARATOR)

    return nick, user, host

def parse_user_modes(modes, current):
    """ Parse user mode changes and return updated mode dictionary. """
    return parse_modes(modes, current, behaviour=USER_MODES_BEHAVIOUR)

def parse_channel_modes(modes, current, behaviour=CHANNEL_MODES_BEHAVIOUR):
    """ Parse channel mode changes and return updated mode dictionary. """
    return parse_modes(modes, current, behaviour=behaviour)

def parse_modes(modes, current, behaviour):
    """ Parse mode change string(s) and return updated dictionary. """
    current = current.copy()

    # Iterate in a somewhat odd way over the list because we want to modify it during iteration.
    i = 0
    while i < len(modes):
        piece = modes[i]
        add = True
        sigiled = False

        for mode in piece:
            # Set mode to addition or deletion of modes.
            if mode == '+':
                add = True
                sigiled = True
                continue
            if mode == '-':
                add = False
                sigiled = True
                continue

            # Find mode behaviour.
            for type, affected in behaviour.items():
                if mode in affected:
                    break
            else:
                # If we don't have a behaviour for this mode, assume it has no parameters...
                type = BEHAVIOUR_NO_PARAMETER

            # Don't parse modes that are meant for list retrieval.
            if type == BEHAVIOUR_LIST and not sigiled:
                continue

            # Do we require a parameter?
            if type in (BEHAVIOUR_PARAMETER, BEHAVIOUR_LIST) or (type == BEHAVIOUR_PARAMETER_ON_SET and add):
                # Do we _have_ a parameter?
                if i + 1 == len(modes):
                    raise ProtocolViolation(
                        'Attempted to parse mode with parameter ({s}{mode}) but no parameters left in mode list.'.format(
                        mode=mode, s='+' if add else '-'), ' '.join(modes))
                param = modes.pop(i + 1)

            # Now update the actual mode dict with our new values.
            if type in (BEHAVIOUR_PARAMETER, BEHAVIOUR_LIST):
                # Add/remove parameter from list.
                if add:
                    if mode not in current:
                        current[mode] = []
                    current[mode].append(param)
                else:
                    if mode in current and param in current[mode]:
                        current[mode].remove(param)
            elif type == BEHAVIOUR_PARAMETER_ON_SET and add:
                # Simply set parameter.
                current[mode] = param
            else:
                # Simply add/remove option.
                if add:
                    current[mode] = True
                else:
                    if mode in current:
                        del current[mode]

        # Onto the next mode.
        i += 1

    return current

<|endoftext|>