<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2010-2018 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 Affero 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/>.

"""Engine: A collection of fundamental functions for initializing and running
calculations."""

import os
import re
import sys
import json
import signal
import traceback
import platform
try:
    from setproctitle import setproctitle
except ImportError:
    def setproctitle(title):
        "Do nothing"
from urllib.request import urlopen, Request
from openquake.baselib.performance import Monitor
from openquake.baselib.python3compat import decode
from openquake.baselib import (
    parallel, general, config, datastore, __version__, zeromq as z)
from openquake.commonlib.oqvalidation import OqParam
from openquake.commonlib import readinput
from openquake.calculators import base, views, export
from openquake.commonlib import logs

OQ_API = 'https://api.openquake.org'
TERMINATE = config.distribution.terminate_workers_on_revoke
OQ_DISTRIBUTE = parallel.oq_distribute()

_PPID = os.getppid()  # the controlling terminal PID

if OQ_DISTRIBUTE == 'zmq':

    def set_concurrent_tasks_default():
        """
        Set the default for concurrent_tasks based on the available
        worker pools .
        """
        num_workers = 0
        w = config.zworkers
        for host, _cores in [hc.split() for hc in w.host_cores.split(',')]:
            url = 'tcp://%s:%s' % (host, w.ctrl_port)
            with z.Socket(url, z.zmq.REQ, 'connect') as sock:
                if not general.socket_ready(url):
                    logs.LOG.warn('%s is not running', host)
                    continue
                num_workers += sock.send('get_num_workers')
        OqParam.concurrent_tasks.default = num_workers * 3
        logs.LOG.info('Using %d zmq workers', num_workers)

elif OQ_DISTRIBUTE.startswith('celery'):
    import celery.task.control

    def set_concurrent_tasks_default():
        """
        Set the default for concurrent_tasks based on the number of available
        celery workers.
        """
        stats = celery.task.control.inspect(timeout=1).stats()
        if not stats:
            sys.exit("No live compute nodes, aborting calculation")
        num_cores = sum(stats[k]['pool']['max-concurrency'] for k in stats)
        OqParam.concurrent_tasks.default = num_cores * 3
        logs.LOG.info(
            'Using %s, %d cores', ', '.join(sorted(stats)), num_cores)

    def celery_cleanup(terminate, task_ids=()):
        """
        Release the resources used by an openquake job.
        In particular revoke the running tasks (if any).

        :param bool terminate: the celery revoke command terminate flag
        :param task_ids: celery task IDs
        """
        # Using the celery API, terminate and revoke and terminate any running
        # tasks associated with the current job.
        if task_ids:
            logs.LOG.warn('Revoking %d tasks', len(task_ids))
        else:  # this is normal when OQ_DISTRIBUTE=no
            logs.LOG.debug('No task to revoke')
        for tid in task_ids:
            celery.task.control.revoke(tid, terminate=terminate)
            logs.LOG.debug('Revoked task %s', tid)

def expose_outputs(dstore):
    """
    Build a correspondence between the outputs in the datastore and the
    ones in the database.

    :param dstore: datastore
    """
    oq = dstore['oqparam']
    exportable = set(ekey[0] for ekey in export.export)
    calcmode = oq.calculation_mode
    dskeys = set(dstore) & exportable  # exportable datastore keys
    dskeys.add('fullreport')
    rlzs = dstore['csm_info'].rlzs
    if len(rlzs) > 1:
        dskeys.add('realizations')
    # expose gmf_data only if < 10 MB
    if oq.ground_motion_fields and calcmode == 'event_based':
        nbytes = dstore['gmf_data'].attrs['nbytes']
        if nbytes < 10 * 1024 ** 2:
            dskeys.add('gmf_data')
    if 'scenario' not in calcmode:  # export sourcegroups.csv
        dskeys.add('sourcegroups')
    hdf5 = dstore.hdf5
    if (len(rlzs) == 1 and 'poes' in hdf5) or 'hcurves' in hdf5:
        dskeys.add('hcurves')
        if oq.uniform_hazard_spectra:
            dskeys.add('uhs')  # export them
        if oq.hazard_maps:
            dskeys.add('hmaps')  # export them
    if 'avg_losses-stats' in dstore or (
            'avg_losses-rlzs' in dstore and len(rlzs)):
        dskeys.add('avg_losses-stats')
    if 'curves-stats' in dstore:
        logs.LOG.warn('loss curves are exportable with oq export')
    if oq.conditional_loss_poes:  # expose loss_maps outputs
        if 'loss_curves-stats' in dstore:
            dskeys.add('loss_maps-stats')
    if 'all_loss_ratios' in dskeys:
        dskeys.remove('all_loss_ratios')  # export only specific IDs
    if 'ruptures' in dskeys and 'scenario' in calcmode:
        exportable.remove('ruptures')  # do not export, as requested by Vitor
    if 'rup_loss_table' in dskeys:  # keep it hidden for the moment
        dskeys.remove('rup_loss_table')
    logs.dbcmd('create_outputs', dstore.calc_id, sorted(dskeys & exportable))

class MasterKilled(KeyboardInterrupt):
    "Exception raised when a job is killed manually"

def raiseMasterKilled(signum, _stack):
    """
    When a SIGTERM is received, raise the MasterKilled
    exception with an appropriate error message.

    :param int signum: the number of the received signal
    :param _stack: the current frame object, ignored
    """
    msg = 'Received a signal %d' % signum

    if signum in (signal.SIGTERM, signal.SIGINT):
        msg = 'The openquake master process was killed manually'

    # kill the calculation only if os.getppid() != _PPID, i.e. the controlling
    # terminal died; in the workers, do nothing
    # NB: there is no SIGHUP on Windows
    if hasattr(signal, 'SIGHUP'):
        if signum == signal.SIGHUP:
            if os.getppid() == _PPID:
                return
            else:
                msg = 'The openquake master lost its controlling terminal'

    parallel.Starmap.shutdown()
    raise MasterKilled(msg)

# register the raiseMasterKilled callback for SIGTERM
# when using the Django development server this module is imported by a thread,
# so one gets a `ValueError: signal only works in main thread` that
# can be safely ignored
try:
    signal.signal(signal.SIGTERM, raiseMasterKilled)
    signal.signal(signal.SIGINT, raiseMasterKilled)
    if hasattr(signal, 'SIGHUP'):
        signal.signal(signal.SIGHUP, raiseMasterKilled)
except ValueError:
    pass

def job_from_file(cfg_file, username, hazard_calculation_id=None):
    """
    Create a full job profile from a job config file.

    :param str cfg_file:
        Path to a job.ini file.
    :param str username:
        The user who will own this job profile and all results
    :param str datadir:
        Data directory of the user
    :param hazard_calculation_id:
        ID of a previous calculation or None
    :returns:
        a pair (job_id, oqparam)
    """
    oq = readinput.get_oqparam(cfg_file, hc_id=hazard_calculation_id)
    job_id = logs.dbcmd('create_job', oq.calculation_mode, oq.description,
                        username, datastore.get_datadir(),
                        hazard_calculation_id)
    return job_id, oq

def run_calc(job_id, oqparam, log_level, log_file, exports,
             hazard_calculation_id=None, **kw):
    """
    Run a calculation.

    :param job_id:
        ID of the current job
    :param oqparam:
        :class:`openquake.commonlib.oqvalidation.OqParam` instance
    :param str log_level:
        The desired logging level. Valid choices are 'debug', 'info',
        'progress', 'warn', 'error', and 'critical'.
    :param str log_file:
        Complete path (including file name) to file where logs will be written.
        If `None`, logging will just be printed to standard output.
    :param exports:
        A comma-separated string of export types.
    """
    setproctitle('oq-job-%d' % job_id)
    monitor = Monitor('total runtime', measuremem=True)
    with logs.handle(job_id, log_level, log_file):  # run the job
        if OQ_DISTRIBUTE.startswith(('celery', 'zmq')):
            set_concurrent_tasks_default()
        msg = check_obsolete_version(oqparam.calculation_mode)
        if msg:
            logs.LOG.warn(msg)
        calc = base.calculators(oqparam, monitor, calc_id=job_id)
        monitor.hdf5path = calc.datastore.hdf5path
        calc.from_engine = True
        tb = 'None\n'
        try:
            logs.dbcmd('set_status', job_id, 'executing')
            _do_run_calc(calc, exports, hazard_calculation_id, **kw)
            duration = monitor.duration
            expose_outputs(calc.datastore)
            monitor.flush()
            records = views.performance_view(calc.datastore)
            logs.dbcmd('save_performance', job_id, records)
            calc.datastore.close()
            logs.LOG.info('Calculation %d finished correctly in %d seconds',
                          job_id, duration)
            logs.dbcmd('finish', job_id, 'complete')
        except BaseException:
            tb = traceback.format_exc()
            try:
                logs.LOG.critical(tb)
                logs.dbcmd('finish', job_id, 'failed')
            except BaseException:  # an OperationalError may always happen
                sys.stderr.write(tb)
            raise
        finally:
            # if there was an error in the calculation, this part may fail;
            # in such a situation, we simply log the cleanup error without
            # taking further action, so that the real error can propagate
            try:
                if OQ_DISTRIBUTE.startswith('celery'):
                    celery_cleanup(TERMINATE, parallel.Starmap.task_ids)
            except BaseException:
                # log the finalization error only if there is no real error
                if tb == 'None\n':
                    logs.LOG.error('finalizing', exc_info=True)
    return calc

def _do_run_calc(calc, exports, hazard_calculation_id, **kw):
    with calc._monitor:
        calc.run(exports=exports, hazard_calculation_id=hazard_calculation_id,
                 close=False, **kw)  # don't close the datastore too soon

def version_triple(tag):
    """
    returns: a triple of integers from a version tag
    """
    groups = re.match(r'v?(\d+)\.(\d+)\.(\d+)', tag).groups()
    return tuple(int(n) for n in groups)

def check_obsolete_version(calculation_mode='WebUI'):
    """
    Check if there is a newer version of the engine.

    :param calculation_mode:
         - the calculation mode when called from the engine
         - an empty string when called from the WebUI
    :returns:
        - a message if the running version of the engine is obsolete
        - the empty string if the engine is updated
        - None if the check could not be performed (i.e. github is down)
    """
    if os.environ.get('JENKINS_URL') or os.environ.get('TRAVIS'):
        # avoid flooding our API server with requests from CI systems
        return

    headers = {'User-Agent': 'OpenQuake Engine %s;%s;%s;%s' %
               (__version__, calculation_mode, platform.platform(),
                config.distribution.oq_distribute)}
    try:
        req = Request(OQ_API + '/engine/latest', headers=headers)
        # NB: a timeout < 1 does not work
        data = urlopen(req, timeout=1).read()  # bytes
        tag_name = json.loads(decode(data))['tag_name']
        current = version_triple(__version__)
        latest = version_triple(tag_name)
    except:  # page not available or wrong version tag
        return
    if current < latest:
        return ('Version %s of the engine is available, but you are '
                'still using version %s' % (tag_name, __version__))
    else:
        return ''
<commit_msg>Inhibit multiple CTRL-C to give time Celery to revoke tasks<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2010-2018 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 Affero 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/>.

"""Engine: A collection of fundamental functions for initializing and running
calculations."""

import os
import re
import sys
import json
import signal
import traceback
import platform
try:
    from setproctitle import setproctitle
except ImportError:
    def setproctitle(title):
        "Do nothing"
from urllib.request import urlopen, Request
from openquake.baselib.performance import Monitor
from openquake.baselib.python3compat import decode
from openquake.baselib import (
    parallel, general, config, datastore, __version__, zeromq as z)
from openquake.commonlib.oqvalidation import OqParam
from openquake.commonlib import readinput
from openquake.calculators import base, views, export
from openquake.commonlib import logs

OQ_API = 'https://api.openquake.org'
TERMINATE = config.distribution.terminate_workers_on_revoke
OQ_DISTRIBUTE = parallel.oq_distribute()

_PPID = os.getppid()  # the controlling terminal PID

if OQ_DISTRIBUTE == 'zmq':

    def set_concurrent_tasks_default():
        """
        Set the default for concurrent_tasks based on the available
        worker pools .
        """
        num_workers = 0
        w = config.zworkers
        for host, _cores in [hc.split() for hc in w.host_cores.split(',')]:
            url = 'tcp://%s:%s' % (host, w.ctrl_port)
            with z.Socket(url, z.zmq.REQ, 'connect') as sock:
                if not general.socket_ready(url):
                    logs.LOG.warn('%s is not running', host)
                    continue
                num_workers += sock.send('get_num_workers')
        OqParam.concurrent_tasks.default = num_workers * 3
        logs.LOG.info('Using %d zmq workers', num_workers)

elif OQ_DISTRIBUTE.startswith('celery'):
    import celery.task.control

    def set_concurrent_tasks_default():
        """
        Set the default for concurrent_tasks based on the number of available
        celery workers.
        """
        stats = celery.task.control.inspect(timeout=1).stats()
        if not stats:
            sys.exit("No live compute nodes, aborting calculation")
        num_cores = sum(stats[k]['pool']['max-concurrency'] for k in stats)
        OqParam.concurrent_tasks.default = num_cores * 3
        logs.LOG.info(
            'Using %s, %d cores', ', '.join(sorted(stats)), num_cores)

    def celery_cleanup(terminate, task_ids=()):
        """
        Release the resources used by an openquake job.
        In particular revoke the running tasks (if any).

        :param bool terminate: the celery revoke command terminate flag
        :param task_ids: celery task IDs
        """
        # Using the celery API, terminate and revoke and terminate any running
        # tasks associated with the current job.
        if task_ids:
            logs.LOG.warn('Revoking %d tasks', len(task_ids))
        else:  # this is normal when OQ_DISTRIBUTE=no
            logs.LOG.debug('No task to revoke')
        for tid in task_ids:
            celery.task.control.revoke(tid, terminate=terminate)
            logs.LOG.debug('Revoked task %s', tid)

def expose_outputs(dstore):
    """
    Build a correspondence between the outputs in the datastore and the
    ones in the database.

    :param dstore: datastore
    """
    oq = dstore['oqparam']
    exportable = set(ekey[0] for ekey in export.export)
    calcmode = oq.calculation_mode
    dskeys = set(dstore) & exportable  # exportable datastore keys
    dskeys.add('fullreport')
    rlzs = dstore['csm_info'].rlzs
    if len(rlzs) > 1:
        dskeys.add('realizations')
    # expose gmf_data only if < 10 MB
    if oq.ground_motion_fields and calcmode == 'event_based':
        nbytes = dstore['gmf_data'].attrs['nbytes']
        if nbytes < 10 * 1024 ** 2:
            dskeys.add('gmf_data')
    if 'scenario' not in calcmode:  # export sourcegroups.csv
        dskeys.add('sourcegroups')
    hdf5 = dstore.hdf5
    if (len(rlzs) == 1 and 'poes' in hdf5) or 'hcurves' in hdf5:
        dskeys.add('hcurves')
        if oq.uniform_hazard_spectra:
            dskeys.add('uhs')  # export them
        if oq.hazard_maps:
            dskeys.add('hmaps')  # export them
    if 'avg_losses-stats' in dstore or (
            'avg_losses-rlzs' in dstore and len(rlzs)):
        dskeys.add('avg_losses-stats')
    if 'curves-stats' in dstore:
        logs.LOG.warn('loss curves are exportable with oq export')
    if oq.conditional_loss_poes:  # expose loss_maps outputs
        if 'loss_curves-stats' in dstore:
            dskeys.add('loss_maps-stats')
    if 'all_loss_ratios' in dskeys:
        dskeys.remove('all_loss_ratios')  # export only specific IDs
    if 'ruptures' in dskeys and 'scenario' in calcmode:
        exportable.remove('ruptures')  # do not export, as requested by Vitor
    if 'rup_loss_table' in dskeys:  # keep it hidden for the moment
        dskeys.remove('rup_loss_table')
    logs.dbcmd('create_outputs', dstore.calc_id, sorted(dskeys & exportable))

class MasterKilled(KeyboardInterrupt):
    "Exception raised when a job is killed manually"

def inhibitSigInt(signum, _stack):
    logs.LOG.warn('Revoking tasks, please wait.')

def raiseMasterKilled(signum, _stack):
    """
    When a SIGTERM is received, raise the MasterKilled
    exception with an appropriate error message.

    :param int signum: the number of the received signal
    :param _stack: the current frame object, ignored
    """
    # Disable further CTRL-C to allow tasks revocation
    signal.signal(signal.SIGINT, inhibitSigInt)

    msg = 'Received a signal %d' % signum
    if signum in (signal.SIGTERM, signal.SIGINT):
        msg = 'The openquake master process was killed manually'

    # kill the calculation only if os.getppid() != _PPID, i.e. the controlling
    # terminal died; in the workers, do nothing
    # NB: there is no SIGHUP on Windows
    if hasattr(signal, 'SIGHUP'):
        if signum == signal.SIGHUP:
            if os.getppid() == _PPID:
                return
            else:
                msg = 'The openquake master lost its controlling terminal'

    parallel.Starmap.shutdown()
    raise MasterKilled(msg)

# register the raiseMasterKilled callback for SIGTERM
# when using the Django development server this module is imported by a thread,
# so one gets a `ValueError: signal only works in main thread` that
# can be safely ignored
try:
    signal.signal(signal.SIGTERM, raiseMasterKilled)
    signal.signal(signal.SIGINT, raiseMasterKilled)
    if hasattr(signal, 'SIGHUP'):
        signal.signal(signal.SIGHUP, raiseMasterKilled)
except ValueError:
    pass

def job_from_file(cfg_file, username, hazard_calculation_id=None):
    """
    Create a full job profile from a job config file.

    :param str cfg_file:
        Path to a job.ini file.
    :param str username:
        The user who will own this job profile and all results
    :param str datadir:
        Data directory of the user
    :param hazard_calculation_id:
        ID of a previous calculation or None
    :returns:
        a pair (job_id, oqparam)
    """
    oq = readinput.get_oqparam(cfg_file, hc_id=hazard_calculation_id)
    job_id = logs.dbcmd('create_job', oq.calculation_mode, oq.description,
                        username, datastore.get_datadir(),
                        hazard_calculation_id)
    return job_id, oq

def run_calc(job_id, oqparam, log_level, log_file, exports,
             hazard_calculation_id=None, **kw):
    """
    Run a calculation.

    :param job_id:
        ID of the current job
    :param oqparam:
        :class:`openquake.commonlib.oqvalidation.OqParam` instance
    :param str log_level:
        The desired logging level. Valid choices are 'debug', 'info',
        'progress', 'warn', 'error', and 'critical'.
    :param str log_file:
        Complete path (including file name) to file where logs will be written.
        If `None`, logging will just be printed to standard output.
    :param exports:
        A comma-separated string of export types.
    """
    setproctitle('oq-job-%d' % job_id)
    monitor = Monitor('total runtime', measuremem=True)
    with logs.handle(job_id, log_level, log_file):  # run the job
        if OQ_DISTRIBUTE.startswith(('celery', 'zmq')):
            set_concurrent_tasks_default()
        msg = check_obsolete_version(oqparam.calculation_mode)
        if msg:
            logs.LOG.warn(msg)
        calc = base.calculators(oqparam, monitor, calc_id=job_id)
        monitor.hdf5path = calc.datastore.hdf5path
        calc.from_engine = True
        tb = 'None\n'
        try:
            logs.dbcmd('set_status', job_id, 'executing')
            _do_run_calc(calc, exports, hazard_calculation_id, **kw)
            duration = monitor.duration
            expose_outputs(calc.datastore)
            monitor.flush()
            records = views.performance_view(calc.datastore)
            logs.dbcmd('save_performance', job_id, records)
            calc.datastore.close()
            logs.LOG.info('Calculation %d finished correctly in %d seconds',
                          job_id, duration)
            logs.dbcmd('finish', job_id, 'complete')
        except BaseException:
            tb = traceback.format_exc()
            try:
                logs.LOG.critical(tb)
                logs.dbcmd('finish', job_id, 'failed')
            except BaseException:  # an OperationalError may always happen
                sys.stderr.write(tb)
            raise
        finally:
            # if there was an error in the calculation, this part may fail;
            # in such a situation, we simply log the cleanup error without
            # taking further action, so that the real error can propagate
            try:
                if OQ_DISTRIBUTE.startswith('celery'):
                    celery_cleanup(TERMINATE, parallel.Starmap.task_ids)
            except BaseException:
                # log the finalization error only if there is no real error
                if tb == 'None\n':
                    logs.LOG.error('finalizing', exc_info=True)
    return calc

def _do_run_calc(calc, exports, hazard_calculation_id, **kw):
    with calc._monitor:
        calc.run(exports=exports, hazard_calculation_id=hazard_calculation_id,
                 close=False, **kw)  # don't close the datastore too soon

def version_triple(tag):
    """
    returns: a triple of integers from a version tag
    """
    groups = re.match(r'v?(\d+)\.(\d+)\.(\d+)', tag).groups()
    return tuple(int(n) for n in groups)

def check_obsolete_version(calculation_mode='WebUI'):
    """
    Check if there is a newer version of the engine.

    :param calculation_mode:
         - the calculation mode when called from the engine
         - an empty string when called from the WebUI
    :returns:
        - a message if the running version of the engine is obsolete
        - the empty string if the engine is updated
        - None if the check could not be performed (i.e. github is down)
    """
    if os.environ.get('JENKINS_URL') or os.environ.get('TRAVIS'):
        # avoid flooding our API server with requests from CI systems
        return

    headers = {'User-Agent': 'OpenQuake Engine %s;%s;%s;%s' %
               (__version__, calculation_mode, platform.platform(),
                config.distribution.oq_distribute)}
    try:
        req = Request(OQ_API + '/engine/latest', headers=headers)
        # NB: a timeout < 1 does not work
        data = urlopen(req, timeout=1).read()  # bytes
        tag_name = json.loads(decode(data))['tag_name']
        current = version_triple(__version__)
        latest = version_triple(tag_name)
    except:  # page not available or wrong version tag
        return
    if current < latest:
        return ('Version %s of the engine is available, but you are '
                'still using version %s' % (tag_name, __version__))
    else:
        return ''
<|endoftext|>


<commit_before>#! /usr/bin/python
# -*- coding: utf8 -*-

from __future__ import absolute_import, division, print_function

import os
import sys
import inspect

from .utils import logger
from .cascadingconfig import CascadingConfig, DefaultFreeArgumentParser

import tldp.doctypes

def isdirectory(d):
    if os.path.exists(d):
        return d
    return None

def collectconfiguration(argv):
    tag = 'ldptool'
    argparser = DefaultFreeArgumentParser()
    argparser.add_argument('--sourcedir', '--source-dir', '--source-directory',
                           '-s',
                           action='append', default=None, type=isdirectory,
                           help='a directory containing LDP source documents')
    argparser.add_argument('--pubdir', '--output', '--outputdir', '--outdir',
                           '-o',
                           default=None, type=isdirectory,
                           help='a directory containing LDP output documents')
    argparser.add_argument('--configfile', '--config-file', '--cfg',
                           '-c',
                           default=None, type=isdirectory,
                           help='a configuration file')

    # -- collect up the fragments of CLI; automate detection?
    #
    tldp.doctypes.linuxdoc.config_fragment(argparser)
    tldp.doctypes.docbooksgml.config_fragment(argparser)
    tldp.doctypes.docbook4xml.config_fragment(argparser)
    tldp.doctypes.docbook5xml.config_fragment(argparser)

    cc = CascadingConfig(tag, argparser, argv)
    return cc.config

def main(argv):
    config = collectconfiguration(argv)
    import pprint
    pprint.pprint(vars(config))
    return 0

#
# -- end of file
<commit_msg>absolute imports (and move isdirectory to utils)<commit_after>#! /usr/bin/python
# -*- coding: utf8 -*-

from __future__ import absolute_import, division, print_function

from tldp.utils import logger, isdirectory
from tldp.cascadingconfig import CascadingConfig, DefaultFreeArgumentParser

import tldp.doctypes

def collectconfiguration(argv):
    tag = 'ldptool'
    argparser = DefaultFreeArgumentParser()
    argparser.add_argument('--sourcedir', '--source-dir', '--source-directory',
                           '-s',
                           action='append', default=None, type=isdirectory,
                           help='a directory containing LDP source documents')
    argparser.add_argument('--pubdir', '--output', '--outputdir', '--outdir',
                           '-o',
                           default=None, type=isdirectory,
                           help='a directory containing LDP output documents')
    argparser.add_argument('--configfile', '--config-file', '--cfg',
                           '-c',
                           default=None, type=isdirectory,
                           help='a configuration file')

    # -- collect up the fragments of CLI; automate detection?
    #
    tldp.doctypes.linuxdoc.config_fragment(argparser)
    tldp.doctypes.docbooksgml.config_fragment(argparser)
    tldp.doctypes.docbook4xml.config_fragment(argparser)
    tldp.doctypes.docbook5xml.config_fragment(argparser)

    cc = CascadingConfig(tag, argparser, argv)
    return cc.config

def main(argv):
    config = collectconfiguration(argv)
    import pprint
    pprint.pprint(vars(config))
    return 0

#
# -- end of file
<|endoftext|>


<commit_before>import Player.py
<commit_msg>changed import code<commit_after>from Player import Player
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2010-2017 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 Affero 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/>.

"""Engine: A collection of fundamental functions for initializing and running
calculations."""

import os
import sys
import signal
import traceback

from openquake.baselib.performance import Monitor
from openquake.hazardlib import valid
from openquake.baselib import parallel
from openquake.commonlib.oqvalidation import OqParam
from openquake.commonlib import datastore, config, readinput
from openquake.calculators import base, views, export
from openquake.commonlib import logs

TERMINATE = valid.boolean(
    config.get('distribution', 'terminate_workers_on_revoke') or 'false')

USE_CELERY = config.get('distribution', 'oq_distribute') == 'celery'

if USE_CELERY:
    import celery.task.control

    def set_concurrent_tasks_default():
        """
        Set the default for concurrent_tasks.
        Returns the number of live celery nodes (i.e. the number of machines).
        """
        stats = celery.task.control.inspect(timeout=1).stats()
        if not stats:
            sys.exit("No live compute nodes, aborting calculation")
        num_cores = sum(stats[k]['pool']['max-concurrency'] for k in stats)
        OqParam.concurrent_tasks.default = num_cores * 5
        logs.LOG.info(
            'Using %s, %d cores', ', '.join(sorted(stats)), num_cores)

    def celery_cleanup(terminate, task_ids=()):
        """
        Release the resources used by an openquake job.
        In particular revoke the running tasks (if any).

        :param bool terminate: the celery revoke command terminate flag
        :param task_ids: celery task IDs
        """
        # Using the celery API, terminate and revoke and terminate any running
        # tasks associated with the current job.
        if task_ids:
            logs.LOG.warn('Revoking %d tasks', len(task_ids))
        else:  # this is normal when OQ_DISTRIBUTE=no
            logs.LOG.debug('No task to revoke')
        for tid in task_ids:
            celery.task.control.revoke(tid, terminate=terminate)
            logs.LOG.debug('Revoked task %s', tid)

def expose_outputs(dstore):
    """
    Build a correspondence between the outputs in the datastore and the
    ones in the database.

    :param dstore: datastore
    """
    oq = dstore['oqparam']
    exportable = set(ekey[0] for ekey in export.export)
    calcmode = oq.calculation_mode
    dskeys = set(dstore) & exportable  # exportable datastore keys
    dskeys.add('fullreport')
    try:
        rlzs = list(dstore['realizations'])
    except KeyError:
        rlzs = []
    # expose gmf_data only if < 10 MB
    if oq.ground_motion_fields and calcmode == 'event_based':
        nbytes = dstore['gmf_data'].attrs['nbytes']
        if nbytes < 10 * 1024 ** 2:
            dskeys.add('gmf_data')
    if 'scenario' not in calcmode:  # export sourcegroups.csv
        dskeys.add('sourcegroups')
    hdf5 = dstore.hdf5
    if 'poes' in hdf5 or 'hcurves' in hdf5:
        dskeys.add('hcurves')
        if oq.uniform_hazard_spectra:
            dskeys.add('uhs')  # export them
        if oq.hazard_maps:
            dskeys.add('hmaps')  # export them
    if 'avg_losses-rlzs' in dstore and rlzs:
        dskeys.add('avg_losses-stats')
    if oq.conditional_loss_poes:  # expose loss_maps outputs
        if 'rcurves-rlzs' in dstore or 'loss_curves-rlzs' in dstore:
            dskeys.add('loss_maps-rlzs')
        if 'rcurves-stats' in dstore or 'loss_curves-stats' in dstore:
            if len(rlzs) > 1:
                dskeys.add('loss_maps-stats')
    if 'all_loss_ratios' in dskeys:
        dskeys.remove('all_loss_ratios')  # export only specific IDs
    if 'realizations' in dskeys and len(rlzs) <= 1:
        dskeys.remove('realizations')  # do not export a single realization
    if 'ruptures' in dskeys and 'scenario' in calcmode:
        exportable.remove('ruptures')  # do not export, as requested by Vitor
    logs.dbcmd('create_outputs', dstore.calc_id, sorted(dskeys & exportable))

class MasterKilled(KeyboardInterrupt):
    "Exception raised when a job is killed manually"

def raiseMasterKilled(signum, _stack):
    """
    When a SIGTERM is received, raise the MasterKilled
    exception with an appropriate error message.

    :param int signum: the number of the received signal
    :param _stack: the current frame object, ignored
    """
    if signum == signal.SIGTERM:
        msg = 'The openquake master process was killed manually'
    else:
        msg = 'Received a signal %d' % signum
    raise MasterKilled(msg)

# register the raiseMasterKilled callback for SIGTERM
# when using the Django development server this module is imported by a thread,
# so one gets a `ValueError: signal only works in main thread` that
# can be safely ignored
try:
    signal.signal(signal.SIGTERM, raiseMasterKilled)
except ValueError:
    pass

def job_from_file(cfg_file, username, hazard_calculation_id=None):
    """
    Create a full job profile from a job config file.

    :param str cfg_file:
        Path to a job.ini file.
    :param str username:
        The user who will own this job profile and all results
    :param str datadir:
        Data directory of the user
    :param hazard_calculation_id:
        ID of a previous calculation or None
    :returns:
        a pair (job_id, oqparam)
    """
    oq = readinput.get_oqparam(cfg_file)
    job_id = logs.dbcmd('create_job', oq.calculation_mode, oq.description,
                        username, datastore.DATADIR, hazard_calculation_id)
    return job_id, oq

def run_calc(job_id, oqparam, log_level, log_file, exports,
             hazard_calculation_id=None, **kw):
    """
    Run a calculation.

    :param job_id:
        ID of the current job
    :param oqparam:
        :class:`openquake.commonlib.oqvalidation.OqParam` instance
    :param str log_level:
        The desired logging level. Valid choices are 'debug', 'info',
        'progress', 'warn', 'error', and 'critical'.
    :param str log_file:
        Complete path (including file name) to file where logs will be written.
        If `None`, logging will just be printed to standard output.
    :param exports:
        A comma-separated string of export types.
    """
    monitor = Monitor('total runtime', measuremem=True)
    with logs.handle(job_id, log_level, log_file):  # run the job
        if USE_CELERY and os.environ.get('OQ_DISTRIBUTE') == 'celery':
            set_concurrent_tasks_default()
        calc = base.calculators(oqparam, monitor, calc_id=job_id)
        calc.from_engine = True
        tb = 'None\n'
        try:
            logs.dbcmd('set_status', job_id, 'executing')
            _do_run_calc(calc, exports, hazard_calculation_id, **kw)
            expose_outputs(calc.datastore)
            records = views.performance_view(calc.datastore)
            logs.dbcmd('save_performance', job_id, records)
            calc.datastore.close()
            logs.LOG.info('Calculation %d finished correctly in %d seconds',
                          job_id, calc._monitor.duration)
            logs.dbcmd('finish', job_id, 'complete')
        except:
            tb = traceback.format_exc()
            try:
                logs.LOG.critical(tb)
                logs.dbcmd('finish', job_id, 'failed')
            except:  # an OperationalError may always happen
                sys.stderr.write(tb)
            raise
        finally:
            # if there was an error in the calculation, this part may fail;
            # in such a situation, we simply log the cleanup error without
            # taking further action, so that the real error can propagate
            try:
                if USE_CELERY:
                    celery_cleanup(TERMINATE, parallel.Starmap.task_ids)
            except:
                # log the finalization error only if there is no real error
                if tb == 'None\n':
                    logs.LOG.error('finalizing', exc_info=True)
    return calc

def _do_run_calc(calc, exports, hazard_calculation_id, **kw):
    with calc._monitor:
        calc.run(exports=exports, hazard_calculation_id=hazard_calculation_id,
                 close=False, **kw)  # don't close the datastore too soon
<commit_msg>Stored the total calculation duration<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2010-2017 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 Affero 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/>.

"""Engine: A collection of fundamental functions for initializing and running
calculations."""

import os
import sys
import signal
import traceback

from openquake.baselib.performance import Monitor
from openquake.hazardlib import valid
from openquake.baselib import parallel
from openquake.commonlib.oqvalidation import OqParam
from openquake.commonlib import datastore, config, readinput
from openquake.calculators import base, views, export
from openquake.commonlib import logs

TERMINATE = valid.boolean(
    config.get('distribution', 'terminate_workers_on_revoke') or 'false')

USE_CELERY = config.get('distribution', 'oq_distribute') == 'celery'

if USE_CELERY:
    import celery.task.control

    def set_concurrent_tasks_default():
        """
        Set the default for concurrent_tasks.
        Returns the number of live celery nodes (i.e. the number of machines).
        """
        stats = celery.task.control.inspect(timeout=1).stats()
        if not stats:
            sys.exit("No live compute nodes, aborting calculation")
        num_cores = sum(stats[k]['pool']['max-concurrency'] for k in stats)
        OqParam.concurrent_tasks.default = num_cores * 5
        logs.LOG.info(
            'Using %s, %d cores', ', '.join(sorted(stats)), num_cores)

    def celery_cleanup(terminate, task_ids=()):
        """
        Release the resources used by an openquake job.
        In particular revoke the running tasks (if any).

        :param bool terminate: the celery revoke command terminate flag
        :param task_ids: celery task IDs
        """
        # Using the celery API, terminate and revoke and terminate any running
        # tasks associated with the current job.
        if task_ids:
            logs.LOG.warn('Revoking %d tasks', len(task_ids))
        else:  # this is normal when OQ_DISTRIBUTE=no
            logs.LOG.debug('No task to revoke')
        for tid in task_ids:
            celery.task.control.revoke(tid, terminate=terminate)
            logs.LOG.debug('Revoked task %s', tid)

def expose_outputs(dstore):
    """
    Build a correspondence between the outputs in the datastore and the
    ones in the database.

    :param dstore: datastore
    """
    oq = dstore['oqparam']
    exportable = set(ekey[0] for ekey in export.export)
    calcmode = oq.calculation_mode
    dskeys = set(dstore) & exportable  # exportable datastore keys
    dskeys.add('fullreport')
    try:
        rlzs = list(dstore['realizations'])
    except KeyError:
        rlzs = []
    # expose gmf_data only if < 10 MB
    if oq.ground_motion_fields and calcmode == 'event_based':
        nbytes = dstore['gmf_data'].attrs['nbytes']
        if nbytes < 10 * 1024 ** 2:
            dskeys.add('gmf_data')
    if 'scenario' not in calcmode:  # export sourcegroups.csv
        dskeys.add('sourcegroups')
    hdf5 = dstore.hdf5
    if 'poes' in hdf5 or 'hcurves' in hdf5:
        dskeys.add('hcurves')
        if oq.uniform_hazard_spectra:
            dskeys.add('uhs')  # export them
        if oq.hazard_maps:
            dskeys.add('hmaps')  # export them
    if 'avg_losses-rlzs' in dstore and rlzs:
        dskeys.add('avg_losses-stats')
    if oq.conditional_loss_poes:  # expose loss_maps outputs
        if 'rcurves-rlzs' in dstore or 'loss_curves-rlzs' in dstore:
            dskeys.add('loss_maps-rlzs')
        if 'rcurves-stats' in dstore or 'loss_curves-stats' in dstore:
            if len(rlzs) > 1:
                dskeys.add('loss_maps-stats')
    if 'all_loss_ratios' in dskeys:
        dskeys.remove('all_loss_ratios')  # export only specific IDs
    if 'realizations' in dskeys and len(rlzs) <= 1:
        dskeys.remove('realizations')  # do not export a single realization
    if 'ruptures' in dskeys and 'scenario' in calcmode:
        exportable.remove('ruptures')  # do not export, as requested by Vitor
    logs.dbcmd('create_outputs', dstore.calc_id, sorted(dskeys & exportable))

class MasterKilled(KeyboardInterrupt):
    "Exception raised when a job is killed manually"

def raiseMasterKilled(signum, _stack):
    """
    When a SIGTERM is received, raise the MasterKilled
    exception with an appropriate error message.

    :param int signum: the number of the received signal
    :param _stack: the current frame object, ignored
    """
    if signum == signal.SIGTERM:
        msg = 'The openquake master process was killed manually'
    else:
        msg = 'Received a signal %d' % signum
    raise MasterKilled(msg)

# register the raiseMasterKilled callback for SIGTERM
# when using the Django development server this module is imported by a thread,
# so one gets a `ValueError: signal only works in main thread` that
# can be safely ignored
try:
    signal.signal(signal.SIGTERM, raiseMasterKilled)
except ValueError:
    pass

def job_from_file(cfg_file, username, hazard_calculation_id=None):
    """
    Create a full job profile from a job config file.

    :param str cfg_file:
        Path to a job.ini file.
    :param str username:
        The user who will own this job profile and all results
    :param str datadir:
        Data directory of the user
    :param hazard_calculation_id:
        ID of a previous calculation or None
    :returns:
        a pair (job_id, oqparam)
    """
    oq = readinput.get_oqparam(cfg_file)
    job_id = logs.dbcmd('create_job', oq.calculation_mode, oq.description,
                        username, datastore.DATADIR, hazard_calculation_id)
    return job_id, oq

def run_calc(job_id, oqparam, log_level, log_file, exports,
             hazard_calculation_id=None, **kw):
    """
    Run a calculation.

    :param job_id:
        ID of the current job
    :param oqparam:
        :class:`openquake.commonlib.oqvalidation.OqParam` instance
    :param str log_level:
        The desired logging level. Valid choices are 'debug', 'info',
        'progress', 'warn', 'error', and 'critical'.
    :param str log_file:
        Complete path (including file name) to file where logs will be written.
        If `None`, logging will just be printed to standard output.
    :param exports:
        A comma-separated string of export types.
    """
    monitor = Monitor('total runtime', measuremem=True)
    with logs.handle(job_id, log_level, log_file):  # run the job
        if USE_CELERY and os.environ.get('OQ_DISTRIBUTE') == 'celery':
            set_concurrent_tasks_default()
        calc = base.calculators(oqparam, monitor, calc_id=job_id)
        monitor.hdf5path = calc.datastore.hdf5path
        calc.from_engine = True
        tb = 'None\n'
        try:
            logs.dbcmd('set_status', job_id, 'executing')
            _do_run_calc(calc, exports, hazard_calculation_id, **kw)
            duration = monitor.duration
            expose_outputs(calc.datastore)
            monitor.flush()
            records = views.performance_view(calc.datastore)
            logs.dbcmd('save_performance', job_id, records)
            calc.datastore.close()
            logs.LOG.info('Calculation %d finished correctly in %d seconds',
                          job_id, duration)
            logs.dbcmd('finish', job_id, 'complete')
        except:
            tb = traceback.format_exc()
            try:
                logs.LOG.critical(tb)
                logs.dbcmd('finish', job_id, 'failed')
            except:  # an OperationalError may always happen
                sys.stderr.write(tb)
            raise
        finally:
            # if there was an error in the calculation, this part may fail;
            # in such a situation, we simply log the cleanup error without
            # taking further action, so that the real error can propagate
            try:
                if USE_CELERY:
                    celery_cleanup(TERMINATE, parallel.Starmap.task_ids)
            except:
                # log the finalization error only if there is no real error
                if tb == 'None\n':
                    logs.LOG.error('finalizing', exc_info=True)
    return calc

def _do_run_calc(calc, exports, hazard_calculation_id, **kw):
    with calc._monitor:
        calc.run(exports=exports, hazard_calculation_id=hazard_calculation_id,
                 close=False, **kw)  # don't close the datastore too soon
<|endoftext|>


<commit_before>from __future__ import print_function
import torchaudio
import torch.utils.data as data
from torch import stack
import os
import errno
import random
import shutil
import json
import math
from itertools import accumulate

import requests
import bs4
import chardet

# heavily inspired by:
# https://github.com/patyork/python-voxforge-download/blob/master/python-voxforge-download.ipynb

class VOXFORGE(data.Dataset):
    """`Voxforge <http://voxforge.org>`_ Dataset.

    Args:
        TODO: update documentation
        basedir (string): Root directory of dataset.
        randomize (bool, optional): Gets a random selection of files for download.
        dev_mode(bool, optional): if true, clean up is not performed on downloaded
            files.  Useful to keep raw audio and transcriptions.
    """

    URLS = {
        "de": 'http://www.repository.voxforge1.org/downloads/de/Trunk/Audio/Main/16kHz_16bit/',
        "en": 'http://www.repository.voxforge1.org/downloads/SpeechCorpus/Trunk/Audio/Main/16kHz_16bit/',
        "fr": 'http://www.repository.voxforge1.org/downloads/fr/Trunk/Audio/Main/16kHz_16bit/',
        "es": 'http://www.repository.voxforge1.org/downloads/es/Trunk/Audio/Main/16kHz_16bit/',
        "it": 'http://www.repository.voxforge1.org/downloads/it/Trunk/Audio/Main/16kHz_16bit/',
    }

    LABELS = dict([[x, i] for i, x in enumerate(sorted(list(URLS.keys())))])

    SPLITS = ["train", "valid", "test"]

    # set random seed
    random.seed(12345)

    def __init__(self, basedir, transform=None, target_transform=None,
                 langs=["de", "en"], label_type="lang", download=False, num_zips=10,
                 ratios=[0.7, 0.1, 0.2], split="train", use_cache=False,
                 randomize=False, dev_mode=False):
        _make_dir_iff(basedir)
        self.basedir = basedir
        self.rand = randomize
        self.use_cache = use_cache
        self.dev_mode = dev_mode
        self.num_zips = num_zips
        self.langs = langs
        self.split = split
        self.transform = transform
        self.target_transform = target_transform
        self.allow_anon = False
        self.processeddir = os.path.join(basedir, "processed")
        _make_dir_iff(self.processeddir)
        _make_dir_iff(os.path.join(self.processeddir, "audio"))
        _make_dir_iff(os.path.join(self.processeddir, "prompts"))

        if download:
            self.batch_download(num_zips, self.rand, self.allow_anon)
            self.extract_all()

        audiodir = os.path.join(self.processeddir, "audio")

        audiomanifest = [os.path.join(audiodir, fp) for fp in os.listdir(audiodir)]

        if randomize:
            random.shuffle(audiomanifest)

        if label_type == "lang":
            audiolabels = [os.path.basename(l).split("__", 1)[0] for l in audiomanifest]
        else:
            with open(os.path.join(self.processeddir, "prompts", "prompts.json"), "r") as json_f:
                prompts = json.load(json_f)
            audiolabels = []
            for l in audiomanifest:
                l = os.path.basename(l)
                lang, spkr, fp_a = l.split("__")
                fp_a = fp_a.rsplit(".", 1)[0]
                audiolabels.append(prompts[lang][spkr][fp_a])

        num_files = len(audiomanifest)
        split_pts = list(accumulate([math.floor(num_files * r) for r in ratios]))
        split_pts = [range(0, split_pts[0]),
                     range(split_pts[0], split_pts[1]),
                     range(split_pts[1], num_files)]
        self.splits = {spn: spi for spn, spi in zip(self.SPLITS, split_pts)}
        self.data, self.labels = audiomanifest, audiolabels
        self.cache = {}

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (audio, label) where target is index of the target class.
        """
        idx_split = self.splits[self.split][index]
        audio_path = self.data[idx_split]
        if self.use_cache and audio_path in self.cache:
            audio, target = self.cache[audio_path]
        else:
            audio, sr = torchaudio.load(audio_path, normalization=True)
            target = self.labels[idx_split]
            assert sr == 16000

            if self.transform is not None:
                audio = self.transform(audio)

            if self.target_transform is not None:
                target = self.target_transform(target)

            if self.use_cache:
                self.cache[audio_path] = (audio, target)

        return audio, target

    def __len__(self):
        return len(self.splits[self.split])

    def set_split(self, s):
        self.split = s

    def find_max_len(self):
        self.maxlen = 0
        for fp in self.data:
            sig, sr = torchaudio.load(fp)
            self.maxlen = sig.size(0) if sig.size(0) > self.maxlen else self.maxlen

    def _set_lang(self, lang):
        try:
            self.url = self.URLS[lang]
            self.lang = lang
            self.langdir = os.path.join(self.basedir, "raw", lang)
            self.langzipsdir = os.path.join(self.basedir, "raw", "zips", lang)
            self.langfilelist = os.path.join(self.langzipsdir, "filelist")
            # make dirs
            _make_dir_iff(self.langzipsdir)
        except:
            print("{} is not a valid language selection".format(lang))

    def _get_source(self, url):
        response = requests.get(url)

        base_url = url.split('?')[0]

        if response.status_code == 200: return (base_url, response.status_code, response.text.replace('', '') + '')
        return (base_url, response.status_code, None)

    def _extract_links(self, source):
        soup = bs4.BeautifulSoup(source, "html5lib")

        links = soup.find_all('a')

        links = [a.get('href') for a in links][12:-1] # The beauty of static headers and footers
        if self.rand:
            random.shuffle(links)
        return links

    def _validate(self, links, allow_anon):
        # If a database of already downloaded files exists
        #  ..we just want to download new ones
        if allow_anon:
            return links
        else:
            return [link for link in links if not 'anonymous' in link]

    def _acquire(self, base_url, file_name):

        if os.path.isfile(os.path.join(self.langzipsdir, file_name)):
            return 1

        response = requests.get(base_url + file_name)

        if response.status_code == 200:
            # Download and save
            with open(os.path.join(self.langzipsdir, file_name), 'wb') as f:
                f.write(response.content)

            return 1
        return 0

    def batch_download(self, maximum=10, rand=False, allow_anon=False):
        # imports that are only needed to download

        for lang in self.langs:
            self._set_lang(lang)
            print(self.url)
            base_url, status, source = self._get_source(self.url)
            if status != 200:
                print('Could not connect to', self.url)
                break

            links = self._extract_links(source)
            if not len(links) > 0:
                print('No links found at url', url)

            to_download = self._validate(links, allow_anon)
            if not len(to_download) > 0:
                print('All available files have already been downloaded from', url)

            # Lets download
            counter = 0
            for link in to_download:
                result = self._acquire(base_url, link)

                counter += result

                if result==0: print('Unable to download', link)
                elif counter == maximum: break

    def extract_all(self):
        # import needed only for extracting zips
        import tarfile

        prompts = {}
        for lang in self.langs:
            prompts[lang] = {}
            self._set_lang(lang)
            existing = [f for f in os.listdir(self.langzipsdir) if os.path.isfile(os.path.join(self.langzipsdir, f))]
            for fp in existing:
                fp_noext = fp.rsplit(".", 1)[0]
                tardir = os.path.join(self.langzipsdir, fp_noext)
                if not os.path.isdir(tardir):
                    full_path = os.path.join(self.langzipsdir, fp)
                    # extract zip file
                    with tarfile.open(full_path) as zip_f:
                        zip_f.extractall(self.langzipsdir)
                # copy prompts
                promptdir = os.path.join(tardir, "etc")
                if os.path.isfile(os.path.join(promptdir, "prompts-original")):
                    promptfile = os.path.join(promptdir, "prompts-original")
                else:
                    promptfile = os.path.join(promptdir, "PROMPTS")
                prompts[lang][fp_noext] = {}
                with open(promptfile, "rb") as detect_f:
                    charset = chardet.detect(detect_f.read())
                    if "ISO" in charset["encoding"]:
                        encoding = charset["encoding"]
                    else:
                        encoding = None
                try:
                    with open(promptfile, "r", encoding=encoding) as prompt_f:
                        tmp_prompts = [line.strip().split(" ", 1) for line in prompt_f if len(line.strip()) > 0]
                        prompts[lang][fp_noext].update(dict(tmp_prompts))
                except Exception as e:
                    print(e)
                    print(tmp_prompts)
                # move audio files to destination folder
                if os.path.isdir(os.path.join(tardir, "flac")):
                    audiodir = os.path.join(tardir, "flac")
                elif os.path.isdir(os.path.join(tardir, "wav")):
                    audiodir = os.path.join(tardir, "wav")
                for audio_fp in os.listdir(audiodir):
                    audio_fp_name = os.path.basename(audio_fp)
                    src_file_path = os.path.join(audiodir, audio_fp_name)
                    dest_file_path = os.path.join(self.processeddir,
                                                  "audio",
                                                  "{}__{}__{}".format(self.lang,
                                                                    fp_noext,
                                                                    audio_fp_name)
                    )
                    audio_fp_noext = audio_fp_name.rsplit(".", 1)[0]
                    if not os.path.isfile(dest_file_path) and audio_fp_noext in prompts[lang][fp_noext]:
                        shutil.copyfile(src_file_path, dest_file_path)

        prompts = json.dumps(prompts, ensure_ascii=False, indent=4, sort_keys=True)
        dest_prompts_file = os.path.join(self.processeddir, "prompts", "prompts.json")
        with open(dest_prompts_file, "w", encoding='utf8') as json_f:
            json_f.write(prompts)

def _make_dir_iff(d):
    try:
        os.makedirs(os.path.join(d))
        print("{} created".format(d))
    except OSError as e:
        if e.errno == errno.EEXIST:
            pass
        else:
            raise

def basic_collate(batch):
    """Puts batch of inputs into a tensor and labels into a list
       Args:
         batch: (list) [inputs, labels].  In this simple example, I'm just
            assuming the inputs are tensors and labels are strings
       Output:
         minibatch: (Tensor)
         targets: (list[str])
    """

    minibatch, targets = zip(*[(a, b) for (a,b) in batch])
    minibatch = stack(minibatch, dim=0)
    return minibatch, targets
<commit_msg>noise mixer<commit_after>from __future__ import print_function
import torchaudio
import torch.utils.data as data
from torch import stack
import os
import errno
import random
import shutil
import json
import math
from itertools import accumulate

import requests
import bs4
import chardet

# heavily inspired by:
# https://github.com/patyork/python-voxforge-download/blob/master/python-voxforge-download.ipynb

class VOXFORGE(data.Dataset):
    """`Voxforge <http://voxforge.org>`_ Dataset.

    Args:
        TODO: update documentation
        basedir (string): Root directory of dataset.
        randomize (bool, optional): Gets a random selection of files for download.
        dev_mode(bool, optional): if true, clean up is not performed on downloaded
            files.  Useful to keep raw audio and transcriptions.
    """

    URLS = {
        "de": 'http://www.repository.voxforge1.org/downloads/de/Trunk/Audio/Main/16kHz_16bit/',
        "en": 'http://www.repository.voxforge1.org/downloads/SpeechCorpus/Trunk/Audio/Main/16kHz_16bit/',
        "fr": 'http://www.repository.voxforge1.org/downloads/fr/Trunk/Audio/Main/16kHz_16bit/',
        "es": 'http://www.repository.voxforge1.org/downloads/es/Trunk/Audio/Main/16kHz_16bit/',
        "it": 'http://www.repository.voxforge1.org/downloads/it/Trunk/Audio/Main/16kHz_16bit/',
    }

    NOISES = ("noises", 'http://web.cse.ohio-state.edu/pnl/corpus/HuNonspeech/Nonspeech.zip')

    LABELS = dict([[x, i] for i, x in enumerate(sorted(list(URLS.keys())))])

    SPLITS = ["train", "valid", "test"]

    # set random seed
    random.seed(12345)

    def __init__(self, basedir, transform=None, target_transform=None,
                 langs=["de", "en"], label_type="lang", download=False, num_zips=10,
                 ratios=[0.7, 0.1, 0.2], split="train", use_cache=False,
                 mix_noise=True, mix_prob=.5, mix_vol=0.1,
                 randomize=False, dev_mode=False):
        _make_dir_iff(basedir)
        self.basedir = basedir
        self.rand = randomize
        self.use_cache = use_cache
        self.dev_mode = dev_mode
        self.num_zips = num_zips
        self.langs = langs
        self.split = split
        self.transform = transform
        self.target_transform = target_transform
        self.allow_anon = False
        self.processeddir = os.path.join(basedir, "processed")
        _make_dir_iff(self.processeddir)
        _make_dir_iff(os.path.join(self.processeddir, "audio"))
        _make_dir_iff(os.path.join(self.processeddir, "noises"))
        _make_dir_iff(os.path.join(self.processeddir, "prompts"))
        self.mix_noise = mix_noise
        self.mix_prob = mix_prob
        self.mix_vol = lambda: random.uniform(-1, 1) * 0.3 * mix_vol + mix_vol

        if download:
            self.batch_download(num_zips, self.rand, self.allow_anon)
            self.extract_all()

        audiodir = os.path.join(self.processeddir, "audio")

        audiomanifest = [os.path.join(audiodir, fp) for fp in os.listdir(audiodir)]

        if randomize:
            random.shuffle(audiomanifest)

        if label_type == "lang":
            audiolabels = [os.path.basename(l).split("__", 1)[0] for l in audiomanifest]
        else:
            with open(os.path.join(self.processeddir, "prompts", "prompts.json"), "r") as json_f:
                prompts = json.load(json_f)
            audiolabels = []
            for l in audiomanifest:
                l = os.path.basename(l)
                lang, spkr, fp_a = l.split("__")
                fp_a = fp_a.rsplit(".", 1)[0]
                audiolabels.append(prompts[lang][spkr][fp_a])

        num_files = len(audiomanifest)
        split_pts = list(accumulate([math.floor(num_files * r) for r in ratios]))
        split_pts = [range(0, split_pts[0]),
                     range(split_pts[0], split_pts[1]),
                     range(split_pts[1], num_files)]
        self.splits = {spn: spi for spn, spi in zip(self.SPLITS, split_pts)}
        self.data, self.labels = audiomanifest, audiolabels

        # noise manifest
        self.noises = [os.path.join(os.path.join(self.processeddir, "audio"), x) for x in os.listdir(os.path.join(self.processeddir, "audio"))]

        self.cache = {}

        if self.use_cache:
            self.cache.update({ap: (torchaudio.load(ap, normalization=True)[0], None) for ap in self.noises})
            self.cache.update({ap: (torchaudio.load(ap, normalization=True)[0], al) for ap, al in zip(audiomanifest, audiolabels)})

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (audio, label) where target is index of the target class.
        """
        idx_split = self.splits[self.split][index]
        audio_path = self.data[idx_split]
        if self.use_cache and audio_path in self.cache:
            audio, target = self.cache[audio_path]
        else:
            audio, sr = torchaudio.load(audio_path, normalization=True)
            target = self.labels[idx_split]
            assert sr == 16000

        if self.split == "train" and self.mix_noise and self.mix_prob > random.random():
            noise_path = random.choice(self.noises)
            if noise_path in self.cache:
                noise_sig, _ = self.cache[noise_path]
            else:
                noise_sig, _ = torchaudio.load(noise_path, normalization=True)
            diff = audio.size(0) - noise_sig.size(0)
            if diff > 0: # audio longer than noise
                st = random.randrange(0, diff)
                end = audio.size(0) - diff + st
                audio[st:end] += noise_sig * self.mix_vol()
            elif diff < 0:  # noise longer than audio
                st = random.randrange(0, -diff)
                end = st + audio.size(0)
                audio += noise_sig[st:end] * self.mix_vol()
            else:
                audio += noise_sig * self.mix_vol()

        if self.transform is not None:
            audio = self.transform(audio)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return audio, target

    def __len__(self):
        return len(self.splits[self.split])

    def set_split(self, s):
        self.split = s

    def find_max_len(self):
        self.maxlen = 0
        for fp in self.data:
            sig, sr = torchaudio.load(fp)
            self.maxlen = sig.size(0) if sig.size(0) > self.maxlen else self.maxlen

    def _set_lang(self, lang):
        try:
            self.url = self.URLS[lang]
            self.lang = lang
            self.langdir = os.path.join(self.basedir, "raw", lang)
            self.langzipsdir = os.path.join(self.basedir, "raw", "zips", lang)
            self.langfilelist = os.path.join(self.langzipsdir, "filelist")
            # make dirs
            _make_dir_iff(self.langzipsdir)
        except:
            print("{} is not a valid language selection".format(lang))

    def _get_source(self, url):
        response = requests.get(url)

        base_url = url.split('?')[0]

        if response.status_code == 200: return (base_url, response.status_code, response.text.replace('', '') + '')
        return (base_url, response.status_code, None)

    def _extract_links(self, source):
        soup = bs4.BeautifulSoup(source, "html5lib")

        links = soup.find_all('a')

        links = [a.get('href') for a in links][12:-1] # The beauty of static headers and footers
        if self.rand:
            random.shuffle(links)
        return links

    def _validate(self, links, allow_anon):
        # If a database of already downloaded files exists
        #  ..we just want to download new ones
        if allow_anon:
            return links
        else:
            return [link for link in links if not 'anonymous' in link]

    def _acquire(self, base_url, file_name):

        if os.path.isfile(os.path.join(self.langzipsdir, file_name)):
            return 1

        response = requests.get(base_url + file_name)

        if response.status_code == 200:
            # Download and save
            with open(os.path.join(self.langzipsdir, file_name), 'wb') as f:
                f.write(response.content)

            return 1
        return 0

    def batch_download(self, maximum=10, rand=False, allow_anon=False):
        # imports that are only needed to download

        for lang in self.langs:
            self._set_lang(lang)
            print(self.url)
            base_url, status, source = self._get_source(self.url)
            if status != 200:
                print('Could not connect to', self.url)
                break

            links = self._extract_links(source)
            if not len(links) > 0:
                print('No links found at url', url)

            to_download = self._validate(links, allow_anon)
            if not len(to_download) > 0:
                print('All available files have already been downloaded from', url)

            # Lets download
            counter = 0
            for link in to_download:
                result = self._acquire(base_url, link)

                counter += result

                if result==0: print('Unable to download', link)
                elif counter == maximum: break

        # download noises

        nlabel, nuri = self.NOISES
        noiseszipdir = os.path.join(self.basedir, "raw", "zips", nlabel)
        _make_dir_iff(noiseszipdir)
        file_name = os.path.basename(nuri)
        zip_path = os.path.join(noiseszipdir, file_name)
        if os.path.exists(zip_path):
            print("Noises already downloaded")
        else:
            resp = requests.get(nuri)
            if resp.status_code == 200:
                # Download and save
                with open(zip_path, 'wb') as f:
                    f.write(resp.content)

    def extract_all(self):
        # import needed only for extracting zips
        import tarfile

        # extract noises
        nlabel, nuri = self.NOISES
        noiseszipdir = os.path.join(self.basedir, "raw", "zips", nlabel)
        file_name = os.path.basename(nuri)
        zip_path = os.path.join(noiseszipdir, file_name)
        if not os.path.isdir(os.path.join(noiseszipdir, "Nonspeech")):
            import zipfile
            with zipfile.ZipFile(zip_path, "r") as zip_f:
                zip_f.extractall(noiseszipdir)

        audiodir = os.path.join(noiseszipdir, "Nonspeech")
        for audio_fp in os.listdir(audiodir):
            audio_fp_name = os.path.basename(audio_fp)
            src_file_path = os.path.join(audiodir, audio_fp_name)
            dest_file_path = os.path.join(self.processeddir,
                                          "noises",
                                          audio_fp_name)
            if not os.path.isfile(dest_file_path):
                # resample
                sox_cmd = "sox {} -r 16000 {}".format(src_file_path, dest_file_path)
                os.system(sox_cmd)

        # extract languages
        prompts = {}
        for lang in self.langs:
            prompts[lang] = {}
            self._set_lang(lang)
            existing = [f for f in os.listdir(self.langzipsdir) if os.path.isfile(os.path.join(self.langzipsdir, f))]
            for fp in existing:
                fp_noext = fp.rsplit(".", 1)[0]
                tardir = os.path.join(self.langzipsdir, fp_noext)
                if not os.path.isdir(tardir):
                    full_path = os.path.join(self.langzipsdir, fp)
                    # extract zip file
                    with tarfile.open(full_path) as zip_f:
                        zip_f.extractall(self.langzipsdir)
                # copy prompts
                promptdir = os.path.join(tardir, "etc")
                if os.path.isfile(os.path.join(promptdir, "prompts-original")):
                    promptfile = os.path.join(promptdir, "prompts-original")
                else:
                    promptfile = os.path.join(promptdir, "PROMPTS")
                prompts[lang][fp_noext] = {}
                with open(promptfile, "rb") as detect_f:
                    charset = chardet.detect(detect_f.read())
                    if "ISO" in charset["encoding"]:
                        encoding = charset["encoding"]
                    else:
                        encoding = None
                try:
                    with open(promptfile, "r", encoding=encoding) as prompt_f:
                        tmp_prompts = [line.strip().split(" ", 1) for line in prompt_f if len(line.strip()) > 0]
                        prompts[lang][fp_noext].update(dict(tmp_prompts))
                except Exception as e:
                    print(e)
                    print(tmp_prompts)
                # move audio files to destination folder
                if os.path.isdir(os.path.join(tardir, "flac")):
                    audiodir = os.path.join(tardir, "flac")
                elif os.path.isdir(os.path.join(tardir, "wav")):
                    audiodir = os.path.join(tardir, "wav")
                for audio_fp in os.listdir(audiodir):
                    audio_fp_name = os.path.basename(audio_fp)
                    src_file_path = os.path.join(audiodir, audio_fp_name)
                    dest_file_path = os.path.join(self.processeddir,
                                                  "audio",
                                                  "{}__{}__{}".format(self.lang,
                                                                    fp_noext,
                                                                    audio_fp_name)
                    )
                    audio_fp_noext = audio_fp_name.rsplit(".", 1)[0]
                    if not os.path.isfile(dest_file_path) and audio_fp_noext in prompts[lang][fp_noext]:
                        shutil.copyfile(src_file_path, dest_file_path)

        prompts = json.dumps(prompts, ensure_ascii=False, indent=4, sort_keys=True)
        dest_prompts_file = os.path.join(self.processeddir, "prompts", "prompts.json")
        with open(dest_prompts_file, "w", encoding='utf8') as json_f:
            json_f.write(prompts)

def _make_dir_iff(d):
    try:
        os.makedirs(os.path.join(d))
        print("{} created".format(d))
    except OSError as e:
        if e.errno == errno.EEXIST:
            pass
        else:
            raise

def basic_collate(batch):
    """Puts batch of inputs into a tensor and labels into a list
       Args:
         batch: (list) [inputs, labels].  In this simple example, I'm just
            assuming the inputs are tensors and labels are strings
       Output:
         minibatch: (Tensor)
         targets: (list[str])
    """

    minibatch, targets = zip(*[(a, b) for (a,b) in batch])
    minibatch = stack(minibatch, dim=0)
    return minibatch, targets
<|endoftext|>


<commit_before>import re
import os
from natsort import natsorted
import tmt
from tmt.image import is_image_file
from tmt.illumstats import Illumstats
from tmt.image import IntensityImage, MaskImage

class Project(object):
    '''
    Utility class for a project.

    A project represents the directory that holds the image folder
    and potentially additional folders for segmentation, shift, and
    illumination statistics files.
    The project folder may either correspond to an "experiment"
    or a "subexperiment".
    '''

    def __init__(self, project_dir, cfg):
        '''
        Initialize Project class.

        Parameters
        ----------
        project_dir: str
            absolute path to project folder
        cfg: Dict[str, str]
            configuration settings

        Raises
        ------
        IOError
            when `project_dir` is not a directory
        '''
        self.project_dir = project_dir
        if not os.path.isdir(project_dir):
            raise IOError('Failed to initialize class "Project" because '
                          '"%s" is not a directory.'
                          % self.project_dir)
        self.cfg = cfg
        self._experiment_dir = None
        self._experiment = None
        self._subexperiment = None
        self.is_MPcycle = False
        self._image_dir = None
        self._image_files = None
        self._segmentation_dir = None
        self._segmentation_files = None
        self._shift_dir = None
        self._shift_file = None
        self._stats_dir = None
        self._stats_files = None

    @property
    def experiment_dir(self):
        '''
        Returns
        -------
        str
            path to experiment directory, determined from `project_dir`

        Raises
        ------
        OSError
            when determined directory does not exist
        '''
        if self._experiment_dir is None:
            if self.cfg['SUBEXPERIMENTS_EXIST']:
                levels = 1
            else:
                levels = 0
            exp_dir = os.path.realpath(os.path.join(self.project_dir,
                                                    * ['..'] * levels))
            if not os.path.exists(exp_dir):
                raise OSError('Experiment directory could not be determined '
                              'correctly. The determined directory does not '
                              'exist: %s' % exp_dir)
            self._experiment_dir = exp_dir
        return self._experiment_dir

    @property
    def experiment(self):
        '''
        Returns
        -------
        str
            name of the experiment
        '''
        return os.path.basename(self.experiment_dir)

    @property
    def subexperiment(self):
        '''
        Returns
        -------
        str
            name of the subexperiment
            (empty string if the project is no subexperiment)
        '''
        if self._subexperiment is None:
            if self.cfg['SUBEXPERIMENTS_EXIST']:
                self._subexperiment = os.path.basename(self.project_dir)
            else:
                self._subexperiment = ''
        return self._subexperiment

    @property
    def image_dir(self):
        '''
        Returns
        -------
        str
            path to directory holding image files
        '''
        if self._image_dir is None:
            folder = self.cfg['IMAGE_FOLDER_LOCATION'].format(
                            experiment_dir=self.experiment_dir,
                            subexperiment=self.subexperiment)
            self._image_dir = folder
        return self._image_dir

    @property
    def image_files(self):
        '''
        Returns
        -------
        List[tmt.image.IntensityImage]
            image files (intensity images)

        Raises
        ------
        OSError
            when image directory does not exist
        IOError
            when no image files are found
        '''
        if self._image_files is None:
            if not os.path.exists(self.image_dir):
                raise OSError('Image directory does not exist: %s'
                              % self.image_dir)
            files = [os.path.join(self.image_dir, f)
                     for f in os.listdir(self.image_dir) if is_image_file(f)]
            files = natsorted(files)
            if not files:
                raise IOError('No image files found in "%s"'
                              % self.image_dir)
            self._image_files = [IntensityImage(f, self.cfg) for f in files]
        return self._image_files

    @property
    def segmentation_dir(self):
        '''
        Returns
        -------
        str
            path to directory holding segmentation image files
        '''
        if self._segmentation_dir is None:
            try:
                # In case of MPCycle experiments the segmentation directory
                # is defined in the JSON shift descriptor file. The path
                # is stored for use in Jterator pipelines. Therefore, we
                # have to remove all relative parts of the path.
                folder = os.path.join(
                            self.project_dir,
                            re.sub(r'../', '',
                                   self.shift_file.description.segmentationDir))
            except:
                folder = self.cfg['SEGMENTATION_FOLDER_LOCATION'].format(
                                    experiment_dir=self.experiment_dir,
                                    subexperiment=self.subexperiment)
            self._segmentation_dir = folder
        return self._segmentation_dir

    @property
    def segmentation_files(self):
        '''
        Returns
        -------
        List[tmt.image.MaskImage]
            segmentation image files (mask images)

        Raises
        ------
        OSError
            when segmentation directory does not exist
        IOError
            when no segmentation files are found
        '''
        if self._segmentation_files is None:
            if not os.path.exists(self.segmentation_dir):
                raise OSError('Segmentation directory does not exist: %s'
                              % self.segmentation_dir)
            files = [os.path.join(self.segmentation_dir, f)
                     for f in os.listdir(self.segmentation_dir)
                     if is_image_file(f)]
            files = natsorted(files)
            if not files:
                raise IOError('No image files found in "%s"'
                              % self.segmentation_dir)
            self._segmentation_files = [MaskImage(f, self.cfg) for f in files]
        return self._segmentation_files

    @property
    def stats_dir(self):
        '''
        Returns
        -------
        str
            path to directory holding illumination statistic files
        '''
        if self._stats_dir is None:
            folder = self.cfg['STATS_FOLDER_LOCATION'].format(
                                experiment_dir=self.experiment_dir,
                                subexperiment=self.subexperiment)
            self._stats_dir = folder
        return self._stats_dir

    @property
    def stats_files(self):
        '''
        Returns
        -------
        List[tmt.illumstats.Illumstats]
            files containing calculated illumination statistics

        Raises
        ------
        OSError
            when stats directory does not exist
        IOError
            when no illumination statistic files are found
        '''
        if self._stats_files is None:
            stats_pattern = self.cfg['STATS_FILE_FORMAT'].format(channel='\d+')
            stats_pattern = re.compile(self.stats_dir)
            if not os.path.exists(self.stats_dir):
                raise OSError('Stats directory does not exist: %s'
                              % self.stats_dir)
            files = [os.path.join(self.stats_dir, f)
                     for f in os.listdir(self.stats_dir)
                     if re.search(stats_pattern, f)]
            files = natsorted(files)
            if not files:
                raise IOError('No illumination statistic files found in "%s"'
                              % self.stats_dir)
            self._stats_files = [Illumstats(f, self.cfg) for f in files]
        return self._stats_files

    @property
    def shift_dir(self):
        '''
        Returns
        -------
        str
            path to directory holding shift descriptor files
        '''
        if self._shift_dir is None:
            folder = self.cfg['SHIFT_FOLDER_LOCATION'].format(
                                experiment_dir=self.experiment_dir,
                                subexperiment=self.subexperiment)
            self._shift_dir = folder
        return self._shift_dir

    @property
    def shift_file(self):
        '''
        Returns
        -------
        Namespacified
            description (content of a shift descriptor JSON file)
            and filename

        Raises
        ------
        OSError
            when shift directory does not exist
        IOError
            when no shift descriptor file is found
        '''
        if self._shift_file is None:
            shift_pattern = self.cfg['SHIFT_FILE_FORMAT']
            shift_pattern = re.compile(shift_pattern)
            if not os.path.exists(self.shift_dir):
                raise OSError('Shift directory does not exist: %s'
                              % self.shift_dir)
            files = [os.path.join(self.shift_dir, f)
                     for f in os.listdir(self.shift_dir)
                     if re.search(shift_pattern, f)]
            # there should only be one file, but in case there are more take
            # the first one
            files = natsorted(files)
            if not files:
                raise IOError('No shift descriptor file found in "%s" '
                              'that matches the provided pattern "%s".\n'
                              'Check your configuration settings!'
                              % (self.shift_dir,
                                 self.cfg['SHIFT_FILE_FORMAT']))
            content = tmt.util.load_shift_descriptor(files[0])
            self._shift_file.description = tmt.util.Namespacified(content)
            self._shift_file.filename = files[0]
        return self._shift_file
<commit_msg>Update according to changes in config<commit_after>import re
import os
from natsort import natsorted
from tmt.image import is_image_file
from tmt.illumstats import Illumstats
from tmt.image import IntensityImage, MaskImage
from tmt.shift import ShiftDescriptor

class Project(object):
    '''
    Utility class for a project.

    A project represents the directory that holds the image folder
    and potentially additional folders for segmentation, shift, and
    illumination statistics files.
    The project folder may either correspond to an "experiment"
    or a "subexperiment".
    '''

    def __init__(self, project_dir, cfg):
        '''
        Initialize Project class.

        Parameters
        ----------
        project_dir: str
            absolute path to project folder
        cfg: Dict[str, str]
            configuration settings

        Raises
        ------
        IOError
            when `project_dir` is not a directory
        '''
        self.project_dir = project_dir
        if not os.path.isdir(project_dir):
            raise IOError('Failed to initialize class "Project" because '
                          '"%s" is not a directory.'
                          % self.project_dir)
        self.cfg = cfg
        self._experiment_dir = None
        self._experiment = None
        self._subexperiment = None
        self.is_MPcycle = False
        self._image_dir = None
        self._image_files = None
        self._segmentation_dir = None
        self._segmentation_files = None
        self._shift_dir = None
        self._shift_file = None
        self._stats_dir = None
        self._stats_files = None

    @property
    def experiment_dir(self):
        '''
        Returns
        -------
        str
            path to experiment directory, determined from `project_dir`

        Raises
        ------
        OSError
            when determined directory does not exist
        '''
        if self._experiment_dir is None:
            if self.cfg['SUBEXPERIMENTS_EXIST']:
                levels = 1
            else:
                levels = 0
            exp_dir = os.path.realpath(os.path.join(self.project_dir,
                                                    * ['..'] * levels))
            if not os.path.exists(exp_dir):
                raise OSError('Experiment directory could not be determined '
                              'correctly. The determined directory does not '
                              'exist: %s' % exp_dir)
            self._experiment_dir = exp_dir
        return self._experiment_dir

    @property
    def experiment(self):
        '''
        Returns
        -------
        str
            name of the experiment
        '''
        return os.path.basename(self.experiment_dir)

    @property
    def subexperiment(self):
        '''
        Returns
        -------
        str
            name of the subexperiment
            (empty string if the project is no subexperiment)
        '''
        if self._subexperiment is None:
            if self.cfg['SUBEXPERIMENTS_EXIST']:
                self._subexperiment = os.path.basename(self.project_dir)
            else:
                self._subexperiment = ''
        return self._subexperiment

    @property
    def image_dir(self):
        '''
        Returns
        -------
        str
            path to directory holding image files
        '''
        if self._image_dir is None:
            folder = self.cfg['IMAGE_FOLDER_LOCATION'].format(
                            experiment_dir=self.experiment_dir,
                            subexperiment=self.subexperiment,
                            sep=os.path.sep)
            self._image_dir = folder
        return self._image_dir

    @property
    def image_files(self):
        '''
        Returns
        -------
        List[tmt.image.IntensityImage]
            image files (intensity images)

        Raises
        ------
        OSError
            when image directory does not exist
        IOError
            when no image files are found
        '''
        if self._image_files is None:
            if not os.path.exists(self.image_dir):
                raise OSError('Image directory does not exist: %s'
                              % self.image_dir)
            files = [os.path.join(self.image_dir, f)
                     for f in os.listdir(self.image_dir) if is_image_file(f)]
            files = natsorted(files)
            if not files:
                raise IOError('No image files found in "%s"'
                              % self.image_dir)
            self._image_files = [IntensityImage(f, self.cfg) for f in files]
        return self._image_files

    @property
    def segmentation_dir(self):
        '''
        Returns
        -------
        str
            path to directory holding segmentation image files
        '''
        if self._segmentation_dir is None:
            try:
                # In case of MPCycle experiments the segmentation directory
                # is defined in the JSON shift descriptor file. The path
                # is stored for use in Jterator pipelines. Therefore, we
                # have to remove all relative parts of the path.
                folder = os.path.join(
                            self.project_dir,
                            re.sub(r'../', '',
                                   self.shift_file.description.segmentationDir))
            except:
                folder = self.cfg['SEGMENTATION_FOLDER_LOCATION'].format(
                                    experiment_dir=self.experiment_dir,
                                    subexperiment=self.subexperiment,
                                    sep=os.path.sep)
            self._segmentation_dir = folder
        return self._segmentation_dir

    @property
    def segmentation_files(self):
        '''
        Returns
        -------
        List[tmt.image.MaskImage]
            segmentation image files (mask images)

        Raises
        ------
        OSError
            when segmentation directory does not exist
        IOError
            when no segmentation files are found
        '''
        if self._segmentation_files is None:
            if not os.path.exists(self.segmentation_dir):
                raise OSError('Segmentation directory does not exist: %s'
                              % self.segmentation_dir)
            files = [os.path.join(self.segmentation_dir, f)
                     for f in os.listdir(self.segmentation_dir)
                     if is_image_file(f)]
            files = natsorted(files)
            if not files:
                raise IOError('No image files found in "%s"'
                              % self.segmentation_dir)
            self._segmentation_files = [MaskImage(f, self.cfg) for f in files]
        return self._segmentation_files

    @property
    def stats_dir(self):
        '''
        Returns
        -------
        str
            path to directory holding illumination statistic files
        '''
        if self._stats_dir is None:
            folder = self.cfg['STATS_FOLDER_LOCATION'].format(
                                experiment_dir=self.experiment_dir,
                                subexperiment=self.subexperiment,
                                sep=os.path.sep)
            self._stats_dir = folder
        return self._stats_dir

    @property
    def stats_files(self):
        '''
        Returns
        -------
        List[tmt.illumstats.Illumstats]
            files containing calculated illumination statistics

        Raises
        ------
        OSError
            when stats directory does not exist
        IOError
            when no illumination statistic files are found
        '''
        if self._stats_files is None:
            stats_pattern = self.cfg['STATS_FILE_FORMAT'].format(channel='\d+')
            stats_pattern = re.compile(self.stats_dir)
            if not os.path.exists(self.stats_dir):
                raise OSError('Stats directory does not exist: %s'
                              % self.stats_dir)
            files = [os.path.join(self.stats_dir, f)
                     for f in os.listdir(self.stats_dir)
                     if re.search(stats_pattern, f)]
            files = natsorted(files)
            if not files:
                raise IOError('No illumination statistic files found in "%s"'
                              % self.stats_dir)
            self._stats_files = [Illumstats(f, self.cfg) for f in files]
        return self._stats_files

    @property
    def shift_dir(self):
        '''
        Returns
        -------
        str
            path to directory holding shift descriptor files
        '''
        if self._shift_dir is None:
            folder = self.cfg['SHIFT_FOLDER_LOCATION'].format(
                                experiment_dir=self.experiment_dir,
                                subexperiment=self.subexperiment,
                                sep=os.path.sep)
            self._shift_dir = folder
        return self._shift_dir

    @property
    def shift_file(self):
        '''
        Returns
        -------
        Namespacified
            description (content of a shift descriptor JSON file)
            and filename

        Raises
        ------
        OSError
            when shift directory does not exist
        IOError
            when no shift descriptor file is found
        '''
        if self._shift_file is None:
            shift_pattern = self.cfg['SHIFT_FILE_FORMAT']
            shift_pattern = re.compile(shift_pattern)
            if not os.path.exists(self.shift_dir):
                raise OSError('Shift directory does not exist: %s'
                              % self.shift_dir)
            files = [os.path.join(self.shift_dir, f)
                     for f in os.listdir(self.shift_dir)
                     if re.search(shift_pattern, f)]
            # there should only be one file, but in case there are more take
            # the first one
            files = natsorted(files)
            if not files:
                raise IOError('No shift descriptor file found in "%s" '
                              'that matches the provided pattern "%s".\n'
                              'Check your configuration settings!'
                              % (self.shift_dir,
                                 self.cfg['SHIFT_FILE_FORMAT']))
            self._shift_file = ShiftDescriptor(files[0], self.cfg)
        return self._shift_file
<|endoftext|>


<commit_before>import asyncio
import logging

from cached_property import cached_property

from .dispatchers import LoaferDispatcher
from .exceptions import ConfigurationError
from .routes import Route
from .runners import LoaferRunner

logger = logging.getLogger(__name__)

class LoaferManager:
    def __init__(self, routes, runner=None):
        self.runner = runner or LoaferRunner(on_stop_callback=self.on_loop__stop)
        self.routes = routes

    @cached_property
    def dispatcher(self):
        if not (self.routes and all(isinstance(r, Route) for r in self.routes)):
            raise ConfigurationError('invalid routes to dispatch, routes={}'.format(self.routes))

        return LoaferDispatcher(self.routes)

    def run(self, forever=True):
        loop = self.runner.loop
        self._future = asyncio.gather(
            self.dispatcher.dispatch_providers(loop, forever=forever),
            loop=loop,
        )
        self._future.add_done_callback(self.on_future__errors)
        self.runner.start(self._future, run_forever=forever)

    #
    # Callbacks
    #

    def on_future__errors(self, future):
        exc = future.exception()
        # Unhandled errors crashes the event loop execution
        if isinstance(exc, BaseException):
            logger.critical('fatal error caught: {!r}'.format(exc))
            self.runner.stop()

    def on_loop__stop(self, *args, **kwargs):
        logger.info('cancel schedulled operations ...')
        for task in asyncio.Task.all_tasks(self.runner.loop):
            logger.debug('cancelling {}'.format(task))
            task.cancel()

        if hasattr(self, '_future'):
            self._future.cancel()

        self.dispatcher.stop()
<commit_msg>Fix: wrong use of asyncio.gather<commit_after>import asyncio
import logging

from cached_property import cached_property

from .dispatchers import LoaferDispatcher
from .exceptions import ConfigurationError
from .routes import Route
from .runners import LoaferRunner

logger = logging.getLogger(__name__)

class LoaferManager:
    def __init__(self, routes, runner=None):
        self.runner = runner or LoaferRunner(on_stop_callback=self.on_loop__stop)
        self.routes = routes

    @cached_property
    def dispatcher(self):
        if not (self.routes and all(isinstance(r, Route) for r in self.routes)):
            raise ConfigurationError('invalid routes to dispatch, routes={}'.format(self.routes))

        return LoaferDispatcher(self.routes)

    def run(self, forever=True):
        loop = self.runner.loop
        self._future = asyncio.ensure_future(
            self.dispatcher.dispatch_providers(loop, forever=forever),
            loop=loop,
        )
        self._future.add_done_callback(self.on_future__errors)
        self.runner.start(self._future, run_forever=forever)

    #
    # Callbacks
    #

    def on_future__errors(self, future):
        exc = future.exception()
        # Unhandled errors crashes the event loop execution
        if isinstance(exc, BaseException):
            logger.critical('fatal error caught: {!r}'.format(exc))
            self.runner.stop()

    def on_loop__stop(self, *args, **kwargs):
        logger.info('cancel schedulled operations ...')
        for task in asyncio.Task.all_tasks(self.runner.loop):
            logger.debug('cancelling {}'.format(task))
            task.cancel()

        if hasattr(self, '_future'):
            self._future.cancel()

        self.dispatcher.stop()
<|endoftext|>