<commit_before>from __future__ import unicode_literals

import re

from django.utils.encoding import python_2_unicode_compatible

from django.utils.encoding import force_text
from django.db.models import signals
from django.db import models
from django.utils.translation import ugettext_lazy as _

from unidecode import unidecode

from south.modelsinspector import add_introspection_rules

import autoslug

from .settings import *

__all__ = ['Country', 'Region', 'City', 'CONTINENT_CHOICES', 'to_search',
    'to_ascii']

ALPHA_REGEXP = re.compile('[\W_]+', re.UNICODE)

CONTINENT_CHOICES = (
    ('OC', _('Oceania')),
    ('EU', _('Europe')),
    ('AF', _('Africa')),
    ('NA', _('North America')),
    ('AN', _('Antarctica')),
    ('SA', _('South America')),
    ('AS', _('Asia')),
)

def to_ascii(value):
    if isinstance(value, str):
        value = force_text(value)

    return unidecode(value)

def to_search(value):
    """
    Convert a string value into a string that is usable against
    City.search_names.

    For example, 'Paris Texas' would become 'paristexas'.
    """

    return ALPHA_REGEXP.sub('', to_ascii(value)).lower()

def set_name_ascii(sender, instance=None, **kwargs):
    """
    Signal reciever that sets instance.name_ascii from instance.name.

    Ascii versions of names are often useful for autocompletes and search.
    """
    name_ascii = to_ascii(instance.name)

    if not name_ascii.strip():
        return

    if name_ascii and not instance.name_ascii:
        instance.name_ascii = to_ascii(instance.name)

def set_display_name(sender, instance=None, **kwargs):
    """
    Set instance.display_name to instance.get_display_name(), avoid spawning
    queries during __str__().
    """
    instance.display_name = instance.get_display_name()

@python_2_unicode_compatible
class Base(models.Model):
    """
    Base model with boilerplate for all models.
    """

    name_ascii = models.CharField(max_length=200, blank=True, db_index=True)
    slug = autoslug.AutoSlugField(populate_from='name_ascii')
    geoname_id = models.IntegerField(null=True, blank=True, unique=True)
    alternate_names = models.TextField(null=True, blank=True, default='')

    class Meta:
        abstract = True
        ordering = ['name']

    def __str__(self):
        display_name = getattr(self, 'display_name', None)
        if display_name:
            return display_name
        return self.name

@python_2_unicode_compatible
class Country(Base):
    """
    Country model.
    """

    name = models.CharField(max_length=200, unique=True)

    code2 = models.CharField(max_length=2, null=True, blank=True, unique=True)
    code3 = models.CharField(max_length=3, null=True, blank=True, unique=True)
    continent = models.CharField(max_length=2, db_index=True,
        choices=CONTINENT_CHOICES)
    tld = models.CharField(max_length=5, blank=True, db_index=True)

    class Meta:
        verbose_name_plural = _('countries')
signals.pre_save.connect(set_name_ascii, sender=Country)

@python_2_unicode_compatible
class Region(Base):
    """
    Region/State model.
    """

    name = models.CharField(max_length=200, db_index=True)
    display_name = models.CharField(max_length=200)
    geoname_code = models.CharField(max_length=50, null=True, blank=True,
        db_index=True)

    country = models.ForeignKey(Country)

    class Meta:
        unique_together = (('country', 'name'), )
        verbose_name = _('region/state')
        verbose_name_plural = _('regions/states')

    def get_display_name(self):
        return '%s, %s' % (self.name, self.country.name)

signals.pre_save.connect(set_name_ascii, sender=Region)
signals.pre_save.connect(set_display_name, sender=Region)

class ToSearchTextField(models.TextField):
    """
    Trivial TextField subclass that passes values through to_search
    automatically.
    """
    def get_prep_lookup(self, lookup_type, value):
        """
        Return the value passed through to_search().
        """
        value = super(ToSearchTextField, self).get_prep_lookup(lookup_type,
            value)
        return to_search(value)

    def south_field_triple(self):
        "Returns a suitable description of this field for South."
        from south.modelsinspector import introspector
        field_class = self.__class__.__module__ + "." + self.__class__.__name__
        args, kwargs = introspector(self)
        # That's our definition!
        return (field_class, args, kwargs)

@python_2_unicode_compatible
class City(Base):
    """
    City model.
    """

    name = models.CharField(max_length=200, db_index=True)
    display_name = models.CharField(max_length=200)

    search_names = ToSearchTextField(max_length=4000,
        db_index=INDEX_SEARCH_NAMES, blank=True, default='')

    latitude = models.DecimalField(max_digits=8, decimal_places=5,
        null=True, blank=True)
    longitude = models.DecimalField(max_digits=8, decimal_places=5,
        null=True, blank=True)

    region = models.ForeignKey(Region, blank=True, null=True)
    country = models.ForeignKey(Country)

    class Meta:
        unique_together = (('region', 'name'),)
        verbose_name_plural = _('cities')

    def get_display_name(self):
        if self.region_id:
            return '%s, %s, %s' % (self.name, self.region.name,
                self.country.name)
        else:
            return '%s, %s' % (self.name, self.country.name)
signals.pre_save.connect(set_name_ascii, sender=City)
signals.pre_save.connect(set_display_name, sender=City)

def city_country(sender, instance, **kwargs):
    if instance.region_id and not instance.country_id:
        instance.country = instance.region.country
signals.pre_save.connect(city_country, sender=City)

def city_search_names(sender, instance, **kwargs):
    search_names = []

    country_names = [instance.country.name]
    if instance.country.alternate_names:
        country_names += instance.country.alternate_names.split(',')

    city_names = [instance.name]
    if instance.alternate_names:
        city_names += instance.alternate_names.split(',')

    if instance.region_id:
        region_names = [instance.region.name]
        if instance.region.alternate_names:
            region_names += instance.region.alternate_names.split(',')

    for city_name in city_names:
        for country_name in country_names:
            name = to_search(city_name + country_name)
            if name not in search_names:
                search_names.append(name)

            if instance.region_id:
                for region_name in region_names:
                    name = to_search(city_name + region_name + country_name)
                    if name not in search_names:
                        search_names.append(name)

    instance.search_names = ' '.join(search_names)
signals.pre_save.connect(city_search_names, sender=City)
<commit_msg>Use bytestring regexp against ascii bytestring<commit_after>from __future__ import unicode_literals

import re

from django.utils.encoding import python_2_unicode_compatible

from django.utils.encoding import force_text
from django.db.models import signals
from django.db import models
from django.utils.translation import ugettext_lazy as _

from unidecode import unidecode

from south.modelsinspector import add_introspection_rules

import autoslug

from .settings import *

__all__ = ['Country', 'Region', 'City', 'CONTINENT_CHOICES', 'to_search',
    'to_ascii']

ALPHA_REGEXP = re.compile(b'[\W_]+')

CONTINENT_CHOICES = (
    ('OC', _('Oceania')),
    ('EU', _('Europe')),
    ('AF', _('Africa')),
    ('NA', _('North America')),
    ('AN', _('Antarctica')),
    ('SA', _('South America')),
    ('AS', _('Asia')),
)

def to_ascii(value):
    if isinstance(value, str):
        value = force_text(value)

    return unidecode(value)

def to_search(value):
    """
    Convert a string value into a string that is usable against
    City.search_names.

    For example, 'Paris Texas' would become 'paristexas'.
    """

    return ALPHA_REGEXP.sub('', to_ascii(value)).lower()

def set_name_ascii(sender, instance=None, **kwargs):
    """
    Signal reciever that sets instance.name_ascii from instance.name.

    Ascii versions of names are often useful for autocompletes and search.
    """
    name_ascii = to_ascii(instance.name)

    if not name_ascii.strip():
        return

    if name_ascii and not instance.name_ascii:
        instance.name_ascii = to_ascii(instance.name)

def set_display_name(sender, instance=None, **kwargs):
    """
    Set instance.display_name to instance.get_display_name(), avoid spawning
    queries during __str__().
    """
    instance.display_name = instance.get_display_name()

@python_2_unicode_compatible
class Base(models.Model):
    """
    Base model with boilerplate for all models.
    """

    name_ascii = models.CharField(max_length=200, blank=True, db_index=True)
    slug = autoslug.AutoSlugField(populate_from='name_ascii')
    geoname_id = models.IntegerField(null=True, blank=True, unique=True)
    alternate_names = models.TextField(null=True, blank=True, default='')

    class Meta:
        abstract = True
        ordering = ['name']

    def __str__(self):
        display_name = getattr(self, 'display_name', None)
        if display_name:
            return display_name
        return self.name

@python_2_unicode_compatible
class Country(Base):
    """
    Country model.
    """

    name = models.CharField(max_length=200, unique=True)

    code2 = models.CharField(max_length=2, null=True, blank=True, unique=True)
    code3 = models.CharField(max_length=3, null=True, blank=True, unique=True)
    continent = models.CharField(max_length=2, db_index=True,
        choices=CONTINENT_CHOICES)
    tld = models.CharField(max_length=5, blank=True, db_index=True)

    class Meta:
        verbose_name_plural = _('countries')
signals.pre_save.connect(set_name_ascii, sender=Country)

@python_2_unicode_compatible
class Region(Base):
    """
    Region/State model.
    """

    name = models.CharField(max_length=200, db_index=True)
    display_name = models.CharField(max_length=200)
    geoname_code = models.CharField(max_length=50, null=True, blank=True,
        db_index=True)

    country = models.ForeignKey(Country)

    class Meta:
        unique_together = (('country', 'name'), )
        verbose_name = _('region/state')
        verbose_name_plural = _('regions/states')

    def get_display_name(self):
        return '%s, %s' % (self.name, self.country.name)

signals.pre_save.connect(set_name_ascii, sender=Region)
signals.pre_save.connect(set_display_name, sender=Region)

class ToSearchTextField(models.TextField):
    """
    Trivial TextField subclass that passes values through to_search
    automatically.
    """
    def get_prep_lookup(self, lookup_type, value):
        """
        Return the value passed through to_search().
        """
        value = super(ToSearchTextField, self).get_prep_lookup(lookup_type,
            value)
        return to_search(value)

    def south_field_triple(self):
        "Returns a suitable description of this field for South."
        from south.modelsinspector import introspector
        field_class = self.__class__.__module__ + "." + self.__class__.__name__
        args, kwargs = introspector(self)
        # That's our definition!
        return (field_class, args, kwargs)

@python_2_unicode_compatible
class City(Base):
    """
    City model.
    """

    name = models.CharField(max_length=200, db_index=True)
    display_name = models.CharField(max_length=200)

    search_names = ToSearchTextField(max_length=4000,
        db_index=INDEX_SEARCH_NAMES, blank=True, default='')

    latitude = models.DecimalField(max_digits=8, decimal_places=5,
        null=True, blank=True)
    longitude = models.DecimalField(max_digits=8, decimal_places=5,
        null=True, blank=True)

    region = models.ForeignKey(Region, blank=True, null=True)
    country = models.ForeignKey(Country)

    class Meta:
        unique_together = (('region', 'name'),)
        verbose_name_plural = _('cities')

    def get_display_name(self):
        if self.region_id:
            return '%s, %s, %s' % (self.name, self.region.name,
                self.country.name)
        else:
            return '%s, %s' % (self.name, self.country.name)
signals.pre_save.connect(set_name_ascii, sender=City)
signals.pre_save.connect(set_display_name, sender=City)

def city_country(sender, instance, **kwargs):
    if instance.region_id and not instance.country_id:
        instance.country = instance.region.country
signals.pre_save.connect(city_country, sender=City)

def city_search_names(sender, instance, **kwargs):
    search_names = []

    country_names = [instance.country.name]
    if instance.country.alternate_names:
        country_names += instance.country.alternate_names.split(',')

    city_names = [instance.name]
    if instance.alternate_names:
        city_names += instance.alternate_names.split(',')

    if instance.region_id:
        region_names = [instance.region.name]
        if instance.region.alternate_names:
            region_names += instance.region.alternate_names.split(',')

    for city_name in city_names:
        for country_name in country_names:
            name = to_search(city_name + country_name)
            if name not in search_names:
                search_names.append(name)

            if instance.region_id:
                for region_name in region_names:
                    name = to_search(city_name + region_name + country_name)
                    if name not in search_names:
                        search_names.append(name)

    instance.search_names = ' '.join(search_names)
signals.pre_save.connect(city_search_names, sender=City)
<|endoftext|>


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

import os

from setuptools import find_packages
from setuptools import setup

PACKAGE_ROOT = os.path.abspath(os.path.dirname(__file__))

with open(os.path.join(PACKAGE_ROOT, 'README.rst')) as file_obj:
    README = file_obj.read()

# NOTE: This is duplicated throughout and we should try to
#       consolidate.
SETUP_BASE = {
    'author': 'Google Cloud Platform',
    'author_email': 'jjg+google-cloud-python@google.com',
    'scripts': [],
    'url': 'https://github.com/GoogleCloudPlatform/google-cloud-python',
    'license': 'Apache 2.0',
    'platforms': 'Posix; MacOS X; Windows',
    'include_package_data': True,
    'zip_safe': False,
    'classifiers': [
        'Development Status :: 4 - Beta',
        'Intended Audience :: Developers',
        'License :: OSI Approved :: Apache Software License',
        'Operating System :: OS Independent',
        'Programming Language :: Python :: 2',
        'Programming Language :: Python :: 2.7',
        'Programming Language :: Python :: 3',
        'Programming Language :: Python :: 3.4',
        'Programming Language :: Python :: 3.5',
        'Programming Language :: Python :: 3.6',
        'Topic :: Internet',
    ],
}

REQUIREMENTS = [
    'enum34',
    'google-cloud-core >= 0.24.0, < 0.25dev',
    'gapic-google-cloud-vision-v1 >= 0.90.3, < 0.91dev',
]

setup(
    name='google-cloud-vision',
    version='0.24.0',
    description='Python Client for Google Cloud Vision',
    long_description=README,
    namespace_packages=[
        'google',
        'google.cloud',
    ],
    packages=find_packages(exclude=('tests*',)),
    install_requires=REQUIREMENTS,
    **SETUP_BASE
)
<commit_msg>Making enum34 dependency conditional on Python version. (#3402)<commit_after># Copyright 2016 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os

from setuptools import find_packages
from setuptools import setup

PACKAGE_ROOT = os.path.abspath(os.path.dirname(__file__))

with open(os.path.join(PACKAGE_ROOT, 'README.rst')) as file_obj:
    README = file_obj.read()

# NOTE: This is duplicated throughout and we should try to
#       consolidate.
SETUP_BASE = {
    'author': 'Google Cloud Platform',
    'author_email': 'jjg+google-cloud-python@google.com',
    'scripts': [],
    'url': 'https://github.com/GoogleCloudPlatform/google-cloud-python',
    'license': 'Apache 2.0',
    'platforms': 'Posix; MacOS X; Windows',
    'include_package_data': True,
    'zip_safe': False,
    'classifiers': [
        'Development Status :: 4 - Beta',
        'Intended Audience :: Developers',
        'License :: OSI Approved :: Apache Software License',
        'Operating System :: OS Independent',
        'Programming Language :: Python :: 2',
        'Programming Language :: Python :: 2.7',
        'Programming Language :: Python :: 3',
        'Programming Language :: Python :: 3.4',
        'Programming Language :: Python :: 3.5',
        'Programming Language :: Python :: 3.6',
        'Topic :: Internet',
    ],
}

REQUIREMENTS = [
    'google-cloud-core >= 0.24.0, < 0.25dev',
    'gapic-google-cloud-vision-v1 >= 0.90.3, < 0.91dev',
]
EXTRAS_REQUIRE = {
    ':python_version<"3.4"': ['enum34'],
}

setup(
    name='google-cloud-vision',
    version='0.24.0',
    description='Python Client for Google Cloud Vision',
    long_description=README,
    namespace_packages=[
        'google',
        'google.cloud',
    ],
    packages=find_packages(exclude=('tests*',)),
    install_requires=REQUIREMENTS,
    extras_require=EXTRAS_REQUIRE,
    **SETUP_BASE
)
<|endoftext|>


<commit_before>"""All necessary calculations in pre_simulate()
"""
import os
import logging

from energy_demand.read_write import data_loader
from energy_demand.basic import basic_functions
from energy_demand.scripts import init_scripts
from energy_demand.read_write import write_data
from energy_demand.assumptions import general_assumptions
from energy_demand.validation import lad_validation
from energy_demand.scripts import s_disaggregation
from energy_demand.basic import demand_supply_interaction
from energy_demand.read_write import read_data
from energy_demand.technologies import tech_related
from energy_demand.scripts import s_scenario_param
from energy_demand.geography import weather_region
from energy_demand.basic import testing_functions

def load_data_before_simulation(
        data,
        sim_yrs,
        config,
        curr_yr
    ):
    # ---------
    # Configuration
    # -----------
    base_yr = config['CONFIG']['base_yr']
    weather_yr_scenario = config['CONFIG']['weather_yr_scenario']
    path_new_scenario = data['path_new_scenario']

    data['weather_station_count_nr'] = [] # Default value is '[]' to use all stations
    data['data_path'] = os.path.normpath(config['PATHS']['path_local_data'])
    data['processed_path'] = os.path.normpath(config['PATHS']['path_processed_data'])
    data['result_path'] = os.path.normpath(config['PATHS']['path_result_data'])
    data['paths'] = data_loader.load_paths(config['PATHS']['path_config_data'])

    # Downloaded (FTP) data
    data['local_paths'] = data_loader.get_local_paths(
        data['data_path'])

    # ------------------------------------------------
    # Load Inputs
    # ------------------------------------------------
    data['enduses'], data['sectors'], data['fuels'], lookup_enduses, lookup_sector_enduses = data_loader.load_fuels(
        data['paths'])

    # ------------------------------------------------
    # Load Assumptions
    # ------------------------------------------------
    data['assumptions'] = general_assumptions.Assumptions(
        lookup_enduses=lookup_enduses,
        lookup_sector_enduses=lookup_sector_enduses,
        base_yr=base_yr,
        weather_by=config['CONFIG']['user_defined_weather_by'],
        simulation_end_yr=config['CONFIG']['user_defined_simulation_end_yr'],
        curr_yr=curr_yr,
        sim_yrs=sim_yrs,
        paths=data['paths'],
        enduses=data['enduses'],
        sectors=data['sectors'],
        reg_nrs=len(data['regions']))

    # --------------------------------------------
    # Make selection of weather stations and data
    # --------------------------------------------
    weather_stations_selection = {}
    temp_data_selection = {}
    if data['weather_station_count_nr'] != []:
        for year in sim_yrs:
            weather_stations_selection, station_id = weather_region.get_weather_station_selection(
                data['weather_stations'],
                counter=data['weather_station_count_nr'],
                weather_yr=weather_yr_scenario)
            temp_data_selection[year] = data['temp_data'][year][station_id]

            if year == weather_yr_scenario:
                data['simulation_name'] = str(weather_yr_scenario) + "__" + str(station_id)
    else:
        for year in sim_yrs:
            weather_stations_selection = data['weather_stations']
            temp_data_selection[year] = data['temp_data'][year]

            if year == weather_yr_scenario:
                data['simulation_name'] = str(weather_yr_scenario) + "__" + "all_stations"

    # Replace weather station with selection
    data['weather_stations'] = weather_stations_selection
    data['temp_data'] = temp_data_selection

    # ------------------------------------------
    # Make selection of regions to model
    # ------------------------------------------
    if config['CRITERIA']['reg_selection']:

        region_selection = read_data.get_region_selection(
            os.path.join(data['local_paths']['local_path_datafolder'],
            "region_definitions",
            config['CRITERIA']['reg_selection_csv_name']))
        #region_selection = ['E02003237', 'E02003238']
        setattr(data['assumptions'], 'reg_nrs', len(region_selection))
    else:
        region_selection = data['regions']

    # Create .ini file with simulation parameter
    write_data.write_simulation_inifile(
        data['path_new_scenario'], data, region_selection)

    # -------------------------------------------
    # Weather year specific initialisations
    # -------------------------------------------
    path_folder_weather_yr = os.path.join(
        os.path.join(path_new_scenario, data['simulation_name']))

    data['result_paths'] = data_loader.get_result_paths(path_folder_weather_yr)

    folders_to_create = [
        path_folder_weather_yr,
        data['result_paths']['data_results'],
        data['result_paths']['data_results_PDF'],
        data['result_paths']['data_results_validation'],
        data['result_paths']['data_results_model_runs']]
    for folder in folders_to_create:
        basic_functions.create_folder(folder)

    # ------------------------------------------------
    # Load load profiles of technologies
    # ------------------------------------------------
    data['tech_lp'] = data_loader.load_data_profiles(
        data['paths'],
        data['local_paths'],
        data['assumptions'].model_yeardays,
        data['assumptions'].model_yeardays_daytype)

    # Obtain population data for disaggregation
    if config['CRITERIA']['msoa_crit']:
        name_population_dataset = data['local_paths']['path_population_data_for_disaggregation_MSOA']
    else:
        name_population_dataset = data['local_paths']['path_population_data_for_disaggregation_LAD']
    data['pop_for_disag'] = data_loader.read_scenario_data(name_population_dataset)

    # ------------------------------------------------
    # Load building related data
    # ------------------------------------------------
    if config['CRITERIA']['virtual_building_stock_criteria']:
        data['scenario_data']['floor_area']['rs_floorarea'], data['scenario_data']['floor_area']['ss_floorarea'], data['service_building_count'], rs_regions_without_floorarea, ss_regions_without_floorarea = data_loader.floor_area_virtual_dw(
            data['regions'],
            data['sectors'],
            data['local_paths'],
            data['scenario_data']['population'][data['assumptions'].base_yr],
            base_yr=data['assumptions'].base_yr)

        # Add all areas with no floor area data
        data['assumptions'].update("rs_regions_without_floorarea", rs_regions_without_floorarea)
        data['assumptions'].update("ss_regions_without_floorarea", ss_regions_without_floorarea)
    else:
        # ------------------------------------------------
        # Load floor area directly from Newcastle
        # ------------------------------------------------
        #rs_floorarea = defaultdict(dict)
        #ss_floorarea = defaultdict(dict)
        pass

    return data

def before_simulation(
        data,
        config,
        sim_yrs,
        pop_density,
        service_switches,
        fuel_switches,
        capacity_switches
    ):
    """
    """
    # ------------------------------------------------
    # Disaggregate national energy demand to regional demands
    # ------------------------------------------------
    fuel_disagg = s_disaggregation.disaggr_demand(
        data,
        config['CRITERIA']['crit_temp_min_max'],
        spatial_calibration=config['CRITERIA']['spatial_calibration'])

    # ------------------------------------------------
    # Calculate spatial diffusion factors
    #
    # Here the real values used for the spatial disaggregation (speec_con_max)
    # need to be defined. If not population density is used,
    # this needs to be replaced by any other values which are loaded from
    # a csv file in the form of: {{region_name: value}}
    # ------------------------------------------------
    f_reg, f_reg_norm, f_reg_norm_abs, crit_all_the_same = init_scripts.create_spatial_diffusion_factors(
        narrative_spatial_explicit_diffusion=data['assumptions'].strategy_vars['spatial_explicit_diffusion'],
        fuel_disagg=fuel_disagg,
        regions=data['regions'],
        real_values=pop_density,
        narrative_speed_con_max=data['assumptions'].strategy_vars['speed_con_max'])

    # ------------------------------------------------
    # Calculate parameter values for every region
    # ------------------------------------------------
    regional_vars = init_scripts.spatial_explicit_modelling_strategy_vars(
        data['assumptions'].strategy_vars,
        data['assumptions'].spatially_modelled_vars,
        data['regions'],
        fuel_disagg,
        f_reg,
        f_reg_norm,
        f_reg_norm_abs)
    data['assumptions'].update('strategy_vars', regional_vars)

    # ------------------------------------------------
    # Calculate parameter values for every simulated year
    # based on narratives. Also calculate annual parameters for
    # technologies diffused by switches.
    # ------------------------------------------------
    regional_vars, non_regional_vars = s_scenario_param.generate_annual_param_vals(
        data['regions'],
        data['assumptions'].strategy_vars,
        sim_yrs)

    # ------------------------------------------------
    # Switches calculations
    # ------------------------------------------------
    # Update assumptions
    crit_switch_happening = testing_functions.switch_testing(
        fuel_switches=fuel_switches,
        service_switches=service_switches,
        capacity_switches=capacity_switches)
    setattr(data['assumptions'],'crit_switch_happening', crit_switch_happening)

    annual_tech_diff_params = init_scripts.switch_calculations(
        sim_yrs,
        data,
        f_reg,
        f_reg_norm,
        f_reg_norm_abs,
        crit_all_the_same,
        service_switches=service_switches,
        fuel_switches=fuel_switches,
        capacity_switches=capacity_switches)

    for region in data['regions']:
        regional_vars[region]['annual_tech_diff_params'] = annual_tech_diff_params[region]

    return regional_vars, non_regional_vars, fuel_disagg, crit_switch_happening

def write_user_defined_results(
        criterias,
        result_paths,
        sim_obj,
        data,
        curr_yr,
        region_selection
    ):
    """
    Write annual results to files
    """

    logging.info("... Start writing results to file")
    if criterias['write_txt_additional_results']:
        # Write full results (Note: Results in very large data written to file)
        ##write_data.write_full_results(
        ##    data_handle.current_timestep,
        ##    data['result_paths']['data_results_model_runs'],
        ##    sim_obj.ed_submodel_enduse_fueltype_regs_yh, #TODO CHANGED FORMAT
        ##    "out_enduse_specific")
        write_data.write_supply_results(
            curr_yr,
            "ed_fueltype_regs_yh",
            result_paths['data_results_model_runs'],
            sim_obj.ed_fueltype_regs_yh,
            "result_tot_submodels_fueltypes")
        write_data.write_enduse_specific(
            curr_yr,
            result_paths['data_results_model_runs'],
            sim_obj.tot_fuel_y_enduse_specific_yh,
            "out_enduse_specific")
        write_data.write_lf(
            result_paths['data_results_model_runs'], "result_reg_load_factor_y",
            [curr_yr], sim_obj.reg_load_factor_y, 'reg_load_factor_y')
        write_data.write_lf(
            result_paths['data_results_model_runs'], "result_reg_load_factor_yd",
            [curr_yr], sim_obj.reg_load_factor_yd, 'reg_load_factor_yd')

    # ----------------------------------------------------------------------------------------
    # Write out national demand for every fueltype (used for first sending of demand data)
    # ----------------------------------------------------------------------------------------
    if criterias['write_out_national']:

        # Write out gas
        demand_supply_interaction.write_national_results(
            path_folder=result_paths,
            results_unconstrained=sim_obj.results_unconstrained,
            enduse_specific_results=sim_obj.tot_fuel_y_enduse_specific_yh,
            fueltype_str='gas',
            fuelype_nr=tech_related.get_fueltype_int('gas'),
            year=curr_yr,
            submodels_names=data['assumptions'].submodels_names)

        # Write out elec
        demand_supply_interaction.write_national_results(
            path_folder=result_paths,
            results_unconstrained=sim_obj.results_unconstrained,
            enduse_specific_results=sim_obj.tot_fuel_y_enduse_specific_yh,
            fueltype_str='electricity',
            fuelype_nr=tech_related.get_fueltype_int('electricity'),
            year=curr_yr,
            submodels_names=data['assumptions'].submodels_names)

    # ------------------------------------------------
    # Temporal Validation
    # ------------------------------------------------
    if (criterias['validation_criteria'] == True) and (
        curr_yr == data['assumptions'].base_yr) and (
           ['cluster_calc'] != True):
        lad_validation.spatio_temporal_val(
            sim_obj.ed_fueltype_national_yh,
            sim_obj.ed_fueltype_regs_yh,
            result_paths,
            data['paths'],
            region_selection,
            data['assumptions'].seasons,
            data['assumptions'].model_yeardays_daytype,
            plot_crit=False)

def load_gva_sector(
        data_handle,
        regions,
        sectors_to_load,
        msoa_crit,
        simulate=False
    ):
    """Load sector specific GVA

    Arguments
    ---------
    data_handle : object
        Data handler
    pop_array : array
        Population
    regions : list
        Regions
    sectors_to_load : list
        Sectors which are loaded
    msoa_crit : bool
        Criteria wheter modelled on MSOA level or LAD level
    simulate : bool
        Criteria wheter run in simulate() or not
    """
    sector_data = {}
    if msoa_crit:
        logging.info("Don't load sector GVA {}")
    else:
        for gva_sector_nr in sectors_to_load:
            try:
                logging.info("... Loading GVA data for sector_Nr {}".format(gva_sector_nr))
                if simulate:
                    gva_sector_data = data_handle.get_data(
                        'gva_per_head_sector__{}'.format(gva_sector_nr))
                else:
                    gva_sector_data = data_handle.get_base_timestep_data(
                        'gva_per_head_sector__{}'.format(gva_sector_nr))

                sector_data[gva_sector_nr] = basic_functions.assign_array_to_dict(
                    gva_sector_data, regions)
            except KeyError:
                # In case no data could be loaded, generate constant dummy data
                raise Exception("Could not load data %s", 'gva_per_head_sector__{}'.format(gva_sector_nr))

    return sector_data

def plots(
        data,
        curr_yr,
        fuel_disagg,
        config
    ):
    """
    """
    # Spatial validation
    if (config['CRITERIA']['validation_criteria'] == True) and (
        curr_yr == data['assumptions'].base_yr) and (
            config['CRITERIA']['cluster_calc'] != True):
        lad_validation.spatial_validation_lad_level(
            fuel_disagg,
            data['result_paths'],
            data['paths'],
            data['regions'],
            data['reg_coord'],
            plot_crit=False)

    # Plot map with weather station
    if config['CRITERIA']['cluster_calc'] != True:
        data_loader.create_weather_station_map(
            data['weather_stations'],
            os.path.join(data['result_path'], 'weatherst_distr_weathyr_{}.pdf'.format(
                config['CONFIG']['weather_yr_scenario'])),
            path_shapefile=data['local_paths']['lad_shapefile'])
<commit_msg>WIP - results folder dynamically created from modelrun name<commit_after>"""All necessary calculations in pre_simulate()
"""
import os
import logging

from energy_demand.read_write import data_loader
from energy_demand.basic import basic_functions
from energy_demand.scripts import init_scripts
from energy_demand.read_write import write_data
from energy_demand.assumptions import general_assumptions
from energy_demand.validation import lad_validation
from energy_demand.scripts import s_disaggregation
from energy_demand.basic import demand_supply_interaction
from energy_demand.read_write import read_data
from energy_demand.technologies import tech_related
from energy_demand.scripts import s_scenario_param
from energy_demand.geography import weather_region
from energy_demand.basic import testing_functions

def load_data_before_simulation(
        data,
        sim_yrs,
        config,
        curr_yr
    ):
    # ---------
    # Configuration
    # -----------
    base_yr = config['CONFIG']['base_yr']
    weather_yr_scenario = config['CONFIG']['weather_yr_scenario']
    path_new_scenario = data['path_new_scenario']

    data['weather_station_count_nr'] = [] # Default value is '[]' to use all stations
    data['data_path'] = os.path.normpath(config['PATHS']['path_local_data'])
    data['processed_path'] = os.path.normpath(config['PATHS']['path_processed_data'])
    data['result_path'] = os.path.normpath(config['PATHS']['path_result_data'])
    data['paths'] = config['CONFIG_DATA']

    # Downloaded (FTP) data
    data['local_paths'] = config['DATA_PATHS']

    # ------------------------------------------------
    # Load Inputs
    # ------------------------------------------------
    data['enduses'], data['sectors'], data['fuels'], lookup_enduses, \
        lookup_sector_enduses = data_loader.load_fuels(data['paths'])

    # ------------------------------------------------
    # Load Assumptions
    # ------------------------------------------------
    data['assumptions'] = general_assumptions.Assumptions(
        lookup_enduses=lookup_enduses,
        lookup_sector_enduses=lookup_sector_enduses,
        base_yr=base_yr,
        weather_by=config['CONFIG']['user_defined_weather_by'],
        simulation_end_yr=config['CONFIG']['user_defined_simulation_end_yr'],
        curr_yr=curr_yr,
        sim_yrs=sim_yrs,
        paths=data['paths'],
        enduses=data['enduses'],
        sectors=data['sectors'],
        reg_nrs=len(data['regions']))

    # --------------------------------------------
    # Make selection of weather stations and data
    # --------------------------------------------
    weather_stations_selection = {}
    temp_data_selection = {}
    if data['weather_station_count_nr'] != []:
        for year in sim_yrs:
            weather_stations_selection, station_id = weather_region.get_weather_station_selection(
                data['weather_stations'],
                counter=data['weather_station_count_nr'],
                weather_yr=weather_yr_scenario)
            temp_data_selection[year] = data['temp_data'][year][station_id]

            if year == weather_yr_scenario:
                data['simulation_name'] = str(weather_yr_scenario) + "__" + str(station_id)
    else:
        for year in sim_yrs:
            weather_stations_selection = data['weather_stations']
            temp_data_selection[year] = data['temp_data'][year]

            if year == weather_yr_scenario:
                data['simulation_name'] = str(weather_yr_scenario) + "__" + "all_stations"

    # Replace weather station with selection
    data['weather_stations'] = weather_stations_selection
    data['temp_data'] = temp_data_selection

    # ------------------------------------------
    # Make selection of regions to model
    # ------------------------------------------
    if config['CRITERIA']['reg_selection']:

        region_selection = read_data.get_region_selection(
            os.path.join(data['local_paths']['local_path_datafolder'],
            "region_definitions",
            config['CRITERIA']['reg_selection_csv_name']))
        #region_selection = ['E02003237', 'E02003238']
        setattr(data['assumptions'], 'reg_nrs', len(region_selection))
    else:
        region_selection = data['regions']

    # Create .ini file with simulation parameter
    write_data.write_simulation_inifile(
        data['path_new_scenario'], data, region_selection)

    # -------------------------------------------
    # Weather year specific initialisations
    # -------------------------------------------
    path_folder_weather_yr = os.path.join(
        os.path.join(path_new_scenario, data['simulation_name']))

    data['result_paths'] = data_loader.get_result_paths(path_folder_weather_yr)

    folders_to_create = [
        path_folder_weather_yr,
        data['result_paths']['data_results'],
        data['result_paths']['data_results_PDF'],
        data['result_paths']['data_results_validation'],
        data['result_paths']['data_results_model_runs']]
    for folder in folders_to_create:
        basic_functions.create_folder(folder)

    # ------------------------------------------------
    # Load load profiles of technologies
    # ------------------------------------------------
    data['tech_lp'] = data_loader.load_data_profiles(
        data['paths'],
        data['local_paths'],
        data['assumptions'].model_yeardays,
        data['assumptions'].model_yeardays_daytype)

    # Obtain population data for disaggregation
    if config['CRITERIA']['msoa_crit']:
        name_population_dataset = data['local_paths']['path_population_data_for_disaggregation_MSOA']
    else:
        name_population_dataset = data['local_paths']['path_population_data_for_disaggregation_LAD']
    data['pop_for_disag'] = data_loader.read_scenario_data(name_population_dataset)

    # ------------------------------------------------
    # Load building related data
    # ------------------------------------------------
    if config['CRITERIA']['virtual_building_stock_criteria']:
        data['scenario_data']['floor_area']['rs_floorarea'], data['scenario_data']['floor_area']['ss_floorarea'], data['service_building_count'], rs_regions_without_floorarea, ss_regions_without_floorarea = data_loader.floor_area_virtual_dw(
            data['regions'],
            data['sectors'],
            data['local_paths'],
            data['scenario_data']['population'][data['assumptions'].base_yr],
            base_yr=data['assumptions'].base_yr)

        # Add all areas with no floor area data
        data['assumptions'].update("rs_regions_without_floorarea", rs_regions_without_floorarea)
        data['assumptions'].update("ss_regions_without_floorarea", ss_regions_without_floorarea)
    else:
        # ------------------------------------------------
        # Load floor area directly from Newcastle
        # ------------------------------------------------
        #rs_floorarea = defaultdict(dict)
        #ss_floorarea = defaultdict(dict)
        pass

    return data

def before_simulation(
        data,
        config,
        sim_yrs,
        pop_density,
        service_switches,
        fuel_switches,
        capacity_switches
    ):
    """
    """
    # ------------------------------------------------
    # Disaggregate national energy demand to regional demands
    # ------------------------------------------------
    fuel_disagg = s_disaggregation.disaggr_demand(
        data,
        config['CRITERIA']['crit_temp_min_max'],
        spatial_calibration=config['CRITERIA']['spatial_calibration'])

    # ------------------------------------------------
    # Calculate spatial diffusion factors
    #
    # Here the real values used for the spatial disaggregation (speec_con_max)
    # need to be defined. If not population density is used,
    # this needs to be replaced by any other values which are loaded from
    # a csv file in the form of: {{region_name: value}}
    # ------------------------------------------------
    f_reg, f_reg_norm, f_reg_norm_abs, crit_all_the_same = init_scripts.create_spatial_diffusion_factors(
        narrative_spatial_explicit_diffusion=data['assumptions'].strategy_vars['spatial_explicit_diffusion'],
        fuel_disagg=fuel_disagg,
        regions=data['regions'],
        real_values=pop_density,
        narrative_speed_con_max=data['assumptions'].strategy_vars['speed_con_max'])

    # ------------------------------------------------
    # Calculate parameter values for every region
    # ------------------------------------------------
    regional_vars = init_scripts.spatial_explicit_modelling_strategy_vars(
        data['assumptions'].strategy_vars,
        data['assumptions'].spatially_modelled_vars,
        data['regions'],
        fuel_disagg,
        f_reg,
        f_reg_norm,
        f_reg_norm_abs)
    data['assumptions'].update('strategy_vars', regional_vars)

    # ------------------------------------------------
    # Calculate parameter values for every simulated year
    # based on narratives. Also calculate annual parameters for
    # technologies diffused by switches.
    # ------------------------------------------------
    regional_vars, non_regional_vars = s_scenario_param.generate_annual_param_vals(
        data['regions'],
        data['assumptions'].strategy_vars,
        sim_yrs)

    # ------------------------------------------------
    # Switches calculations
    # ------------------------------------------------
    # Update assumptions
    crit_switch_happening = testing_functions.switch_testing(
        fuel_switches=fuel_switches,
        service_switches=service_switches,
        capacity_switches=capacity_switches)
    setattr(data['assumptions'],'crit_switch_happening', crit_switch_happening)

    annual_tech_diff_params = init_scripts.switch_calculations(
        sim_yrs,
        data,
        f_reg,
        f_reg_norm,
        f_reg_norm_abs,
        crit_all_the_same,
        service_switches=service_switches,
        fuel_switches=fuel_switches,
        capacity_switches=capacity_switches)

    for region in data['regions']:
        regional_vars[region]['annual_tech_diff_params'] = annual_tech_diff_params[region]

    return regional_vars, non_regional_vars, fuel_disagg, crit_switch_happening

def write_user_defined_results(
        criterias,
        result_paths,
        sim_obj,
        data,
        curr_yr,
        region_selection
    ):
    """
    Write annual results to files
    """

    logging.info("... Start writing results to file")
    if criterias['write_txt_additional_results']:
        # Write full results (Note: Results in very large data written to file)
        ##write_data.write_full_results(
        ##    data_handle.current_timestep,
        ##    data['result_paths']['data_results_model_runs'],
        ##    sim_obj.ed_submodel_enduse_fueltype_regs_yh, #TODO CHANGED FORMAT
        ##    "out_enduse_specific")
        write_data.write_supply_results(
            curr_yr,
            "ed_fueltype_regs_yh",
            result_paths['data_results_model_runs'],
            sim_obj.ed_fueltype_regs_yh,
            "result_tot_submodels_fueltypes")
        write_data.write_enduse_specific(
            curr_yr,
            result_paths['data_results_model_runs'],
            sim_obj.tot_fuel_y_enduse_specific_yh,
            "out_enduse_specific")
        write_data.write_lf(
            result_paths['data_results_model_runs'], "result_reg_load_factor_y",
            [curr_yr], sim_obj.reg_load_factor_y, 'reg_load_factor_y')
        write_data.write_lf(
            result_paths['data_results_model_runs'], "result_reg_load_factor_yd",
            [curr_yr], sim_obj.reg_load_factor_yd, 'reg_load_factor_yd')

    # ----------------------------------------------------------------------------------------
    # Write out national demand for every fueltype (used for first sending of demand data)
    # ----------------------------------------------------------------------------------------
    if criterias['write_out_national']:

        # Write out gas
        demand_supply_interaction.write_national_results(
            path_folder=result_paths,
            results_unconstrained=sim_obj.results_unconstrained,
            enduse_specific_results=sim_obj.tot_fuel_y_enduse_specific_yh,
            fueltype_str='gas',
            fuelype_nr=tech_related.get_fueltype_int('gas'),
            year=curr_yr,
            submodels_names=data['assumptions'].submodels_names)

        # Write out elec
        demand_supply_interaction.write_national_results(
            path_folder=result_paths,
            results_unconstrained=sim_obj.results_unconstrained,
            enduse_specific_results=sim_obj.tot_fuel_y_enduse_specific_yh,
            fueltype_str='electricity',
            fuelype_nr=tech_related.get_fueltype_int('electricity'),
            year=curr_yr,
            submodels_names=data['assumptions'].submodels_names)

    # ------------------------------------------------
    # Temporal Validation
    # ------------------------------------------------
    if (criterias['validation_criteria'] == True) and (
        curr_yr == data['assumptions'].base_yr) and (
           ['cluster_calc'] != True):
        lad_validation.spatio_temporal_val(
            sim_obj.ed_fueltype_national_yh,
            sim_obj.ed_fueltype_regs_yh,
            result_paths,
            data['paths'],
            region_selection,
            data['assumptions'].seasons,
            data['assumptions'].model_yeardays_daytype,
            plot_crit=False)

def load_gva_sector(
        data_handle,
        regions,
        sectors_to_load,
        msoa_crit,
        simulate=False
    ):
    """Load sector specific GVA

    Arguments
    ---------
    data_handle : object
        Data handler
    pop_array : array
        Population
    regions : list
        Regions
    sectors_to_load : list
        Sectors which are loaded
    msoa_crit : bool
        Criteria wheter modelled on MSOA level or LAD level
    simulate : bool
        Criteria wheter run in simulate() or not
    """
    sector_data = {}
    if msoa_crit:
        logging.info("Don't load sector GVA {}")
    else:
        for gva_sector_nr in sectors_to_load:
            try:
                logging.info("... Loading GVA data for sector_Nr {}".format(gva_sector_nr))
                if simulate:
                    gva_sector_data = data_handle.get_data(
                        'gva_per_head_sector__{}'.format(gva_sector_nr))
                else:
                    gva_sector_data = data_handle.get_base_timestep_data(
                        'gva_per_head_sector__{}'.format(gva_sector_nr))

                sector_data[gva_sector_nr] = basic_functions.assign_array_to_dict(
                    gva_sector_data, regions)
            except KeyError:
                # In case no data could be loaded, generate constant dummy data
                raise Exception("Could not load data %s", 'gva_per_head_sector__{}'.format(gva_sector_nr))

    return sector_data

def plots(
        data,
        curr_yr,
        fuel_disagg,
        config
    ):
    """
    """
    # Spatial validation
    if (config['CRITERIA']['validation_criteria'] == True) and (
        curr_yr == data['assumptions'].base_yr) and (
            config['CRITERIA']['cluster_calc'] != True):
        lad_validation.spatial_validation_lad_level(
            fuel_disagg,
            data['result_paths'],
            data['paths'],
            data['regions'],
            data['reg_coord'],
            plot_crit=False)

    # Plot map with weather station
    if config['CRITERIA']['cluster_calc'] != True:
        data_loader.create_weather_station_map(
            data['weather_stations'],
            os.path.join(data['result_path'], 'weatherst_distr_weathyr_{}.pdf'.format(
                config['CONFIG']['weather_yr_scenario'])),
            path_shapefile=data['local_paths']['lad_shapefile'])
<|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 sys
import pysam

from utils import logger

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

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

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

class bamUtils(object):
    """
    Tool for aligning sequence reads to a genome using BWA
    """

    def __init__(self):
        """
        Initialise the tool with its configuration.

        Parameters
        ----------
        configuration : dict
            a dictionary containing parameters that define how the operation
            should be carried out, which are specific to each Tool.
        """
        logger.info("BAM Utils")

    @staticmethod
    def bam_sort(bam_file):
        """
        Wrapper for the pysam SAMtools sort function

        Parameters
        ----------
        bam_file : str
            Location of the bam file to sort
        """
        try:
            pysam.sort("-o", bam_file, "-T", bam_file + "_sort", bam_file)
        except IOError:
            return False
        return True

    @staticmethod
    def bam_merge(bam_file_1, bam_file_2):
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file_1 : str
            Location of the bam file to merge into
        bam_file_2 : str
            Location of the bam file that is to get merged into bam_file_1
        """
        logger.info("Merging: " + bam_file_1 + " - " + bam_file_2)
        pysam.merge("-f", bam_file_1 + "_merge.bam", bam_file_1, bam_file_2)

        try:
            with open(bam_file_1 + "_merge.bam", "rb") as f_in:
                with open(bam_file_1, "wb") as f_out:
                    f_out.write(f_in.read())
        except IOError:
            return False

        os.remove(bam_file_1 + "_merge.bam")

        return True

    @staticmethod
    def bam_copy(bam_in, bam_out):
        """
        Wrapper function to copy from one bam file to another

        Parameters
        ----------
        bam_in : str
            Location of the input bam file
        bam_out : str
            Location of the output bam file
        """
        logger.info("Copying: " + bam_in + " - " + bam_out)

        try:
            with open(bam_in, "rb") as f_in:
                with open(bam_out, "wb") as f_out:
                    f_out.write(f_in.read())
        except IOError:
            return False

        return True

    @staticmethod
    def bam_index(bam_file, bam_idx_file):
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file : str
            Location of the bam file that is to be indexed
        bam_idx_file : str
            Location of the bam index file (.bai)
        """
        pysam.index(bam_file, bam_file + "_tmp.bai")

        try:
            with open(bam_file + "_tmp.bai", "rb") as f_in:
                with open(bam_idx_file, "wb") as f_out:
                    f_out.write(f_in.read())
        except IOError:
            return False

        return True

    @staticmethod
    def check_header(bam_file):
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file_1 : str
            Location of the bam file to merge into
        bam_file_2 : str
            Location of the bam file that is to get merged into bam_file_1
        """
        output = True

        bam_file_handle = pysam.AlignmentFile(bam_file, "rb")
        if ("SO" not in bam_file_handle.header["HD"] or
                bam_file_handle.header["HD"]["SO"] == "unsorted"):
            output = False
        bam_file_handle.close()

        return output

class bamUtilsTask(object):
    """
    Wrappers so that the function above can be used as part of a @task within
    COMPSs avoiding the files being copied around the infrstructure too many
    times
    """

    def __init__(self):
        """
        Init function
        """
        logger.info("BAM @task Utils")

    @task(bam_file=FILE_INOUT)
    def bam_sort(self, bam_file):  # pylint: disable=no-self-use
        """
        Wrapper for the pysam SAMtools sort function

        Parameters
        ----------
        bam_file : str
            Location of the bam file to sort
        """
        bam_handle = bamUtils()
        return bam_handle.bam_sort(bam_file)

    @task(bam_file_1=FILE_INOUT, bam_file_2=FILE_IN)
    def bam_merge(self, bam_file_1, bam_file_2):  # pylint: disable=no-self-use
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file_1 : str
            Location of the bam file to merge into
        bam_file_2 : str
            Location of the bam file that is to get merged into bam_file_1
        """
        bam_handle = bamUtils()
        return bam_handle.bam_merge(bam_file_1, bam_file_2)

    @task(bam_in=FILE_IN, bam_out=FILE_OUT)
    def bam_copy(self, bam_in, bam_out):  # pylint: disable=no-self-use
        """
        Wrapper function to copy from one bam file to another

        Parameters
        ----------
        bam_in : str
            Location of the input bam file
        bam_out : str
            Location of the output bam file
        """
        bam_handle = bamUtils()
        return bam_handle.bam_copy(bam_in, bam_out)

    @task(bam_file=FILE_IN, bam_idx_file=FILE_OUT)
    def bam_index(self, bam_file, bam_idx_file):  # pylint: disable=no-self-use
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file : str
            Location of the bam file that is to be indexed
        bam_idx_file : str
            Location of the bam index file (.bai)
        """
        bam_handle = bamUtils()
        return bam_handle.bam_index(bam_file, bam_idx_file)

    @task(bam_file=FILE_IN)
    def check_header(self, bam_file):  # pylint: disable=no-self-use
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file_1 : str
            Location of the bam file to merge into
        bam_file_2 : str
            Location of the bam file that is to get merged into bam_file_1
        """
        bam_handle = bamUtils()
        return bam_handle.check_header(bam_file)
<commit_msg>Update to the bam_merge so that it is able to handle a list of bam files or a set of strings<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 sys
import pysam

from utils import logger

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

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

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

class bamUtils(object):
    """
    Tool for aligning sequence reads to a genome using BWA
    """

    def __init__(self):
        """
        Initialise the tool with its configuration.

        Parameters
        ----------
        configuration : dict
            a dictionary containing parameters that define how the operation
            should be carried out, which are specific to each Tool.
        """
        logger.info("BAM Utils")

    @staticmethod
    def bam_sort(bam_file):
        """
        Wrapper for the pysam SAMtools sort function

        Parameters
        ----------
        bam_file : str
            Location of the bam file to sort
        """
        try:
            pysam.sort("-o", bam_file, "-T", bam_file + "_sort", bam_file)
        except IOError:
            return False
        return True

    @staticmethod
    def bam_merge(*args):
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file_1 : str
            Location of the bam file to merge into
        bam_file_2 : str
            Location of the bam file that is to get merged into bam_file_1
        """
        logger.info("Merging:")

        if isinstance(args[0], list):
            final_bam = args[0][0]
            tmp_bam = final_bam + "_merge.bam"
            pysam.merge("-f", tmp_bam, *args[0])
        else:
            final_bam = args[0]
            tmp_bam = final_bam + "_merge.bam"
            pysam.merge("-f", tmp_bam, *args)

        try:
            with open(tmp_bam, "rb") as f_in:
                with open(final_bam, "wb") as f_out:
                    f_out.write(f_in.read())
        except IOError:
            return False

        os.remove(tmp_bam)

        return True

    @staticmethod
    def bam_copy(bam_in, bam_out):
        """
        Wrapper function to copy from one bam file to another

        Parameters
        ----------
        bam_in : str
            Location of the input bam file
        bam_out : str
            Location of the output bam file
        """
        logger.info("Copying: " + bam_in + " - " + bam_out)

        try:
            with open(bam_in, "rb") as f_in:
                with open(bam_out, "wb") as f_out:
                    f_out.write(f_in.read())
        except IOError:
            return False

        return True

    @staticmethod
    def bam_index(bam_file, bam_idx_file):
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file : str
            Location of the bam file that is to be indexed
        bam_idx_file : str
            Location of the bam index file (.bai)
        """
        pysam.index(bam_file, bam_file + "_tmp.bai")

        try:
            with open(bam_file + "_tmp.bai", "rb") as f_in:
                with open(bam_idx_file, "wb") as f_out:
                    f_out.write(f_in.read())
        except IOError:
            return False

        return True

    @staticmethod
    def check_header(bam_file):
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file_1 : str
            Location of the bam file to merge into
        bam_file_2 : str
            Location of the bam file that is to get merged into bam_file_1
        """
        output = True

        bam_file_handle = pysam.AlignmentFile(bam_file, "rb")
        if ("SO" not in bam_file_handle.header["HD"] or
                bam_file_handle.header["HD"]["SO"] == "unsorted"):
            output = False
        bam_file_handle.close()

        return output

class bamUtilsTask(object):
    """
    Wrappers so that the function above can be used as part of a @task within
    COMPSs avoiding the files being copied around the infrstructure too many
    times
    """

    def __init__(self):
        """
        Init function
        """
        logger.info("BAM @task Utils")

    @task(bam_file=FILE_INOUT)
    def bam_sort(self, bam_file):  # pylint: disable=no-self-use
        """
        Wrapper for the pysam SAMtools sort function

        Parameters
        ----------
        bam_file : str
            Location of the bam file to sort
        """
        bam_handle = bamUtils()
        return bam_handle.bam_sort(bam_file)

    @task(bam_file_1=FILE_INOUT, bam_file_2=FILE_IN)
    def bam_merge(self, bam_file_1, bam_file_2):  # pylint: disable=no-self-use
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file_1 : str
            Location of the bam file to merge into
        bam_file_2 : str
            Location of the bam file that is to get merged into bam_file_1
        """
        bam_handle = bamUtils()
        return bam_handle.bam_merge(bam_file_1, bam_file_2)

    @task(bam_in=FILE_IN, bam_out=FILE_OUT)
    def bam_copy(self, bam_in, bam_out):  # pylint: disable=no-self-use
        """
        Wrapper function to copy from one bam file to another

        Parameters
        ----------
        bam_in : str
            Location of the input bam file
        bam_out : str
            Location of the output bam file
        """
        bam_handle = bamUtils()
        return bam_handle.bam_copy(bam_in, bam_out)

    @task(bam_file=FILE_IN, bam_idx_file=FILE_OUT)
    def bam_index(self, bam_file, bam_idx_file):  # pylint: disable=no-self-use
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file : str
            Location of the bam file that is to be indexed
        bam_idx_file : str
            Location of the bam index file (.bai)
        """
        bam_handle = bamUtils()
        return bam_handle.bam_index(bam_file, bam_idx_file)

    @task(bam_file=FILE_IN)
    def check_header(self, bam_file):  # pylint: disable=no-self-use
        """
        Wrapper for the pysam SAMtools merge function

        Parameters
        ----------
        bam_file_1 : str
            Location of the bam file to merge into
        bam_file_2 : str
            Location of the bam file that is to get merged into bam_file_1
        """
        bam_handle = bamUtils()
        return bam_handle.check_header(bam_file)
<|endoftext|>


<commit_before>#------------------------------------------------------------------------------
#
#  Copyright (c) 2005, Enthought, Inc.
#  All rights reserved.
#
#  This software is provided without warranty under the terms of the BSD
#  license included in enthought/LICENSE.txt and may be redistributed only
#  under the conditions described in the aforementioned license.  The license
#  is also available online at http://www.enthought.com/licenses/BSD.txt
#  Thanks for using Enthought open source!
#
#  Author: David C. Morrill
#  Date:   06/21/2002
#
#  Refactored into a separate module: 07/04/2003
#
#------------------------------------------------------------------------------

""" Defines common, low-level capabilities needed by the Traits package.
"""

#-------------------------------------------------------------------------------
#  Imports:
#-------------------------------------------------------------------------------

import os
import sys

from os \
    import getcwd, makedirs

from os.path \
    import isdir, dirname, exists, join

from string \
    import lowercase, uppercase

from types \
    import ListType, TupleType, DictType, StringType, UnicodeType, IntType, \
           LongType, FloatType, ComplexType, ClassType, TypeType

try:
    from enthought.ets.api import ETS
except:
    # If the ETS package is not available, fake it:
    class ETS ( object ):

        def _get_application_data ( self ):
            """ Initializes the (default) application data directory. """

            environment_variable = 'HOME'
            directory_name       = '.enthought'
            if sys.platform == 'win32':
                environment_variable = 'APPDATA'
                directory_name       = 'Enthought'

            # Lookup the environment variable:
            parent_directory = os.environ.get( environment_variable, None )
            if parent_directory is None:
                raise ValueError(
                    'Environment variable "%s" not set' % environment_variable )

            application_data = join( parent_directory, directory_name )

            # If a file already exists with this name then make sure that it is
            # a directory!
            if exists( application_data ):
                if not isdir( application_data ):
                    raise ValueError( 'File "%s" already exists' %
                                      application_data )

            # Otherwise, create the directory:
            else:
                makedirs( application_data )

            return application_data

        application_data = property( _get_application_data )

    ETS = ETS()

#-------------------------------------------------------------------------------
#  Provide Python 2.3+ compatible definitions (if necessary):
#-------------------------------------------------------------------------------

try:
    from types import BooleanType
except ImportError:
    BooleanType = IntType

def _enumerate ( seq ):
    for i in xrange( len( seq) ):
        yield i, seq[i]
try:
    enumerate = enumerate
except:
    enumerate = _enumerate
del _enumerate

#-------------------------------------------------------------------------------
#  Constants:
#-------------------------------------------------------------------------------

ClassTypes    = ( ClassType, TypeType )

SequenceTypes = ( ListType, TupleType )

ComplexTypes  = ( float, int )

TypeTypes     = ( StringType,  UnicodeType, IntType,   LongType, FloatType,
                  ComplexType, ListType,    TupleType, DictType, BooleanType )

TraitNotifier = '__trait_notifier__'

#-------------------------------------------------------------------------------
#  Singleton 'Undefined' object (used as undefined trait name and/or value):
#-------------------------------------------------------------------------------

class _Undefined ( object ):

   def __repr__ ( self ):
       return '<undefined>'

# Singleton object that indicates that a trait attribute has not yet had a
# value set (i.e., its value is undefined. This object is used instead of
# None, because None is often has other meanings, such as that a value
# is not used. When a trait attribute is first assigned a value, and its
# associated trait notification handlers are called, Undefined is passed
# as the *old* parameter, to indicate that the attribute previously had no
# value.
Undefined = _Undefined()

# Tell the C-base code about the 'Undefined' objects:
import ctraits
ctraits._undefined( Undefined )

#-------------------------------------------------------------------------------
#  Singleton 'Missing' object (used as missing method argument marker):
#-------------------------------------------------------------------------------

class _Missing ( object ):

   def __repr__ ( self ):
       return '<missing>'

# Singleton object that indicates that a method argument is missing from a
# type-checked method signature.
Missing = _Missing()

#-------------------------------------------------------------------------------
#  Singleton 'Self' object (used as object reference to current 'object'):
#-------------------------------------------------------------------------------

class _Self ( object ):

   def __repr__ ( self ):
       return '<self>'

# Singleton object that references the current 'object'.
Self = _Self()

#-------------------------------------------------------------------------------
#  Define a special 'string' coercion function:
#-------------------------------------------------------------------------------

def strx ( arg ):
    """ Wraps the built-in str() function to raise a TypeError if the
    argument is not of a type in StringTypes.
    """
    if type( arg ) in StringTypes:
       return str( arg )
    raise TypeError

#-------------------------------------------------------------------------------
#  Constants:
#-------------------------------------------------------------------------------

StringTypes = ( StringType, UnicodeType, IntType, LongType, FloatType,
                ComplexType )

#-------------------------------------------------------------------------------
#  Define a mapping of coercable types:
#-------------------------------------------------------------------------------

# Mapping of coercable types.
CoercableTypes = {
    LongType:    ( 11, long, int ),
    FloatType:   ( 11, float, int ),
    ComplexType: ( 11, complex, float, int ),
    UnicodeType: ( 11, unicode, str )
}

#-------------------------------------------------------------------------------
#  Return a string containing the class name of an object with the correct
#  article (a or an) preceding it (e.g. 'an Image', 'a PlotValue'):
#-------------------------------------------------------------------------------

def class_of ( object ):
    """ Returns a string containing the class name of an object with the
    correct indefinite article ('a' or 'an') preceding it (e.g., 'an Image',
    'a PlotValue').
    """
    if isinstance( object, basestring ):
       return add_article( object )

    return add_article( object.__class__.__name__ )

#-------------------------------------------------------------------------------
#  Return a string containing the right article (i.e. 'a' or 'an') prefixed to
#  a specified string:
#-------------------------------------------------------------------------------

def add_article ( name ):
    """ Returns a string containing the correct indefinite article ('a' or 'an')
    prefixed to the specified string.
    """
    if name[:1].lower() in 'aeiou':
       return 'an ' + name

    return 'a ' + name

#----------------------------------------------------------------------------
#  Return a 'user-friendly' name for a specified trait:
#----------------------------------------------------------------------------

def user_name_for ( name ):
    """ Returns a "user-friendly" version of a string, with the first letter
    capitalized and with underscore characters replaced by spaces. For example,
    ``user_name_for('user_name_for')`` returns ``'User name for'``.
    """
    name       = name.replace( '_', ' ' ).capitalize()
    result     = ''
    last_lower = 0

    for c in name:
        if (c in uppercase) and last_lower:
           result += ' '
        last_lower = (c in lowercase)
        result    += c

    return result

#-------------------------------------------------------------------------------
#  Gets the path to the traits home directory:
#-------------------------------------------------------------------------------

_traits_home = None

def traits_home ( ):
    """ Gets the path to the Traits home directory.
    """
    global _traits_home

    if _traits_home is None:
        _traits_home = _verify_path( join( ETS.application_data, 'traits' ) )

    return _traits_home

#-------------------------------------------------------------------------------
#  Verify that a specified path exists, and try to create it if it doesn't:
#-------------------------------------------------------------------------------

def _verify_path ( path ):
    """ Verify that a specified path exists, and try to create it if it
        does not exist.
    """
    if not exists( path ):
        try:
            os.mkdir( path )
        except:
            pass

    return path

#-------------------------------------------------------------------------------
#  Returns the name of the module the caller's caller is located in:
#-------------------------------------------------------------------------------

def get_module_name ( level = 2 ):
    """ Returns the name of the module that the caller's caller is located in.
    """
    return sys._getframe( level ).f_globals.get( '__name__', '__main__' )

#-------------------------------------------------------------------------------
#  Returns a resource path calculated from the caller's stack:
#-------------------------------------------------------------------------------

def get_resource_path ( level = 2 ):
    """Returns a resource path calculated from the caller's stack.
    """
    module = sys._getframe( level ).f_globals.get( '__name__', '__main__' )

    if module != '__main__':
        # Return the path to the module:
        try:
            return dirname( getattr( sys.modules.get( module ), '__file__' ) )
        except:
            # Apparently 'module' is not a registered module...treat it like
            # '__main__':
            pass

    # '__main__' is not a real module, so we need a work around:
    for path in [ dirname( sys.argv[0] ), getcwd() ]:
        if exists( path ):
            break

    return path

<commit_msg>changed ETS to app_data_locator, the package name changed a couple of weeks ago but this code had not been updated<commit_after>#------------------------------------------------------------------------------
#
#  Copyright (c) 2005, Enthought, Inc.
#  All rights reserved.
#
#  This software is provided without warranty under the terms of the BSD
#  license included in enthought/LICENSE.txt and may be redistributed only
#  under the conditions described in the aforementioned license.  The license
#  is also available online at http://www.enthought.com/licenses/BSD.txt
#  Thanks for using Enthought open source!
#
#  Author: David C. Morrill
#  Date:   06/21/2002
#
#  Refactored into a separate module: 07/04/2003
#
#------------------------------------------------------------------------------

""" Defines common, low-level capabilities needed by the Traits package.
"""

#-------------------------------------------------------------------------------
#  Imports:
#-------------------------------------------------------------------------------

import os
import sys

from os \
    import getcwd, makedirs

from os.path \
    import isdir, dirname, exists, join

from string \
    import lowercase, uppercase

from types \
    import ListType, TupleType, DictType, StringType, UnicodeType, IntType, \
           LongType, FloatType, ComplexType, ClassType, TypeType

try:
    from enthought.app_data_locator.api import AppDataLocator
except:
    # If the AppDataLocator package is not available, fake it:
    class AppDataLocator ( object ):

        def _get_application_data ( self ):
            """ Initializes the (default) application data directory. """

            environment_variable = 'HOME'
            directory_name       = '.enthought'
            if sys.platform == 'win32':
                environment_variable = 'APPDATA'
                directory_name       = 'Enthought'

            # Lookup the environment variable:
            parent_directory = os.environ.get( environment_variable, None )
            if parent_directory is None:
                raise ValueError(
                    'Environment variable "%s" not set' % environment_variable )

            application_data = join( parent_directory, directory_name )

            # If a file already exists with this name then make sure that it is
            # a directory!
            if exists( application_data ):
                if not isdir( application_data ):
                    raise ValueError( 'File "%s" already exists' %
                                      application_data )

            # Otherwise, create the directory:
            else:
                makedirs( application_data )

            return application_data

        application_data = property( _get_application_data )

    AppDataLocator = AppDataLocator()

#-------------------------------------------------------------------------------
#  Provide Python 2.3+ compatible definitions (if necessary):
#-------------------------------------------------------------------------------

try:
    from types import BooleanType
except ImportError:
    BooleanType = IntType

def _enumerate ( seq ):
    for i in xrange( len( seq) ):
        yield i, seq[i]
try:
    enumerate = enumerate
except:
    enumerate = _enumerate
del _enumerate

#-------------------------------------------------------------------------------
#  Constants:
#-------------------------------------------------------------------------------

ClassTypes    = ( ClassType, TypeType )

SequenceTypes = ( ListType, TupleType )

ComplexTypes  = ( float, int )

TypeTypes     = ( StringType,  UnicodeType, IntType,   LongType, FloatType,
                  ComplexType, ListType,    TupleType, DictType, BooleanType )

TraitNotifier = '__trait_notifier__'

#-------------------------------------------------------------------------------
#  Singleton 'Undefined' object (used as undefined trait name and/or value):
#-------------------------------------------------------------------------------

class _Undefined ( object ):

   def __repr__ ( self ):
       return '<undefined>'

# Singleton object that indicates that a trait attribute has not yet had a
# value set (i.e., its value is undefined. This object is used instead of
# None, because None is often has other meanings, such as that a value
# is not used. When a trait attribute is first assigned a value, and its
# associated trait notification handlers are called, Undefined is passed
# as the *old* parameter, to indicate that the attribute previously had no
# value.
Undefined = _Undefined()

# Tell the C-base code about the 'Undefined' objects:
import ctraits
ctraits._undefined( Undefined )

#-------------------------------------------------------------------------------
#  Singleton 'Missing' object (used as missing method argument marker):
#-------------------------------------------------------------------------------

class _Missing ( object ):

   def __repr__ ( self ):
       return '<missing>'

# Singleton object that indicates that a method argument is missing from a
# type-checked method signature.
Missing = _Missing()

#-------------------------------------------------------------------------------
#  Singleton 'Self' object (used as object reference to current 'object'):
#-------------------------------------------------------------------------------

class _Self ( object ):

   def __repr__ ( self ):
       return '<self>'

# Singleton object that references the current 'object'.
Self = _Self()

#-------------------------------------------------------------------------------
#  Define a special 'string' coercion function:
#-------------------------------------------------------------------------------

def strx ( arg ):
    """ Wraps the built-in str() function to raise a TypeError if the
    argument is not of a type in StringTypes.
    """
    if type( arg ) in StringTypes:
       return str( arg )
    raise TypeError

#-------------------------------------------------------------------------------
#  Constants:
#-------------------------------------------------------------------------------

StringTypes = ( StringType, UnicodeType, IntType, LongType, FloatType,
                ComplexType )

#-------------------------------------------------------------------------------
#  Define a mapping of coercable types:
#-------------------------------------------------------------------------------

# Mapping of coercable types.
CoercableTypes = {
    LongType:    ( 11, long, int ),
    FloatType:   ( 11, float, int ),
    ComplexType: ( 11, complex, float, int ),
    UnicodeType: ( 11, unicode, str )
}

#-------------------------------------------------------------------------------
#  Return a string containing the class name of an object with the correct
#  article (a or an) preceding it (e.g. 'an Image', 'a PlotValue'):
#-------------------------------------------------------------------------------

def class_of ( object ):
    """ Returns a string containing the class name of an object with the
    correct indefinite article ('a' or 'an') preceding it (e.g., 'an Image',
    'a PlotValue').
    """
    if isinstance( object, basestring ):
       return add_article( object )

    return add_article( object.__class__.__name__ )

#-------------------------------------------------------------------------------
#  Return a string containing the right article (i.e. 'a' or 'an') prefixed to
#  a specified string:
#-------------------------------------------------------------------------------

def add_article ( name ):
    """ Returns a string containing the correct indefinite article ('a' or 'an')
    prefixed to the specified string.
    """
    if name[:1].lower() in 'aeiou':
       return 'an ' + name

    return 'a ' + name

#----------------------------------------------------------------------------
#  Return a 'user-friendly' name for a specified trait:
#----------------------------------------------------------------------------

def user_name_for ( name ):
    """ Returns a "user-friendly" version of a string, with the first letter
    capitalized and with underscore characters replaced by spaces. For example,
    ``user_name_for('user_name_for')`` returns ``'User name for'``.
    """
    name       = name.replace( '_', ' ' ).capitalize()
    result     = ''
    last_lower = 0

    for c in name:
        if (c in uppercase) and last_lower:
           result += ' '
        last_lower = (c in lowercase)
        result    += c

    return result

#-------------------------------------------------------------------------------
#  Gets the path to the traits home directory:
#-------------------------------------------------------------------------------

_traits_home = None

def traits_home ( ):
    """ Gets the path to the Traits home directory.
    """
    global _traits_home

    if _traits_home is None:
        _traits_home = _verify_path( join( AppDataLocator.application_data, 'traits' ) )

    return _traits_home

#-------------------------------------------------------------------------------
#  Verify that a specified path exists, and try to create it if it doesn't:
#-------------------------------------------------------------------------------

def _verify_path ( path ):
    """ Verify that a specified path exists, and try to create it if it
        does not exist.
    """
    if not exists( path ):
        try:
            os.mkdir( path )
        except:
            pass

    return path

#-------------------------------------------------------------------------------
#  Returns the name of the module the caller's caller is located in:
#-------------------------------------------------------------------------------

def get_module_name ( level = 2 ):
    """ Returns the name of the module that the caller's caller is located in.
    """
    return sys._getframe( level ).f_globals.get( '__name__', '__main__' )

#-------------------------------------------------------------------------------
#  Returns a resource path calculated from the caller's stack:
#-------------------------------------------------------------------------------

def get_resource_path ( level = 2 ):
    """Returns a resource path calculated from the caller's stack.
    """
    module = sys._getframe( level ).f_globals.get( '__name__', '__main__' )

    if module != '__main__':
        # Return the path to the module:
        try:
            return dirname( getattr( sys.modules.get( module ), '__file__' ) )
        except:
            # Apparently 'module' is not a registered module...treat it like
            # '__main__':
            pass

    # '__main__' is not a real module, so we need a work around:
    for path in [ dirname( sys.argv[0] ), getcwd() ]:
        if exists( path ):
            break

    return path

<|endoftext|>


<commit_before># Copyright (c) 2014 - 2017, Carlos Millett
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

#
# Import
#
import os
import re

from shutil import move
from itertools import zip_longest

#
# Class
#
class LocalPath():
    def __init__(self, filename):
        self.dirName = os.path.dirname(filename)
        self.curFileName = os.path.basename(filename)
        self.fileNameExt = os.path.splitext(filename)[-1]
        self._newFileName = None

    def rename(self):
        cur = os.path.join(self.dirName, self.curFileName)
        new = os.path.join(self.dirName, self._newFileName)

        if cur == new:
            raise SameFileError("{} and {} are the same file.".format(cur, new))

        move(cur, new)

    @property
    def newFileName(self):
         return self._newFileName

    @newFileName.setter
    def newFileName(self, newName):
        new = self._sanitizeName(newName)
        self._newFileName = "{}{}".format(new, self.fileNameExt)

    def _sanitizeName(self, name):
        table = {
                 ord("á"): "a",
                 ord("à"): "a",
                 ord("ã"): "a",
                 ord("â"): "a",
                 ord("é"): "e",
                 ord("è"): "e",
                 ord("ẽ"): "e",
                 ord("ê"): "e",
                 ord("í"): "i",
                 ord("ì"): "i",
                 ord("ĩ"): "i",
                 ord("î"): "i",
                 ord("ó"): "o",
                 ord("ò"): "o",
                 ord("õ"): "o",
                 ord("ô"): "o",
                 ord("ú"): "u",
                 ord("ù"): "u",
                 ord("ũ"): "u",
                 ord("û"): "u",
                 ord("ç"): "c",
                 ord("ñ"): "n",
                 ord(":"): " -",
                 ord(">"): None,
                 ord("<"): None,
                 ord("?"): None,
                 ord("!"): None,
                 ord("*"): None,
                 ord("/"): None,
                 ord("\\"): None,
                 ord("\""): None,
                 ord("\'"): None
                }

        return name.translate(table)

    def _formatName(self, name):
        sep = re.compile("[\W]+")
        year = re.compile("[12][0-9]{3}")
        country = re.compile("[A-Z]{2}")

        formated = year.sub("", country.sub("", name))
        formated = sep.sub(" ", formated).strip()

        return formated.title()

class SerieFile(LocalPath):
    def __init__(self, filename):
        super().__init__(filename)
        regex1 = re.compile("\.S(\d{2})((E\d{2})+)\.")
        regex2 = re.compile("\.(\d{4}.)?(\d{3,})\.")

        if regex1.search(self.curFileName):
            show, season, info = regex1.split(self.curFileName)[:3]
            eps = re.findall("\d{2}", info)

        elif regex2.search(self.curFileName):
            show, _, info = regex2.split(self.curFileName)[:3]

            info_lst = list(info)
            info_lst.reverse()
            info_itr = [ iter(info_lst) ] * 2
            info_grp = zip_longest(*info_itr, fillvalue="0")
            info_lst = list(info_grp)
            s_num, s_dec = info_lst.pop()
            info_lst.reverse()

            season = "{}{}".format(s_dec, s_num)
            eps = [ "{}{}".format(e_dec, e_num) for e_num, e_dec in info_lst ]

        else:
            strerror = "Can't find show pattern for {}".format(self.curFileName)
            raise MatchNotFoundError(strerror)

        self.show = super()._formatName(show)
        self.season = "{:0>2}".format(season)
        self.episodes = eps

class MovieFile(LocalPath):
    def __init__(self, filename):
        super().__init__(filename)
        regex = re.compile("([a-z0-9.-]+).([12][0-9]{3})", re.I)

        if not regex.search(self.curFileName):
            strerror = "Can't find movie pattern for {}".format(self.curFileName)
            raise MatchNotFoundError(strerror)

        splitName = regex.findall(self.curFileName)
        movie, year = splitName[0]

        self.title = super()._formatName(movie)
        self.year = year

#
# Exception
#
class MatchNotFoundError(Exception):
    pass

class SameFileError(Exception):
    pass
<commit_msg>update table of forbidden characters on a filename<commit_after># Copyright (c) 2014 - 2017, Carlos Millett
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

#
# Import
#
import os
import re

from shutil import move
from itertools import zip_longest

#
# Class
#
class LocalPath():
    def __init__(self, filename):
        self.dirName = os.path.dirname(filename)
        self.curFileName = os.path.basename(filename)
        self.fileNameExt = os.path.splitext(filename)[-1]
        self._newFileName = None

    def rename(self):
        cur = os.path.join(self.dirName, self.curFileName)
        new = os.path.join(self.dirName, self._newFileName)

        if cur == new:
            raise SameFileError("{} and {} are the same file.".format(cur, new))

        move(cur, new)

    @property
    def newFileName(self):
         return self._newFileName

    @newFileName.setter
    def newFileName(self, newName):
        new = self._sanitizeName(newName)
        self._newFileName = "{}{}".format(new, self.fileNameExt)

    def _sanitizeName(self, name):
        table = {
                 ord("á"): "a",
                 ord("à"): "a",
                 ord("ã"): "a",
                 ord("â"): "a",
                 ord("é"): "e",
                 ord("è"): "e",
                 ord("ẽ"): "e",
                 ord("ê"): "e",
                 ord("í"): "i",
                 ord("ì"): "i",
                 ord("ĩ"): "i",
                 ord("î"): "i",
                 ord("ó"): "o",
                 ord("ò"): "o",
                 ord("õ"): "o",
                 ord("ô"): "o",
                 ord("ú"): "u",
                 ord("ù"): "u",
                 ord("ũ"): "u",
                 ord("û"): "u",
                 ord("ç"): "c",
                 ord("ñ"): "n",
                 ord(":"): " -",
                 ord(">"): None,
                 ord("<"): None,
                 ord("?"): None,
                 ord("!"): None,
                 ord("*"): None,
                 ord("#"): None,
                 ord("/"): None,
                 ord("\\"): None,
                 ord("\""): None,
                 ord("\'"): None
                }

        return name.translate(table)

    def _formatName(self, name):
        sep = re.compile("[\W]+")
        year = re.compile("[12][0-9]{3}")
        country = re.compile("[A-Z]{2}")

        formated = year.sub("", country.sub("", name))
        formated = sep.sub(" ", formated).strip()

        return formated.title()

class SerieFile(LocalPath):
    def __init__(self, filename):
        super().__init__(filename)
        regex1 = re.compile("\.S(\d{2})((E\d{2})+)\.")
        regex2 = re.compile("\.(\d{4}.)?(\d{3,})\.")

        if regex1.search(self.curFileName):
            show, season, info = regex1.split(self.curFileName)[:3]
            eps = re.findall("\d{2}", info)

        elif regex2.search(self.curFileName):
            show, _, info = regex2.split(self.curFileName)[:3]

            info_lst = list(info)
            info_lst.reverse()
            info_itr = [ iter(info_lst) ] * 2
            info_grp = zip_longest(*info_itr, fillvalue="0")
            info_lst = list(info_grp)
            s_num, s_dec = info_lst.pop()
            info_lst.reverse()

            season = "{}{}".format(s_dec, s_num)
            eps = [ "{}{}".format(e_dec, e_num) for e_num, e_dec in info_lst ]

        else:
            strerror = "Can't find show pattern for {}".format(self.curFileName)
            raise MatchNotFoundError(strerror)

        self.show = super()._formatName(show)
        self.season = "{:0>2}".format(season)
        self.episodes = eps

class MovieFile(LocalPath):
    def __init__(self, filename):
        super().__init__(filename)
        regex = re.compile("([a-z0-9.-]+).([12][0-9]{3})", re.I)

        if not regex.search(self.curFileName):
            strerror = "Can't find movie pattern for {}".format(self.curFileName)
            raise MatchNotFoundError(strerror)

        splitName = regex.findall(self.curFileName)
        movie, year = splitName[0]

        self.title = super()._formatName(movie)
        self.year = year

#
# Exception
#
class MatchNotFoundError(Exception):
    pass

class SameFileError(Exception):
    pass
<|endoftext|>