<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#-----------------------
# Name: util.py    Assorted utilities used in tmdb_api
# Python Library
# Author: Raymond Wagner
#-----------------------

from copy import copy
from locales import get_locale

class Poller( object ):
    """
    Wrapper for an optional callable to populate an Element derived class
    with raw data, or data from a Request.
    """
    def __init__(self, func, lookup, inst=None):
        self.func = func
        self.lookup = lookup
        self.inst = inst
        if func:
            # with function, this allows polling data from the API
            self.__doc__ = func.__doc__
            self.__name__ = func.__name__
            self.__module__ = func.__module__
        else:
            # without function, this is just a dummy poller used for applying
            # raw data to a new Element class with the lookup table
            self.__name__ = '_populate'

    def __get__(self, inst, owner):
        # normal decorator stuff
        # return self for a class
        # return instantiated copy of self for an object
        if inst is None:
            return self
        func = None
        if self.func:
            func = self.func.__get__(inst, owner)
        return self.__class__(func, self.lookup, inst)

    def __call__(self):
        # retrieve data from callable function, and apply
        if not callable(self.func):
            raise RuntimeError('Poller object called without a source function')
        req = self.func()
        if (('language' in req._kwargs) or ('country' in req._kwargs)) \
                and self.inst._locale.fallthrough:
            # request specifies a locale filter, and fallthrough is enabled
            # run a first pass with specified filter
            if not self.apply(req.readJSON(), False):
                return
            # if first pass results in missed data, run a second pass to
            # fill in the gaps
            self.apply(req.new(language=None, country=None).readJSON())
            # re-apply the filtered first pass data over top the second
            # unfiltered set. this is to work around the issue that the
            # properties have no way of knowing when they should or
            # should not overwrite existing data. the cache engine will
            # take care of the duplicate query
        self.apply(req.readJSON())

    def apply(self, data, set_nones=True):
        # apply data directly, bypassing callable function
        unfilled = False
        for k,v in self.lookup.items():
            if (k in data) and \
                    ((data[k] is not None) if callable(self.func) else True):
                # argument received data, populate it
                setattr(self.inst, v, data[k])
            elif set_nones:
                # argument did not receive data, so fill it with None
                # to indicate such and prevent a repeat scan
                setattr(self.inst, v, None)
            else:
                # argument does not need data, so ignore it allowing it to
                # trigger a later poll. this is intended for use when
                # initializing a class with raw data, or when performing a
                # first pass through when performing locale fall through
                unfilled = True
        return unfilled

class Data( object ):
    """
    Basic response definition class
    This maps to a single key in a JSON dictionary received from the API
    """
    def __init__(self, field, initarg=None, handler=None, poller=None,
                 raw=True, default=u'', lang=False):
        """
        This defines how the dictionary value is to be processed by the poller
            field   -- defines the dictionary key that filters what data this uses
            initarg -- (optional) specifies that this field must be supplied when
                       creating a new instance of the Element class this definition
                       is mapped to. Takes an integer for the order it should be
                       used in the input arguments
            handler -- (optional) callable used to process the received value
                       before being stored in the Element object.
            poller  -- (optional) callable to be used if data is requested and
                       this value has not yet been defined. the callable should
                       return a dictionary of data from a JSON query. many
                       definitions may share a single poller, which will be
                       and the data used to populate all referenced definitions
                       based off their defined field
            raw     -- (optional) if the specified handler is an Element class,
                       the data will be passed into it using the 'raw' keyword
                       attribute.  setting this to false will force the data to
                       instead be passed in as the first argument
        """
        self.field = field
        self.initarg = initarg
        self.poller = poller
        self.raw = raw
        self.default = default
        self.sethandler(handler)

    def __get__(self, inst, owner):
        if inst is None:
            return self
        if self.field not in inst._data:
            if self.poller is None:
                return None
            self.poller.__get__(inst, owner)()
        return inst._data[self.field]

    def __set__(self, inst, value):
        if value:
            value = self.handler(value)
        else:
            value = self.default
        if isinstance(value, Element):
            value._locale = inst._locale
        inst._data[self.field] = value

    def sethandler(self, handler):
        # ensure handler is always callable, even for passthrough data
        if handler is None:
            self.handler = lambda x: x
        elif isinstance(handler, ElementType) and self.raw:
            self.handler = lambda x: handler(raw=x)
        else:
            self.handler = lambda x: handler(x)

class Datapoint( Data ):
    pass

class Datalist( Data ):
    """
    Response definition class for list data
    This maps to a key in a JSON dictionary storing a list of data
    """
    def __init__(self, field, handler=None, poller=None, sort=None, raw=True):
        """
        This defines how the dictionary value is to be processed by the poller
            field   -- defines the dictionary key that filters what data this uses
            handler -- (optional) callable used to process the received value
                       before being stored in the Element object.
            poller  -- (optional) callable to be used if data is requested and
                       this value has not yet been defined. the callable should
                       return a dictionary of data from a JSON query. many
                       definitions may share a single poller, which will be
                       and the data used to populate all referenced definitions
                       based off their defined field
            sort    -- (optional) name of attribute in resultant data to be used
                       to sort the list after processing. this effectively
                       a handler be defined to process the data into something
                       that has attributes
            raw     -- (optional) if the specified handler is an Element class,
                       the data will be passed into it using the 'raw' keyword
                       attribute.  setting this to false will force the data to
                       instead be passed in as the first argument
        """
        super(Datalist, self).__init__(field, None, handler, poller, raw)
        self.sort = sort
    def __set__(self, inst, value):
        data = []
        if value:
            for val in value:
                val = self.handler(val)
                if isinstance(val, Element):
                    val._locale = inst._locale
                data.append(val)
            if self.sort:
                data.sort(key=lambda x: getattr(x, self.sort))
        inst._data[self.field] = data

class Datadict( Data ):
    """
    Response definition class for dictionary data
    This maps to a key in a JSON dictionary storing a dictionary of data
    """
    def __init__(self, field, handler=None, poller=None, raw=True, key=None, attr=None):
        """
        This defines how the dictionary value is to be processed by the poller
            field   -- defines the dictionary key that filters what data this uses
            handler -- (optional) callable used to process the received value
                       before being stored in the Element object.
            poller  -- (optional) callable to be used if data is requested and
                       this value has not yet been defined. the callable should
                       return a dictionary of data from a JSON query. many
                       definitions may share a single poller, which will be
                       and the data used to populate all referenced definitions
                       based off their defined field
            key     -- (optional) name of key in resultant data to be used as
                       the key in the stored dictionary. if this is not the
                       field name from the source data is used instead
            attr    -- (optional) name of attribute in resultant data to be used
                       as the key in the stored dictionary. if this is not
                       the field name from the source data is used instead
            raw     -- (optional) if the specified handler is an Element class,
                       the data will be passed into it using the 'raw' keyword
                       attribute.  setting this to false will force the data to
                       instead be passed in as the first argument
        """
        if key and attr:
            raise TypeError("`key` and `attr` cannot both be defined")
        super(Datadict, self).__init__(field, None, handler, poller, raw)
        if key:
            self.getkey = lambda x: x[key]
        elif attr:
            self.getkey = lambda x: getattr(x, attr)
        else:
            raise TypeError("Datadict requires `key` or `attr` be defined "+\
                            "for populating the dictionary")
    def __set__(self, inst, value):
        data = {}
        if value:
            for val in value:
                val = self.handler(val)
                if isinstance(val, Element):
                    val._locale = inst._locale
                data[self.getkey(val)] = val
        inst._data[self.field] = data

class ElementType( type ):
    """
    MetaClass used to pre-process Element-derived classes and set up the
    Data definitions
    """
    def __new__(mcs, name, bases, attrs):
        # any Data or Poller object defined in parent classes must be cloned
        # and processed in this class to function properly
        # scan through available bases for all such definitions and insert
        # a copy into this class's attributes
        # run in reverse order so higher priority values overwrite lower ones
        data = {}
        pollers = {'_populate':None}

        for base in reversed(bases):
            if isinstance(base, mcs):
                for k, attr in base.__dict__.items():
                    if isinstance(attr, Data):
                        # extract copies of each defined Data element from
                        # parent classes
                        attr = copy(attr)
                        attr.poller = attr.poller.func
                        data[k] = attr
                    elif isinstance(attr, Poller):
                        # extract copies of each defined Poller function
                        # from parent classes
                        pollers[k] = attr.func
        for k,attr in attrs.items():
            if isinstance(attr, Data):
                data[k] = attr
        if '_populate' in attrs:
            pollers['_populate'] = attrs['_populate']

        # process all defined Data attribues, testing for use as an initial
        # argument, and building a list of what Pollers are used to populate
        # which Data points
        pollermap = dict([(k,[]) for k in pollers])
        initargs = []
        for k,v in data.items():
            v.name = k
            if v.initarg:
                initargs.append(v)
            if v.poller:
                pn = v.poller.__name__
                if pn not in pollermap:
                    pollermap[pn] = []
                if pn not in pollers:
                    pollers[pn] = v.poller
                pollermap[pn].append(v)
            else:
                pollermap['_populate'].append(v)

        # wrap each used poller function with a Poller class, and push into
        # the new class attributes
        for k,v in pollermap.items():
            if len(v) == 0:
                continue
            lookup = dict([(attr.field, attr.name) for attr in v])
            poller = Poller(pollers[k], lookup)
            attrs[k] = poller
            # backfill wrapped Poller into each mapped Data object, and ensure
            # the data elements are defined for this new class
            for attr in v:
                attr.poller = poller
                attrs[attr.name] = attr

        # build sorted list of arguments used for intialization
        attrs['_InitArgs'] = tuple([a.name for a in \
                                sorted(initargs, key=lambda x: x.initarg)])
        return type.__new__(mcs, name, bases, attrs)

    def __call__(cls, *args, **kwargs):
        obj = cls.__new__(cls)
        if 'locale' in kwargs:
            obj._locale = kwargs['locale']
        else:
            obj._locale = get_locale()

        obj._data = {}
        if 'raw' in kwargs:
            # if 'raw' keyword is supplied, create populate object manually
            if len(args) != 0:
                raise TypeError('__init__() takes exactly 2 arguments (1 given)')
            obj._populate.apply(kwargs['raw'], False)
        else:
            # if not, the number of input arguments must exactly match that
            # defined by the Data definitions
            if len(args) != len(cls._InitArgs):
                raise TypeError('__init__() takes exactly {0} arguments ({1} given)'\
                            .format(len(cls._InitArgs)+1, len(args)+1))
            for a,v in zip(cls._InitArgs, args):
                setattr(obj, a, v)

        obj.__init__()
        return obj

class Element( object ):
    __metaclass__ = ElementType
    _lang = 'en'

<commit_msg>Allow Datalists to be sorted natively without supplying a key.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#-----------------------
# Name: util.py    Assorted utilities used in tmdb_api
# Python Library
# Author: Raymond Wagner
#-----------------------

from copy import copy
from locales import get_locale

class Poller( object ):
    """
    Wrapper for an optional callable to populate an Element derived class
    with raw data, or data from a Request.
    """
    def __init__(self, func, lookup, inst=None):
        self.func = func
        self.lookup = lookup
        self.inst = inst
        if func:
            # with function, this allows polling data from the API
            self.__doc__ = func.__doc__
            self.__name__ = func.__name__
            self.__module__ = func.__module__
        else:
            # without function, this is just a dummy poller used for applying
            # raw data to a new Element class with the lookup table
            self.__name__ = '_populate'

    def __get__(self, inst, owner):
        # normal decorator stuff
        # return self for a class
        # return instantiated copy of self for an object
        if inst is None:
            return self
        func = None
        if self.func:
            func = self.func.__get__(inst, owner)
        return self.__class__(func, self.lookup, inst)

    def __call__(self):
        # retrieve data from callable function, and apply
        if not callable(self.func):
            raise RuntimeError('Poller object called without a source function')
        req = self.func()
        if (('language' in req._kwargs) or ('country' in req._kwargs)) \
                and self.inst._locale.fallthrough:
            # request specifies a locale filter, and fallthrough is enabled
            # run a first pass with specified filter
            if not self.apply(req.readJSON(), False):
                return
            # if first pass results in missed data, run a second pass to
            # fill in the gaps
            self.apply(req.new(language=None, country=None).readJSON())
            # re-apply the filtered first pass data over top the second
            # unfiltered set. this is to work around the issue that the
            # properties have no way of knowing when they should or
            # should not overwrite existing data. the cache engine will
            # take care of the duplicate query
        self.apply(req.readJSON())

    def apply(self, data, set_nones=True):
        # apply data directly, bypassing callable function
        unfilled = False
        for k,v in self.lookup.items():
            if (k in data) and \
                    ((data[k] is not None) if callable(self.func) else True):
                # argument received data, populate it
                setattr(self.inst, v, data[k])
            elif set_nones:
                # argument did not receive data, so fill it with None
                # to indicate such and prevent a repeat scan
                setattr(self.inst, v, None)
            else:
                # argument does not need data, so ignore it allowing it to
                # trigger a later poll. this is intended for use when
                # initializing a class with raw data, or when performing a
                # first pass through when performing locale fall through
                unfilled = True
        return unfilled

class Data( object ):
    """
    Basic response definition class
    This maps to a single key in a JSON dictionary received from the API
    """
    def __init__(self, field, initarg=None, handler=None, poller=None,
                 raw=True, default=u'', lang=False):
        """
        This defines how the dictionary value is to be processed by the poller
            field   -- defines the dictionary key that filters what data this uses
            initarg -- (optional) specifies that this field must be supplied when
                       creating a new instance of the Element class this definition
                       is mapped to. Takes an integer for the order it should be
                       used in the input arguments
            handler -- (optional) callable used to process the received value
                       before being stored in the Element object.
            poller  -- (optional) callable to be used if data is requested and
                       this value has not yet been defined. the callable should
                       return a dictionary of data from a JSON query. many
                       definitions may share a single poller, which will be
                       and the data used to populate all referenced definitions
                       based off their defined field
            raw     -- (optional) if the specified handler is an Element class,
                       the data will be passed into it using the 'raw' keyword
                       attribute.  setting this to false will force the data to
                       instead be passed in as the first argument
        """
        self.field = field
        self.initarg = initarg
        self.poller = poller
        self.raw = raw
        self.default = default
        self.sethandler(handler)

    def __get__(self, inst, owner):
        if inst is None:
            return self
        if self.field not in inst._data:
            if self.poller is None:
                return None
            self.poller.__get__(inst, owner)()
        return inst._data[self.field]

    def __set__(self, inst, value):
        if value:
            value = self.handler(value)
        else:
            value = self.default
        if isinstance(value, Element):
            value._locale = inst._locale
        inst._data[self.field] = value

    def sethandler(self, handler):
        # ensure handler is always callable, even for passthrough data
        if handler is None:
            self.handler = lambda x: x
        elif isinstance(handler, ElementType) and self.raw:
            self.handler = lambda x: handler(raw=x)
        else:
            self.handler = lambda x: handler(x)

class Datapoint( Data ):
    pass

class Datalist( Data ):
    """
    Response definition class for list data
    This maps to a key in a JSON dictionary storing a list of data
    """
    def __init__(self, field, handler=None, poller=None, sort=None, raw=True):
        """
        This defines how the dictionary value is to be processed by the poller
            field   -- defines the dictionary key that filters what data this uses
            handler -- (optional) callable used to process the received value
                       before being stored in the Element object.
            poller  -- (optional) callable to be used if data is requested and
                       this value has not yet been defined. the callable should
                       return a dictionary of data from a JSON query. many
                       definitions may share a single poller, which will be
                       and the data used to populate all referenced definitions
                       based off their defined field
            sort    -- (optional) name of attribute in resultant data to be used
                       to sort the list after processing. this effectively
                       a handler be defined to process the data into something
                       that has attributes
            raw     -- (optional) if the specified handler is an Element class,
                       the data will be passed into it using the 'raw' keyword
                       attribute.  setting this to false will force the data to
                       instead be passed in as the first argument
        """
        super(Datalist, self).__init__(field, None, handler, poller, raw)
        self.sort = sort
    def __set__(self, inst, value):
        data = []
        if value:
            for val in value:
                val = self.handler(val)
                if isinstance(val, Element):
                    val._locale = inst._locale
                data.append(val)
            if self.sort:
                if self.sort is True:
                    data.sort()
                else:
                    data.sort(key=lambda x: getattr(x, self.sort))
        inst._data[self.field] = data

class Datadict( Data ):
    """
    Response definition class for dictionary data
    This maps to a key in a JSON dictionary storing a dictionary of data
    """
    def __init__(self, field, handler=None, poller=None, raw=True, key=None, attr=None):
        """
        This defines how the dictionary value is to be processed by the poller
            field   -- defines the dictionary key that filters what data this uses
            handler -- (optional) callable used to process the received value
                       before being stored in the Element object.
            poller  -- (optional) callable to be used if data is requested and
                       this value has not yet been defined. the callable should
                       return a dictionary of data from a JSON query. many
                       definitions may share a single poller, which will be
                       and the data used to populate all referenced definitions
                       based off their defined field
            key     -- (optional) name of key in resultant data to be used as
                       the key in the stored dictionary. if this is not the
                       field name from the source data is used instead
            attr    -- (optional) name of attribute in resultant data to be used
                       as the key in the stored dictionary. if this is not
                       the field name from the source data is used instead
            raw     -- (optional) if the specified handler is an Element class,
                       the data will be passed into it using the 'raw' keyword
                       attribute.  setting this to false will force the data to
                       instead be passed in as the first argument
        """
        if key and attr:
            raise TypeError("`key` and `attr` cannot both be defined")
        super(Datadict, self).__init__(field, None, handler, poller, raw)
        if key:
            self.getkey = lambda x: x[key]
        elif attr:
            self.getkey = lambda x: getattr(x, attr)
        else:
            raise TypeError("Datadict requires `key` or `attr` be defined "+\
                            "for populating the dictionary")
    def __set__(self, inst, value):
        data = {}
        if value:
            for val in value:
                val = self.handler(val)
                if isinstance(val, Element):
                    val._locale = inst._locale
                data[self.getkey(val)] = val
        inst._data[self.field] = data

class ElementType( type ):
    """
    MetaClass used to pre-process Element-derived classes and set up the
    Data definitions
    """
    def __new__(mcs, name, bases, attrs):
        # any Data or Poller object defined in parent classes must be cloned
        # and processed in this class to function properly
        # scan through available bases for all such definitions and insert
        # a copy into this class's attributes
        # run in reverse order so higher priority values overwrite lower ones
        data = {}
        pollers = {'_populate':None}

        for base in reversed(bases):
            if isinstance(base, mcs):
                for k, attr in base.__dict__.items():
                    if isinstance(attr, Data):
                        # extract copies of each defined Data element from
                        # parent classes
                        attr = copy(attr)
                        attr.poller = attr.poller.func
                        data[k] = attr
                    elif isinstance(attr, Poller):
                        # extract copies of each defined Poller function
                        # from parent classes
                        pollers[k] = attr.func
        for k,attr in attrs.items():
            if isinstance(attr, Data):
                data[k] = attr
        if '_populate' in attrs:
            pollers['_populate'] = attrs['_populate']

        # process all defined Data attribues, testing for use as an initial
        # argument, and building a list of what Pollers are used to populate
        # which Data points
        pollermap = dict([(k,[]) for k in pollers])
        initargs = []
        for k,v in data.items():
            v.name = k
            if v.initarg:
                initargs.append(v)
            if v.poller:
                pn = v.poller.__name__
                if pn not in pollermap:
                    pollermap[pn] = []
                if pn not in pollers:
                    pollers[pn] = v.poller
                pollermap[pn].append(v)
            else:
                pollermap['_populate'].append(v)

        # wrap each used poller function with a Poller class, and push into
        # the new class attributes
        for k,v in pollermap.items():
            if len(v) == 0:
                continue
            lookup = dict([(attr.field, attr.name) for attr in v])
            poller = Poller(pollers[k], lookup)
            attrs[k] = poller
            # backfill wrapped Poller into each mapped Data object, and ensure
            # the data elements are defined for this new class
            for attr in v:
                attr.poller = poller
                attrs[attr.name] = attr

        # build sorted list of arguments used for intialization
        attrs['_InitArgs'] = tuple([a.name for a in \
                                sorted(initargs, key=lambda x: x.initarg)])
        return type.__new__(mcs, name, bases, attrs)

    def __call__(cls, *args, **kwargs):
        obj = cls.__new__(cls)
        if 'locale' in kwargs:
            obj._locale = kwargs['locale']
        else:
            obj._locale = get_locale()

        obj._data = {}
        if 'raw' in kwargs:
            # if 'raw' keyword is supplied, create populate object manually
            if len(args) != 0:
                raise TypeError('__init__() takes exactly 2 arguments (1 given)')
            obj._populate.apply(kwargs['raw'], False)
        else:
            # if not, the number of input arguments must exactly match that
            # defined by the Data definitions
            if len(args) != len(cls._InitArgs):
                raise TypeError('__init__() takes exactly {0} arguments ({1} given)'\
                            .format(len(cls._InitArgs)+1, len(args)+1))
            for a,v in zip(cls._InitArgs, args):
                setattr(obj, a, v)

        obj.__init__()
        return obj

class Element( object ):
    __metaclass__ = ElementType
    _lang = 'en'

<|endoftext|>


<commit_before>"""
.. See the NOTICE file distributed with this work for additional information
   regarding copyright ownership.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
"""
from __future__ import print_function

import os
import shlex
import subprocess
import sys

from utils import logger

try:
    if hasattr(sys, '_run_from_cmdl') is True:
        raise ImportError
    from pycompss.api.parameter import FILE_IN, FILE_OUT, IN
    from pycompss.api.task import task
    from pycompss.api.api import compss_wait_on
except ImportError:
    logger.warn("[Warning] Cannot import \"pycompss\" API packages.")
    logger.warn("          Using mock decorators.")

    from utils.dummy_pycompss import FILE_IN, FILE_OUT, IN  # pylint: disable=ungrouped-imports
    from utils.dummy_pycompss import task
    from utils.dummy_pycompss import compss_wait_on

from basic_modules.tool import Tool
from basic_modules.metadata import Metadata

# ------------------------------------------------------------------------------

class idearTool(Tool):
    """
    Tool for peak calling for MNase-seq data
    """

    def __init__(self, configuration=None):
        """
        Init function
        """
        print("iDEAR Peak Caller")
        Tool.__init__(self)

        if configuration is None:
            configuration = {}

        self.configuration.update(configuration)

    @task(
        returns=int,
        sample_name=IN, bg_name=IN, sample_bam_file_1=FILE_IN, sample_bam_file_2=FILE_IN,
        bg_bam_file_1=FILE_IN, bg_bam_file_2=FILE_IN, species=IN, assembly=IN, bsgenome=FILE_IN,
        peak_bw=FILE_OUT, isModifier=False)
    def idear_peak_calling(
            self, sample_name, bg_name, sample_bam_file_1, sample_bam_file_2,
            bg_bam_file_1, bg_bam_file_2, common_species_name, assembly,
            bsgenome, peak_bw):
        """
        Make iDamID-seq peak calls. These are saved as bed files That can then
        get displayed on genome browsers. Uses an R script that wraps teh iDEAR
        protocol.

        Parameters
        ----------
        sample_name : str
        bg_name : str
        sample_bam_file_1 : str
            Location of the aligned sequences in bam format
        sample_bam_file_2 : str
            Location of the aligned sequences in bam format
        bg_bam_file_1 : str
            Location of the aligned background sequences in bam format
        bg_bam_file_2 : str
            Location of the aligned background sequences in bam format
        species : str
            Species name for the alignments
        assembly : str
            Assembly used for teh aligned sequences
        peak_bed : str
            Location of the peak bed file

        Returns
        -------
        peak_bed : str
            Location of the collated bed file
        """

        # mkdir tmp_R_lib
        # R CMD INSTALL -l ${PWD}/tmp_R_lib BSgenome.human.grch38_1.4.2.tar.gz

        if not os.path.exists("tmp_R_lib"):
            os.makedirs("tmp_R_lib")

        args = shlex.split("R CMD INSTALL -l tmp_R_lib " + bsgenome)
        process = subprocess.Popen(args)
        process.wait()

        rscript = os.path.join(os.path.dirname(__file__), "../scripts/idear.R")

        args = [
            'Rscript', rscript,
            '--sample_name', sample_name,
            '--background_name', bg_name,
            '--file1', sample_bam_file_1,
            '--file2', sample_bam_file_2,
            '--file3', bg_bam_file_1,
            '--file4', bg_bam_file_2,
            '--species', str(common_species_name),
            '--assembly', assembly,
            '--output', peak_bw + '.tmp']
        logger.info("iDEAR CMD: " + ' '.join(args))

        process = subprocess.Popen(args)
        process.wait()

        try:
            with open(peak_bw + ".tmp", "rb") as f_in:
                with open(peak_bw, "wb") as f_out:
                    f_out.write(f_in.read())
        except IOError:
            logger.fatal("iDEAR failed to generate peak file")
            return False

        return True

    def run(self, input_files, input_metadata, output_files):
        """
        The main function to run iNPS for peak calling over a given BAM file
        and matching background BAM file.

        Parameters
        ----------
        input_files : list
            List of input bam file locations where 0 is the bam data file and 1
            is the matching background bam file
        metadata : dict

        Returns
        -------
        output_files : list
            List of locations for the output files.
        output_metadata : list
            List of matching metadata dict objects

        """

        sample_name = None
        background_name = None

        if "idear_sample_param" in self.configuration:
            sample_name = str(self.configuration["idear_sample_param"])
        if "idear_background_param" in self.configuration:
            background_name = str(self.configuration["idear_background_param"])

        results = self.idear_peak_calling(
            sample_name, background_name,
            input_files["bam_1"], input_files["bam_2"],
            input_files["bg_bam_1"], input_files["bg_bam_2"],
            input_metadata["bam_1"].taxon_id,
            input_metadata["bam_1"].meta_data["assembly"],
            input_files["bsgenome"],
            output_files["bigwig"]
        )
        results = compss_wait_on(results)

        output_metadata = {
            "bigwig": Metadata(
                data_type=input_metadata['bam_1'].data_type,
                file_type="BIGWIG",
                file_path=output_files["bigwig"],
                sources=[
                    input_metadata["bam_1"].file_path,
                    input_metadata["bam_2"].file_path,
                    input_metadata["bg_bam_1"].file_path,
                    input_metadata["bg_bam_2"].file_path,
                    input_metadata["bsgenome"].file_path
                ],
                taxon_id=input_metadata["bam_1"].taxon_id,
                meta_data={
                    "assembly": input_metadata["bam_1"].meta_data["assembly"],
                    "tool": "idear"
                }
            )
        }

        return (output_files, output_metadata)

# ------------------------------------------------------------------------------
<commit_msg>Passing of common name and local R lib dir<commit_after>"""
.. See the NOTICE file distributed with this work for additional information
   regarding copyright ownership.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
"""
from __future__ import print_function

import os
import shlex
import subprocess
import sys

from utils import logger

try:
    if hasattr(sys, '_run_from_cmdl') is True:
        raise ImportError
    from pycompss.api.parameter import FILE_IN, FILE_OUT, IN
    from pycompss.api.task import task
    from pycompss.api.api import compss_wait_on
except ImportError:
    logger.warn("[Warning] Cannot import \"pycompss\" API packages.")
    logger.warn("          Using mock decorators.")

    from utils.dummy_pycompss import FILE_IN, FILE_OUT, IN  # pylint: disable=ungrouped-imports
    from utils.dummy_pycompss import task
    from utils.dummy_pycompss import compss_wait_on

from basic_modules.tool import Tool
from basic_modules.metadata import Metadata

# ------------------------------------------------------------------------------

class idearTool(Tool):
    """
    Tool for peak calling for MNase-seq data
    """

    def __init__(self, configuration=None):
        """
        Init function
        """
        print("iDEAR Peak Caller")
        Tool.__init__(self)

        if configuration is None:
            configuration = {}

        self.configuration.update(configuration)

    @task(
        returns=int,
        sample_name=IN, bg_name=IN, sample_bam_file_1=FILE_IN, sample_bam_file_2=FILE_IN,
        bg_bam_file_1=FILE_IN, bg_bam_file_2=FILE_IN, species=IN, assembly=IN, bsgenome=FILE_IN,
        peak_bw=FILE_OUT, isModifier=False)
    def idear_peak_calling(
            self, sample_name, bg_name, sample_bam_file_1, sample_bam_file_2,
            bg_bam_file_1, bg_bam_file_2, common_species_name, assembly,
            bsgenome, peak_bw):
        """
        Make iDamID-seq peak calls. These are saved as bed files That can then
        get displayed on genome browsers. Uses an R script that wraps teh iDEAR
        protocol.

        Parameters
        ----------
        sample_name : str
        bg_name : str
        sample_bam_file_1 : str
            Location of the aligned sequences in bam format
        sample_bam_file_2 : str
            Location of the aligned sequences in bam format
        bg_bam_file_1 : str
            Location of the aligned background sequences in bam format
        bg_bam_file_2 : str
            Location of the aligned background sequences in bam format
        species : str
            Species name for the alignments
        assembly : str
            Assembly used for teh aligned sequences
        peak_bed : str
            Location of the peak bed file

        Returns
        -------
        peak_bed : str
            Location of the collated bed file
        """

        # mkdir tmp_R_lib
        # R CMD INSTALL -l ${PWD}/tmp_R_lib BSgenome.human.grch38_1.4.2.tar.gz

        rscript = os.path.join(os.path.dirname(__file__), "../scripts/idear.R")
        rlib = os.path.join(os.getcwd(), "tmp_R_lib")

        if not os.path.exists(rlib):
            os.makedirs(rlib)

        args = shlex.split("R CMD INSTALL -l tmp_R_lib " + bsgenome)
        process = subprocess.Popen(args)
        process.wait()

        args = [
            'Rscript', rscript,
            '--sample_name', sample_name,
            '--background_name', bg_name,
            '--file1', sample_bam_file_1,
            '--file2', sample_bam_file_2,
            '--file3', bg_bam_file_1,
            '--file4', bg_bam_file_2,
            '--species', str(common_species_name),
            '--assembly', assembly,
            '--output', peak_bw + '.tmp',
            '--local_lib', rlib]
        logger.info("iDEAR CMD: " + ' '.join(args))

        process = subprocess.Popen(args)
        process.wait()

        try:
            with open(peak_bw + ".tmp", "rb") as f_in:
                with open(peak_bw, "wb") as f_out:
                    f_out.write(f_in.read())
        except IOError:
            logger.fatal("iDEAR failed to generate peak file")
            return False

        return True

    def run(self, input_files, input_metadata, output_files):
        """
        The main function to run iNPS for peak calling over a given BAM file
        and matching background BAM file.

        Parameters
        ----------
        input_files : list
            List of input bam file locations where 0 is the bam data file and 1
            is the matching background bam file
        metadata : dict

        Returns
        -------
        output_files : list
            List of locations for the output files.
        output_metadata : list
            List of matching metadata dict objects

        """

        sample_name = None
        background_name = None
        common_name = None

        if "idear_sample_param" in self.configuration:
            sample_name = str(self.configuration["idear_sample_param"])
        if "idear_background_param" in self.configuration:
            background_name = str(self.configuration["idear_background_param"])
        if "idear_common_name" in self.configuration:
            common_name = str(self.configuration["idear_common_name"])

        results = self.idear_peak_calling(
            sample_name, background_name,
            input_files["bam_1"], input_files["bam_2"],
            input_files["bg_bam_1"], input_files["bg_bam_2"],
            common_name,
            input_metadata["bam_1"].meta_data["assembly"],
            input_files["bsgenome"],
            output_files["bigwig"]
        )
        results = compss_wait_on(results)

        output_metadata = {
            "bigwig": Metadata(
                data_type=input_metadata['bam_1'].data_type,
                file_type="BIGWIG",
                file_path=output_files["bigwig"],
                sources=[
                    input_metadata["bam_1"].file_path,
                    input_metadata["bam_2"].file_path,
                    input_metadata["bg_bam_1"].file_path,
                    input_metadata["bg_bam_2"].file_path,
                    input_metadata["bsgenome"].file_path
                ],
                taxon_id=input_metadata["bam_1"].taxon_id,
                meta_data={
                    "assembly": input_metadata["bam_1"].meta_data["assembly"],
                    "tool": "idear"
                }
            )
        }

        return (output_files, output_metadata)

# ------------------------------------------------------------------------------
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Placeholder"""
# Works only for files where each dataset has the same number of rows, but can be easily modified to do so

import h5py
import numpy as np
import math
import argparse
import sys
from memory_profiler import profile # for memory profiling, call with @profile; myfunc()

@profile
def get_cum_number_of_rows(file_list):
    """
    Returns the cumulative number of rows (axis_0) in a list based on the specified input .h5 files.
    This information is needed for concatenating the .h5 files later on.
    Additionally, the average number of rows for all the input files is calculated, in order to derive a sensible chunksize (optimized for diskspace).
    :param list file_list: list that contains all filepaths of the input files.
    :return: list cum_number_of_rows_list: list that contains the cumulative number of rows (i.e. [0,100,200,300,...] if each file has 100 rows)
    :return: int mean_number_of_rows: specifies the average number of rows (rounded up to int) for the files in the file_list.
    """
    total_number_of_rows = 0
    cum_number_of_rows_list = [0]
    number_of_rows_list = [] # used for approximating the chunksize

    # Get total number of rows for the files in the list, faster than resizing the dataset in each iteration of the file loop in concatenate_h5_files()
    for file_name in file_list:
        f = h5py.File(file_name, 'r')
        # get number of rows from the first folder of the file -> each folder needs to have the same number of rows
        total_number_of_rows += f[f.keys()[0]].shape[0]
        cum_number_of_rows_list.append(total_number_of_rows)
        number_of_rows_list.append(f[f.keys()[0]].shape[0])

        f.close()

    mean_number_of_rows = math.ceil(np.mean(number_of_rows_list))
    return cum_number_of_rows_list, mean_number_of_rows

def parse_input():
    """
    Parses the user input in order to return the most important information:
    1) list of files that should be concatenated 2) the filepath of the output .h5 file 3) use custom chunksize or not
    :return: list file_list: list that contains all filepaths of the input files.
    :return: str output_filepath: specifies the filepath (path+name) of the output .h5 file.
    :return: (bool, int) custom_chunksize: specifies if a custom_chunksize should be used and if yes, which chunksize has been specified. I.e. (True, 1000).
    """
    parser = argparse.ArgumentParser(description='E.g. < python concatenate_h5.py file_1 file_2 /path/to/output.h5 > or '
                                                 '< python concatenate_h5.py --list filepaths.txt /path/to/output.h5 >.\n'
                                                 'Concatenates arrays stored in .h5 files for either multiple direct .h5 inputs or a .txt file of .h5 files (--list option).\n'
                                                 'Outputs a new .h5 file with the concatenated arrays. This output is chunked!\n'
                                                 'Careful: The folders of one file need to have the same number of rows (axis_0)!\n'
                                                 'Make a .txt file with < find /path/to/files -name "file_x-*.h5" | sort --version-sort > listname.list >\n'
                                                 'Chunksize: By default, the chunksize is set as the average number of rows per file, in order to keep the filesize small.\n'
                                                 'If you need fast I/O on the output file, you should consider to set the chunksize according to your use case. \n'
                                                 'This can be done with the optional argument --chunksize.\n'
                                                 'I.e. if you feed the output data to a neural network in batches, you should set chunksize=batchsize.',
                                     formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument('files', metavar='file', type=str, nargs='?', help = 'a file that should be concatenated, minimum of two.')
    parser.add_argument('output_filepath', metavar='output_filepath', type=str, nargs=1, help='filepath and name of the output .h5 file')
    parser.add_argument('-l', '--list', dest='list_file', type=str,
                        help = 'filepath of a .list file that contains all .h5 files that should be concatenated')
    parser.add_argument('-c', '--chunksize', dest='chunksize', type=int,
                        help = 'specify a specific chunksize that should be used instead of the automatic selection.')

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()

    if args.list_file:
        file_list = [line.rstrip('\n') for line in open(args.list_file)]
    else:
        file_list = []
        for filepath in args.files:
            file_list.append(filepath)

    output_filepath = args.output_filepath[0]

    custom_chunksize = (False, None)
    if args.chunksize:
        custom_chunksize = (True, args.chunksize)

    return file_list, output_filepath, custom_chunksize

@profile
def concatenate_h5_files():
    """
    Main code. Concatenates .h5 files with multiple datasets, where each dataset in one file needs to have the same number of rows (axis_0).
    Gets user input with aid of the parse_input() function. By default, the chunksize for the output .h5 file is automatically computed
    based on the average number of rows per file, in order to eliminate padding (wastes disk space).
    For faster I/O, the chunksize should be set by the user depending on the use case.
    In deep learning applications for example, the chunksize should be equal to the batch size that is used later on for reading the data.
    """

    file_list, output_filepath, custom_chunksize = parse_input()
    cum_rows_list, mean_number_of_rows = get_cum_number_of_rows(file_list)

    file_output = h5py.File(output_filepath, 'w')

    for n, input_file_name in enumerate(file_list):

        print 'Processing file ' + file_list[n]

        input_file = h5py.File(input_file_name, 'r')

        for folder_name in input_file:

            folder_data = input_file[folder_name]
            print 'Shape and dtype of dataset ' + folder_name + ': ' + str(folder_data.shape) + ' ; ' + str(folder_data.dtype)

            if n == 0:
                # first file; create the dummy dataset with no max shape
                maxshape = (None,) + folder_data.shape[1:] # change shape of axis zero to None

                if custom_chunksize[0] is True:
                    output_dataset = file_output.create_dataset(folder_name, data=folder_data, maxshape=maxshape,
                                                                chunks=(custom_chunksize[1],)+ folder_data.shape[1:])
                else:
                    output_dataset = file_output.create_dataset(folder_name, data=folder_data, maxshape=maxshape,
                                                                chunks=(mean_number_of_rows,)+ folder_data.shape[1:])
                output_dataset.resize(cum_rows_list[-1], axis=0)

            else:
                file_output[folder_name][cum_rows_list[n]:cum_rows_list[n+1]] = folder_data

        file_output.flush()

    print 'Output information:'
    print '-------------------'
    print 'The output file contains the following datasets:'
    for folder_name in file_output:
        print 'Dataset ' + folder_name + ' with the following shape, dtype and chunks (first argument is the chunksize in axis_0): \n' \
              +  str(file_output[folder_name].shape) + ' ; ' + str(file_output[folder_name].dtype) + ' ; ' + str(file_output[folder_name].chunks)

    file_output.close()

if __name__ == '__main__':
    concatenate_h5_files()
<commit_msg>Minor.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Concatenates .h5 files. Works only for files where each dataset has the same number of rows."""

import h5py
import numpy as np
import math
import argparse
import sys
from memory_profiler import profile # for memory profiling, call with @profile; myfunc()

__author__ = 'Michael Moser'
__license__ = 'AGPL'
__version__= '1.0'
__email__ = 'michael.m.moser@fau.de'
__status__ = 'Production'

@profile
def get_cum_number_of_rows(file_list):
    """
    Returns the cumulative number of rows (axis_0) in a list based on the specified input .h5 files.
    This information is needed for concatenating the .h5 files later on.
    Additionally, the average number of rows for all the input files is calculated, in order to derive a sensible chunksize (optimized for diskspace).
    :param list file_list: list that contains all filepaths of the input files.
    :return: list cum_number_of_rows_list: list that contains the cumulative number of rows (i.e. [0,100,200,300,...] if each file has 100 rows)
    :return: int mean_number_of_rows: specifies the average number of rows (rounded up to int) for the files in the file_list.
    """
    total_number_of_rows = 0
    cum_number_of_rows_list = [0]
    number_of_rows_list = [] # used for approximating the chunksize

    # Get total number of rows for the files in the list, faster than resizing the dataset in each iteration of the file loop in concatenate_h5_files()
    for file_name in file_list:
        f = h5py.File(file_name, 'r')
        # get number of rows from the first folder of the file -> each folder needs to have the same number of rows
        total_number_of_rows += f[f.keys()[0]].shape[0]
        cum_number_of_rows_list.append(total_number_of_rows)
        number_of_rows_list.append(f[f.keys()[0]].shape[0])

        f.close()

    mean_number_of_rows = math.ceil(np.mean(number_of_rows_list))
    return cum_number_of_rows_list, mean_number_of_rows

def parse_input():
    """
    Parses the user input in order to return the most important information:
    1) list of files that should be concatenated 2) the filepath of the output .h5 file 3) use custom chunksize or not
    :return: list file_list: list that contains all filepaths of the input files.
    :return: str output_filepath: specifies the filepath (path+name) of the output .h5 file.
    :return: (bool, int) custom_chunksize: specifies if a custom_chunksize should be used and if yes, which chunksize has been specified. I.e. (True, 1000).
    """
    parser = argparse.ArgumentParser(description='E.g. < python concatenate_h5.py file_1 file_2 /path/to/output.h5 > or '
                                                 '< python concatenate_h5.py --list filepaths.txt /path/to/output.h5 >.\n'
                                                 'Concatenates arrays stored in .h5 files for either multiple direct .h5 inputs or a .txt file of .h5 files (--list option).\n'
                                                 'Outputs a new .h5 file with the concatenated arrays. This output is chunked!\n'
                                                 'Careful: The folders of one file need to have the same number of rows (axis_0)!\n'
                                                 'Make a .txt file with < find /path/to/files -name "file_x-*.h5" | sort --version-sort > listname.list >\n'
                                                 'Chunksize: By default, the chunksize is set as the average number of rows per file, in order to keep the filesize small.\n'
                                                 'If you need fast I/O on the output file, you should consider to set the chunksize according to your use case. \n'
                                                 'This can be done with the optional argument --chunksize.\n'
                                                 'I.e. if you feed the output data to a neural network in batches, you should set chunksize=batchsize.',
                                     formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument('files', metavar='file', type=str, nargs='?', help = 'a file that should be concatenated, minimum of two.')
    parser.add_argument('output_filepath', metavar='output_filepath', type=str, nargs=1, help='filepath and name of the output .h5 file')
    parser.add_argument('-l', '--list', dest='list_file', type=str,
                        help = 'filepath of a .list file that contains all .h5 files that should be concatenated')
    parser.add_argument('-c', '--chunksize', dest='chunksize', type=int,
                        help = 'specify a specific chunksize that should be used instead of the automatic selection.')

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()

    if args.list_file:
        file_list = [line.rstrip('\n') for line in open(args.list_file)]
    else:
        file_list = []
        for filepath in args.files:
            file_list.append(filepath)

    output_filepath = args.output_filepath[0]

    custom_chunksize = (False, None)
    if args.chunksize:
        custom_chunksize = (True, args.chunksize)

    return file_list, output_filepath, custom_chunksize

@profile
def concatenate_h5_files():
    """
    Main code. Concatenates .h5 files with multiple datasets, where each dataset in one file needs to have the same number of rows (axis_0).
    Gets user input with aid of the parse_input() function. By default, the chunksize for the output .h5 file is automatically computed
    based on the average number of rows per file, in order to eliminate padding (wastes disk space).
    For faster I/O, the chunksize should be set by the user depending on the use case.
    In deep learning applications for example, the chunksize should be equal to the batch size that is used later on for reading the data.
    """

    file_list, output_filepath, custom_chunksize = parse_input()
    cum_rows_list, mean_number_of_rows = get_cum_number_of_rows(file_list)

    file_output = h5py.File(output_filepath, 'w')

    for n, input_file_name in enumerate(file_list):

        print 'Processing file ' + file_list[n]

        input_file = h5py.File(input_file_name, 'r')

        for folder_name in input_file:

            folder_data = input_file[folder_name]
            print 'Shape and dtype of dataset ' + folder_name + ': ' + str(folder_data.shape) + ' ; ' + str(folder_data.dtype)

            if n == 0:
                # first file; create the dummy dataset with no max shape
                maxshape = (None,) + folder_data.shape[1:] # change shape of axis zero to None

                if custom_chunksize[0] is True:
                    output_dataset = file_output.create_dataset(folder_name, data=folder_data, maxshape=maxshape,
                                                                chunks=(custom_chunksize[1],)+ folder_data.shape[1:])
                else:
                    output_dataset = file_output.create_dataset(folder_name, data=folder_data, maxshape=maxshape,
                                                                chunks=(mean_number_of_rows,)+ folder_data.shape[1:])
                output_dataset.resize(cum_rows_list[-1], axis=0)

            else:
                file_output[folder_name][cum_rows_list[n]:cum_rows_list[n+1]] = folder_data

        file_output.flush()

    print 'Output information:'
    print '-------------------'
    print 'The output file contains the following datasets:'
    for folder_name in file_output:
        print 'Dataset ' + folder_name + ' with the following shape, dtype and chunks (first argument is the chunksize in axis_0): \n' \
              +  str(file_output[folder_name].shape) + ' ; ' + str(file_output[folder_name].dtype) + ' ; ' + str(file_output[folder_name].chunks)

    file_output.close()

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


<commit_before>"""Tests for distutils.dir_util."""
import unittest
import os
import stat
import shutil

from distutils.dir_util import (mkpath, remove_tree, create_tree, copy_tree,
                                ensure_relative)

from distutils import log
from distutils.tests import support

class DirUtilTestCase(support.TempdirManager, unittest.TestCase):

    def _log(self, msg, *args):
        if len(args) > 0:
            self._logs.append(msg % args)
        else:
            self._logs.append(msg)

    def setUp(self):
        super(DirUtilTestCase, self).setUp()
        self._logs = []
        tmp_dir = self.mkdtemp()
        self.root_target = os.path.join(tmp_dir, 'deep')
        self.target = os.path.join(self.root_target, 'here')
        self.target2 = os.path.join(tmp_dir, 'deep2')
        self.old_log = log.info
        log.info = self._log

    def tearDown(self):
        log.info = self.old_log
        super(DirUtilTestCase, self).tearDown()

    def test_mkpath_remove_tree_verbosity(self):

        mkpath(self.target, verbose=0)
        wanted = []
        self.assertEquals(self._logs, wanted)
        remove_tree(self.root_target, verbose=0)

        mkpath(self.target, verbose=1)
        wanted = ['creating %s' % self.root_target,
                  'creating %s' % self.target]
        self.assertEquals(self._logs, wanted)
        self._logs = []

        remove_tree(self.root_target, verbose=1)
        wanted = ["removing '%s' (and everything under it)" % self.root_target]
        self.assertEquals(self._logs, wanted)

    def test_mkpath_with_custom_mode(self):
        mkpath(self.target, 0o700)
        self.assertEqual(stat.S_IMODE(os.stat(self.target).st_mode), 0o700)
        mkpath(self.target2, 0o555)
        self.assertEqual(stat.S_IMODE(os.stat(self.target2).st_mode), 0o555)

    def test_create_tree_verbosity(self):

        create_tree(self.root_target, ['one', 'two', 'three'], verbose=0)
        self.assertEquals(self._logs, [])
        remove_tree(self.root_target, verbose=0)

        wanted = ['creating %s' % self.root_target]
        create_tree(self.root_target, ['one', 'two', 'three'], verbose=1)
        self.assertEquals(self._logs, wanted)

        remove_tree(self.root_target, verbose=0)

    def test_copy_tree_verbosity(self):

        mkpath(self.target, verbose=0)

        copy_tree(self.target, self.target2, verbose=0)
        self.assertEquals(self._logs, [])

        remove_tree(self.root_target, verbose=0)

        mkpath(self.target, verbose=0)
        a_file = os.path.join(self.target, 'ok.txt')
        f = open(a_file, 'w')
        f.write('some content')
        f.close()

        wanted = ['copying %s -> %s' % (a_file, self.target2)]
        copy_tree(self.target, self.target2, verbose=1)
        self.assertEquals(self._logs, wanted)

        remove_tree(self.root_target, verbose=0)
        remove_tree(self.target2, verbose=0)

    def test_ensure_relative(self):
        if os.sep == '/':
            self.assertEquals(ensure_relative('/home/foo'), 'home/foo')
            self.assertEquals(ensure_relative('some/path'), 'some/path')
        else:   # \\
            self.assertEquals(ensure_relative('c:\\home\\foo'), 'c:home\\foo')
            self.assertEquals(ensure_relative('home\\foo'), 'home\\foo')

def test_suite():
    return unittest.makeSuite(DirUtilTestCase)

if __name__ == "__main__":
    unittest.main(defaultTest="test_suite")
<commit_msg>Skip the distutils mode test on Windows OS.<commit_after>"""Tests for distutils.dir_util."""
import unittest
import os
import stat
import shutil
import sys

from distutils.dir_util import (mkpath, remove_tree, create_tree, copy_tree,
                                ensure_relative)

from distutils import log
from distutils.tests import support

class DirUtilTestCase(support.TempdirManager, unittest.TestCase):

    def _log(self, msg, *args):
        if len(args) > 0:
            self._logs.append(msg % args)
        else:
            self._logs.append(msg)

    def setUp(self):
        super(DirUtilTestCase, self).setUp()
        self._logs = []
        tmp_dir = self.mkdtemp()
        self.root_target = os.path.join(tmp_dir, 'deep')
        self.target = os.path.join(self.root_target, 'here')
        self.target2 = os.path.join(tmp_dir, 'deep2')
        self.old_log = log.info
        log.info = self._log

    def tearDown(self):
        log.info = self.old_log
        super(DirUtilTestCase, self).tearDown()

    def test_mkpath_remove_tree_verbosity(self):

        mkpath(self.target, verbose=0)
        wanted = []
        self.assertEquals(self._logs, wanted)
        remove_tree(self.root_target, verbose=0)

        mkpath(self.target, verbose=1)
        wanted = ['creating %s' % self.root_target,
                  'creating %s' % self.target]
        self.assertEquals(self._logs, wanted)
        self._logs = []

        remove_tree(self.root_target, verbose=1)
        wanted = ["removing '%s' (and everything under it)" % self.root_target]
        self.assertEquals(self._logs, wanted)

    @unittest.skipIf(sys.platform.startswith('win'),
                        "This test is only appropriate for POSIX-like systems.")
    def test_mkpath_with_custom_mode(self):
        mkpath(self.target, 0o700)
        self.assertEqual(stat.S_IMODE(os.stat(self.target).st_mode), 0o700)
        mkpath(self.target2, 0o555)
        self.assertEqual(stat.S_IMODE(os.stat(self.target2).st_mode), 0o555)

    def test_create_tree_verbosity(self):

        create_tree(self.root_target, ['one', 'two', 'three'], verbose=0)
        self.assertEquals(self._logs, [])
        remove_tree(self.root_target, verbose=0)

        wanted = ['creating %s' % self.root_target]
        create_tree(self.root_target, ['one', 'two', 'three'], verbose=1)
        self.assertEquals(self._logs, wanted)

        remove_tree(self.root_target, verbose=0)

    def test_copy_tree_verbosity(self):

        mkpath(self.target, verbose=0)

        copy_tree(self.target, self.target2, verbose=0)
        self.assertEquals(self._logs, [])

        remove_tree(self.root_target, verbose=0)

        mkpath(self.target, verbose=0)
        a_file = os.path.join(self.target, 'ok.txt')
        f = open(a_file, 'w')
        f.write('some content')
        f.close()

        wanted = ['copying %s -> %s' % (a_file, self.target2)]
        copy_tree(self.target, self.target2, verbose=1)
        self.assertEquals(self._logs, wanted)

        remove_tree(self.root_target, verbose=0)
        remove_tree(self.target2, verbose=0)

    def test_ensure_relative(self):
        if os.sep == '/':
            self.assertEquals(ensure_relative('/home/foo'), 'home/foo')
            self.assertEquals(ensure_relative('some/path'), 'some/path')
        else:   # \\
            self.assertEquals(ensure_relative('c:\\home\\foo'), 'c:home\\foo')
            self.assertEquals(ensure_relative('home\\foo'), 'home\\foo')

def test_suite():
    return unittest.makeSuite(DirUtilTestCase)

if __name__ == "__main__":
    unittest.main(defaultTest="test_suite")
<|endoftext|>


<commit_before>import sys
import os

if not os.path.exists('Public'):
    os.makedirs('Public')

if not os.path.exists('External'):
    os.makedirs('External')

allFiles = set()

for root, directories, filenames in os.walk('Sources/'):
    for filename in filenames:

        # Matching only header files
        if not filename.lower().endswith(".h"):
            continue

        # File logic path like im/actor/core/Messenger.h
        path = os.path.join(root[8:], filename)
        # Full directory of file like im/actor/core
        dirName = os.path.split(os.path.join(root, filename))[0][8:] + '/'
        # Source file full path like Sources/actor/core/Messenger.h
        srcFile = os.path.join(root, filename)
        # Converted File path like Sources2/actor/core/Messenger.h
        destFile = os.path.join('Public', path)
        externalFile = os.path.join('External', path)

        allFiles.add(path)

        # Auto Create directory
        if not os.path.exists(os.path.dirname(destFile)):
            os.makedirs(os.path.dirname(destFile))

        if not os.path.exists(os.path.dirname(externalFile)):
            os.makedirs(os.path.dirname(externalFile))

        with open(srcFile, 'r') as f:

            allLines = f.read()
            destLines = ""

            with open(externalFile, 'w') as d:
                d.write(destLines)

            for line in allLines.splitlines():
                if (line.startswith("#include") or line.startswith("#import")) and '\"' in line:
                    start = line.index('\"')
                    end = line.index('\"', start + 1)
                    includedFile = line[start + 1:end]
                    if includedFile != 'objc/runtime.h':
                        localIncludePath = os.path.join("Sources/", includedFile)
                        if os.path.exists(localIncludePath):
                            line = line[0:start] + "<ActorSDK/" + includedFile + ">" + line[end+1:]
                        else:
                            line = line[0:start] + "<j2objc/" + includedFile + ">" + line[end+1:]

                destLines = destLines + line + "\n"

            isUpdated = True

            if os.path.exists(destFile):
                with open(destFile, 'rw') as d:
                    if d.read() == destLines:
                        # print "Not Updated"
                        isUpdated = False

            if isUpdated:
                with open(destFile, 'w') as d:
                    d.write(destLines)

with open('Public/ActorCoreUmbrella.h', 'w') as d:
    for line in allFiles:
        d.write("#import <ActorSDK/" + line + ">\n")

<commit_msg>fix(sdk): Fixing external headers copying<commit_after>import sys
import os

if not os.path.exists('Public'):
    os.makedirs('Public')

if not os.path.exists('External'):
    os.makedirs('External')

allFiles = set()

for root, directories, filenames in os.walk('Sources/'):
    for filename in filenames:

        # Matching only header files
        if not filename.lower().endswith(".h"):
            continue

        # File logic path like im/actor/core/Messenger.h
        path = os.path.join(root[8:], filename)
        # Full directory of file like im/actor/core
        dirName = os.path.split(os.path.join(root, filename))[0][8:] + '/'
        # Source file full path like Sources/actor/core/Messenger.h
        srcFile = os.path.join(root, filename)
        # Converted File path like Sources2/actor/core/Messenger.h
        destFile = os.path.join('Public', path)
        externalFile = os.path.join('External', path)

        allFiles.add(path)

        # Auto Create directory
        if not os.path.exists(os.path.dirname(destFile)):
            os.makedirs(os.path.dirname(destFile))

        if not os.path.exists(os.path.dirname(externalFile)):
            os.makedirs(os.path.dirname(externalFile))

        with open(srcFile, 'r') as f:

            allLines = f.read()
            destLines = ""

            with open(externalFile, 'w') as d:
                d.write(allLines)

            for line in allLines.splitlines():
                if (line.startswith("#include") or line.startswith("#import")) and '\"' in line:
                    start = line.index('\"')
                    end = line.index('\"', start + 1)
                    includedFile = line[start + 1:end]
                    if includedFile != 'objc/runtime.h':
                        localIncludePath = os.path.join("Sources/", includedFile)
                        if os.path.exists(localIncludePath):
                            line = line[0:start] + "<ActorSDK/" + includedFile + ">" + line[end+1:]
                        else:
                            line = line[0:start] + "<j2objc/" + includedFile + ">" + line[end+1:]

                destLines = destLines + line + "\n"

            isUpdated = True

            if os.path.exists(destFile):
                with open(destFile, 'rw') as d:
                    if d.read() == destLines:
                        # print "Not Updated"
                        isUpdated = False

            if isUpdated:
                with open(destFile, 'w') as d:
                    d.write(destLines)

with open('Public/ActorCoreUmbrella.h', 'w') as d:
    for line in allFiles:
        d.write("#import <ActorSDK/" + line + ">\n")

<|endoftext|>


<commit_before>import sqlite3

def create_sql(db_name):
    sql_con = sqlite3.connect(db_name)
    cursor = sql_con.cursor()

    # user choice
    cursor.execute("""CREATE TABLE user_choice
                        (
                            user_id INT PRIMARY KEY,
                            divisions_json TEXT,
                            alias TEXT,
                            division_name TEXT,
                            study_programs_json TEXT,
                            study_level_name TEXT,
                            study_program_combination_name TEXT,
                            admission_year_name TEXT,
                            student_groups_json TEXT,
                            student_group_name TEXT,
                            student_group_id INT,
                            step TEXT DEFAULT 'handle_start' NOT NULL
                        )""")
    sql_con.commit()

    # group data
    cursor.execute("""CREATE TABLE groups_data
                        (
                            id INT PRIMARY KEY NOT NULL,
                            alias TEXT NOT NULL,
                            json_week_data TEXT,
                            interim_attestation TEXT
                        )""")
    sql_con.commit()

    # user data
    cursor.execute("""CREATE TABLE user_data
                        (
                            id INT PRIMARY KEY NOT NULL,
                            alias TEXT,
                            group_id INT NOT NULL,
                            full_place INT DEFAULT 1 NOT NULL,
                            sending INT DEFAULT 0 NOT NULL,
                            delete_skips INT DEFAULT 1 NOT NULL,
                            step TEXT DEFAULT 'main_menu' NOT NULL,
                            rate INT DEFAULT 0 NOT NULL,
                            home_station_code TEXT DEFAULT 'c2' NOT NULL,
                            is_univer INT DEFAULT 1 NOT NULL,
                            CONSTRAINT user_data_groups_data_id_fk
                              FOREIGN KEY (group_id) REFERENCES groups_data (id)
                        )""")
    sql_con.commit()

    # lessons
    cursor.execute("""CREATE TABLE lessons
                        (
                            id INTEGER PRIMARY KEY AUTOINCREMENT,
                            name TEXT NOT NULL,
                            type TEXT NOT NULL,
                            day TEXT NOT NULL DEFAULT 'all',
                            time TEXT NOT NULL DEFAULT 'all',
                            UNIQUE (name, type, day, time)
                        )""")
    sql_con.commit()

    # skips
    cursor.execute("""CREATE TABLE skips
                        (
                            lesson_id INT NOT NULL,
                            user_id INT NOT NULL,
                            CONSTRAINT skips_lesson_id_user_id_pk PRIMARY KEY (lesson_id, user_id),
                            CONSTRAINT skips_user_data_id_fk FOREIGN KEY (user_id) REFERENCES user_data (id),
                            CONSTRAINT skips_lessons_id_fk FOREIGN KEY (lesson_id) REFERENCES lessons (id)
                        )""")
    sql_con.commit()

    # user groups
    cursor.execute("""CREATE TABLE user_groups
                        (
                            group_id INT NOT NULL,
                            user_id INT NOT NULL,
                            CONSTRAINT user_groups_group_id_user_id_pk PRIMARY KEY (group_id, user_id),
                            CONSTRAINT user_groups_user_data_id_fk FOREIGN KEY (user_id) REFERENCES user_data (id),
                            CONSTRAINT user_groups_groups_data_id_fk FOREIGN KEY (group_id) REFERENCES groups_data (id)
                        )""")
    sql_con.commit()

    cursor.close()
    sql_con.close()

if __name__ == '__main__':
    create_sql("Bot_db")
<commit_msg>code style<commit_after>import sqlite3

def create_sql(db_name):
    sql_con = sqlite3.connect(db_name)
    cursor = sql_con.cursor()

    # user choice
    cursor.execute("""CREATE TABLE user_choice
                        (
                            user_id INT PRIMARY KEY,
                            divisions_json TEXT,
                            alias TEXT,
                            division_name TEXT,
                            study_programs_json TEXT,
                            study_level_name TEXT,
                            study_program_combination_name TEXT,
                            admission_year_name TEXT,
                            student_groups_json TEXT,
                            student_group_name TEXT,
                            student_group_id INT,
                            step TEXT DEFAULT 'handle_start' NOT NULL
                        )""")
    sql_con.commit()

    # group data
    cursor.execute("""CREATE TABLE groups_data
                        (
                            id INT PRIMARY KEY NOT NULL,
                            alias TEXT NOT NULL,
                            json_week_data TEXT,
                            interim_attestation TEXT
                        )""")
    sql_con.commit()

    # user data
    cursor.execute("""CREATE TABLE user_data
                        (
                            id INT PRIMARY KEY NOT NULL,
                            alias TEXT,
                            group_id INT NOT NULL,
                            full_place INT DEFAULT 1 NOT NULL,
                            sending INT DEFAULT 0 NOT NULL,
                            delete_skips INT DEFAULT 1 NOT NULL,
                            step TEXT DEFAULT 'main_menu' NOT NULL,
                            rate INT DEFAULT 0 NOT NULL,
                            home_station_code TEXT DEFAULT 'c2' NOT NULL,
                            is_univer INT DEFAULT 1 NOT NULL,
                            CONSTRAINT user_data_groups_data_id_fk
                              FOREIGN KEY (group_id) REFERENCES groups_data (id)
                        )""")
    sql_con.commit()

    # lessons
    cursor.execute("""CREATE TABLE lessons
                        (
                            id INTEGER PRIMARY KEY AUTOINCREMENT,
                            name TEXT NOT NULL,
                            type TEXT NOT NULL,
                            day TEXT NOT NULL DEFAULT 'all',
                            time TEXT NOT NULL DEFAULT 'all',
                            UNIQUE (name, type, day, time)
                        )""")
    sql_con.commit()

    # skips
    cursor.execute("""CREATE TABLE skips
                        (
                            lesson_id INT NOT NULL,
                            user_id INT NOT NULL,
                            CONSTRAINT skips_lesson_id_user_id_pk
                              PRIMARY KEY (lesson_id, user_id),
                            CONSTRAINT skips_user_data_id_fk
                              FOREIGN KEY (user_id) REFERENCES user_data (id),
                            CONSTRAINT skips_lessons_id_fk
                              FOREIGN KEY (lesson_id) REFERENCES lessons (id)
                        )""")
    sql_con.commit()

    # user groups
    cursor.execute("""CREATE TABLE user_groups
                        (
                            group_id INT NOT NULL,
                            user_id INT NOT NULL,
                            CONSTRAINT user_groups_group_id_user_id_pk
                              PRIMARY KEY (group_id, user_id),
                            CONSTRAINT user_groups_user_data_id_fk
                              FOREIGN KEY (user_id) REFERENCES user_data (id),
                            CONSTRAINT user_groups_groups_data_id_fk
                              FOREIGN KEY (group_id) REFERENCES groups_data (id)
                        )""")
    sql_con.commit()

    cursor.close()
    sql_con.close()

if __name__ == '__main__':
    create_sql("Bot_db")
<|endoftext|>


<commit_before># Copyright (c) 2015 Ansible, Inc.
# All Rights Reserved

import logging
from datetime import datetime

from django.core.management.base import NoArgsCommand

from awx.fact.models.fact import * # noqa
from awx.main.socket import Socket

_MODULES = ['packages', 'services', 'files']

logger = logging.getLogger('awx.main.commands.run_fact_cache_receiver')
class FactCacheReceiver(object):
    def __init__(self):
        self.timestamp = None

    def _determine_module(self, facts):
        for x in _MODULES:
            if x in facts:
                return x
        return 'ansible'

    def _extract_module_facts(self, module, facts):
        if module in facts:
            f = facts[module]
            return f
        return facts

    def process_facts(self, facts):
        module = self._determine_module(facts)
        facts = self._extract_module_facts(module, facts)
        return (module, facts)

    def process_fact_message(self, message):
        hostname = message['host']
        facts_data = message['facts']
        date_key = message['date_key']

        # TODO: in ansible < v2 module_setup is emitted for "smart" fact caching.
        # ansible v2 will not emit this message. Thus, this can be removed at that time.
        if 'module_setup' in facts_data and len(facts_data) == 1:
            return

        try:
            host = FactHost.objects.get(hostname=hostname)
        except FactHost.DoesNotExist:
            host = FactHost(hostname=hostname)
            host.save()
        except FactHost.MultipleObjectsReturned:
            query = "db['fact_host'].find(hostname=%s)" % hostname
            print('Database inconsistent. Multiple FactHost "%s" exist. Try the query %s to find the records.' % (hostname, query))
            return

        (module, facts) = self.process_facts(facts_data)
        self.timestamp = datetime.fromtimestamp(date_key, None)

        try:
            # Update existing Fact entry
            version_obj = FactVersion.objects.get(timestamp=self.timestamp, host=host, module=module)
            Fact.objects(id=version_obj.fact.id).update_one(fact=facts)
        except FactVersion.DoesNotExist:
            # Create new Fact entry
            (fact_obj, version_obj) = Fact.add_fact(self.timestamp, facts, host, module)

    def run_receiver(self):
        with Socket('fact_cache', 'r') as facts:
            for message in facts.listen():
                if 'host' not in message or 'facts' not in message or 'date_key' not in message:
                    continue
                self.process_fact_message(message)

class Command(NoArgsCommand):
    '''
    blah blah
    '''
    help = 'Launch the Fact Cache Receiver'

    def handle_noargs(self, **options):
        fcr = FactCacheReceiver()
        try:
            fcr.run_receiver()
        except KeyboardInterrupt:
            pass

<commit_msg>allow user provided fact scan modules<commit_after># Copyright (c) 2015 Ansible, Inc.
# All Rights Reserved

import logging
from datetime import datetime

from django.core.management.base import NoArgsCommand

from awx.fact.models.fact import * # noqa
from awx.main.socket import Socket

logger = logging.getLogger('awx.main.commands.run_fact_cache_receiver')
class FactCacheReceiver(object):
    def __init__(self):
        self.timestamp = None

    def _determine_module(self, facts):
        # Symantically determine the module type
        if len(facts) == 1:
            return facts.iterkeys().next()
        return 'ansible'

    def _extract_module_facts(self, module, facts):
        if module in facts:
            f = facts[module]
            return f
        return facts

    def process_facts(self, facts):
        module = self._determine_module(facts)
        facts = self._extract_module_facts(module, facts)
        return (module, facts)

    def process_fact_message(self, message):
        hostname = message['host']
        facts_data = message['facts']
        date_key = message['date_key']

        # TODO: in ansible < v2 module_setup is emitted for "smart" fact caching.
        # ansible v2 will not emit this message. Thus, this can be removed at that time.
        if 'module_setup' in facts_data and len(facts_data) == 1:
            return

        try:
            host = FactHost.objects.get(hostname=hostname)
        except FactHost.DoesNotExist:
            host = FactHost(hostname=hostname)
            host.save()
        except FactHost.MultipleObjectsReturned:
            query = "db['fact_host'].find(hostname=%s)" % hostname
            print('Database inconsistent. Multiple FactHost "%s" exist. Try the query %s to find the records.' % (hostname, query))
            return

        (module, facts) = self.process_facts(facts_data)
        self.timestamp = datetime.fromtimestamp(date_key, None)

        try:
            # Update existing Fact entry
            version_obj = FactVersion.objects.get(timestamp=self.timestamp, host=host, module=module)
            Fact.objects(id=version_obj.fact.id).update_one(fact=facts)
        except FactVersion.DoesNotExist:
            # Create new Fact entry
            (fact_obj, version_obj) = Fact.add_fact(self.timestamp, facts, host, module)

    def run_receiver(self):
        with Socket('fact_cache', 'r') as facts:
            for message in facts.listen():
                if 'host' not in message or 'facts' not in message or 'date_key' not in message:
                    continue
                self.process_fact_message(message)

class Command(NoArgsCommand):
    '''
    blah blah
    '''
    help = 'Launch the Fact Cache Receiver'

    def handle_noargs(self, **options):
        fcr = FactCacheReceiver()
        try:
            fcr.run_receiver()
        except KeyboardInterrupt:
            pass

<|endoftext|>


<commit_before>
import os
import logging
import warnings
from lofarpipe.support.parset import parameterset
import tkp.utility.accessors
from tkp.database import DataBase, DataSet, Image
from tkp.utility.accessors.dataaccessor import extract_metadata
from tempfile import NamedTemporaryFile
from pyrap.images import image as pyrap_image

logger = logging.getLogger(__name__)

def parse_parset(parset_file):
    """parse a persistence recipe specific parset file
    """
    parset = parameterset(parset_file)
    result = {}
    result['description'] = parset.getString('description', 'TRAP dataset')
    result['dataset_id'] = parset.getInt('dataset_id', -1)
    result['copy_images'] = parset.getBool('copy_images', False)
    result['mongo_host'] = parset.getString('mongo_host', 'localhost')
    result['mongo_port'] = parset.getInt('mongo_port', 27017)
    result['mongo_db'] = parset.getString('mongo_db', 'tkp')
    return result

def image_to_mongodb(filename, hostname, port, db):
    """Copy a file into mongodb"""

    try:
        import pymongo
        import gridfs
    except ImportError:
        msg = "Could not import MongoDB modules"
        logger.error(msg)
        warnings.warn(msg)
        return False

    try:
        connection = pymongo.Connection(host=hostname, port=port)
        gfs = gridfs.GridFS(connection[db])
        if gfs.exists(filename=filename):
            logger.debug("File already in database")

        else:
            # This conversion should work whether the input file
            # is in FITS or CASA format.
            # temp_fits_file is removed automatically when closed.
            temp_fits_file = NamedTemporaryFile()
            i = pyrap_image(filename)
            i.tofits(temp_fits_file.name)
            new_file = gfs.new_file(filename=filename)
            with open(temp_fits_file.name, "r") as f:
                new_file.write(f)
            new_file.close()
    except Exception, e:
        msg = "Failed to save image to MongoDB: %s" % (str(e),)
        logger.error(msg)
        warnings.warn(msg)
        return False

    finally:
        connection.close()

        # Only close this if it has been created.
        if "temp_fits_file" in locals():
            temp_fits_file.close()

    return True

def create_dataset(dataset_id, description):
    """ Creates a dataset if it doesn't exists
    Note: Should only be used in a master recipe
    Returns:
      the database ID of this dataset
    """
    database = DataBase()
    if dataset_id == -1:
        dataset = DataSet({'description': description}, database)
    else:
        dataset = DataSet(id=dataset_id, database=database)
    return dataset.id

def extract_metadatas(images):
    results = []
    for image in images:
        accessor = tkp.utility.accessors.open(image)
        results.append(extract_metadata(accessor))
    return results

def store_images(images_metadata, dataset_id):
    """ Add images to database.
    Note that all images in one dataset should be inserted in one go, since the
    order is very important here. If you don't add them all in once, you should
    make sure they are added in the correct order e.g. sorted by observation
    time.

    Note: Should only be used in a master recipe

    Args:
        images_metadata: list of dicts containing image metadata
        dataset_id: dataset id to be used. don't use value from parset file
                    since this can be -1 (trap way of setting auto increment)
    Returns:
        the database ID of this dataset
    """
    database = DataBase()
    dataset = DataSet(id=dataset_id, database=database)
    image_ids = []

    # sort images by timestamp
    images_metadata.sort(key=lambda m: m['taustart_ts'])

    for metadata in images_metadata:
        filename = metadata['url']
        db_image = Image(data=metadata, dataset=dataset)
        image_ids.append(db_image.id)
        logger.info("stored %s with ID %s" % (os.path.basename(filename), db_image.id))
    return image_ids

def node_steps(images, parset_file):
    """
    this function executes all persistence steps that should be executed on a node.
    Note: Should only be used in a node recipe
    """
    persistence_parset = parse_parset(parset_file)
    mongohost = persistence_parset['mongo_host']
    mongoport = persistence_parset['mongo_port']
    mongodb = persistence_parset['mongo_db']
    copy_images = persistence_parset['copy_images']

    if copy_images:
        logger.info("Copying %s images to mongodb" % len(images))
        for image in images:
            image_to_mongodb(image, mongohost, mongoport, mongodb)
            logger.info("Saved local copy of %s on %s" \
                        % (os.path.basename(image), mongohost))
    else:
        logger.info("Not copying the %s images to mongodb" % len(images))

    logger.info("Extracting metadata from images")
    metadatas = extract_metadatas(images)
    return metadatas

def master_steps(metadatas, parset_file):
    """this function executes all persistence steps that should be executed on
        a master.
    Args:
        parset_file: path to a parset file containig persistence settings
        metadatas: a list of dicts containing info from Image Accessors. This
                   is returned by the node recipe
    """
    persistence_parset = parse_parset(parset_file)
    logger.info("creating dataset in database ...")
    dataset_id = create_dataset(
        persistence_parset['dataset_id'],
        persistence_parset['description'])
    logger.info("added dataset with ID %s" % dataset_id)

    logger.info("Storing images")
    image_ids = store_images(metadatas, dataset_id)
    return dataset_id, image_ids

def all(images, parset_file):
    """
    execute node and then master code, should be used in local run only!
    """
    metadatas = node_steps(images, parset_file)
    return master_steps(metadatas, parset_file)
<commit_msg>more useful logging<commit_after>
import os
import logging
import warnings
from lofarpipe.support.parset import parameterset
import tkp.utility.accessors
from tkp.database import DataBase, DataSet, Image
from tkp.utility.accessors.dataaccessor import extract_metadata
from tempfile import NamedTemporaryFile
from pyrap.images import image as pyrap_image

logger = logging.getLogger(__name__)

def parse_parset(parset_file):
    """parse a persistence recipe specific parset file
    """
    parset = parameterset(parset_file)
    result = {}
    result['description'] = parset.getString('description', 'TRAP dataset')
    result['dataset_id'] = parset.getInt('dataset_id', -1)
    result['copy_images'] = parset.getBool('copy_images', False)
    result['mongo_host'] = parset.getString('mongo_host', 'localhost')
    result['mongo_port'] = parset.getInt('mongo_port', 27017)
    result['mongo_db'] = parset.getString('mongo_db', 'tkp')
    return result

def image_to_mongodb(filename, hostname, port, db):
    """Copy a file into mongodb"""

    try:
        import pymongo
        import gridfs
    except ImportError:
        msg = "Could not import MongoDB modules"
        logger.error(msg)
        warnings.warn(msg)
        return False

    try:
        connection = pymongo.Connection(host=hostname, port=port)
        gfs = gridfs.GridFS(connection[db])
        if gfs.exists(filename=filename):
            logger.debug("File already in database")

        else:
            # This conversion should work whether the input file
            # is in FITS or CASA format.
            # temp_fits_file is removed automatically when closed.
            temp_fits_file = NamedTemporaryFile()
            i = pyrap_image(filename)
            i.tofits(temp_fits_file.name)
            new_file = gfs.new_file(filename=filename)
            with open(temp_fits_file.name, "r") as f:
                new_file.write(f)
            new_file.close()
            logger.info("Saved local copy of %s on %s"\
                        % (os.path.basename(filename), hostname))
    except Exception, e:
        msg = "Failed to save image to MongoDB: %s" % (str(e),)
        logger.error(msg)
        warnings.warn(msg)
        return False

    finally:
        connection.close()

        # Only close this if it has been created.
        if "temp_fits_file" in locals():
            temp_fits_file.close()

    return True

def create_dataset(dataset_id, description):
    """ Creates a dataset if it doesn't exists
    Note: Should only be used in a master recipe
    Returns:
      the database ID of this dataset
    """
    database = DataBase()
    if dataset_id == -1:
        dataset = DataSet({'description': description}, database)
    else:
        dataset = DataSet(id=dataset_id, database=database)
    return dataset.id

def extract_metadatas(images):
    results = []
    for image in images:
        logger.info("Extracting metadata from %s" % image)
        accessor = tkp.utility.accessors.open(image)
        results.append(extract_metadata(accessor))
    return results

def store_images(images_metadata, dataset_id):
    """ Add images to database.
    Note that all images in one dataset should be inserted in one go, since the
    order is very important here. If you don't add them all in once, you should
    make sure they are added in the correct order e.g. sorted by observation
    time.

    Note: Should only be used in a master recipe

    Args:
        images_metadata: list of dicts containing image metadata
        dataset_id: dataset id to be used. don't use value from parset file
                    since this can be -1 (trap way of setting auto increment)
    Returns:
        the database ID of this dataset
    """
    database = DataBase()
    dataset = DataSet(id=dataset_id, database=database)
    image_ids = []

    # sort images by timestamp
    images_metadata.sort(key=lambda m: m['taustart_ts'])

    for metadata in images_metadata:
        filename = metadata['url']
        db_image = Image(data=metadata, dataset=dataset)
        image_ids.append(db_image.id)
        logger.info("stored %s with ID %s" % (os.path.basename(filename), db_image.id))
    return image_ids

def node_steps(images, parset_file):
    """
    this function executes all persistence steps that should be executed on a node.
    Note: Should only be used in a node recipe
    """
    persistence_parset = parse_parset(parset_file)
    mongohost = persistence_parset['mongo_host']
    mongoport = persistence_parset['mongo_port']
    mongodb = persistence_parset['mongo_db']
    copy_images = persistence_parset['copy_images']

    if copy_images:
        for image in images:
            image_to_mongodb(image, mongohost, mongoport, mongodb)
    else:
        logger.info("Not copying images to mongodb")

    metadatas = extract_metadatas(images)
    return metadatas

def master_steps(metadatas, parset_file):
    """this function executes all persistence steps that should be executed on
        a master.
    Args:
        parset_file: path to a parset file containig persistence settings
        metadatas: a list of dicts containing info from Image Accessors. This
                   is returned by the node recipe
    """
    persistence_parset = parse_parset(parset_file)
    logger.info("creating dataset in database ...")
    dataset_id = create_dataset(
        persistence_parset['dataset_id'],
        persistence_parset['description'])
    logger.info("added dataset with ID %s" % dataset_id)

    logger.info("Storing images")
    image_ids = store_images(metadatas, dataset_id)
    return dataset_id, image_ids

def all(images, parset_file):
    """
    execute node and then master code, should be used in local run only!
    """
    metadatas = node_steps(images, parset_file)
    return master_steps(metadatas, parset_file)
<|endoftext|>