<commit_before>import argparse

parent_parser = argparse.ArgumentParser(add_help=False)

config_group = parent_parser.add_argument_group(title="Lampost Engine Configuration")
config_group.add_argument('-cid', '--config_id', help="database configuration id", default='lampost')
config_group.add_argument('-a', '--app_id', help="application module", default='lampmud.lpmud')

log_group = parent_parser.add_argument_group(title="Logging Configuration")
log_group.add_argument('-l', '--log', help="default logging level", default='info', dest='log_level')
log_group.add_argument('-lf', '--log_file', help="log output file", default=None)
log_group.add_argument('-lm', '--log_mode', help="log file open mode", default='w')

db_group = parent_parser.add_argument_group(title="Redis Database Configuration")
db_group.add_argument('-db_host', help="database server host name", default='localhost')
db_group.add_argument('-db_port', help="database server port", type=int, default=6379)
db_group.add_argument('-db_num', help="Redis database number", type=int, default=0)
db_group.add_argument('-db_pw', help="Redis database password", default=None)

main_parser = argparse.ArgumentParser(parents=[parent_parser], formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                      description="lampost -- The Lampost Game Engine  https://github.com/genzgd/Lampost-Mud")
web_group = main_parser.add_argument_group(title="Web Server Configuration")
web_group.add_argument('-p', '--port', help="web server port", type=int, default=2500)
web_group.add_argument('-si', help="web server network interface", default='127.0.0.1', metavar="NW_INT", dest="server_interface")

config_group.add_argument('-cd', '--config_dir', help="yaml configuration directory", default='conf')

create_parser = argparse.ArgumentParser(parents=[parent_parser], formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                      description="lampost_setup -- Creation utility for the Lampost Game Engine https://github.com/genzgd/Lampost-Mud")

create_parser.add_argument('-flush', help="Flush (remove all keys) form existing Redis database", const=True, default=False, action='store_const')

imm_group = create_parser.add_argument_group(title="Immortal User Configuration")
imm_group.add_argument('-ia', '--imm_account', help="root account name", default='root')
imm_group.add_argument('-im', '--imm_name', help="name of the root account primary primary", required=True)
imm_group.add_argument('-ipw', '--imm_password', help="root account password", default='lampost')

tools_parser = argparse.ArgumentParser(parents=[parent_parser], formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                       description="lampost_tools -- Some command line tools for lampost")
tools_parser.add_argument('op', metavar="OPERATION", help="Tools operation")

<commit_msg>Make client file path configurable.<commit_after>import argparse

parent_parser = argparse.ArgumentParser(add_help=False)

config_group = parent_parser.add_argument_group(title="Lampost Engine Configuration")
config_group.add_argument('-cid', '--config_id', help="database configuration id", default='lampost')
config_group.add_argument('-a', '--app_id', help="application module", default='lampmud.lpmud')

log_group = parent_parser.add_argument_group(title="Logging Configuration")
log_group.add_argument('-ll', '--log_level', help="default logging level", default='info', dest='log_level')
log_group.add_argument('-lf', '--log_file', help="log output file", default=None)
log_group.add_argument('-lm', '--log_mode', help="log file open mode", default='w')

db_group = parent_parser.add_argument_group(title="Redis Database Configuration")
db_group.add_argument('-db_host', help="database server host name", default='localhost')
db_group.add_argument('-db_port', help="database server port", type=int, default=6379)
db_group.add_argument('-db_num', help="Redis database number", type=int, default=0)
db_group.add_argument('-db_pw', help="Redis database password", default=None)

main_parser = argparse.ArgumentParser(parents=[parent_parser], formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                      description="lampost -- The Lampost Game Engine  https://github.com/genzgd/Lampost-Mud")
web_group = main_parser.add_argument_group(title="Web Server Configuration")
web_group.add_argument('-p', '--port', help="web server port", type=int, default=2500)
web_group.add_argument('-si', help="web server network interface", default='127.0.0.1', metavar="NW_INT", dest="server_interface")
web_group.add_argument('-wf', '--web_files', help="file system path for static web client", default='webclient')

config_group.add_argument('-cd', '--config_dir', help="yaml configuration directory", default='conf')

create_parser = argparse.ArgumentParser(parents=[parent_parser], formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                      description="lampost_setup -- Creation utility for the Lampost Game Engine https://github.com/genzgd/Lampost-Mud")

create_parser.add_argument('-flush', help="Flush (remove all keys) form existing Redis database", const=True, default=False, action='store_const')

imm_group = create_parser.add_argument_group(title="Immortal User Configuration")
imm_group.add_argument('-ia', '--imm_account', help="root account name", default='root')
imm_group.add_argument('-im', '--imm_name', help="name of the root account primary primary", required=True)
imm_group.add_argument('-ipw', '--imm_password', help="root account password", default='lampost')

tools_parser = argparse.ArgumentParser(parents=[parent_parser], formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                       description="lampost_tools -- Some command line tools for lampost")
tools_parser.add_argument('op', metavar="OPERATION", help="Tools operation")

<|endoftext|>


<commit_before># Copyright (c) 2010-2014, GEM Foundation.
#
# OpenQuake 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.
#
# OpenQuake 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 General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Functions for exporting risk artifacts from the database.
"""

import os
import csv

from openquake.engine.db import models
from openquake.engine.export import core
from openquake.engine.utils import FileWrapper
from openquake.commonlib import risk_writers, writers
from openquake.baselib.general import block_splitter

LOSS_CURVE_FILENAME_FMT = 'loss-curves-%(loss_curve_id)s.xml'
LOSS_MAP_FILENAME_FMT = 'loss-maps-%(loss_map_id)s.%(file_ext)s'
LOSS_FRACTION_FILENAME_FMT = 'loss-fractions-%(loss_fraction_id)s.xml'
AGGREGATE_LOSS_FILENAME_FMT = 'aggregate-loss-%s.csv'
BCR_FILENAME_FMT = 'bcr-distribution-%(bcr_distribution_id)s.xml'
EVENT_LOSS_FILENAME_FMT = 'event-loss-%s.csv'
EVENT_LOSS_ASSET_FILENAME_FMT = 'event-loss-asset-%s.csv'
DMG_PER_ASSET = 'dmg-per-asset-%s.csv'

DAMAGE = dict(
    collapse_map=models.DmgDistPerAsset,
    dmg_dist_per_asset=models.DmgDistPerAsset,
    dmg_dist_per_taxonomy=models.DmgDistPerTaxonomy,
    dmg_dist_total=models.DmgDistTotal)

def _get_result_export_dest(target, output, file_ext='xml'):
    """
    Get the full absolute path (including file name) for a given ``result``.

    As a side effect, intermediate directories are created such that the file
    can be created and written to immediately.

    :param target:
        Destination directory location for exported files OR a file-like
        object. If file-like, we just simply return it.
    :param output:
        :class:`~openquake.engine.db.models.Output`.
    :param file_ext:
        Desired file extension for the output file.
        Defaults to 'xml'.

    :returns:
        Full path (including filename) to the destination export file.
        If the ``target`` is a file-like, we don't do anything special
        and simply return it.
    """
    if not isinstance(target, basestring):
        # It's not a file path. In this case, we expect a file-like object.
        # Just return it.
        return target

    output_type = output.output_type

    filename = None

    if output_type in ('loss_curve', 'agg_loss_curve', 'event_loss_curve'):
        filename = LOSS_CURVE_FILENAME_FMT % dict(
            loss_curve_id=output.loss_curve.id,
        )
    elif output_type == 'loss_map':
        filename = LOSS_MAP_FILENAME_FMT % dict(
            loss_map_id=output.loss_map.id,
            file_ext=file_ext,
        )
    elif output_type == 'loss_fraction':
        filename = LOSS_FRACTION_FILENAME_FMT % dict(
            loss_fraction_id=output.loss_fraction.id,
        )
    elif output_type == 'bcr_distribution':
        filename = BCR_FILENAME_FMT % dict(
            bcr_distribution_id=output.bcr_distribution.id,
        )
    elif output_type == 'event_loss':
        filename = EVENT_LOSS_FILENAME_FMT % output.id
    elif output_type == 'event_loss_asset':
        filename = EVENT_LOSS_ASSET_FILENAME_FMT % output.id
    elif output_type == 'aggregate_loss':
        filename = AGGREGATE_LOSS_FILENAME_FMT % (output.aggregate_loss.id)
    elif output_type in ('dmg_dist_per_asset', 'dmg_dist_per_taxonomy',
                         'dmg_dist_total', 'collapse_map'):
        filename = '%(output_type)s-%(job_id)s.%(file_ext)s' % dict(
            output_type=output_type,
            job_id=output.oq_job.id,
            file_ext=file_ext,
        )
    elif output_type == 'dmg_per_asset':
        filename = DMG_PER_ASSET % output.id
    return os.path.abspath(os.path.join(target, filename))

def _export_common(output, loss_type):
    """
    Returns a dict containing the output metadata which are serialized
    by nrml writers before actually writing the `output` data.
    """
    metadata = output.hazard_metadata
    if metadata.sm_path is not None:
        source_model_tree_path = core.LT_PATH_JOIN_TOKEN.join(
            metadata.sm_path)
    else:
        source_model_tree_path = None
    if metadata.gsim_path is not None:
        gsim_tree_path = core.LT_PATH_JOIN_TOKEN.join(metadata.gsim_path)
    else:
        gsim_tree_path = None

    return dict(investigation_time=metadata.investigation_time,
                statistics=metadata.statistics,
                quantile_value=metadata.quantile,
                source_model_tree_path=source_model_tree_path,
                gsim_tree_path=gsim_tree_path,
                unit=output.oq_job.exposure_model.unit(loss_type),
                loss_type=loss_type)

@core.export_output.add(('agg_loss_curve', 'xml'))
def export_agg_loss_curve_xml(key, output, target):
    """
    Export `output` to `target` by using a nrml loss curves
    serializers
    """
    args = _export_common(output, output.loss_curve.loss_type)
    dest = _get_result_export_dest(target, output)
    risk_writers.AggregateLossCurveXMLWriter(dest, **args).serialize(
        output.loss_curve.aggregatelosscurvedata)
    return dest

@core.export_output.add(('agg_loss_curve', 'csv'))
def export_agg_loss_curve_csv(key, output, target):
    """
    Export `output` to `target` in csv format
    """
    dest = _get_result_export_dest(target, output)[:-3] + 'csv'
    row = output.loss_curve.aggregatelosscurvedata
    data = ('aggregate', row.losses, row.poes, row.average_loss,
            row.stddev_loss)
    return writers.save_csv(dest, [data])

@core.export_output.add(('loss_curve', 'xml'), ('event_loss_curve', 'xml'))
def export_loss_curve_xml(key, output, target):
    """
    Export `output` to `target` by using a nrml loss curves
    serializer
    """
    args = _export_common(output, output.loss_curve.loss_type)
    dest = _get_result_export_dest(target, output)
    args['insured'] = output.loss_curve.insured

    data = output.loss_curve.losscurvedata_set.all().order_by('asset_ref')

    risk_writers.LossCurveXMLWriter(dest, **args).serialize(data)
    return dest

@core.export_output.add(('loss_curve', 'csv'))
def export_avgloss_csv(key, output, target):
    """
    Export `output` to `target` in csv format for a given loss type
    """
    dest = _get_result_export_dest(target, output)[:-3] + 'csv'
    data = output.loss_curve.losscurvedata_set.all().order_by('asset_ref')
    header = ['lon', 'lat', 'asset_ref', 'asset_value', 'average_loss',
              'stddev_loss', 'loss_type']
    rows = [(c.location.x, c.location.y, c.asset_ref, c.asset_value,
             c.average_loss, c.stddev_loss or '', c.loss_curve.loss_type)
            for c in data]
    writers.save_csv(dest, [header] + rows)
    return dest

@core.export_output.add(('loss_map', 'xml'), ('loss_map', 'geojson'))
def export_loss_map(key, output, target):
    """
    General loss map export code.
    """
    args = _export_common(output, output.loss_map.loss_type)
    dest = _get_result_export_dest(target, output, file_ext=key[1])
    args.update(dict(
        poe=output.loss_map.poe,
        loss_category=output.oq_job.exposure_model.category))

    writercls = (risk_writers.LossMapXMLWriter if key[1] == 'xml'
                 else risk_writers.LossMapGeoJSONWriter)
    writercls(dest, **args).serialize(
        models.order_by_location(
            output.loss_map.lossmapdata_set.all().order_by('asset_ref')))
    return dest

@core.export_output.add(('loss_fraction', 'xml'))
def export_loss_fraction_xml(key, output, target):
    """
    Export `output` to `target` by using a nrml loss fractions
    serializer
    """
    args = _export_common(output, output.loss_fraction.loss_type)
    hazard_metadata = models.Output.HazardMetadata(
        investigation_time=args['investigation_time'],
        statistics=args['statistics'],
        quantile=args.get('quantile_value'),
        sm_path=args['source_model_tree_path'],
        gsim_path=args['gsim_tree_path'])
    dest = _get_result_export_dest(target, output)
    poe = output.loss_fraction.poe
    variable = output.loss_fraction.variable
    loss_category = output.oq_job.exposure_model.category
    risk_writers.LossFractionsWriter(
        dest, variable, args['unit'], args['loss_type'],
        loss_category, hazard_metadata, poe).serialize(
        output.loss_fraction.total_fractions(),
        output.loss_fraction)
    return dest

@core.export_output.add(('bcr_distribution', 'xml'))
def export_bcr_distribution_xml(key, output, target):
    """
    Export `output` to `target` by using a nrml bcr distribution
    serializer
    """
    risk_calculation = output.oq_job.get_oqparam()
    args = _export_common(output, output.bcr_distribution.loss_type)

    dest = _get_result_export_dest(target, output)
    args.update(
        dict(interest_rate=risk_calculation.interest_rate,
             asset_life_expectancy=risk_calculation.asset_life_expectancy))
    del args['investigation_time']

    risk_writers.BCRMapXMLWriter(dest, **args).serialize(
        output.bcr_distribution.bcrdistributiondata_set.all().order_by(
            'asset_ref'))
    return dest

# clearly this is not a good approach for large exposures
@core.export_output.add(('collapse_map', 'xml'),
                        ('dmg_dist_per_taxonomy', 'xml'),
                        ('dmg_dist_per_asset', 'xml'),
                        ('dmg_dist_total', 'xml'))
def export_dmg_dist(key, output, target):
    """
    Export the damage distribution identified
    by the given output to the `target`.

    :param output:
        DB output record which identifies the distribution. A
        :class:`openquake.engine.db.models.Output` object.
    :param target:
        Destination directory of the exported file, or a file-like object
    """
    job = output.oq_job
    dest = _get_result_export_dest(target, output)
    damage_states = list(models.DmgState.objects.filter(
        risk_calculation__id=job.id).order_by('lsi'))
    writer = risk_writers.DamageWriter(damage_states)
    damagecls = DAMAGE[key[0]]
    if key[0] == 'collapse_map':
        data = damagecls.objects.filter(dmg_state=damage_states[-1])
    else:
        data = damagecls.objects.filter(dmg_state__risk_calculation__id=job.id)
    writer.to_nrml(key[0], data, dest)
    return dest

@core.export_output.add(('dmg_per_asset', 'csv'))
def export_dmg_per_asset_csv(key, output, target):
    """
    Classical Damage Per Asset in CSV format
    """
    dest = _get_result_export_dest(target, output)

    damage_states = list(models.DmgState.objects.filter(
        risk_calculation=output.oq_job).order_by('lsi'))
    data = block_splitter(
        models.DamageData.objects.filter(
            dmg_state__risk_calculation=output.oq_job).order_by(
            'exposure_data', 'dmg_state'),
        len(damage_states))

    with FileWrapper(dest, mode='wb') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['asset_ref'] + [ds.dmg_state for ds in damage_states])
        for row in data:
            asset = row[0].exposure_data
            fractions = [rec.fraction for rec in row]
            writer.writerow(
                [asset.asset_ref] + map(writers.scientificformat, fractions))
    return dest

@core.export_output.add(('aggregate_loss', 'csv'))
def export_aggregate_loss_csv(key, output, target):
    """
    Export aggregate losses in CSV
    """
    dest = _get_result_export_dest(target, output)

    with FileWrapper(dest, mode='wb') as csvfile:
        writer = csv.writer(csvfile, delimiter='|')
        writer.writerow(['Mean', 'Standard Deviation'])
        writer.writerow([output.aggregate_loss.mean,
                        output.aggregate_loss.std_dev])
    return dest

@core.export_output.add(('event_loss', 'csv'))
def export_event_loss_csv(key, output, target):
    """
    Export Event Loss Table in CSV format
    """

    dest = _get_result_export_dest(target, output)

    with FileWrapper(dest, mode='wb') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['Rupture', 'Magnitude', 'Aggregate Loss'])

        for event_loss in models.EventLossData.objects.filter(
                event_loss__output=output).select_related().order_by(
                '-aggregate_loss'):
            writer.writerow([event_loss.rupture.tag,
                             "%.07f" % event_loss.rupture.rupture.magnitude,
                             "%.07f" % event_loss.aggregate_loss])
    return dest

@core.export_output.add(('event_loss_asset', 'csv'))
def export_event_loss_asset_csv(key, output, target):
    """
    Export Event Loss Per Asset in CSV format
    """
    dest = _get_result_export_dest(target, output)
    rows = []
    for event_loss in models.EventLossAsset.objects.filter(
            event_loss__output=output).select_related().order_by(
            'rupture__tag', 'asset__asset_ref'):
        rows.append([event_loss.rupture.tag,
                     event_loss.asset.asset_ref,
                     event_loss.loss])
    writers.save_csv(dest, rows)
    return dest
<commit_msg>Reduced the precision when exporting the loss curves in CSV<commit_after># Copyright (c) 2010-2014, GEM Foundation.
#
# OpenQuake 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.
#
# OpenQuake 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 General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Functions for exporting risk artifacts from the database.
"""

import os
import csv

from openquake.engine.db import models
from openquake.engine.export import core
from openquake.engine.utils import FileWrapper
from openquake.commonlib import risk_writers, writers
from openquake.baselib.general import block_splitter

LOSS_CURVE_FILENAME_FMT = 'loss-curves-%(loss_curve_id)s.xml'
LOSS_MAP_FILENAME_FMT = 'loss-maps-%(loss_map_id)s.%(file_ext)s'
LOSS_FRACTION_FILENAME_FMT = 'loss-fractions-%(loss_fraction_id)s.xml'
AGGREGATE_LOSS_FILENAME_FMT = 'aggregate-loss-%s.csv'
BCR_FILENAME_FMT = 'bcr-distribution-%(bcr_distribution_id)s.xml'
EVENT_LOSS_FILENAME_FMT = 'event-loss-%s.csv'
EVENT_LOSS_ASSET_FILENAME_FMT = 'event-loss-asset-%s.csv'
DMG_PER_ASSET = 'dmg-per-asset-%s.csv'

DAMAGE = dict(
    collapse_map=models.DmgDistPerAsset,
    dmg_dist_per_asset=models.DmgDistPerAsset,
    dmg_dist_per_taxonomy=models.DmgDistPerTaxonomy,
    dmg_dist_total=models.DmgDistTotal)

def _get_result_export_dest(target, output, file_ext='xml'):
    """
    Get the full absolute path (including file name) for a given ``result``.

    As a side effect, intermediate directories are created such that the file
    can be created and written to immediately.

    :param target:
        Destination directory location for exported files OR a file-like
        object. If file-like, we just simply return it.
    :param output:
        :class:`~openquake.engine.db.models.Output`.
    :param file_ext:
        Desired file extension for the output file.
        Defaults to 'xml'.

    :returns:
        Full path (including filename) to the destination export file.
        If the ``target`` is a file-like, we don't do anything special
        and simply return it.
    """
    if not isinstance(target, basestring):
        # It's not a file path. In this case, we expect a file-like object.
        # Just return it.
        return target

    output_type = output.output_type

    filename = None

    if output_type in ('loss_curve', 'agg_loss_curve', 'event_loss_curve'):
        filename = LOSS_CURVE_FILENAME_FMT % dict(
            loss_curve_id=output.loss_curve.id,
        )
    elif output_type == 'loss_map':
        filename = LOSS_MAP_FILENAME_FMT % dict(
            loss_map_id=output.loss_map.id,
            file_ext=file_ext,
        )
    elif output_type == 'loss_fraction':
        filename = LOSS_FRACTION_FILENAME_FMT % dict(
            loss_fraction_id=output.loss_fraction.id,
        )
    elif output_type == 'bcr_distribution':
        filename = BCR_FILENAME_FMT % dict(
            bcr_distribution_id=output.bcr_distribution.id,
        )
    elif output_type == 'event_loss':
        filename = EVENT_LOSS_FILENAME_FMT % output.id
    elif output_type == 'event_loss_asset':
        filename = EVENT_LOSS_ASSET_FILENAME_FMT % output.id
    elif output_type == 'aggregate_loss':
        filename = AGGREGATE_LOSS_FILENAME_FMT % (output.aggregate_loss.id)
    elif output_type in ('dmg_dist_per_asset', 'dmg_dist_per_taxonomy',
                         'dmg_dist_total', 'collapse_map'):
        filename = '%(output_type)s-%(job_id)s.%(file_ext)s' % dict(
            output_type=output_type,
            job_id=output.oq_job.id,
            file_ext=file_ext,
        )
    elif output_type == 'dmg_per_asset':
        filename = DMG_PER_ASSET % output.id
    return os.path.abspath(os.path.join(target, filename))

def _export_common(output, loss_type):
    """
    Returns a dict containing the output metadata which are serialized
    by nrml writers before actually writing the `output` data.
    """
    metadata = output.hazard_metadata
    if metadata.sm_path is not None:
        source_model_tree_path = core.LT_PATH_JOIN_TOKEN.join(
            metadata.sm_path)
    else:
        source_model_tree_path = None
    if metadata.gsim_path is not None:
        gsim_tree_path = core.LT_PATH_JOIN_TOKEN.join(metadata.gsim_path)
    else:
        gsim_tree_path = None

    return dict(investigation_time=metadata.investigation_time,
                statistics=metadata.statistics,
                quantile_value=metadata.quantile,
                source_model_tree_path=source_model_tree_path,
                gsim_tree_path=gsim_tree_path,
                unit=output.oq_job.exposure_model.unit(loss_type),
                loss_type=loss_type)

@core.export_output.add(('agg_loss_curve', 'xml'))
def export_agg_loss_curve_xml(key, output, target):
    """
    Export `output` to `target` by using a nrml loss curves
    serializers
    """
    args = _export_common(output, output.loss_curve.loss_type)
    dest = _get_result_export_dest(target, output)
    risk_writers.AggregateLossCurveXMLWriter(dest, **args).serialize(
        output.loss_curve.aggregatelosscurvedata)
    return dest

@core.export_output.add(('agg_loss_curve', 'csv'))
def export_agg_loss_curve_csv(key, output, target):
    """
    Export `output` to `target` in csv format
    """
    dest = _get_result_export_dest(target, output)[:-3] + 'csv'
    row = output.loss_curve.aggregatelosscurvedata
    data = ('aggregate', row.losses, row.poes, row.average_loss,
            row.stddev_loss)
    return writers.save_csv(dest, [data], fmt='%10.6E')

@core.export_output.add(('loss_curve', 'xml'), ('event_loss_curve', 'xml'))
def export_loss_curve_xml(key, output, target):
    """
    Export `output` to `target` by using a nrml loss curves
    serializer
    """
    args = _export_common(output, output.loss_curve.loss_type)
    dest = _get_result_export_dest(target, output)
    args['insured'] = output.loss_curve.insured

    data = output.loss_curve.losscurvedata_set.all().order_by('asset_ref')

    risk_writers.LossCurveXMLWriter(dest, **args).serialize(data)
    return dest

@core.export_output.add(('loss_curve', 'csv'))
def export_avgloss_csv(key, output, target):
    """
    Export `output` to `target` in csv format for a given loss type
    """
    dest = _get_result_export_dest(target, output)[:-3] + 'csv'
    data = output.loss_curve.losscurvedata_set.all().order_by('asset_ref')
    header = ['lon', 'lat', 'asset_ref', 'asset_value', 'average_loss',
              'stddev_loss', 'loss_type']
    rows = [(c.location.x, c.location.y, c.asset_ref, c.asset_value,
             c.average_loss, c.stddev_loss or '', c.loss_curve.loss_type)
            for c in data]
    writers.save_csv(dest, [header] + rows)
    return dest

@core.export_output.add(('loss_map', 'xml'), ('loss_map', 'geojson'))
def export_loss_map(key, output, target):
    """
    General loss map export code.
    """
    args = _export_common(output, output.loss_map.loss_type)
    dest = _get_result_export_dest(target, output, file_ext=key[1])
    args.update(dict(
        poe=output.loss_map.poe,
        loss_category=output.oq_job.exposure_model.category))

    writercls = (risk_writers.LossMapXMLWriter if key[1] == 'xml'
                 else risk_writers.LossMapGeoJSONWriter)
    writercls(dest, **args).serialize(
        models.order_by_location(
            output.loss_map.lossmapdata_set.all().order_by('asset_ref')))
    return dest

@core.export_output.add(('loss_fraction', 'xml'))
def export_loss_fraction_xml(key, output, target):
    """
    Export `output` to `target` by using a nrml loss fractions
    serializer
    """
    args = _export_common(output, output.loss_fraction.loss_type)
    hazard_metadata = models.Output.HazardMetadata(
        investigation_time=args['investigation_time'],
        statistics=args['statistics'],
        quantile=args.get('quantile_value'),
        sm_path=args['source_model_tree_path'],
        gsim_path=args['gsim_tree_path'])
    dest = _get_result_export_dest(target, output)
    poe = output.loss_fraction.poe
    variable = output.loss_fraction.variable
    loss_category = output.oq_job.exposure_model.category
    risk_writers.LossFractionsWriter(
        dest, variable, args['unit'], args['loss_type'],
        loss_category, hazard_metadata, poe).serialize(
        output.loss_fraction.total_fractions(),
        output.loss_fraction)
    return dest

@core.export_output.add(('bcr_distribution', 'xml'))
def export_bcr_distribution_xml(key, output, target):
    """
    Export `output` to `target` by using a nrml bcr distribution
    serializer
    """
    risk_calculation = output.oq_job.get_oqparam()
    args = _export_common(output, output.bcr_distribution.loss_type)

    dest = _get_result_export_dest(target, output)
    args.update(
        dict(interest_rate=risk_calculation.interest_rate,
             asset_life_expectancy=risk_calculation.asset_life_expectancy))
    del args['investigation_time']

    risk_writers.BCRMapXMLWriter(dest, **args).serialize(
        output.bcr_distribution.bcrdistributiondata_set.all().order_by(
            'asset_ref'))
    return dest

# clearly this is not a good approach for large exposures
@core.export_output.add(('collapse_map', 'xml'),
                        ('dmg_dist_per_taxonomy', 'xml'),
                        ('dmg_dist_per_asset', 'xml'),
                        ('dmg_dist_total', 'xml'))
def export_dmg_dist(key, output, target):
    """
    Export the damage distribution identified
    by the given output to the `target`.

    :param output:
        DB output record which identifies the distribution. A
        :class:`openquake.engine.db.models.Output` object.
    :param target:
        Destination directory of the exported file, or a file-like object
    """
    job = output.oq_job
    dest = _get_result_export_dest(target, output)
    damage_states = list(models.DmgState.objects.filter(
        risk_calculation__id=job.id).order_by('lsi'))
    writer = risk_writers.DamageWriter(damage_states)
    damagecls = DAMAGE[key[0]]
    if key[0] == 'collapse_map':
        data = damagecls.objects.filter(dmg_state=damage_states[-1])
    else:
        data = damagecls.objects.filter(dmg_state__risk_calculation__id=job.id)
    writer.to_nrml(key[0], data, dest)
    return dest

@core.export_output.add(('dmg_per_asset', 'csv'))
def export_dmg_per_asset_csv(key, output, target):
    """
    Classical Damage Per Asset in CSV format
    """
    dest = _get_result_export_dest(target, output)

    damage_states = list(models.DmgState.objects.filter(
        risk_calculation=output.oq_job).order_by('lsi'))
    data = block_splitter(
        models.DamageData.objects.filter(
            dmg_state__risk_calculation=output.oq_job).order_by(
            'exposure_data', 'dmg_state'),
        len(damage_states))

    with FileWrapper(dest, mode='wb') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['asset_ref'] + [ds.dmg_state for ds in damage_states])
        for row in data:
            asset = row[0].exposure_data
            fractions = [rec.fraction for rec in row]
            writer.writerow(
                [asset.asset_ref] + map(writers.scientificformat, fractions))
    return dest

@core.export_output.add(('aggregate_loss', 'csv'))
def export_aggregate_loss_csv(key, output, target):
    """
    Export aggregate losses in CSV
    """
    dest = _get_result_export_dest(target, output)

    with FileWrapper(dest, mode='wb') as csvfile:
        writer = csv.writer(csvfile, delimiter='|')
        writer.writerow(['Mean', 'Standard Deviation'])
        writer.writerow([output.aggregate_loss.mean,
                        output.aggregate_loss.std_dev])
    return dest

@core.export_output.add(('event_loss', 'csv'))
def export_event_loss_csv(key, output, target):
    """
    Export Event Loss Table in CSV format
    """

    dest = _get_result_export_dest(target, output)

    with FileWrapper(dest, mode='wb') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(['Rupture', 'Magnitude', 'Aggregate Loss'])

        for event_loss in models.EventLossData.objects.filter(
                event_loss__output=output).select_related().order_by(
                '-aggregate_loss'):
            writer.writerow([event_loss.rupture.tag,
                             "%.07f" % event_loss.rupture.rupture.magnitude,
                             "%.07f" % event_loss.aggregate_loss])
    return dest

@core.export_output.add(('event_loss_asset', 'csv'))
def export_event_loss_asset_csv(key, output, target):
    """
    Export Event Loss Per Asset in CSV format
    """
    dest = _get_result_export_dest(target, output)
    rows = []
    for event_loss in models.EventLossAsset.objects.filter(
            event_loss__output=output).select_related().order_by(
            'rupture__tag', 'asset__asset_ref'):
        rows.append([event_loss.rupture.tag,
                     event_loss.asset.asset_ref,
                     event_loss.loss])
    writers.save_csv(dest, rows)
    return dest
<|endoftext|>


<commit_before>__author__ = 'mongolrgata'

import sys
import re
import binascii
import json

def rotr8(int8, shift):
    """
    :param int8:
    :type int8: int
    :param shift:
    :type shift: int
    :return:
    :rtype: int
    """

    return (int8 >> shift) | (int8 << (8 - shift) & 0xff)

def shift_decode(string):
    """
    :param string:
    :return:
    :rtype: str
    """

    result = ''

    for char_code in string:
        result += chr(rotr8(char_code, 2))

    return result

def hex_decode(string):
    """
    :param string:
    :return:
    :rtype: str
    """

    return binascii.hexlify(string).decode()

def convert(ws2_filename, json_filename):
    """
    :param ws2_filename:
    :type ws2_filename: str
    :param json_filename
    :type json_filename: str
    :return:
    """

    with open(ws2_filename, 'rb') as f:
        text = f.read()

    pattern = re.compile(
        b'\x54(\x94\x31\x0D(?P<name>.*?))?\x00\x50(?P<id>..)'
        b'\x00\x00\x8D\xA1\x85\xC9\x00(?P<line>.*?)\x94\x2D\x94\x41\x00',
        re.DOTALL
    )

    result = []

    for match in pattern.finditer(text):
        result.append({
            'name': shift_decode(match.group('name') or ''),
            'id': hex_decode(match.group('id')),
            'line': {
                'en': shift_decode(match.group('line'))
            }
        })

    with open(json_filename, 'w+') as f:
        f.write(json.dumps(result, sort_keys=True, indent=4))

def main():
    convert(sys.argv[1], sys.argv[1] + '.json')

if __name__ == '__main__':
    main()
<commit_msg>Изменение интерфейса функции shift_decode<commit_after>__author__ = 'mongolrgata'

import sys
import re
import binascii
import json

def rotr8(int8, shift):
    """
    :param int8:
    :type int8: int
    :param shift:
    :type shift: int
    :return:
    :rtype: int
    """

    return (int8 >> shift) | (int8 << (8 - shift) & 0xff)

def shift_decode(string):
    """
    :param string:
    :type string: bytes
    :return:
    :rtype: bytes
    """

    result = bytearray()

    for char_code in string:
        result.append(rotr8(char_code, 2))

    return bytes(result)

def hex_decode(string):
    """
    :param string:
    :return:
    :rtype: str
    """

    return binascii.hexlify(string).decode()

def convert(ws2_filename, json_filename):
    """
    :param ws2_filename:
    :type ws2_filename: str
    :param json_filename
    :type json_filename: str
    :return:
    """

    with open(ws2_filename, 'rb') as f:
        text = f.read()

    pattern = re.compile(
        b'\x54(\x94\x31\x0D(?P<name>.*?))?\x00\x50(?P<id>..)'
        b'\x00\x00\x8D\xA1\x85\xC9\x00(?P<line>.*?)\x94\x2D\x94\x41\x00',
        re.DOTALL
    )

    result = []

    for match in pattern.finditer(text):
        result.append({
            'name': shift_decode(match.group('name') or '').decode(),
            'id': hex_decode(match.group('id')),
            'line': {
                'en': shift_decode(match.group('line')).decode()
            }
        })

    with open(json_filename, 'w+') as f:
        f.write(json.dumps(result, sort_keys=True, indent=4))

def main():
    convert(sys.argv[1], sys.argv[1] + '.json')

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright (c) 2010-2014, GEM Foundation.
#
# OpenQuake 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.
#
# OpenQuake 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 General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""Utility functions related to splitting work into tasks."""

import operator

from celery.result import ResultSet
from celery.app import current_app
from celery.task import task

from openquake.baselib.general import split_in_blocks, AccumDict
from openquake.commonlib.parallel import \
    TaskManager, safely_call, check_mem_usage, pickle_sequence, no_distribute
from openquake.engine import logs
from openquake.engine.db import models
from openquake.engine.utils import config
from openquake.engine.writer import CacheInserter
from openquake.engine.performance import EnginePerformanceMonitor

CONCURRENT_TASKS = int(config.get('celery', 'concurrent_tasks'))
SOFT_MEM_LIMIT = int(config.get('memory', 'soft_mem_limit'))
HARD_MEM_LIMIT = int(config.get('memory', 'hard_mem_limit'))

class JobNotRunning(Exception):
    pass

class OqTaskManager(TaskManager):
    """
    A celery-based task manager. The usage is::

      oqm = OqTaskManager(do_something, logs.LOG.progress)
      oqm.send(arg1, arg2)
      oqm.send(arg3, arg4)
      print oqm.aggregate_results(agg, acc)

    Progress report is built-in.
    """
    def submit(self, *args):
        """
        Submit an oqtask with the given arguments to celery and return
        an AsyncResult. If the variable OQ_NO_DISTRIBUTE is set, the
        task function is run in process and the result is returned.
        """
        # log a warning if too much memory is used
        check_mem_usage(SOFT_MEM_LIMIT, HARD_MEM_LIMIT)
        if no_distribute():
            res = safely_call(self.oqtask.task_func, args)
        else:
            piks = pickle_sequence(args)
            self.sent += sum(len(p) for p in piks)
            res = self.oqtask.delay(*piks)
        self.results.append(res)

    def aggregate_result_set(self, agg, acc):
        """
        Loop on a set of celery AsyncResults and update the accumulator
        by using the aggregation function.

        :param agg: the aggregation function, (acc, val) -> new acc
        :param acc: the initial value of the accumulator
        :returns: the final value of the accumulator
        """
        if not self.results:
            return acc
        backend = current_app().backend
        rset = ResultSet(self.results)
        for task_id, result_dict in rset.iter_native():
            # log a warning if too much memory is used
            check_mem_usage(SOFT_MEM_LIMIT, HARD_MEM_LIMIT)
            result = result_dict['result']
            if isinstance(result, BaseException):
                raise result
            self.received += len(result)
            acc = agg(acc, result.unpickle())
            del backend._cache[task_id]  # work around a celery bug
        return acc

# a convenient alias
starmap = OqTaskManager.starmap

def apply_reduce(task, task_args,
                 agg=operator.add,
                 acc=None,
                 concurrent_tasks=CONCURRENT_TASKS,
                 weight=lambda item: 1,
                 key=lambda item: 'Unspecified',
                 name=None):
    """
    Apply a task to a tuple of the form (job_id, data, *args)
    by splitting the data in chunks and reduce the results with an
    aggregation function.

    :param task: an oqtask
    :param task_args: the arguments to be passed to the task function
    :param agg: the aggregation function
    :param acc: initial value of the accumulator
    :param concurrent_tasks: hint about how many tasks to generate
    :param weight: function to extract the weight of an item in data
    :param key: function to extract the kind of an item in data
    """
    if acc is None:
        acc = AccumDict()
    job_id = task_args[0]
    data = task_args[1]
    args = task_args[2:]
    if not data:
        return acc
    elif len(data) == 1 or not concurrent_tasks:
        return agg(acc, task.task_func(job_id, data, *args))
    blocks = split_in_blocks(data, concurrent_tasks, weight, key)
    task_args = [(job_id, block) + args for block in blocks]
    return starmap(task, task_args, logs.LOG.progress, name).reduce(agg, acc)

def oqtask(task_func):
    """
    Task function decorator which sets up logging and catches (and logs) any
    errors which occur inside the task. Also checks to make sure the job is
    actually still running. If it is not running, the task doesn't get
    executed, so we don't do useless computation.

    :param task_func: the function to decorate
    """
    def wrapped(*args):
        """
        Initialize logs, make sure the job is still running, and run the task
        code surrounded by a try-except. If any error occurs, log it as a
        critical failure.
        """
        # job_id is always assumed to be the first argument
        job_id = args[0]
        job = models.OqJob.objects.get(id=job_id)
        if job.is_running is False:
            # the job was killed, it is useless to run the task
            raise JobNotRunning(job_id)

        # it is important to save the task id soon, so that
        # the revoke functionality can work
        EnginePerformanceMonitor.store_task_id(job_id, tsk)

        with (EnginePerformanceMonitor(
                'total ' + task_func.__name__, job_id, tsk, flush=True),
              logs.handle(job)):
            try:
                # log a warning if too much memory is used
                check_mem_usage(SOFT_MEM_LIMIT, HARD_MEM_LIMIT)
                # run the task
                return task_func(*args)
            finally:
                # save on the db
                CacheInserter.flushall()
                # the task finished, we can remove from the performance
                # table the associated row 'storing task id'
                models.Performance.objects.filter(
                    oq_job=job,
                    operation='storing task id',
                    task_id=tsk.request.id).delete()
    celery_queue = config.get('amqp', 'celery_queue')
    f = lambda *args: safely_call(wrapped, args, pickle=True)
    f.__name__ = task_func.__name__
    tsk = task(f, queue=celery_queue)
    tsk.task_func = task_func
    return tsk
<commit_msg>Removed parens breaking a with statement<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright (c) 2010-2014, GEM Foundation.
#
# OpenQuake 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.
#
# OpenQuake 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 General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""Utility functions related to splitting work into tasks."""

import operator

from celery.result import ResultSet
from celery.app import current_app
from celery.task import task

from openquake.baselib.general import split_in_blocks, AccumDict
from openquake.commonlib.parallel import \
    TaskManager, safely_call, check_mem_usage, pickle_sequence, no_distribute
from openquake.engine import logs
from openquake.engine.db import models
from openquake.engine.utils import config
from openquake.engine.writer import CacheInserter
from openquake.engine.performance import EnginePerformanceMonitor

CONCURRENT_TASKS = int(config.get('celery', 'concurrent_tasks'))
SOFT_MEM_LIMIT = int(config.get('memory', 'soft_mem_limit'))
HARD_MEM_LIMIT = int(config.get('memory', 'hard_mem_limit'))

class JobNotRunning(Exception):
    pass

class OqTaskManager(TaskManager):
    """
    A celery-based task manager. The usage is::

      oqm = OqTaskManager(do_something, logs.LOG.progress)
      oqm.send(arg1, arg2)
      oqm.send(arg3, arg4)
      print oqm.aggregate_results(agg, acc)

    Progress report is built-in.
    """
    def submit(self, *args):
        """
        Submit an oqtask with the given arguments to celery and return
        an AsyncResult. If the variable OQ_NO_DISTRIBUTE is set, the
        task function is run in process and the result is returned.
        """
        # log a warning if too much memory is used
        check_mem_usage(SOFT_MEM_LIMIT, HARD_MEM_LIMIT)
        if no_distribute():
            res = safely_call(self.oqtask.task_func, args)
        else:
            piks = pickle_sequence(args)
            self.sent += sum(len(p) for p in piks)
            res = self.oqtask.delay(*piks)
        self.results.append(res)

    def aggregate_result_set(self, agg, acc):
        """
        Loop on a set of celery AsyncResults and update the accumulator
        by using the aggregation function.

        :param agg: the aggregation function, (acc, val) -> new acc
        :param acc: the initial value of the accumulator
        :returns: the final value of the accumulator
        """
        if not self.results:
            return acc
        backend = current_app().backend
        rset = ResultSet(self.results)
        for task_id, result_dict in rset.iter_native():
            # log a warning if too much memory is used
            check_mem_usage(SOFT_MEM_LIMIT, HARD_MEM_LIMIT)
            result = result_dict['result']
            if isinstance(result, BaseException):
                raise result
            self.received += len(result)
            acc = agg(acc, result.unpickle())
            del backend._cache[task_id]  # work around a celery bug
        return acc

# a convenient alias
starmap = OqTaskManager.starmap

def apply_reduce(task, task_args,
                 agg=operator.add,
                 acc=None,
                 concurrent_tasks=CONCURRENT_TASKS,
                 weight=lambda item: 1,
                 key=lambda item: 'Unspecified',
                 name=None):
    """
    Apply a task to a tuple of the form (job_id, data, *args)
    by splitting the data in chunks and reduce the results with an
    aggregation function.

    :param task: an oqtask
    :param task_args: the arguments to be passed to the task function
    :param agg: the aggregation function
    :param acc: initial value of the accumulator
    :param concurrent_tasks: hint about how many tasks to generate
    :param weight: function to extract the weight of an item in data
    :param key: function to extract the kind of an item in data
    """
    if acc is None:
        acc = AccumDict()
    job_id = task_args[0]
    data = task_args[1]
    args = task_args[2:]
    if not data:
        return acc
    elif len(data) == 1 or not concurrent_tasks:
        return agg(acc, task.task_func(job_id, data, *args))
    blocks = split_in_blocks(data, concurrent_tasks, weight, key)
    task_args = [(job_id, block) + args for block in blocks]
    return starmap(task, task_args, logs.LOG.progress, name).reduce(agg, acc)

def oqtask(task_func):
    """
    Task function decorator which sets up logging and catches (and logs) any
    errors which occur inside the task. Also checks to make sure the job is
    actually still running. If it is not running, the task doesn't get
    executed, so we don't do useless computation.

    :param task_func: the function to decorate
    """
    def wrapped(*args):
        """
        Initialize logs, make sure the job is still running, and run the task
        code surrounded by a try-except. If any error occurs, log it as a
        critical failure.
        """
        # job_id is always assumed to be the first argument
        job_id = args[0]
        job = models.OqJob.objects.get(id=job_id)
        if job.is_running is False:
            # the job was killed, it is useless to run the task
            raise JobNotRunning(job_id)

        # it is important to save the task id soon, so that
        # the revoke functionality can work
        EnginePerformanceMonitor.store_task_id(job_id, tsk)

        with EnginePerformanceMonitor(
                'total ' + task_func.__name__, job_id, tsk, flush=True), \
                logs.handle(job):
            try:
                # log a warning if too much memory is used
                check_mem_usage(SOFT_MEM_LIMIT, HARD_MEM_LIMIT)
                # run the task
                return task_func(*args)
            finally:
                # save on the db
                CacheInserter.flushall()
                # the task finished, we can remove from the performance
                # table the associated row 'storing task id'
                models.Performance.objects.filter(
                    oq_job=job,
                    operation='storing task id',
                    task_id=tsk.request.id).delete()
    celery_queue = config.get('amqp', 'celery_queue')
    f = lambda *args: safely_call(wrapped, args, pickle=True)
    f.__name__ = task_func.__name__
    tsk = task(f, queue=celery_queue)
    tsk.task_func = task_func
    return tsk
<|endoftext|>


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

from ws4py.messaging import TextMessage, BinaryMessage, CloseControlMessage,\
     PingControlMessage, PongControlMessage
from ws4py.framing import Frame, OPCODE_CONTINUATION, OPCODE_TEXT, \
     OPCODE_BINARY, OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG
from ws4py.exc import FrameTooLargeException, ProtocolException, InvalidBytesError,\
     TextFrameEncodingException, UnsupportedFrameTypeException, StreamClosed

class Stream(object):
    def __init__(self):
        """
        Represents a websocket stream of bytes flowing in and out.

        The stream doesn't know about the data provider itself and
        doesn't even know about sockets. Instead the stream simply
        yields for more bytes whenever it requires it. The stream owner
        is responsible to provide the stream with those bytes until
        a frame can be interpreted.

        >>> s = Stream()
        >>> s.parser.send(BYTES)
        >>> s.has_messages
        False
        >>> s.parser.send(MORE_BYTES)
        >>> s.has_messages
        True
        >>> s.messages.pop()
        <TextMessage ... >

        """
        self.message = None
        """
        Parsed test or binary messages. Whenever the parser
        reads more bytes from a fragment message, those bytes
        are appended to the most recent message.
        """

        self.pings = []
        """
        Parsed ping control messages. They are instances of
        messaging.PingControlMessage
        """

        self.pongs = []
        """
        Parsed pong control messages. They are instances of
        messaging.PongControlMessage
        """

        self.closing = None
        """
        Parsed close control messsage. Instance of
        messaging.CloseControlMessage
        """

        self.errors = []
        """
        Detected errors while parsing. Instances of
        messaging.CloseControlMessage
        """

        self.parser = self.receiver()
        """
        Parser in charge to process bytes it is fed with.
        """

        # Python generators must be initialized once.
        self.parser.next()

    def text_message(self, text):
        """
        Returns a messaging.TextMessage instance
        ready to be built. Convenience method so
        that the caller doesn't need to import the
        TextMessage class itself.

        @param text: data to be carried by the message
        """
        return TextMessage(text=text)

    def binary_message(self, bytes):
        """
        Returns a messaging.BinaryMessage instance
        ready to be built. Convenience method so
        that the caller doesn't need to import the
        BinaryMessage class itself.

        @param text: data to be carried by the message
        """
        return BinaryMessage(bytes)

    @property
    def has_message(self):
        """
        Checks if the stream has received any message
        which, if fragmented, is completed.
        """
        if self.message is not None:
            return self.message.completed

        return False

    def close(self, code=1000, reason=''):
        """
        Returns a close control message built from
        a messaging.CloseControlMessage instance.

        @param code: closing status code
        @param reason: status message
        @return: bytes representing a close control single framed message
        """
        return CloseControlMessage(code=code, reason=reason).single()

    def ping(self, data=''):
        """
        Returns a ping control message built from
        a messaging.PingControlMessage instance.

        @param data: ping data
        @return: bytes representing a ping single framed message
        """
        return PingControlMessage(data).single()

    def pong(self, data=''):
        """
        Returns a ping control message built from
        a messaging.PongControlMessage instance.

        @param data: pong data
        @return: bytes representing a pong single framed message
        """
        return PongControlMessage(data).single()

    def receiver(self):
        """
        Parser that keeps trying to interpret bytes it is fed with as
        incoming frames part of a message.

        Control message are single frames only while data messages, like text
        and binary, may be fragmented accross frames.

        The way it works is by instanciating a framing.Frame object,
        then running its parser generator which yields how much bytes
        it requires to performs its task. The stream parser yields this value
        to its caller and feeds the frame parser.

        When the frame parser raises StopIteration, the stream parser
        tries to make sense of the parsed frame. It dispatches the frame's bytes
        to the most appropriate message type based on the frame's opcode.

        Overall this makes the stream parser totally agonstic to
        the data provider.
        """
        running = True
        while running:
            frame = Frame()
            while True:
                try:
                    bytes = (yield frame.parser.next())
                    if bytes is None:
                        raise InvalidBytesError()

                    frame.parser.send(bytes)
                except StopIteration:
                    bytes = frame.body or ''
                    if frame.masking_key and bytes:
                        bytes = frame.unmask(bytes)

                    if frame.opcode == OPCODE_TEXT:
                        if self.message and not self.message.completed:
                            # We got a text frame before we completed the previous one
                            self.errors.append(CloseControlMessage(code=1002))
                            break

                        try:
                            m = TextMessage(bytes.decode("utf-8"))
                            m.completed = (frame.fin == 1)
                            self.message = m
                        except UnicodeDecodeError:
                            self.errors.append(CloseControlMessage(code=1007))
                            break

                    elif frame.opcode == OPCODE_BINARY:
                        m = BinaryMessage(bytes)
                        m.completed = (frame.fin == 1)
                        self.message = m

                    elif frame.opcode == OPCODE_CONTINUATION:
                        m = self.message
                        if m is None:
                            self.errors.append(CloseControlMessage(code=1002))
                            break

                        m.completed = (frame.fin == 1)
                        if m.opcode == OPCODE_TEXT:
                            try:
                                m.extend(bytes.decode("utf-8"))
                            except UnicodeDecodeError:
                                self.errors.append(CloseControlMessage(code=1007))
                                break
                        else:
                            m.extend(bytes)

                    elif frame.opcode == OPCODE_CLOSE:
                        code = 1000
                        reason = ""
                        bytes = str(bytes)
                        if len(bytes) > 1:
                            code = struct.unpack("!H", bytes[0:2])[0]
                            if len(bytes) > 2:
                                reason = bytes[2:].decode("utf-8", "replace")
                        self.closing = CloseControlMessage(code=code, reason=reason)

                    elif frame.opcode == OPCODE_PING:
                        self.pings.append(PingControlMessage(bytes))

                    elif frame.opcode == OPCODE_PONG:
                        self.pongs.append(PongControlMessage(bytes))

                    else:
                        self.errors.append(CloseControlMessage(code=1003))

                    # When the frame's payload is empty, we must yield
                    # once more so that the caller is properly aligned
                    if not bytes:
                        yield 0

                    break

                except ProtocolException:
                    self.errors.append(CloseControlMessage(code=1002))
                except FrameTooLargeException:
                    self.errors.append(CloseControlMessage(code=1004))
                except StreamClosed:
                    running = False
                    break

            frame.parser.close()
<commit_msg>fixes for various closing frames use cases<commit_after># -*- coding: utf-8 -*-
import struct

from ws4py.messaging import TextMessage, BinaryMessage, CloseControlMessage,\
     PingControlMessage, PongControlMessage
from ws4py.framing import Frame, OPCODE_CONTINUATION, OPCODE_TEXT, \
     OPCODE_BINARY, OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG
from ws4py.exc import FrameTooLargeException, ProtocolException, InvalidBytesError,\
     TextFrameEncodingException, UnsupportedFrameTypeException, StreamClosed

class Stream(object):
    def __init__(self):
        """
        Represents a websocket stream of bytes flowing in and out.

        The stream doesn't know about the data provider itself and
        doesn't even know about sockets. Instead the stream simply
        yields for more bytes whenever it requires it. The stream owner
        is responsible to provide the stream with those bytes until
        a frame can be interpreted.

        >>> s = Stream()
        >>> s.parser.send(BYTES)
        >>> s.has_messages
        False
        >>> s.parser.send(MORE_BYTES)
        >>> s.has_messages
        True
        >>> s.messages.pop()
        <TextMessage ... >

        """
        self.message = None
        """
        Parsed test or binary messages. Whenever the parser
        reads more bytes from a fragment message, those bytes
        are appended to the most recent message.
        """

        self.pings = []
        """
        Parsed ping control messages. They are instances of
        messaging.PingControlMessage
        """

        self.pongs = []
        """
        Parsed pong control messages. They are instances of
        messaging.PongControlMessage
        """

        self.closing = None
        """
        Parsed close control messsage. Instance of
        messaging.CloseControlMessage
        """

        self.errors = []
        """
        Detected errors while parsing. Instances of
        messaging.CloseControlMessage
        """

        self.parser = self.receiver()
        """
        Parser in charge to process bytes it is fed with.
        """

        # Python generators must be initialized once.
        self.parser.next()

    def text_message(self, text):
        """
        Returns a messaging.TextMessage instance
        ready to be built. Convenience method so
        that the caller doesn't need to import the
        TextMessage class itself.

        @param text: data to be carried by the message
        """
        return TextMessage(text=text)

    def binary_message(self, bytes):
        """
        Returns a messaging.BinaryMessage instance
        ready to be built. Convenience method so
        that the caller doesn't need to import the
        BinaryMessage class itself.

        @param text: data to be carried by the message
        """
        return BinaryMessage(bytes)

    @property
    def has_message(self):
        """
        Checks if the stream has received any message
        which, if fragmented, is completed.
        """
        if self.message is not None:
            return self.message.completed

        return False

    def close(self, code=1000, reason=''):
        """
        Returns a close control message built from
        a messaging.CloseControlMessage instance.

        @param code: closing status code
        @param reason: status message
        @return: bytes representing a close control single framed message
        """
        return CloseControlMessage(code=code, reason=reason).single()

    def ping(self, data=''):
        """
        Returns a ping control message built from
        a messaging.PingControlMessage instance.

        @param data: ping data
        @return: bytes representing a ping single framed message
        """
        return PingControlMessage(data).single()

    def pong(self, data=''):
        """
        Returns a ping control message built from
        a messaging.PongControlMessage instance.

        @param data: pong data
        @return: bytes representing a pong single framed message
        """
        return PongControlMessage(data).single()

    def receiver(self):
        """
        Parser that keeps trying to interpret bytes it is fed with as
        incoming frames part of a message.

        Control message are single frames only while data messages, like text
        and binary, may be fragmented accross frames.

        The way it works is by instanciating a framing.Frame object,
        then running its parser generator which yields how much bytes
        it requires to performs its task. The stream parser yields this value
        to its caller and feeds the frame parser.

        When the frame parser raises StopIteration, the stream parser
        tries to make sense of the parsed frame. It dispatches the frame's bytes
        to the most appropriate message type based on the frame's opcode.

        Overall this makes the stream parser totally agonstic to
        the data provider.
        """
        running = True
        while running:
            frame = Frame()
            while True:
                try:
                    bytes = (yield frame.parser.next())
                    if bytes is None:
                        raise InvalidBytesError()

                    frame.parser.send(bytes)
                except StopIteration:
                    bytes = frame.body or ''
                    if frame.masking_key and bytes:
                        bytes = frame.unmask(bytes)

                    if frame.opcode == OPCODE_TEXT:
                        if self.message and not self.message.completed:
                            # We got a text frame before we completed the previous one
                            self.errors.append(CloseControlMessage(code=1002))
                            break

                        try:
                            m = TextMessage(bytes.decode("utf-8"))
                            m.completed = (frame.fin == 1)
                            self.message = m
                        except UnicodeDecodeError:
                            self.errors.append(CloseControlMessage(code=1007))
                            break

                    elif frame.opcode == OPCODE_BINARY:
                        m = BinaryMessage(bytes)
                        m.completed = (frame.fin == 1)
                        self.message = m

                    elif frame.opcode == OPCODE_CONTINUATION:
                        m = self.message
                        if m is None:
                            self.errors.append(CloseControlMessage(code=1002))
                            break

                        m.completed = (frame.fin == 1)
                        if m.opcode == OPCODE_TEXT:
                            try:
                                m.extend(bytes.decode("utf-8"))
                            except UnicodeDecodeError:
                                self.errors.append(CloseControlMessage(code=1007))
                                break
                        else:
                            m.extend(bytes)

                    elif frame.opcode == OPCODE_CLOSE:
                        code = 1000
                        reason = ""
                        if len(bytes) == 0:
                            self.errors.append(CloseControlMessage(code=1000))
                        elif 1 < len(bytes) < 126:
                            code = struct.unpack("!H", str(bytes[0:2]))[0]
                            try:
                                code = int(code)
                            except TypeError:
                                code = 1002
                                reason = 'Invalid Closing Frame Code Type'
                            else:
                                # Those codes are reserved or plainly forbidden
                                if code < 1000 or code in [1004, 1005, 1006, 1012, 1013, 1014, 1015,
                                                           1016, 1100, 2000, 2999, 5000, 65536]:
                                    code = 1002
                                    reason = 'Invalid Closing Frame Code'
                                else:
                                    if len(bytes) > 2:
                                        try:
                                            reason = frame.body[2:].decode("utf-8")
                                        except UnicodeDecodeError:
                                            code = 1007
                                            reason = ''
                            self.closing = CloseControlMessage(code=code, reason=reason)
                        else:
                            self.errors.append(CloseControlMessage(code=1002))

                    elif frame.opcode == OPCODE_PING:
                        self.pings.append(PingControlMessage(bytes))

                    elif frame.opcode == OPCODE_PONG:
                        self.pongs.append(PongControlMessage(bytes))

                    else:
                        self.errors.append(CloseControlMessage(code=1003))

                    # When the frame's payload is empty, we must yield
                    # once more so that the caller is properly aligned
                    if not bytes:
                        yield 0

                    break

                except ProtocolException:
                    self.errors.append(CloseControlMessage(code=1002))
                except FrameTooLargeException:
                    self.errors.append(CloseControlMessage(code=1002))
                except StreamClosed:
                    running = False
                    break

            frame.parser.close()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright (c) 2010-2014, GEM Foundation.
#
# OpenQuake 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.
#
# OpenQuake 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 General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""Utility functions related to splitting work into tasks."""

from celery.result import ResultSet
from celery.app import current_app
from celery.task import task

from openquake.commonlib.general import split_in_blocks
from openquake.commonlib.parallel import \
    TaskManager, safely_call, check_mem_usage, pickle_sequence, no_distribute
from openquake.engine import logs
from openquake.engine.db import models
from openquake.engine.utils import config
from openquake.engine.writer import CacheInserter
from openquake.engine.performance import EnginePerformanceMonitor

CONCURRENT_TASKS = int(config.get('celery', 'concurrent_tasks'))

class JobNotRunning(Exception):
    pass

class OqTaskManager(TaskManager):
    """
    A celery-based task manager. The usage is::

      oqm = OqTaskManager(do_something, logs.LOG.progress)
      oqm.send(arg1, arg2)
      oqm.send(arg3, arg4)
      print oqm.aggregate_results(agg, acc)

    Progress report is built-in.
    """
    def submit(self, *args):
        """
        Submit an oqtask with the given arguments to celery and return
        an AsyncResult. If the variable OQ_NO_DISTRIBUTE is set, the
        task function is run in process and the result is returned.
        """
        check_mem_usage()  # log a warning if too much memory is used
        if no_distribute():
            res = safely_call(self.oqtask.task_func, args)
        else:
            piks = pickle_sequence(args)
            self.sent += sum(len(p) for p in piks)
            res = self.oqtask.delay(*piks)
        self.results.append(res)

    def aggregate_result_set(self, agg, acc):
        """
        Loop on a set of celery AsyncResults and update the accumulator
        by using the aggregation function.

        :param agg: the aggregation function, (acc, val) -> new acc
        :param acc: the initial value of the accumulator
        :returns: the final value of the accumulator
        """
        if not self.results:
            return acc
        backend = current_app().backend
        rset = ResultSet(self.results)
        for task_id, result_dict in rset.iter_native():
            check_mem_usage()  # log a warning if too much memory is used
            result = result_dict['result']
            if isinstance(result, BaseException):
                raise result
            self.received += len(result)
            acc = agg(acc, result.unpickle())
            del backend._cache[task_id]  # work around a celery bug
        return acc

def map_reduce(task, task_args, agg, acc, name=None):
    """
    Given a task and an iterable of positional arguments, apply the
    task function to the arguments in parallel and return an aggregate
    result depending on the initial value of the accumulator
    and on the aggregation function. To save memory, the order is
    not preserved and there is no list with the intermediated results:
    the accumulator is incremented as soon as a task result comes.

    NB: if the environment variable OQ_NO_DISTRIBUTE is set the
    tasks are run sequentially in the current process and then
    map_reduce(task, task_args, agg, acc) is the same as
    reduce(agg, itertools.starmap(task, task_args), acc).
    Users of map_reduce should be aware of the fact that when
    thousands of tasks are spawned and large arguments are passed
    or large results are returned they may incur in memory issue:
    this is way the calculators limit the queue with the
    `concurrent_task` concept.

    :param task: a `celery` task callable.
    :param task_args: an iterable over positional arguments
    :param agg: the aggregation function, (acc, val) -> new acc
    :param acc: the initial value of the accumulator
    :returns: the final value of the accumulator
    """
    oqm = OqTaskManager(task, logs.LOG.progress, name)
    for i, args in enumerate(task_args, 1):
        logs.LOG.info('Submitting task %s #%d', oqm.name, i)
        oqm.submit(*args)
    return oqm.aggregate_results(agg, acc)

def apply_reduce(task, task_args,
                 agg=lambda a, x: x,
                 acc=None,
                 concurrent_tasks=CONCURRENT_TASKS,
                 weight=lambda item: 1,
                 key=lambda item: 'Unspecified',
                 name=None):
    """
    Apply a task to a tuple of the form (job_id, data, *args)
    by splitting the data in chunks and reduce the results with an
    aggregation function.

    :param task: an oqtask
    :param task_args: the arguments to be passed to the task function
    :param agg: the aggregation function
    :param acc: initial value of the accumulator
    :param concurrent_tasks: hint about how many tasks to generate
    :param weight: function to extract the weight of an item in data
    :param key: function to extract the kind of an item in data
    """
    job_id = task_args[0]
    data = task_args[1]
    args = task_args[2:]
    if not data:
        return acc
    elif len(data) == 1 or not concurrent_tasks:
        return agg(acc, task.task_func(job_id, data, *args))
    blocks = split_in_blocks(data, concurrent_tasks, weight, key)
    all_args = [(job_id, block) + args for block in blocks]
    return map_reduce(task, all_args, agg, acc, name)

# used to implement BaseCalculator.parallelize, which takes in account
# the `concurrent_task` concept to avoid filling the Celery queue
def parallelize(task, task_args, side_effect=lambda val: None):
    """
    Given a celery task and an iterable of positional arguments, apply the
    callable to the arguments in parallel. It is possible to pass a
    function side_effect(val) which takes the return value of the
    callable and does something with it (such as saving or printing
    it). Notice that the order is not preserved. parallelize returns None.

    NB: if the environment variable OQ_NO_DISTRIBUTE is set the
    tasks are run sequentially in the current process.

    :param task: a celery task
    :param task_args: an iterable over positional arguments
    :param side_effect: a function val -> None
    """
    map_reduce(task, task_args, lambda acc, val: side_effect(val), None)

def oqtask(task_func):
    """
    Task function decorator which sets up logging and catches (and logs) any
    errors which occur inside the task. Also checks to make sure the job is
    actually still running. If it is not running, the task doesn't get
    executed, so we don't do useless computation.

    :param task_func: the function to decorate
    """
    def wrapped(*args):
        """
        Initialize logs, make sure the job is still running, and run the task
        code surrounded by a try-except. If any error occurs, log it as a
        critical failure.
        """
        # job_id is always assumed to be the first argument
        job_id = args[0]
        job = models.OqJob.objects.get(id=job_id)
        if job.is_running is False:
            # the job was killed, it is useless to run the task
            raise JobNotRunning(job_id)

        # it is important to save the task id soon, so that
        # the revoke functionality can work
        EnginePerformanceMonitor.store_task_id(job_id, tsk)

        with EnginePerformanceMonitor(
                'total ' + task_func.__name__, job_id, tsk, flush=True):
            # tasks write on the celery log file
            logs.set_level(job.log_level)
            check_mem_usage()  # log a warning if too much memory is used
            try:
                # run the task
                return task_func(*args)
            finally:
                # save on the db
                CacheInserter.flushall()
                # the task finished, we can remove from the performance
                # table the associated row 'storing task id'
                models.Performance.objects.filter(
                    oq_job=job,
                    operation='storing task id',
                    task_id=tsk.request.id).delete()
    celery_queue = config.get('amqp', 'celery_queue')
    f = lambda *args: safely_call(wrapped, args, pickle=True)
    f.__name__ = task_func.__name__
    tsk = task(f, queue=celery_queue)
    tsk.task_func = task_func
    return tsk
<commit_msg>Add support for customizable hard and soft memory limits<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright (c) 2010-2014, GEM Foundation.
#
# OpenQuake 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.
#
# OpenQuake 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 General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""Utility functions related to splitting work into tasks."""

from celery.result import ResultSet
from celery.app import current_app
from celery.task import task

from openquake.commonlib.general import split_in_blocks
from openquake.commonlib.parallel import \
    TaskManager, safely_call, check_mem_usage, pickle_sequence, no_distribute
from openquake.engine import logs
from openquake.engine.db import models
from openquake.engine.utils import config
from openquake.engine.writer import CacheInserter
from openquake.engine.performance import EnginePerformanceMonitor

CONCURRENT_TASKS = int(config.get('celery', 'concurrent_tasks'))
SOFT_MEM_LIMIT = int(config.get('memory', 'soft_mem_limit'))
HARD_MEM_LIMIT = int(config.get('memory', 'hard_mem_limit'))

class JobNotRunning(Exception):
    pass

class OqTaskManager(TaskManager):
    """
    A celery-based task manager. The usage is::

      oqm = OqTaskManager(do_something, logs.LOG.progress)
      oqm.send(arg1, arg2)
      oqm.send(arg3, arg4)
      print oqm.aggregate_results(agg, acc)

    Progress report is built-in.
    """
    def submit(self, *args):
        """
        Submit an oqtask with the given arguments to celery and return
        an AsyncResult. If the variable OQ_NO_DISTRIBUTE is set, the
        task function is run in process and the result is returned.
        """
        # log a warning if too much memory is used
        check_mem_usage(SOFT_MEM_LIMIT, HARD_MEM_LIMIT)
        if no_distribute():
            res = safely_call(self.oqtask.task_func, args)
        else:
            piks = pickle_sequence(args)
            self.sent += sum(len(p) for p in piks)
            res = self.oqtask.delay(*piks)
        self.results.append(res)

    def aggregate_result_set(self, agg, acc):
        """
        Loop on a set of celery AsyncResults and update the accumulator
        by using the aggregation function.

        :param agg: the aggregation function, (acc, val) -> new acc
        :param acc: the initial value of the accumulator
        :returns: the final value of the accumulator
        """
        if not self.results:
            return acc
        backend = current_app().backend
        rset = ResultSet(self.results)
        for task_id, result_dict in rset.iter_native():
            # log a warning if too much memory is used
            check_mem_usage(SOFT_MEM_LIMIT, HARD_MEM_LIMIT)
            result = result_dict['result']
            if isinstance(result, BaseException):
                raise result
            self.received += len(result)
            acc = agg(acc, result.unpickle())
            del backend._cache[task_id]  # work around a celery bug
        return acc

def map_reduce(task, task_args, agg, acc, name=None):
    """
    Given a task and an iterable of positional arguments, apply the
    task function to the arguments in parallel and return an aggregate
    result depending on the initial value of the accumulator
    and on the aggregation function. To save memory, the order is
    not preserved and there is no list with the intermediated results:
    the accumulator is incremented as soon as a task result comes.

    NB: if the environment variable OQ_NO_DISTRIBUTE is set the
    tasks are run sequentially in the current process and then
    map_reduce(task, task_args, agg, acc) is the same as
    reduce(agg, itertools.starmap(task, task_args), acc).
    Users of map_reduce should be aware of the fact that when
    thousands of tasks are spawned and large arguments are passed
    or large results are returned they may incur in memory issue:
    this is way the calculators limit the queue with the
    `concurrent_task` concept.

    :param task: a `celery` task callable.
    :param task_args: an iterable over positional arguments
    :param agg: the aggregation function, (acc, val) -> new acc
    :param acc: the initial value of the accumulator
    :returns: the final value of the accumulator
    """
    oqm = OqTaskManager(task, logs.LOG.progress, name)
    for i, args in enumerate(task_args, 1):
        logs.LOG.info('Submitting task %s #%d', oqm.name, i)
        oqm.submit(*args)
    return oqm.aggregate_results(agg, acc)

def apply_reduce(task, task_args,
                 agg=lambda a, x: x,
                 acc=None,
                 concurrent_tasks=CONCURRENT_TASKS,
                 weight=lambda item: 1,
                 key=lambda item: 'Unspecified',
                 name=None):
    """
    Apply a task to a tuple of the form (job_id, data, *args)
    by splitting the data in chunks and reduce the results with an
    aggregation function.

    :param task: an oqtask
    :param task_args: the arguments to be passed to the task function
    :param agg: the aggregation function
    :param acc: initial value of the accumulator
    :param concurrent_tasks: hint about how many tasks to generate
    :param weight: function to extract the weight of an item in data
    :param key: function to extract the kind of an item in data
    """
    job_id = task_args[0]
    data = task_args[1]
    args = task_args[2:]
    if not data:
        return acc
    elif len(data) == 1 or not concurrent_tasks:
        return agg(acc, task.task_func(job_id, data, *args))
    blocks = split_in_blocks(data, concurrent_tasks, weight, key)
    all_args = [(job_id, block) + args for block in blocks]
    return map_reduce(task, all_args, agg, acc, name)

# used to implement BaseCalculator.parallelize, which takes in account
# the `concurrent_task` concept to avoid filling the Celery queue
def parallelize(task, task_args, side_effect=lambda val: None):
    """
    Given a celery task and an iterable of positional arguments, apply the
    callable to the arguments in parallel. It is possible to pass a
    function side_effect(val) which takes the return value of the
    callable and does something with it (such as saving or printing
    it). Notice that the order is not preserved. parallelize returns None.

    NB: if the environment variable OQ_NO_DISTRIBUTE is set the
    tasks are run sequentially in the current process.

    :param task: a celery task
    :param task_args: an iterable over positional arguments
    :param side_effect: a function val -> None
    """
    map_reduce(task, task_args, lambda acc, val: side_effect(val), None)

def oqtask(task_func):
    """
    Task function decorator which sets up logging and catches (and logs) any
    errors which occur inside the task. Also checks to make sure the job is
    actually still running. If it is not running, the task doesn't get
    executed, so we don't do useless computation.

    :param task_func: the function to decorate
    """
    def wrapped(*args):
        """
        Initialize logs, make sure the job is still running, and run the task
        code surrounded by a try-except. If any error occurs, log it as a
        critical failure.
        """
        # job_id is always assumed to be the first argument
        job_id = args[0]
        job = models.OqJob.objects.get(id=job_id)
        if job.is_running is False:
            # the job was killed, it is useless to run the task
            raise JobNotRunning(job_id)

        # it is important to save the task id soon, so that
        # the revoke functionality can work
        EnginePerformanceMonitor.store_task_id(job_id, tsk)

        with EnginePerformanceMonitor(
                'total ' + task_func.__name__, job_id, tsk, flush=True):
            # tasks write on the celery log file
            logs.set_level(job.log_level)
            # log a warning if too much memory is used
            check_mem_usage(SOFT_MEM_LIMIT, HARD_MEM_LIMIT)
            try:
                # run the task
                return task_func(*args)
            finally:
                # save on the db
                CacheInserter.flushall()
                # the task finished, we can remove from the performance
                # table the associated row 'storing task id'
                models.Performance.objects.filter(
                    oq_job=job,
                    operation='storing task id',
                    task_id=tsk.request.id).delete()
    celery_queue = config.get('amqp', 'celery_queue')
    f = lambda *args: safely_call(wrapped, args, pickle=True)
    f.__name__ = task_func.__name__
    tsk = task(f, queue=celery_queue)
    tsk.task_func = task_func
    return tsk
<|endoftext|>


<commit_before>import types
from functools import partial

from scrapy.http import Request
from scrapy.utils.spider import iterate_spider_output

def inline_requests(method):
    """Decorator to enable inline requests"""
    def wrapper(self, response):
        callback = types.MethodType(method, self, self.__class__)
        genwrapper = _RequestGenerator(callback)
        return genwrapper(response)
    return wrapper

class _RequestGenerator(object):

    def __init__(self, callback):
        self.callback = callback

    def __call__(self, response):
        output = iterate_spider_output(self.callback(response))
        if isinstance(output, types.GeneratorType):
            return self._unwindGenerator(output)
        else:
            return output

    def _unwindGenerator(self, generator, _prev=None):
        while True:
            if _prev:
                ret, _prev = _prev, None
            else:
                try:
                    ret = generator.next()
                except StopIteration:
                    break
            if isinstance(ret, Request) and ret.callback is None:
                yield self._wrapRequest(ret, generator)
                break
            else:
                yield ret

    def _wrapRequest(self, request, generator):
        request.callback = partial(self._handleSuccess, generator=generator)
        request.errback = partial(self._handleFailure, generator=generator)
        return request

    def _handleSuccess(self, response, generator):
        try:
            ret = generator.send(response)
        except StopIteration:
            return
        return self._unwindGenerator(generator, ret)

    def _handleFailure(self, failure, generator):
        try:
            ret = failure.throwExceptionIntoGenerator(generator)
        except StopIteration:
            return
        return self._unwindGenerator(generator, ret)

<commit_msg>added more magic: allow to decorate either methods or functions.<commit_after>import inspect
import types
from functools import partial, wraps

from scrapy.http import Request
from scrapy.utils.spider import iterate_spider_output

def inline_requests(method_or_func):
    args = inspect.getargspec(method_or_func).args
    if not args:
        raise TypeError("Function must accept at least one argument.")
    # XXX: hardcoded convention of 'self' as first argument for methods
    if args[0] == 'self':
        def wrapper(self, response, **kwargs):
            callback = types.MethodType(method_or_func, self, self.__class__)
            genwrapper = _RequestGenerator(callback, **kwargs)
            return genwrapper(response)
    else:
        def wrapper(response, **kwargs):
            genwrapper = _RequestGenerator(method_or_func, **kwargs)
            return genwrapper(response)

    return wraps(method_or_func)(wrapper)

class _RequestGenerator(object):

    def __init__(self, callback, **kwargs):
        self.callback = callback
        self.kwargs = kwargs

    def __call__(self, response):
        output = iterate_spider_output(self.callback(response=response, **self.kwargs))
        if isinstance(output, types.GeneratorType):
            return self._unwindGenerator(output)
        else:
            return output

    def _unwindGenerator(self, generator, _prev=None):
        while True:
            if _prev:
                ret, _prev = _prev, None
            else:
                try:
                    ret = generator.next()
                except StopIteration:
                    break
            if isinstance(ret, Request) and ret.callback is None:
                yield self._wrapRequest(ret, generator)
                break
            else:
                yield ret

    def _wrapRequest(self, request, generator):
        request.callback = partial(self._handleSuccess, generator=generator)
        request.errback = partial(self._handleFailure, generator=generator)
        return request

    def _handleSuccess(self, response, generator):
        try:
            ret = generator.send(response)
        except StopIteration:
            return
        return self._unwindGenerator(generator, ret)

    def _handleFailure(self, failure, generator):
        try:
            ret = failure.throwExceptionIntoGenerator(generator)
        except StopIteration:
            return
        return self._unwindGenerator(generator, ret)
<|endoftext|>


<commit_before>"""
This file demonstrates writing tests using the unittest module. These will pass
when you run "manage.py test".

Replace this with more appropriate tests for your application.
"""

from django.test import TestCase

class SimpleTest(TestCase):

    def test_basic_addition(self):
        """
        Tests that 1 + 1 always equals 2.
        """
        self.assertEqual(1 + 1, 2)
<commit_msg>Remove unused tests<commit_after><|endoftext|>