<commit_before>from django.conf import settings
import os.path
import errno
import logging
import time
import datetime
import shutil
import requests
import zipfile

logger = logging.getLogger("kpcc_backroom_handshakes")

class Retriever(object):
    """
    a series of reusable methods we'll need for downloading and moving files

    if you change something in here you're gonna want to change something in the test_utils_files script as well

    """

    list_of_expected_files = []

    def _successful_save_results(self, item):
        """
        can i take the response from url can and write it to a timestamped version of the a file that should work no matter the file. it's  based on the file_ext specified in a config dict
        """
        try:
            response = requests.get(item.source_url, headers=settings.REQUEST_HEADERS, stream=True)
            with open(item.file_name, "wb") as output:
                output.write(response.content)
            file_exists = os.path.isfile(item.file_name)
            file_has_size = os.path.getsize(item.file_name)
            if file_exists == True and file_has_size > 0:
                logger.debug("Success!")
        except Exception, exception:
            logger.error(exception)
            raise

    # REFACTOR FUNCTION TO ACCOUNT FOR MULTIPLE FILE TYPES
    def _found_files_in_zipfile(self, item):
        """
        compare files in a zipfile with a list of expected files
        """
        try:
            with zipfile.ZipFile(item.file_name) as zip:
                files = zipfile.ZipFile.namelist(zip)
                for item in self.list_of_expected_files:
                    if item.strip() in set(files) == True:
                        logger.debug("Success!")
        except Exception, exception:
            logger.error(exception)
            raise

    def _copy_timestamped_file_as_latest(self, item, data_directory):
        """
        create timestamped version of a file deemed latest
        """
        try:
            item.file_latest = "%s%s_latest%s" % (data_directory, item.source_short, item.source_type)
            shutil.copyfile(item.file_name, item.file_latest)
            file_exists = os.path.isfile(item.file_latest)
            if file_exists == True:
                logger.debug("Success!")
        except Exception, exception:
            logger.error(exception)
            raise

    def _create_directory_for_latest_file(self, item, data_directory):
        """
        move latest files to a working directory
        """
        working = "%s%s_latest" % (data_directory, item.source_short)
        try:
            os.makedirs(working)
            logger.debug("Success!")
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise

    def _move_latest_files_to_latest_directory(self, item, data_directory):
        """
        move timestamped file to working directory as latest
        """
        working = "%s%s_latest" % (data_directory, item.source_short)
        latest = os.path.join(working, os.path.basename(item.file_latest))
        shutil.copy(item.file_latest, working)
        os.remove(item.file_latest)
        latest_exists = os.path.isfile(latest)
        working_exists = os.path.isfile(working)
        if latest_exists == True and working_exists == False:
            logger.debug("Success!")

    def _archive_downloaded_file(self, item, data_directory):
        """
        move timestamped zipfile to archives
        """
        archives = "%s_archived_files" % (data_directory)
        try:
            os.makedirs(archives)
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise
        shutil.move(item.file_name, archives)
        file_exists = os.path.isfile(item.file_name)
        if file_exists == False:
            logger.debug("Success!")

    def _unzip_latest_file(self, item, data_directory):
        """
        if the item is a zipfile can I extract the files?
        """
        if item.source_type == ".zip":
            working = "%s%s_latest" % (data_directory, item.source_short)
            file_latest = os.path.join(working, os.path.basename(item.file_latest))
            with zipfile.ZipFile(file_latest) as zip:
                if zipfile.ZipFile.testzip(zip) == None:
                    zip.extract("X14GG510v7.xml", working)
                    zip.extract("X14GG530v7.xml", working)
                os.remove(file_latest)
                file_exists = os.path.isfile(file_latest)
                if file_exists == False:
                    logger.debug("Success!")
<commit_msg>updates to the base file_utils<commit_after>from django.conf import settings
import os.path
import errno
import logging
import time
import datetime
import shutil
import requests
from requests.packages.urllib3.util.retry import Retry
from requests.adapters import HTTPAdapter
import zipfile

logger = logging.getLogger("kpcc_backroom_handshakes")

class Retriever(object):
    """
    a series of reusable methods we'll need for downloading and moving files

    if you change something in here you're gonna want to change something in the test_utils_files script as well

    """

    list_of_expected_files = []

    def _successful_save_results(self, item):
        """
        can i take the response from url can and write it to a timestamped version of the a file that should work no matter the file. it's  based on the file_ext specified in a config dict
        """
        session = requests.Session()
        retries = Retry(
            total=5,
            backoff_factor=1,
            status_forcelist=[500, 502, 503, 504]
        )
        session.mount("http://", HTTPAdapter(max_retries=retries))
        response = session.get(
            item.source_url,
            headers=settings.REQUEST_HEADERS,
            timeout=10,
            allow_redirects=False
        )
        try:
            response.raise_for_status()
            logger.debug("%s: success" % (item.source_url))
        except requests.exceptions.ConnectionError as exception:
            # incorrect domain
            logger.error("%s: %s" % (exception, item.source_url))
            raise
        except requests.exceptions.Timeout as exception:
            # maybe set up for a retry, or continue in a retry loop
            logger.error("%s: %s" % (exception, item.source_url))
            raise
        except requests.exceptions.TooManyRedirects as exception:
            # tell the user their url was bad and try a different one
            logger.error("%s: %s" % (exception, item.source_url))
            raise
        except requests.exceptions.RequestException as exception:
            # catastrophic error and bail
            logger.error("%s: %s" % (exception, item.source_url))
            sys.exit(1)
        with open(item.file_name, "wb") as output:
            output.write(response.content)
        file_exists = os.path.isfile(item.file_name)
        file_has_size = os.path.getsize(item.file_name)
        if file_exists == True:
            if file_has_size > 0:
                logger.debug("Success!")
            else:
                logger.error("Your file has zero data")
                raise Exception
        else:
            logger.error("Your file doesn't exist")
            raise Exception

    def _copy_timestamped_file_as_latest(self, item, data_directory):
        """
        create timestamped version of a file deemed latest
        """
        try:
            item.file_latest = "%s%s_latest%s" % (data_directory, item.source_short, item.source_type)
            shutil.copyfile(item.file_name, item.file_latest)
            file_exists = os.path.isfile(item.file_latest)
            if file_exists == True:
                logger.debug("Success!")
        except Exception, exception:
            logger.error(exception)
            raise

    def _create_directory_for_latest_file(self, item, data_directory):
        """
        move latest files to a working directory
        """
        working = "%s%s_latest" % (data_directory, item.source_short)
        try:
            os.makedirs(working)
            logger.debug("Success!")
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise

    def _move_latest_files_to_latest_directory(self, item, data_directory):
        """
        move timestamped file to working directory as latest
        """
        working = "%s%s_latest" % (data_directory, item.source_short)
        latest = os.path.join(working, os.path.basename(item.file_latest))
        shutil.copy(item.file_latest, working)
        os.remove(item.file_latest)
        latest_exists = os.path.isfile(latest)
        working_exists = os.path.isfile(working)
        if latest_exists == True and working_exists == False:
            logger.debug("Success!")

    def _archive_downloaded_file(self, item, data_directory):
        """
        move timestamped zipfile to archives
        """
        archives = "%s_archived_files" % (data_directory)
        try:
            os.makedirs(archives)
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise
        shutil.move(item.file_name, archives)
        file_exists = os.path.isfile(item.file_name)
        if file_exists == False:
            logger.debug("Success!")

    def _found_required_files(self, item, data_directory):
        """
        compare files in a zipfile with a list of expected files
        """
        working = "%s%s_latest" % (data_directory, item.source_short)
        latest = os.path.join(working, os.path.basename(item.file_latest))
        if item.source_type == ".zip":
            try:
                with zipfile.ZipFile(latest) as zip:
                    files = zipfile.ZipFile.namelist(zip)
                    for file in self.list_of_expected_files:
                        self.assertEquals(file.strip() in set(files), True)
                        logger.debug("Success: %s exists" % (file.strip()))
            except Exception, exception:
                logger.error(exception)
        else:
            try:
                for file in self.list_of_expected_files:
                    self.assertEquals(file.strip(), os.path.basename(item.file_latest))
                    logger.debug("Success: %s exists" % (file.strip()))
            except Exception, exception:
                logger.error(exception)
                raise

    def _unzip_latest_file(self, item, data_directory):
        """
        if the item is a zipfile can I extract the files?
        """
        if item.source_type == ".zip":
            working = "%s%s_latest" % (data_directory, item.source_short)
            file_latest = os.path.join(working, os.path.basename(item.file_latest))
            with zipfile.ZipFile(file_latest) as zip:
                if zipfile.ZipFile.testzip(zip) == None:
                    for file in self.list_of_expected_files:
                        file = file.strip()
                        zip.extract(file, working)
                        file_exists = os.path.isfile(os.path.join(working, file))
                        if file_exists == True:
                            logger.debug("Success!")
                os.remove(file_latest)
                file_exists = os.path.isfile(file_latest)
                if file_exists == False:
                    logger.debug("Success!")
<|endoftext|>


<commit_before>#!/usr/bin/env python3

from bisect import bisect_left
import logging

from .base import SubprocessWidget, protect_handler
from barython.tools import splitted_sleep
from barython.hooks.audio import PulseAudioHook

logger = logging.getLogger("barython")

class PulseAudioWidget(SubprocessWidget):
    """
    Show the current volume

    Requires pamixer to work
    """
    _icon = None
    _volume = 0
    _input_mute = False
    _output_mute = False

    @property
    def icon(self):
        # In case the icon is static
        if isinstance(self._icon, str) or self._icon is None:
            return self._icon
        elif self._output_mute and "ouput_mute" in self._icon:
            return self._icon["ouput_mute"]
        elif "volume" in self._icon:
            volume_icons = sorted(self._icon["volume"], key=lambda k: k[0])
            keys = [i[0] for i in volume_icons]
            return volume_icons[bisect_left(keys, self._volume, lo=1) - 1][1]

    @protect_handler
    def handler(self, event, *args, **kwargs):
        """
        Filter events sent by notifications
        """
        # Only notify if there is something changes in pulseaudio
        event_change_msg = "Event 'change' on destination"
        if event_change_msg in event:
            logger.debug("PA: line \"{}\" catched.".format(event))
            with self._lock_update:
                self.update()
                splitted_sleep(self.refresh, stop=self._stop.is_set)

    def organize_result(self, output, *args, **kwargs):
        """
        Override this method to change the infos to print
        """
        volume, output_mute = output.split()
        self._volume = int(volume)
        self._output_mute = output_mute == "true"
        if self.icon:
            return (
                "{} {}".format(self.icon, self._volume)
                if not self._output_mute else "{}".format(self.icon)
            )
        else:
            return "{}".format(self._volume)

    def __init__(self,
                 cmd=["echo $(pamixer --get-volume) $(pamixer --get-mute)"],
                 shell=True, *args, **kwargs):
        super().__init__(*args, **kwargs, cmd=cmd, infinite=False, shell=shell)

        # Update the widget when PA volume changes
        self.hooks.subscribe(self.handler, PulseAudioHook)
<commit_msg>Change pactl expected output<commit_after>#!/usr/bin/env python3

from bisect import bisect_left
import logging

from .base import SubprocessWidget, protect_handler
from barython.tools import splitted_sleep
from barython.hooks.audio import PulseAudioHook

logger = logging.getLogger("barython")

class PulseAudioWidget(SubprocessWidget):
    """
    Show the current volume

    Requires pamixer to work
    """
    _icon = None
    _volume = 0
    _input_mute = False
    _output_mute = False

    @property
    def icon(self):
        # In case the icon is static
        if isinstance(self._icon, str) or self._icon is None:
            return self._icon
        elif self._output_mute and "ouput_mute" in self._icon:
            return self._icon["ouput_mute"]
        elif "volume" in self._icon:
            volume_icons = sorted(self._icon["volume"], key=lambda k: k[0])
            keys = [i[0] for i in volume_icons]
            return volume_icons[bisect_left(keys, self._volume, lo=1) - 1][1]

    @protect_handler
    def handler(self, event, *args, **kwargs):
        """
        Filter events sent by notifications
        """
        # Only notify if there is something changes in pulseaudio
        event_change_msg = "Event 'change' on sink"
        if event_change_msg in event:
            logger.debug("PA: line \"{}\" catched.".format(event))
            with self._lock_update:
                self.update()
                splitted_sleep(self.refresh, stop=self._stop.is_set)

    def organize_result(self, output, *args, **kwargs):
        """
        Override this method to change the infos to print
        """
        volume, output_mute = output.split()
        self._volume = int(volume)
        self._output_mute = output_mute == "true"
        if self.icon:
            return (
                "{} {}".format(self.icon, self._volume)
                if not self._output_mute else "{}".format(self.icon)
            )
        else:
            return "{}".format(self._volume)

    def __init__(self,
                 cmd=["echo $(pamixer --get-volume) $(pamixer --get-mute)"],
                 shell=True, *args, **kwargs):
        super().__init__(*args, **kwargs, cmd=cmd, infinite=False, shell=shell)

        # Update the widget when PA volume changes
        self.hooks.subscribe(self.handler, PulseAudioHook)
<|endoftext|>


<commit_before># This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

import inspect
from flask import g
from flask import json
import relengapi.app
from nose.tools import make_decorator

class TestContext(object):

    _known_options = set([
        'databases',
        'reuse_app',
        'app_setup',
        'db_setup',
        'db_teardown',
        'config',
        'actions',
    ])

    def __init__(self, **options):
        self._validate(options)
        self.options = options
        self._app = None

    def specialize(self, **options):
        new_options = self.options.copy()
        new_options.update(options)
        return TestContext(**new_options)

    def _validate(self, options):
        unknown = set(options) - self._known_options
        if unknown:
            raise ValueError("unknown options %s" % (', '.join(unknown)))
        # verify that all actions exist
        if options.get('actions'):
            for action in options['actions']:
                assert action.exists()

    def _make_app(self):
        if self.options.get('reuse_app') and self._app:
            return self._app
        config = self.options.get('config', {}).copy()
        config['TESTING'] = True
        config['SECRET_KEY'] = 'test'
        config['SQLALCHEMY_DATABASE_URIS'] = uris = {}
        dbnames = self.options.get('databases', [])
        for dbname in dbnames:
            uris[dbname] = 'sqlite://'
        app = relengapi.app.create_app(test_config=config)

        # set up actions
        if self.options.get('actions') is not None:
            @app.before_request
            def set_actions():
                g.identity.provides.update(set(self.options['actions']))

        # set up the requested DBs
        for dbname in dbnames:
            meta = app.db.metadata[dbname]
            engine = app.db.engine(dbname)
            meta.create_all(bind=engine)
        self._app = app
        if 'app_setup' in self.options:
            self.options['app_setup'](app)
        return app

    def __call__(self, func):
        arginfo = inspect.getargspec(func)
        args = set((arginfo.args if arginfo.args else []) +
                   (arginfo.keywords if arginfo.keywords else []))

        @make_decorator(func)
        def wrap(**kwargs):
            def post_json(path, data):
                return kwargs['client'].post(
                    path, data=json.dumps(data),
                    headers=[('Content-Type', 'application/json')])
            app = self._make_app()
            if 'app' in args:
                kwargs['app'] = app
            if 'client' in args:
                kwargs['client'] = app.test_client()
                kwargs['client'].post_json = post_json
            if 'db_setup' in self.options:
                self.options['db_setup'](app)
            try:
                func(**kwargs)
            finally:
                if 'db_teardown' in self.options:
                    self.options['db_teardown'](app)
        return wrap
<commit_msg>use wrapt for the TestContext decorator<commit_after># This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

import wrapt
import inspect
from flask import g
from flask import json
import relengapi.app

class TestContext(object):

    _known_options = set([
        'databases',
        'reuse_app',
        'app_setup',
        'db_setup',
        'db_teardown',
        'config',
        'actions',
    ])

    def __init__(self, **options):
        self._validate(options)
        self.options = options
        self._app = None

    def specialize(self, **options):
        new_options = self.options.copy()
        new_options.update(options)
        return TestContext(**new_options)

    def _validate(self, options):
        unknown = set(options) - self._known_options
        if unknown:
            raise ValueError("unknown options %s" % (', '.join(unknown)))
        # verify that all actions exist
        if options.get('actions'):
            for action in options['actions']:
                assert action.exists()

    def _make_app(self):
        if self.options.get('reuse_app') and self._app:
            return self._app
        config = self.options.get('config', {}).copy()
        config['TESTING'] = True
        config['SECRET_KEY'] = 'test'
        config['SQLALCHEMY_DATABASE_URIS'] = uris = {}
        dbnames = self.options.get('databases', [])
        for dbname in dbnames:
            uris[dbname] = 'sqlite://'
        app = relengapi.app.create_app(test_config=config)

        # set up actions
        if self.options.get('actions') is not None:
            @app.before_request
            def set_actions():
                g.identity.provides.update(set(self.options['actions']))

        # set up the requested DBs
        for dbname in dbnames:
            meta = app.db.metadata[dbname]
            engine = app.db.engine(dbname)
            meta.create_all(bind=engine)
        self._app = app
        if 'app_setup' in self.options:
            self.options['app_setup'](app)
        return app

    @wrapt.decorator
    def __call__(self, wrapped, instance, _, kwargs):
        arginfo = inspect.getargspec(wrapped)
        args = set((arginfo.args if arginfo.args else []) +
                   (arginfo.keywords if arginfo.keywords else []))

        def post_json(path, data):
            return kwargs['client'].post(
                path, data=json.dumps(data),
                headers=[('Content-Type', 'application/json')])
        app = self._make_app()
        if 'app' in args:
            kwargs['app'] = app
        if 'client' in args:
            kwargs['client'] = app.test_client()
            kwargs['client'].post_json = post_json
        if 'db_setup' in self.options:
            self.options['db_setup'](app)
        try:
            wrapped(**kwargs)
        finally:
            if 'db_teardown' in self.options:
                self.options['db_teardown'](app)
<|endoftext|>


<commit_before>from __future__ import division, print_function
import os, re, time, csv, sys
from io_util import make_dir, remove_dir, write_json
# from io_util import myopen, make_dir, remove_dir, tree_to_json, write_json
# from collections import defaultdict
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import generic_dna
from Bio.SeqFeature import FeatureLocation
from Bio.Data.CodonTable import TranslationError
from Bio.Seq import CodonTable
import numpy as np
from seq_util import pad_nucleotide_sequences, nuc_alpha, aa_alpha
from datetime import datetime
import json
from pdb import set_trace
from utils import fix_names, num_date, parse_date
from pprint import pprint
from Bio import AlignIO
from Bio.Align import MultipleSeqAlignment
import random

TINY = 1e-10

def calc_af(aln, alpha):
    aln_array = np.array(aln)
    af = np.zeros((len(alpha), aln_array.shape[1]))
    for ai, state in enumerate(alpha):
        af[ai] += (aln_array==state).mean(axis=0)
    af[-1] = 1.0 - af[:-1].sum(axis=0)
    return af

def safe_translate(sequence, report_exceptions=False):
    """Returns an amino acid translation of the given nucleotide sequence accounting
    for gaps in the given sequence.

    Optionally, returns a tuple of the translated sequence and whether an
    exception was raised during initial translation.

    >>> safe_translate("ATG")
    'M'
    >>> safe_translate("ATGGT-")
    'MX'
    >>> safe_translate("ATG---")
    'M-'
    >>> safe_translate("ATGTAG")
    'M*'
    >>> safe_translate("")
    ''
    >>> safe_translate("ATGT")
    'M'
    >>> safe_translate("ATG", report_exceptions=True)
    ('M', False)
    >>> safe_translate("ATGA-G", report_exceptions=True)
    ('MX', True)
    """
    translation_exception = False

    try:
        # Attempt translation by extracting the sequence according to the
        # BioPhython SeqFeature in frame gaps of three will translate as '-'
        translated_sequence = str(Seq(sequence).translate(gap='-'))
    except TranslationError:
        translation_exception = True

        # Any other codon like '-AA' or 'NNT' etc will fail. Translate codons
        # one by one.
        codon_table  = CodonTable.ambiguous_dna_by_name['Standard'].forward_table
        str_seq = str(sequence)
        codons = np.fromstring(str_seq[:len(str_seq) - len(str_seq) % 3], dtype='S3')
        assert len(codons) > 0
        aas = []

        for c in codons:
            # Parse result of single codon translation, add amino acids as
            # appropriate.
            try:
                aa = codon_table.get(c)
                if aa is None:
                    if c == '---':
                        aas.append('-')
                    else:
                        aas.append('X')
                else:
                    aas.append(aa)
            except (TranslationError, ValueError):
                aas.append('X')

        translated_sequence = "".join(aas)

    if report_exceptions:
        return translated_sequence, translation_exception
    else:
        return translated_sequence

class sequence_set(object):

    def __init__(self, logger, sequences, reference, dateFormat):
        super(sequence_set, self).__init__()
        self.log = logger

        # load sequences from the (parsed) JSON - don't forget to sort out dates
        self.seqs = {}
        for name, data in sequences.iteritems():
            self.seqs[name] = SeqRecord(Seq(data["seq"], generic_dna),
                   id=name, name=name, description=name)
            self.seqs[name].attributes = data["attributes"]
            # tidy up dates
            date_struc = parse_date(self.seqs[name].attributes["raw_date"], dateFormat)
            self.seqs[name].attributes["num_date"] = date_struc[1]
            self.seqs[name].attributes["date"] = date_struc[2]

        # if the reference is to be analysed it'll already be in the (filtered & subsampled)
        # sequences, so no need to add it here, and no need to care about attributes etc
        # we do, however, need it for alignment
        self.reference_in_dataset = reference["included"]
        name = reference["strain"]
        self.reference_seq = SeqRecord(Seq(reference["seq"], generic_dna),
               id=name, name=name, description=name)
        if "genes" in reference and len(reference["genes"]):
            self.proteins = {k:FeatureLocation(start=v["start"], end=v["end"], strand=v["strand"]) for k, v in reference["genes"].iteritems()}
        else:
            self.proteins = None

        # other things:
        self.run_dir = '_'.join(['temp', time.strftime('%Y%m%d-%H%M%S',time.gmtime()), str(random.randint(0,1000000))])
        self.nthreads = 2 # should load from config file

    def convert_trait_to_numerical_date(self, trait, dateFormat):
        for name, seq in self.seqs.iteritems():
            try:
                date_struc = parse_date(seq.attributes[trait], dateFormat)
                seq.attributes[trait] = date_struc[1]
            except KeyError:
                self.log.warn("Attribute {} not found for sequence {}. Ignoring".format(trait, seq.name))

    def codon_align(self):
        self.log.fatal("Codon align not yet implemented")

    def align(self, verbose, debug=False):
        '''
        align sequences using mafft

        side-effects:
            self.aln {MultipleSeqAlignment} reference not present if not in self.seqs
            self.reference_aln {SeqRecord} always set, even if the reference is subsequently discarded
            self.sequence_lookup {dict} map linking seq.id to the alignment, potentialy without the reference
            saves the alignment (always including reference) to fname
        '''
        make_dir(self.run_dir)
        os.chdir(self.run_dir)
        if self.reference_in_dataset:
            out_seqs = self.seqs.values()
        else:
            self.log.notify("Adding reference for alignment step")
            out_seqs = self.seqs.values() + [self.reference_seq]

        SeqIO.write(out_seqs, "temp_in.fasta", "fasta")
        self.log.notify("Running alignment")
        if verbose == 0:
            os.system("mafft --anysymbol --thread " + str(self.nthreads) + " temp_in.fasta 1> temp_out.fasta 2>mafft_stderr")
        else:
            os.system("mafft --anysymbol --thread " + str(self.nthreads) + " temp_in.fasta 1> temp_out.fasta")
        self.aln = AlignIO.read('temp_out.fasta', 'fasta')
        os.chdir("..")
        if not debug: remove_dir(self.run_dir)

        self.set_reference_alignment() #make reference_aln object while reference still in alignment
        self.set_sequence_lookup()
        self.add_attributes_to_aln()

    def remove_reference_from_alignment(self):
        count = len(self.aln)
        self.aln = MultipleSeqAlignment([s for s in self.aln if s.name!=self.reference_seq.name])
        assert(count == (len(self.aln)+1))

    def set_reference_alignment(self):
        self.reference_aln = [x for x in list(self.aln) if x.name==self.reference_seq.name][0]

    def set_sequence_lookup(self):
        self.sequence_lookup = {seq.id:seq for seq in self.aln}

    def add_attributes_to_aln(self):
        for seqid, seq in self.seqs.iteritems():
            self.sequence_lookup[seqid].attributes = seq.attributes

    def try_restore_align_from_disk(self, fname):
        try:
            self.aln = AlignIO.read(fname, "fasta")
        except IOError:
            return
        except Exception as e:
            self.log.notify("Error restoring from alignment... re-doing")
            print(e)
            return

        try:
            self.set_reference_alignment()
        except IndexError:
            self.log.notify("Reference not found in alignment... ok on reload")
            # del self.aln
            # return

        if not self.reference_in_dataset:
            self.remove_reference_from_alignment()

        if len({x.id for x in self.aln} ^ set(self.seqs.keys())) != 0:
            self.log.notify("Alignment on disk had different sequences... re-doing")
            del self.aln
            del self.reference_aln
            return

        # at this stage we are happy with the alignment
        self.set_sequence_lookup()
        self.add_attributes_to_aln()
        self.log.notify("Alignment restored from disk")

    def strip_non_reference(self):
        '''
        remove insertions relative to the reference from the alignment
        '''
        ungapped = np.array(self.reference_aln)!='-'
        for seq in self.aln:
            seq.seq = Seq("".join(np.array(seq)[ungapped]))

    def make_gaps_ambiguous(self):
        '''
        replace all gaps by 'N' in all sequences in the alignment. TreeTime will treat them
        as fully ambiguous and replace then with the most likely state
        '''
        for seq in self.aln:
            seq_array = np.array(seq)
            gaps = seq_array=='-'
            seq_array[gaps]='N'
            seq.seq = Seq("".join(seq_array))

    def make_terminal_gaps_ambiguous(self):
        '''
        replace all gaps at the end of sequences by 'N' in the alignment. TreeTime will treat them
        as fully ambiguous and replace then with the most likely state
        '''
        for seq in self.aln:
            str_seq = str(seq.seq)
            lgaps = len(str_seq) - len(str_seq.lstrip('-'))
            rgaps = len(str_seq) - len(str_seq.rstrip('-'))
            str_seq = 'N'*lgaps + str_seq.strip('-') + 'N'*rgaps
            seq.seq = Seq(str_seq)

    def translate(self):
        '''
        make alignments of translations.
        '''
        from Bio.Seq import CodonTable
        codon_table  = CodonTable.ambiguous_dna_by_name['Standard'].forward_table
        self.translations={}
        if not hasattr(self, "proteins"): # ensure dictionary to hold annotation
            self.proteins={}

        # add a default translation of the entire sequence unless otherwise specified
        if len(self.proteins)==0:
            self.proteins.update({'cds':FeatureLocation(start=0,
                end=self.aln.get_alignment_length(), strand=1)})

        # loop over all proteins and create one MSA for each
        for prot in self.proteins:
            aa_seqs = []
            for seq in self.aln:
                tmpseq = self.proteins[prot].extract(seq)
                translated_seq, translation_exception = safe_translate(str(tmpseq.seq), report_exceptions=True)

                # Strip stop codons from translated sequences.
                if translated_seq.endswith("*"):
                    translated_seq = translated_seq[:-1]

                if translation_exception:
                    self.log.notify("Trouble translating because of invalid codons %s" % seq.id)

                tmpseq.seq = Seq(translated_seq)

                # copy attributes
                tmpseq.attributes = seq.attributes
                aa_seqs.append(tmpseq)
            self.translations[prot] = MultipleSeqAlignment(aa_seqs)

    def clock_filter(self, root_seq=None, n_iqd=3, max_gaps = 1.0, plot=False):
        '''
        remove sequences form the set that are that evolve much faster or slower
        compared the majority. Regions with predominantly gaps can be removed since
        this can skew the evolutionary rates.
        '''
        if root_seq is None: # use consensus
            af = calc_af(self.aln, nuc_alpha)
            root_seq = np.fromstring(nuc_alpha, 'S1')[af.argmax(axis=0)]
        if type(root_seq)==str and root_seq in self.sequence_lookup:
            root_seq = np.array(self.sequence_lookup[root_seq])
        if max_gaps<1.0:
            af=calc_af(self.aln, nuc_alpha)
            good_pos = af[nuc_alpha.index('-')]<max_gaps
        else:
            good_pos = np.ones(self.aln.get_alignment_length(), dtype=bool)
        date_vs_distance = {}
        # self.reference_aln = None already set at alignment step
        for seq in self.aln:
            date_vs_distance[seq.id] = (seq.attributes['num_date'],
                np.mean((np.array(seq)!=root_seq)[(np.array(seq)!='-')&(root_seq!='-')&good_pos]))
            # if seq.id==self.reference.id:
            #     self.reference_aln = seq
        date_vs_distance_array=np.array(date_vs_distance.values())
        from scipy.stats import linregress, scoreatpercentile
        slope, intercept, rval, pval, stderr = linregress(date_vs_distance_array[:,0], date_vs_distance_array[:,1])
        print("distance vs time regression:",slope)
        residuals = (intercept + slope*date_vs_distance_array[:,0]) - date_vs_distance_array[:,1]
        IQD = scoreatpercentile(residuals, 75) - scoreatpercentile(residuals,25)
        if plot:
            import matplotlib.pyplot as plt
            plt.ion()
            plt.scatter(date_vs_distance_array[:,0], date_vs_distance_array[:,1], c='g')
            bad_points = abs(intercept+slope*date_vs_distance_array[:,0] - date_vs_distance_array[:,1])>n_iqd*IQD
            plt.scatter(date_vs_distance_array[bad_points,0], date_vs_distance_array[bad_points,1], c='r')

        print("before clock filter:",len(self.aln))
        tmp = {seq.id:seq for seq in self.aln
                if abs(intercept+slope*date_vs_distance[seq.id][0] - date_vs_distance[seq.id][1])<n_iqd*IQD}
        if self.reference.id not in tmp and self.reference.reference_in_dataset:
            self.log.notify('adding reference again after clock filter')
            tmp[self.reference.id] = self.reference_aln
        self.aln = MultipleSeqAlignment(tmp.values())
        print("after clock filter:",len(self.aln))

    def diversity_statistics(self):
        ''' calculate alignment entropy of nucleotide and optionally protein alignments '''
        if not hasattr(self, "aln"):
            self.log.fatal("Diversity statistics calculated before alignment generated.")
            return
        aln_array = np.array(self.aln)
        self.af = {'nuc': calc_af(self.aln, nuc_alpha)}
        tmp_af = self.af['nuc'][:-2]/self.af['nuc'][:-2].sum(axis=0)
        self.entropy ={'nuc': -(tmp_af*np.log(tmp_af+TINY)).sum(axis=0)}

        if hasattr(self, "translations"):
            for prot, aln in self.translations.iteritems():
                self.af[prot] = calc_af(aln, aa_alpha)
                tmp_af = self.af[prot][:-2]/self.af[prot][:-2].sum(axis=0)
                self.entropy[prot] = -(tmp_af*np.log(tmp_af+TINY)).sum(axis=0)

    def export_diversity(self, fname = 'entropy.json', indent=None):
        '''
        write the alignment entropy of each alignment (nucleotide and translations) to file
        '''
        if not hasattr(self, "entropy"):
            self.diversity_statistics()
        entropy_json = {}
        for feat in self.entropy:
            S = [max(0,round(x,4)) for x in self.entropy[feat]]
            n = len(S)
            if feat=='nuc':
                entropy_json[feat] = {'pos':range(0,n), 'codon':[x//3 for x in range(0,n)], 'val':S}
            else:
                entropy_json[feat] = {'pos':[x for x in self.proteins[feat]][::3],
                                      'codon':[(x-self.proteins[feat].start)//3 for x in self.proteins[feat]][::3], 'val':S}
        write_json(entropy_json, fname, indent=indent)
<commit_msg>Trim stop codon column from MSA<commit_after>from __future__ import division, print_function
import os, re, time, csv, sys
from io_util import make_dir, remove_dir, write_json
# from io_util import myopen, make_dir, remove_dir, tree_to_json, write_json
# from collections import defaultdict
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import generic_dna
from Bio.SeqFeature import FeatureLocation
from Bio.Data.CodonTable import TranslationError
from Bio.Seq import CodonTable
import numpy as np
from seq_util import pad_nucleotide_sequences, nuc_alpha, aa_alpha
from datetime import datetime
import json
from pdb import set_trace
from utils import fix_names, num_date, parse_date
from pprint import pprint
from Bio import AlignIO
from Bio.Align import MultipleSeqAlignment
import random

TINY = 1e-10

def calc_af(aln, alpha):
    aln_array = np.array(aln)
    af = np.zeros((len(alpha), aln_array.shape[1]))
    for ai, state in enumerate(alpha):
        af[ai] += (aln_array==state).mean(axis=0)
    af[-1] = 1.0 - af[:-1].sum(axis=0)
    return af

def safe_translate(sequence, report_exceptions=False):
    """Returns an amino acid translation of the given nucleotide sequence accounting
    for gaps in the given sequence.

    Optionally, returns a tuple of the translated sequence and whether an
    exception was raised during initial translation.

    >>> safe_translate("ATG")
    'M'
    >>> safe_translate("ATGGT-")
    'MX'
    >>> safe_translate("ATG---")
    'M-'
    >>> safe_translate("ATGTAG")
    'M*'
    >>> safe_translate("")
    ''
    >>> safe_translate("ATGT")
    'M'
    >>> safe_translate("ATG", report_exceptions=True)
    ('M', False)
    >>> safe_translate("ATGA-G", report_exceptions=True)
    ('MX', True)
    """
    translation_exception = False

    try:
        # Attempt translation by extracting the sequence according to the
        # BioPhython SeqFeature in frame gaps of three will translate as '-'
        translated_sequence = str(Seq(sequence).translate(gap='-'))
    except TranslationError:
        translation_exception = True

        # Any other codon like '-AA' or 'NNT' etc will fail. Translate codons
        # one by one.
        codon_table  = CodonTable.ambiguous_dna_by_name['Standard'].forward_table
        str_seq = str(sequence)
        codons = np.fromstring(str_seq[:len(str_seq) - len(str_seq) % 3], dtype='S3')
        assert len(codons) > 0
        aas = []

        for c in codons:
            # Parse result of single codon translation, add amino acids as
            # appropriate.
            try:
                aa = codon_table.get(c)
                if aa is None:
                    if c == '---':
                        aas.append('-')
                    else:
                        aas.append('X')
                else:
                    aas.append(aa)
            except (TranslationError, ValueError):
                aas.append('X')

        translated_sequence = "".join(aas)

    if report_exceptions:
        return translated_sequence, translation_exception
    else:
        return translated_sequence

class sequence_set(object):

    def __init__(self, logger, sequences, reference, dateFormat):
        super(sequence_set, self).__init__()
        self.log = logger

        # load sequences from the (parsed) JSON - don't forget to sort out dates
        self.seqs = {}
        for name, data in sequences.iteritems():
            self.seqs[name] = SeqRecord(Seq(data["seq"], generic_dna),
                   id=name, name=name, description=name)
            self.seqs[name].attributes = data["attributes"]
            # tidy up dates
            date_struc = parse_date(self.seqs[name].attributes["raw_date"], dateFormat)
            self.seqs[name].attributes["num_date"] = date_struc[1]
            self.seqs[name].attributes["date"] = date_struc[2]

        # if the reference is to be analysed it'll already be in the (filtered & subsampled)
        # sequences, so no need to add it here, and no need to care about attributes etc
        # we do, however, need it for alignment
        self.reference_in_dataset = reference["included"]
        name = reference["strain"]
        self.reference_seq = SeqRecord(Seq(reference["seq"], generic_dna),
               id=name, name=name, description=name)
        if "genes" in reference and len(reference["genes"]):
            self.proteins = {k:FeatureLocation(start=v["start"], end=v["end"], strand=v["strand"]) for k, v in reference["genes"].iteritems()}
        else:
            self.proteins = None

        # other things:
        self.run_dir = '_'.join(['temp', time.strftime('%Y%m%d-%H%M%S',time.gmtime()), str(random.randint(0,1000000))])
        self.nthreads = 2 # should load from config file

    def convert_trait_to_numerical_date(self, trait, dateFormat):
        for name, seq in self.seqs.iteritems():
            try:
                date_struc = parse_date(seq.attributes[trait], dateFormat)
                seq.attributes[trait] = date_struc[1]
            except KeyError:
                self.log.warn("Attribute {} not found for sequence {}. Ignoring".format(trait, seq.name))

    def codon_align(self):
        self.log.fatal("Codon align not yet implemented")

    def align(self, verbose, debug=False):
        '''
        align sequences using mafft

        side-effects:
            self.aln {MultipleSeqAlignment} reference not present if not in self.seqs
            self.reference_aln {SeqRecord} always set, even if the reference is subsequently discarded
            self.sequence_lookup {dict} map linking seq.id to the alignment, potentialy without the reference
            saves the alignment (always including reference) to fname
        '''
        make_dir(self.run_dir)
        os.chdir(self.run_dir)
        if self.reference_in_dataset:
            out_seqs = self.seqs.values()
        else:
            self.log.notify("Adding reference for alignment step")
            out_seqs = self.seqs.values() + [self.reference_seq]

        SeqIO.write(out_seqs, "temp_in.fasta", "fasta")
        self.log.notify("Running alignment")
        if verbose == 0:
            os.system("mafft --anysymbol --thread " + str(self.nthreads) + " temp_in.fasta 1> temp_out.fasta 2>mafft_stderr")
        else:
            os.system("mafft --anysymbol --thread " + str(self.nthreads) + " temp_in.fasta 1> temp_out.fasta")
        self.aln = AlignIO.read('temp_out.fasta', 'fasta')
        os.chdir("..")
        if not debug: remove_dir(self.run_dir)

        self.set_reference_alignment() #make reference_aln object while reference still in alignment
        self.set_sequence_lookup()
        self.add_attributes_to_aln()

    def remove_reference_from_alignment(self):
        count = len(self.aln)
        self.aln = MultipleSeqAlignment([s for s in self.aln if s.name!=self.reference_seq.name])
        assert(count == (len(self.aln)+1))

    def set_reference_alignment(self):
        self.reference_aln = [x for x in list(self.aln) if x.name==self.reference_seq.name][0]

    def set_sequence_lookup(self):
        self.sequence_lookup = {seq.id:seq for seq in self.aln}

    def add_attributes_to_aln(self):
        for seqid, seq in self.seqs.iteritems():
            self.sequence_lookup[seqid].attributes = seq.attributes

    def try_restore_align_from_disk(self, fname):
        try:
            self.aln = AlignIO.read(fname, "fasta")
        except IOError:
            return
        except Exception as e:
            self.log.notify("Error restoring from alignment... re-doing")
            print(e)
            return

        try:
            self.set_reference_alignment()
        except IndexError:
            self.log.notify("Reference not found in alignment... ok on reload")
            # del self.aln
            # return

        if not self.reference_in_dataset:
            self.remove_reference_from_alignment()

        if len({x.id for x in self.aln} ^ set(self.seqs.keys())) != 0:
            self.log.notify("Alignment on disk had different sequences... re-doing")
            del self.aln
            del self.reference_aln
            return

        # at this stage we are happy with the alignment
        self.set_sequence_lookup()
        self.add_attributes_to_aln()
        self.log.notify("Alignment restored from disk")

    def strip_non_reference(self):
        '''
        remove insertions relative to the reference from the alignment
        '''
        ungapped = np.array(self.reference_aln)!='-'
        for seq in self.aln:
            seq.seq = Seq("".join(np.array(seq)[ungapped]))

    def make_gaps_ambiguous(self):
        '''
        replace all gaps by 'N' in all sequences in the alignment. TreeTime will treat them
        as fully ambiguous and replace then with the most likely state
        '''
        for seq in self.aln:
            seq_array = np.array(seq)
            gaps = seq_array=='-'
            seq_array[gaps]='N'
            seq.seq = Seq("".join(seq_array))

    def make_terminal_gaps_ambiguous(self):
        '''
        replace all gaps at the end of sequences by 'N' in the alignment. TreeTime will treat them
        as fully ambiguous and replace then with the most likely state
        '''
        for seq in self.aln:
            str_seq = str(seq.seq)
            lgaps = len(str_seq) - len(str_seq.lstrip('-'))
            rgaps = len(str_seq) - len(str_seq.rstrip('-'))
            str_seq = 'N'*lgaps + str_seq.strip('-') + 'N'*rgaps
            seq.seq = Seq(str_seq)

    def translate(self):
        '''
        make alignments of translations.
        '''
        from Bio.Seq import CodonTable
        codon_table  = CodonTable.ambiguous_dna_by_name['Standard'].forward_table
        self.translations={}
        if not hasattr(self, "proteins"): # ensure dictionary to hold annotation
            self.proteins={}

        # add a default translation of the entire sequence unless otherwise specified
        if len(self.proteins)==0:
            self.proteins.update({'cds':FeatureLocation(start=0,
                end=self.aln.get_alignment_length(), strand=1)})

        # loop over all proteins and create one MSA for each
        for prot in self.proteins:
            aa_seqs = []
            trim_end = False
            for seq in self.aln:
                tmpseq = self.proteins[prot].extract(seq)
                translated_seq, translation_exception = safe_translate(str(tmpseq.seq), report_exceptions=True)

                if translation_exception:
                    self.log.notify("Trouble translating because of invalid codons %s" % seq.id)

                if translated_seq.endswith("*"):
                    trim_end = True

                tmpseq.seq = Seq(translated_seq)

                # copy attributes
                tmpseq.attributes = seq.attributes
                aa_seqs.append(tmpseq)

            # strip stop codons from translated sequences, some sequences may end in NNNNN
            if trim_end:
                for aa_seq in aa_seqs:
                    aa_seq = aa_seq[:-1]

            self.translations[prot] = MultipleSeqAlignment(aa_seqs)

    def clock_filter(self, root_seq=None, n_iqd=3, max_gaps = 1.0, plot=False):
        '''
        remove sequences form the set that are that evolve much faster or slower
        compared the majority. Regions with predominantly gaps can be removed since
        this can skew the evolutionary rates.
        '''
        if root_seq is None: # use consensus
            af = calc_af(self.aln, nuc_alpha)
            root_seq = np.fromstring(nuc_alpha, 'S1')[af.argmax(axis=0)]
        if type(root_seq)==str and root_seq in self.sequence_lookup:
            root_seq = np.array(self.sequence_lookup[root_seq])
        if max_gaps<1.0:
            af=calc_af(self.aln, nuc_alpha)
            good_pos = af[nuc_alpha.index('-')]<max_gaps
        else:
            good_pos = np.ones(self.aln.get_alignment_length(), dtype=bool)
        date_vs_distance = {}
        # self.reference_aln = None already set at alignment step
        for seq in self.aln:
            date_vs_distance[seq.id] = (seq.attributes['num_date'],
                np.mean((np.array(seq)!=root_seq)[(np.array(seq)!='-')&(root_seq!='-')&good_pos]))
            # if seq.id==self.reference.id:
            #     self.reference_aln = seq
        date_vs_distance_array=np.array(date_vs_distance.values())
        from scipy.stats import linregress, scoreatpercentile
        slope, intercept, rval, pval, stderr = linregress(date_vs_distance_array[:,0], date_vs_distance_array[:,1])
        print("distance vs time regression:",slope)
        residuals = (intercept + slope*date_vs_distance_array[:,0]) - date_vs_distance_array[:,1]
        IQD = scoreatpercentile(residuals, 75) - scoreatpercentile(residuals,25)
        if plot:
            import matplotlib.pyplot as plt
            plt.ion()
            plt.scatter(date_vs_distance_array[:,0], date_vs_distance_array[:,1], c='g')
            bad_points = abs(intercept+slope*date_vs_distance_array[:,0] - date_vs_distance_array[:,1])>n_iqd*IQD
            plt.scatter(date_vs_distance_array[bad_points,0], date_vs_distance_array[bad_points,1], c='r')

        print("before clock filter:",len(self.aln))
        tmp = {seq.id:seq for seq in self.aln
                if abs(intercept+slope*date_vs_distance[seq.id][0] - date_vs_distance[seq.id][1])<n_iqd*IQD}
        if self.reference.id not in tmp and self.reference.reference_in_dataset:
            self.log.notify('adding reference again after clock filter')
            tmp[self.reference.id] = self.reference_aln
        self.aln = MultipleSeqAlignment(tmp.values())
        print("after clock filter:",len(self.aln))

    def diversity_statistics(self):
        ''' calculate alignment entropy of nucleotide and optionally protein alignments '''
        if not hasattr(self, "aln"):
            self.log.fatal("Diversity statistics calculated before alignment generated.")
            return
        aln_array = np.array(self.aln)
        self.af = {'nuc': calc_af(self.aln, nuc_alpha)}
        tmp_af = self.af['nuc'][:-2]/self.af['nuc'][:-2].sum(axis=0)
        self.entropy ={'nuc': -(tmp_af*np.log(tmp_af+TINY)).sum(axis=0)}

        if hasattr(self, "translations"):
            for prot, aln in self.translations.iteritems():
                self.af[prot] = calc_af(aln, aa_alpha)
                tmp_af = self.af[prot][:-2]/self.af[prot][:-2].sum(axis=0)
                self.entropy[prot] = -(tmp_af*np.log(tmp_af+TINY)).sum(axis=0)

    def export_diversity(self, fname = 'entropy.json', indent=None):
        '''
        write the alignment entropy of each alignment (nucleotide and translations) to file
        '''
        if not hasattr(self, "entropy"):
            self.diversity_statistics()
        entropy_json = {}
        for feat in self.entropy:
            S = [max(0,round(x,4)) for x in self.entropy[feat]]
            n = len(S)
            if feat=='nuc':
                entropy_json[feat] = {'pos':range(0,n), 'codon':[x//3 for x in range(0,n)], 'val':S}
            else:
                entropy_json[feat] = {'pos':[x for x in self.proteins[feat]][::3],
                                      'codon':[(x-self.proteins[feat].start)//3 for x in self.proteins[feat]][::3], 'val':S}
        write_json(entropy_json, fname, indent=indent)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#############################################################################
# Copyright (c) 2017 SiteWare Corp. All right reserved
#############################################################################

import os
import sys

import sd2
from .file_rewriter import FileRewriter

g_ssh_config_path = os.path.join(os.getenv('HOME'), '.ssh', 'config')

container_entry_template = '''

host {cont[name]}
    HostName {cont[name]}
    User {host[User]}
    Port {sshport}
    StrictHostKeyChecking no
    UserKnownHostsFile /dev/null
    ServerAliveInterval 60
    ConnectTimeout 5
'''

ssh_option_names = [
    'HostName',
    'Port',
    'User',
    'IdentityFile',
    'ServerAliveInterval',
    'StrictHostKeyChecking',
    'UserKnownHostsFile',
    'ProxyCommand',
    'ConnectTimeout',
    'UseKeychain',
    'AddKeysToAgent',
    'ForwardAgent',
    "PKCS11Provider"
]

def generate_for_host(host):
    from . import util
    rr = ''
    rr += '''\n########## GENERATED DO NOT MODIFY #####################\n'''
    sshport = 22 if util.is_localhost(host['name']) else 2222
    if not util.is_localhost(host['name']):
        if host.get('match'):
            matches = host.get('match')
            for hostname in [host['name'], host['name'] + '-ports']:
                for match in matches:
                    rr += 'Match originalhost {hostname} exec "{match[condition]}"\n'.format(
                        hostname=hostname, match=match)
                    for key in match.keys():
                        if not key in ssh_option_names:
                            continue
                        rr += '    {key} {value}\n'.format(key=key,
                                                           value=match[key])
                rr += '\n'

        rr += 'host {}\n'.format(host['name'])
        for key, val in host.iteritems():
            if not key in ssh_option_names:
                continue
            if not isinstance(val, (list, tuple)):
                val = [val]
            for vv in val:
                rr += '    {key} {value}\n'.format(key=key, value=vv)
        rr += '\n'

        if host.get('containers'):
            rr += 'host {}-ports\n'.format(host['name'])
            if not 'HostName' in host.keys():
                host['HostName'] = host['name']
            for key, val in host.iteritems():
                if not key in ssh_option_names + ['LocalForward']:
                    continue
                if not isinstance(val, (list, tuple)):
                    val = [val]
                for vv in val:
                    rr += '    {key} {value}\n'.format(key=key, value=vv)
                    # rr += '    LocalForward {}-local:2375 localhost:2375\n'.format(host['name'])
            for cont in host.get('containers', []):
                ports = cont['image'].get('ports', [])
                for port in ports + ["{}:22".format(sshport)]:
                    (p1, p2) = port.split(':')
                    rr += (
                        "    LocalForward  {0}:{1} {2}:{1}\n".format(
                            cont['name'], p1, cont['ip']))
            rr += '\n'

    for cont in host.get('containers', []):
        rr += container_entry_template.format(**locals())
    return rr

def get_our_ssh_config():
    rr = ''
    for host in sd2.get_hosts():
        if not host.get("User"):
            host['User'] = os.getenv('USER')
        try:
            rr += generate_for_host(host)
        except Exception as ex:
            sys.stderr.write("ERROR: Processing host {}\n".format(host['name']))
            raise

    return rr

def gen_ssh_config():
    fr = FileRewriter(g_ssh_config_path)
    before, after = fr.read_config()
    rr = get_our_ssh_config()
    fr.write_config(
        before,
        rr.split('\n'),
        after
    )<commit_msg>Fix: Pass the ssh credentials of the host to the container entry so that we can ssh into the container<commit_after>#!/usr/bin/env python
#############################################################################
# Copyright (c) 2017 SiteWare Corp. All right reserved
#############################################################################

import os
import sys

import sd2
from .file_rewriter import FileRewriter

g_ssh_config_path = os.path.join(os.getenv('HOME'), '.ssh', 'config')

container_entry_template = '''

host {cont[name]}
    HostName {cont[name]}
    User {host[User]}
    Port {sshport}
    StrictHostKeyChecking no
    UserKnownHostsFile /dev/null
    ServerAliveInterval 60
    ConnectTimeout 5
'''

ssh_option_names = [
    'HostName',
    'Port',
    'User',
    'IdentityFile',
    'ServerAliveInterval',
    'StrictHostKeyChecking',
    'UserKnownHostsFile',
    'ProxyCommand',
    'ConnectTimeout',
    'UseKeychain',
    'AddKeysToAgent',
    'ForwardAgent',
    "PKCS11Provider"
]

def generate_for_host(host):
    from . import util
    rr = ''
    rr += '''\n########## GENERATED DO NOT MODIFY #####################\n'''
    sshport = 22 if util.is_localhost(host['name']) else 2222
    if not util.is_localhost(host['name']):
        if host.get('match'):
            matches = host.get('match')
            for hostname in [host['name'], host['name'] + '-ports']:
                for match in matches:
                    rr += 'Match originalhost {hostname} exec "{match[condition]}"\n'.format(
                        hostname=hostname, match=match)
                    for key in match.keys():
                        if not key in ssh_option_names:
                            continue
                        rr += '    {key} {value}\n'.format(key=key,
                                                           value=match[key])
                rr += '\n'

        rr += 'host {}\n'.format(host['name'])
        for key, val in host.iteritems():
            if not key in ssh_option_names:
                continue
            if not isinstance(val, (list, tuple)):
                val = [val]
            for vv in val:
                rr += '    {key} {value}\n'.format(key=key, value=vv)
        rr += '\n'

        if host.get('containers'):
            rr += 'host {}-ports\n'.format(host['name'])
            if not 'HostName' in host.keys():
                host['HostName'] = host['name']
            for key, val in host.iteritems():
                if not key in ssh_option_names + ['LocalForward']:
                    continue
                if not isinstance(val, (list, tuple)):
                    val = [val]
                for vv in val:
                    rr += '    {key} {value}\n'.format(key=key, value=vv)
                    # rr += '    LocalForward {}-local:2375 localhost:2375\n'.format(host['name'])
            for cont in host.get('containers', []):
                ports = cont['image'].get('ports', [])
                for port in ports + ["{}:22".format(sshport)]:
                    (p1, p2) = port.split(':')
                    rr += (
                        "    LocalForward  {0}:{1} {2}:{1}\n".format(
                            cont['name'], p1, cont['ip']))
            rr += '\n'

    for cont in host.get('containers', []):
        rr += container_entry_template.format(**locals())
        for key, val in host.iteritems():
            if key in ['IdentityFile','ProxyCommad','PKCS11Provider']:
                rr += '    {} {}\n'.format(key, val)
    return rr

def get_our_ssh_config():
    rr = ''
    for host in sd2.get_hosts():
        if not host.get("User"):
            host['User'] = os.getenv('USER')
        try:
            rr += generate_for_host(host)
        except Exception as ex:
            sys.stderr.write("ERROR: Processing host {}\n".format(host['name']))
            raise

    return rr

def gen_ssh_config():
    fr = FileRewriter(g_ssh_config_path)
    before, after = fr.read_config()
    rr = get_our_ssh_config()
    fr.write_config(
        before,
        rr.split('\n'),
        after
    )<|endoftext|>


<commit_before># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
# pylint: disable=C,R,W

import enum

from croniter import croniter
from flask import flash, g
from flask_appbuilder import expose
from flask_appbuilder.models.sqla.interface import SQLAInterface
from flask_appbuilder.security.decorators import has_access
from flask_babel import gettext as __
from flask_babel import lazy_gettext as _
import simplejson as json
from wtforms import BooleanField, StringField

from superset import app, appbuilder, db, security_manager
from superset.exceptions import SupersetException
from superset.models.core import Dashboard, Slice
from superset.models.schedules import (
    DashboardEmailSchedule,
    ScheduleType,
    SliceEmailSchedule,
)
from superset.tasks.schedules import schedule_email_report
from superset.utils.core import (
    get_email_address_list,
    json_iso_dttm_ser,
)
from superset.views.core import json_success
from .base import DeleteMixin, SupersetModelView

class EmailScheduleView(SupersetModelView, DeleteMixin):
    _extra_data = {
        'test_email': False,
        'test_email_recipients': None,
    }
    schedule_type = None
    schedule_type_model = None

    page_size = 20

    add_exclude_columns = [
        'user',
        'created_on',
        'changed_on',
        'created_by',
        'changed_by',
    ]

    edit_exclude_columns = add_exclude_columns

    description_columns = {
        'deliver_as_group': 'If enabled, send a single email to all '
        'recipients (in email/To: field)',
        'crontab': 'Unix style crontab schedule to deliver emails. '
                   'Changes to schedules reflect in one hour.',
        'delivery_type': 'Indicates how the rendered content is delivered',
    }

    add_form_extra_fields = {
        'test_email': BooleanField(
            'Send Test Email',
            default=False,
            description='If enabled, we send a test mail on create / update',
        ),
        'test_email_recipients': StringField(
            'Test Email Recipients',
            default=None,
            description='List of recipients to send test email to. '
                        'If empty, we send it to the original recipients',
        ),
    }

    edit_form_extra_fields = add_form_extra_fields

    def process_form(self, form, is_created):
        recipients = form.test_email_recipients.data.strip() or None
        self._extra_data['test_email'] = form.test_email.data
        self._extra_data['test_email_recipients'] = recipients

    def pre_add(self, obj):
        try:
            recipients = get_email_address_list(obj.recipients)
            obj.recipients = ', '.join(recipients)
        except Exception:
            raise SupersetException('Invalid email list')

        obj.user = obj.user or g.user
        if not croniter.is_valid(obj.crontab):
            raise SupersetException('Invalid crontab format')

    def pre_update(self, obj):
        self.pre_add(obj)

    def post_add(self, obj):
        # Schedule a test mail if the user requested for it.
        if self._extra_data['test_email']:
            recipients = self._extra_data['test_email_recipients']
            args = (self.schedule_type, obj.id)
            kwargs = dict(recipients=recipients)
            schedule_email_report.apply_async(args=args, kwargs=kwargs)

        # Notify the user that schedule changes will be activate only in the
        # next hour
        if obj.active:
            flash('Schedule changes will get applied in one hour', 'warning')

    def post_update(self, obj):
        self.post_add(obj)

    @has_access
    @expose('/fetch/<int:item_id>/', methods=['GET'])
    def fetch_schedules(self, item_id):

        query = db.session.query(self.datamodel.obj)
        query = query.join(self.schedule_type_model).filter(
            self.schedule_type_model.id == item_id)

        schedules = []
        for schedule in query.all():
            info = {'schedule': schedule.id}

            for col in self.list_columns + self.add_exclude_columns:
                info[col] = getattr(schedule, col)

                if isinstance(info[col], enum.Enum):
                    info[col] = info[col].name
                elif isinstance(info[col], security_manager.user_model):
                    info[col] = info[col].username

            info['user'] = schedule.user.username
            info[self.schedule_type] = getattr(schedule, self.schedule_type).id
            schedules.append(info)

        return json_success(json.dumps(schedules, default=json_iso_dttm_ser))

class DashboardEmailScheduleView(EmailScheduleView):
    schedule_type = ScheduleType.dashboard.name
    schedule_type_model = Dashboard

    add_title = _('Schedule Email Reports for Dashboards')
    edit_title = add_title
    list_title = _('Manage Email Reports for Dashboards')

    datamodel = SQLAInterface(DashboardEmailSchedule)
    order_columns = ['user', 'dashboard', 'created_on']

    list_columns = [
        'dashboard',
        'active',
        'crontab',
        'user',
        'deliver_as_group',
        'delivery_type',
    ]

    add_columns = [
        'dashboard',
        'active',
        'crontab',
        'recipients',
        'deliver_as_group',
        'delivery_type',
        'test_email',
        'test_email_recipients',
    ]

    edit_columns = add_columns

    search_columns = [
        'dashboard',
        'active',
        'user',
        'deliver_as_group',
        'delivery_type',
    ]

    label_columns = {
        'dashboard': _('Dashboard'),
        'created_on': _('Created On'),
        'changed_on': _('Changed On'),
        'user': _('User'),
        'active': _('Active'),
        'crontab': _('Crontab'),
        'recipients': _('Recipients'),
        'deliver_as_group': _('Deliver As Group'),
        'delivery_type': _('Delivery Type'),
    }

    def pre_add(self, obj):
        if obj.dashboard is None:
            raise SupersetException('Dashboard is mandatory')
        super(DashboardEmailScheduleView, self).pre_add(obj)

class SliceEmailScheduleView(EmailScheduleView):
    schedule_type = ScheduleType.slice.name
    schedule_type_model = Slice
    add_title = _('Schedule Email Reports for Charts')
    edit_title = add_title
    list_title = _('Manage Email Reports for Charts')

    datamodel = SQLAInterface(SliceEmailSchedule)
    order_columns = ['user', 'slice', 'created_on']
    list_columns = [
        'slice',
        'active',
        'crontab',
        'user',
        'deliver_as_group',
        'delivery_type',
        'email_format',
    ]

    add_columns = [
        'slice',
        'active',
        'crontab',
        'recipients',
        'deliver_as_group',
        'delivery_type',
        'email_format',
        'test_email',
        'test_email_recipients',
    ]

    edit_columns = add_columns

    search_columns = [
        'slice',
        'active',
        'user',
        'deliver_as_group',
        'delivery_type',
        'email_format',
    ]

    label_columns = {
        'slice': _('Chart'),
        'created_on': _('Created On'),
        'changed_on': _('Changed On'),
        'user': _('User'),
        'active': _('Active'),
        'crontab': _('Crontab'),
        'recipients': _('Recipients'),
        'deliver_as_group': _('Deliver As Group'),
        'delivery_type': _('Delivery Type'),
        'email_format': _('Email Format'),
    }

    def pre_add(self, obj):
        if obj.slice is None:
            raise SupersetException('Slice is mandatory')
        super(SliceEmailScheduleView, self).pre_add(obj)

def _register_schedule_menus():
    appbuilder.add_separator('Manage')

    appbuilder.add_view(
        DashboardEmailScheduleView,
        'Dashboard Email Schedules',
        label=__('Dashboard Emails'),
        category='Manage',
        category_label=__('Manage'),
        icon='fa-search')

    appbuilder.add_view(
        SliceEmailScheduleView,
        'Chart Emails',
        label=__('Chart Email Schedules'),
        category='Manage',
        category_label=__('Manage'),
        icon='fa-search')

if app.config.get('ENABLE_SCHEDULED_EMAIL_REPORTS'):
    _register_schedule_menus()
<commit_msg>Fix NoneType bug & fill the test recipients with original recipients if empty (#7365)<commit_after># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
# pylint: disable=C,R,W

import enum

from croniter import croniter
from flask import flash, g
from flask_appbuilder import expose
from flask_appbuilder.models.sqla.interface import SQLAInterface
from flask_appbuilder.security.decorators import has_access
from flask_babel import gettext as __
from flask_babel import lazy_gettext as _
import simplejson as json
from wtforms import BooleanField, StringField

from superset import app, appbuilder, db, security_manager
from superset.exceptions import SupersetException
from superset.models.core import Dashboard, Slice
from superset.models.schedules import (
    DashboardEmailSchedule,
    ScheduleType,
    SliceEmailSchedule,
)
from superset.tasks.schedules import schedule_email_report
from superset.utils.core import (
    get_email_address_list,
    json_iso_dttm_ser,
)
from superset.views.core import json_success
from .base import DeleteMixin, SupersetModelView

class EmailScheduleView(SupersetModelView, DeleteMixin):
    _extra_data = {
        'test_email': False,
        'test_email_recipients': None,
    }
    schedule_type = None
    schedule_type_model = None

    page_size = 20

    add_exclude_columns = [
        'user',
        'created_on',
        'changed_on',
        'created_by',
        'changed_by',
    ]

    edit_exclude_columns = add_exclude_columns

    description_columns = {
        'deliver_as_group': 'If enabled, send a single email to all '
        'recipients (in email/To: field)',
        'crontab': 'Unix style crontab schedule to deliver emails. '
                   'Changes to schedules reflect in one hour.',
        'delivery_type': 'Indicates how the rendered content is delivered',
    }

    add_form_extra_fields = {
        'test_email': BooleanField(
            'Send Test Email',
            default=False,
            description='If enabled, we send a test mail on create / update',
        ),
        'test_email_recipients': StringField(
            'Test Email Recipients',
            default=None,
            description='List of recipients to send test email to. '
                        'If empty, we send it to the original recipients',
        ),
    }

    edit_form_extra_fields = add_form_extra_fields

    def process_form(self, form, is_created):
        if form.test_email_recipients.data:
            test_email_recipients = form.test_email_recipients.data.strip()
        else:
            test_email_recipients = None
        self._extra_data['test_email'] = form.test_email.data
        self._extra_data['test_email_recipients'] = test_email_recipients

    def pre_add(self, obj):
        try:
            recipients = get_email_address_list(obj.recipients)
            obj.recipients = ', '.join(recipients)
        except Exception:
            raise SupersetException('Invalid email list')

        obj.user = obj.user or g.user
        if not croniter.is_valid(obj.crontab):
            raise SupersetException('Invalid crontab format')

    def pre_update(self, obj):
        self.pre_add(obj)

    def post_add(self, obj):
        # Schedule a test mail if the user requested for it.
        if self._extra_data['test_email']:
            recipients = self._extra_data['test_email_recipients'] or obj.recipients
            args = (self.schedule_type, obj.id)
            kwargs = dict(recipients=recipients)
            schedule_email_report.apply_async(args=args, kwargs=kwargs)

        # Notify the user that schedule changes will be activate only in the
        # next hour
        if obj.active:
            flash('Schedule changes will get applied in one hour', 'warning')

    def post_update(self, obj):
        self.post_add(obj)

    @has_access
    @expose('/fetch/<int:item_id>/', methods=['GET'])
    def fetch_schedules(self, item_id):

        query = db.session.query(self.datamodel.obj)
        query = query.join(self.schedule_type_model).filter(
            self.schedule_type_model.id == item_id)

        schedules = []
        for schedule in query.all():
            info = {'schedule': schedule.id}

            for col in self.list_columns + self.add_exclude_columns:
                info[col] = getattr(schedule, col)

                if isinstance(info[col], enum.Enum):
                    info[col] = info[col].name
                elif isinstance(info[col], security_manager.user_model):
                    info[col] = info[col].username

            info['user'] = schedule.user.username
            info[self.schedule_type] = getattr(schedule, self.schedule_type).id
            schedules.append(info)

        return json_success(json.dumps(schedules, default=json_iso_dttm_ser))

class DashboardEmailScheduleView(EmailScheduleView):
    schedule_type = ScheduleType.dashboard.name
    schedule_type_model = Dashboard

    add_title = _('Schedule Email Reports for Dashboards')
    edit_title = add_title
    list_title = _('Manage Email Reports for Dashboards')

    datamodel = SQLAInterface(DashboardEmailSchedule)
    order_columns = ['user', 'dashboard', 'created_on']

    list_columns = [
        'dashboard',
        'active',
        'crontab',
        'user',
        'deliver_as_group',
        'delivery_type',
    ]

    add_columns = [
        'dashboard',
        'active',
        'crontab',
        'recipients',
        'deliver_as_group',
        'delivery_type',
        'test_email',
        'test_email_recipients',
    ]

    edit_columns = add_columns

    search_columns = [
        'dashboard',
        'active',
        'user',
        'deliver_as_group',
        'delivery_type',
    ]

    label_columns = {
        'dashboard': _('Dashboard'),
        'created_on': _('Created On'),
        'changed_on': _('Changed On'),
        'user': _('User'),
        'active': _('Active'),
        'crontab': _('Crontab'),
        'recipients': _('Recipients'),
        'deliver_as_group': _('Deliver As Group'),
        'delivery_type': _('Delivery Type'),
    }

    def pre_add(self, obj):
        if obj.dashboard is None:
            raise SupersetException('Dashboard is mandatory')
        super(DashboardEmailScheduleView, self).pre_add(obj)

class SliceEmailScheduleView(EmailScheduleView):
    schedule_type = ScheduleType.slice.name
    schedule_type_model = Slice
    add_title = _('Schedule Email Reports for Charts')
    edit_title = add_title
    list_title = _('Manage Email Reports for Charts')

    datamodel = SQLAInterface(SliceEmailSchedule)
    order_columns = ['user', 'slice', 'created_on']
    list_columns = [
        'slice',
        'active',
        'crontab',
        'user',
        'deliver_as_group',
        'delivery_type',
        'email_format',
    ]

    add_columns = [
        'slice',
        'active',
        'crontab',
        'recipients',
        'deliver_as_group',
        'delivery_type',
        'email_format',
        'test_email',
        'test_email_recipients',
    ]

    edit_columns = add_columns

    search_columns = [
        'slice',
        'active',
        'user',
        'deliver_as_group',
        'delivery_type',
        'email_format',
    ]

    label_columns = {
        'slice': _('Chart'),
        'created_on': _('Created On'),
        'changed_on': _('Changed On'),
        'user': _('User'),
        'active': _('Active'),
        'crontab': _('Crontab'),
        'recipients': _('Recipients'),
        'deliver_as_group': _('Deliver As Group'),
        'delivery_type': _('Delivery Type'),
        'email_format': _('Email Format'),
    }

    def pre_add(self, obj):
        if obj.slice is None:
            raise SupersetException('Slice is mandatory')
        super(SliceEmailScheduleView, self).pre_add(obj)

def _register_schedule_menus():
    appbuilder.add_separator('Manage')

    appbuilder.add_view(
        DashboardEmailScheduleView,
        'Dashboard Email Schedules',
        label=__('Dashboard Emails'),
        category='Manage',
        category_label=__('Manage'),
        icon='fa-search')

    appbuilder.add_view(
        SliceEmailScheduleView,
        'Chart Emails',
        label=__('Chart Email Schedules'),
        category='Manage',
        category_label=__('Manage'),
        icon='fa-search')

if app.config.get('ENABLE_SCHEDULED_EMAIL_REPORTS'):
    _register_schedule_menus()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Make an iOS project transportable so that it can be zipped and
# sent to another machine
#

import os, sys, shutil, codecs, glob
template_dir = os.path.abspath(os.path.dirname(sys._getframe(0).f_code.co_filename))

sys.path.append(os.path.abspath(os.path.dirname(template_dir)))
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(template_dir),'module')))
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(template_dir),'common')))

from tiapp import *
from module import ModuleDetector
from pbxproj import PBXProj
from localecompiler import LocaleCompiler
from projector import Projector
from tools import *
from compiler import Compiler

def find_sdk(version):
		sdks = [os.path.join(os.path.expanduser("~/Library/Application Support/Titanium"),"mobilesdk","osx",version),
				os.path.join("/Library","Application Support","Titanium","mobilesdk","osx",version)]

		for sdk in sdks:
			if os.path.exists(sdk):
				return sdk
		print "[ERROR] Is Titanium installed? I can't find it"
		sys.exit(1)

def info(msg):
		print "[INFO] %s" % msg
		sys.stdout.flush()

def main(args):
		if len(args) < 2:
				print "Usage: %s <project_directory> [sdk_verison]" % os.path.basename(args[0])
				sys.exit(1)

		# What needs to be done in order to perform a "true" export?
		# ---
		# Wipe the build dir
		# Migrate resources
		# Migrate tiapp.xml (required for scripts)
		# Generate project from template
		# Populate Info.plist
		# Compile/migrate i18n
		# Migrate scripts for compiling JSS files (and i18n)
		# Modify xcodeproj build steps to call the JSS compiler
		# Then... Share and Enjoy.

		project_dir = os.path.abspath(args[1])
		build_dir = os.path.join(project_dir,'build','iphone')
		titanium_local = os.path.join(build_dir,'titanium')
		buildMode = 'production'

		if len(args) > 2:
			version = args[2]
			sdk_dir = find_sdk(version)
		else:
			sdk_dir = os.path.abspath(os.path.dirname(template_dir))
			version = os.path.basename(sdk_dir)

		if len(args) == 4:
			buildMode = args[3]

		tiappxml = os.path.join(project_dir, 'tiapp.xml')
		tiapp = TiAppXML(tiappxml)

		app_id = tiapp.properties['id']
		app_name = tiapp.properties['name']

		info("Creating project in %s mode" % buildMode)

		if app_id is None or app_name is None:
			info("Your tiapp.xml is malformed - please specify an app name and id")
			sys.exit(1)

		# Clean build dir (if it exists), error otherwise (no iphone support)
		info("Cleaning build...")
		if os.path.exists(build_dir):
			for f in os.listdir(build_dir):
				path = os.path.join(build_dir,f)
				if os.path.isfile(path):
					os.remove(path)
				else:
					shutil.rmtree(path)
		else:
			info("Your project is not configured to be built for iphone.")
			exit(1)

		# Migrate Resources
		info("Migrating resources...")
		project_resources = os.path.join(project_dir, 'Resources')
		resources_dir = os.path.join(build_dir, 'Resources')

		shutil.copytree(project_resources,resources_dir)

		# Migrate platform/iphone contents into Resources.
		info("Migrating platform/iphone to Resources...")
		project_platform = os.path.join(project_dir,'platform','iphone')

		if os.path.isdir(project_platform):
			contents = os.listdir(project_platform)
			for file in contents:
				path = os.path.join(project_platform,file)
				if os.path.isdir(path):
					shutil.copytree(path, os.path.join(resources_dir,file))
				else:
					shutil.copy(path, os.path.join(resources_dir,file))

		# Migrate tiapp.xml
		info("Migrating tiapp.xml...")
		shutil.copy(tiappxml, build_dir)

		ti = TiAppXML(tiappxml)
		# target the requested value if provided
		min_ver = 4.0
		if 'min-ios-ver' in ti.ios:
			min_ver = float(ti.ios['min-ios-ver'])
			if min_ver < 4.0:
				print "[INFO] Minimum iOS version %s is lower than 4.0: Using 4.0 as minimum" % min_ver
				min_ver = 4.0

		# Generate project stuff from the template
		info("Generating project from Titanium template...")
		project = Projector(app_name,version,template_dir,project_dir,app_id, min_ver)
		project.create(template_dir,build_dir)

		# Because the debugger.plist is built as part of the required
		# resources, we need to autogen an empty one
		debug_plist = os.path.join(resources_dir,'debugger.plist')
		force_xcode = write_debugger_plist(None, None, None, None, template_dir, debug_plist)

		# Populate Info.plist
		applogo = None
		info("Populating Info.plist...")
		plist_out = os.path.join(build_dir, 'Info.plist')
		create_info_plist(tiapp, template_dir, project_dir, plist_out)
		applogo = tiapp.generate_infoplist(plist_out, app_id, 'iphone', project_dir, None)

		# Run the compiler to autogenerate .m files
		info("Copying classes, creating generated .m files...")
		compiler = Compiler(project_dir,app_id,app_name,buildMode)
		compiler.compileProject(silent=True)

		#... But we still have to nuke the stuff that gets built that we don't want
		# to bundle.
		ios_build = os.path.join(build_dir,'build')
		if os.path.isdir(ios_build):
			shutil.rmtree(os.path.join(build_dir,'build'))

		# Install applogo/splash/etc.
		info("Copying icons and splash...")
		install_logo(tiapp, applogo, project_dir, template_dir, resources_dir)
		install_defaults(project_dir, template_dir, resources_dir)

		# Get Modules
		detector = ModuleDetector(project_dir)
		missing_modules, modules = detector.find_app_modules(tiapp, 'iphone', buildMode)

		if len(missing_modules) != 0:
			for module in missing_modules:
				info("MISSING MODULE: %s ... Project will not build correctly" % module['id'])
			info("Terminating export: Please fix your modules.")
			sys.exit(1)

		module_search_path, module_asset_dirs = locate_modules(modules, project_dir, resources_dir, info)

		lib_dir = os.path.join(build_dir, 'lib')
		if not os.path.exists(lib_dir):
			os.makedirs(lib_dir)

		if len(module_search_path) > 0:
			info("Copying modules...")
			for module in module_search_path:
				module_name, module_path = module
				info("\t%s..." % module_name)
				shutil.copy(os.path.join(module_path, module_name), lib_dir)
				module[1] = os.path.join(lib_dir, module_name)

			info("Copying module metadata...")
			metadata_dir = os.path.join(build_dir, 'metadata')
			for module in modules:
				module_metadata = os.path.join(module.path,'metadata.json')
				if os.path.exists(module_metadata):
					if not os.path.exists(metadata_dir):
						os.makedirs(metadata_dir)
					target = os.path.join(metadata_dir, "%s.json" % module.manifest.moduleid)
					shutil.copyfile(module_metadata, target)

			# Note: The module link information has to be added to
			# the xcodeproj after it's created.
			# We also have to mod the module_search_path to reference
			# the local 'lib' directory instead of the original
			# module install location
			info("Linking modules...")
			local_modules = []
			for module in module_search_path:
				name = module[0]
				newpath = os.path.join('lib',name)
				local_modules.append([name, newpath])
			link_modules(local_modules, app_name, build_dir, relative=True)

		# Copy libraries
		info("Copying libraries...")
		iphone_dir = os.path.join(sdk_dir, 'iphone')
		for lib in glob.iglob(os.path.join(iphone_dir,'lib*')):
			info("\t%s..." % lib)
			shutil.copy(lib, lib_dir)

		# Process i18n files
		info("Processing i18n...")
		locale_compiler = LocaleCompiler(app_name, project_dir, 'ios', buildMode, resources_dir)
		locale_compiler.compile()

		# Migrate compile scripts
		info("Copying custom Titanium compiler scripts...")
		shutil.copytree(os.path.join(sdk_dir,'common'),titanium_local)
		shutil.copy(os.path.join(sdk_dir,'tiapp.py'),titanium_local)

		iphone_script_dir = os.path.join(titanium_local,'iphone')
		os.mkdir(iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','compiler.py'),iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','tools.py'),iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','run.py'),iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','csspacker.py'),iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','jspacker.py'),iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','titanium_prep'),iphone_script_dir)

		# Add compilation to the build script in project
		info("Modifying pre-compile stage...")
		xcodeproj = os.path.join(build_dir,'%s.xcodeproj' % app_name, 'project.pbxproj')
		contents = codecs.open(xcodeproj,'r',encoding='utf-8').read()

		css_compiler = os.path.join('titanium','css','csscompiler.py')
		ti_compiler = os.path.join('titanium','iphone','compiler.py')
		script = """%s . ios Resources
%s . export-build $TARGETED_DEVICE_FAMILY $SDKROOT %s""" % (css_compiler, ti_compiler, version)
		contents = fix_xcode_script(contents,"Pre-Compile",script)

		# write our new project
		f = codecs.open(xcodeproj,'w',encoding='utf-8')
		f.write(contents)
		f.close()

		info("Finished! Share and Enjoy.")

if __name__ == "__main__":
		main(sys.argv)
		sys.exit(0)
		<commit_msg>Updated per cb1kenobi feedback<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Make an iOS project transportable so that it can be zipped and
# sent to another machine
#

import os, sys, shutil, codecs, glob
template_dir = os.path.abspath(os.path.dirname(sys._getframe(0).f_code.co_filename))

sys.path.append(os.path.abspath(os.path.dirname(template_dir)))
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(template_dir),'module')))
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(template_dir),'common')))

from tiapp import *
from module import ModuleDetector
from pbxproj import PBXProj
from localecompiler import LocaleCompiler
from projector import Projector
from tools import *
from compiler import Compiler

def find_sdk(version):
		sdks = [os.path.join(os.path.expanduser("~/Library/Application Support/Titanium"),"mobilesdk","osx",version),
				os.path.join("/Library","Application Support","Titanium","mobilesdk","osx",version)]

		for sdk in sdks:
			if os.path.exists(sdk):
				return sdk
		print "[ERROR] Is Titanium installed? I can't find it"
		sys.exit(1)

def info(msg):
		print "[INFO] %s" % msg
		sys.stdout.flush()

def main(args):
		if len(args) < 2:
				print "Usage: %s <project_directory> [sdk_verison]" % os.path.basename(args[0])
				sys.exit(1)

		# What needs to be done in order to perform a "true" export?
		# ---
		# Wipe the build dir
		# Migrate resources
		# Migrate tiapp.xml (required for scripts)
		# Generate project from template
		# Populate Info.plist
		# Compile/migrate i18n
		# Migrate scripts for compiling JSS files (and i18n)
		# Modify xcodeproj build steps to call the JSS compiler
		# Then... Share and Enjoy.

		project_dir = os.path.abspath(args[1])
		build_dir = os.path.join(project_dir,'build','iphone')
		titanium_local = os.path.join(build_dir,'titanium')
		buildMode = 'production'

		if len(args) > 2:
			version = args[2]
			sdk_dir = find_sdk(version)
		else:
			sdk_dir = os.path.abspath(os.path.dirname(template_dir))
			version = os.path.basename(sdk_dir)

		if len(args) > 3:
			buildMode = args[3]

		tiappxml = os.path.join(project_dir, 'tiapp.xml')
		tiapp = TiAppXML(tiappxml)

		app_id = tiapp.properties['id']
		app_name = tiapp.properties['name']

		info("Creating project in %s mode" % buildMode)

		if app_id is None or app_name is None:
			info("Your tiapp.xml is malformed - please specify an app name and id")
			sys.exit(1)

		# Clean build dir (if it exists), error otherwise (no iphone support)
		info("Cleaning build...")
		if os.path.exists(build_dir):
			for f in os.listdir(build_dir):
				path = os.path.join(build_dir,f)
				if os.path.isfile(path):
					os.remove(path)
				else:
					shutil.rmtree(path)
		else:
			info("Your project is not configured to be built for iphone.")
			exit(1)

		# Migrate Resources
		info("Migrating resources...")
		project_resources = os.path.join(project_dir, 'Resources')
		resources_dir = os.path.join(build_dir, 'Resources')

		shutil.copytree(project_resources,resources_dir)

		# Migrate platform/iphone contents into Resources.
		info("Migrating platform/iphone to Resources...")
		project_platform = os.path.join(project_dir,'platform','iphone')

		if os.path.isdir(project_platform):
			contents = os.listdir(project_platform)
			for file in contents:
				path = os.path.join(project_platform,file)
				if os.path.isdir(path):
					shutil.copytree(path, os.path.join(resources_dir,file))
				else:
					shutil.copy(path, os.path.join(resources_dir,file))

		# Migrate tiapp.xml
		info("Migrating tiapp.xml...")
		shutil.copy(tiappxml, build_dir)

		ti = TiAppXML(tiappxml)
		# target the requested value if provided
		min_ver = 4.0
		if 'min-ios-ver' in ti.ios:
			min_ver = float(ti.ios['min-ios-ver'])
			if min_ver < 4.0:
				print "[INFO] Minimum iOS version %s is lower than 4.0: Using 4.0 as minimum" % min_ver
				min_ver = 4.0

		# Generate project stuff from the template
		info("Generating project from Titanium template...")
		project = Projector(app_name,version,template_dir,project_dir,app_id, min_ver)
		project.create(template_dir,build_dir)

		# Because the debugger.plist is built as part of the required
		# resources, we need to autogen an empty one
		debug_plist = os.path.join(resources_dir,'debugger.plist')
		force_xcode = write_debugger_plist(None, None, None, None, template_dir, debug_plist)

		# Populate Info.plist
		applogo = None
		info("Populating Info.plist...")
		plist_out = os.path.join(build_dir, 'Info.plist')
		create_info_plist(tiapp, template_dir, project_dir, plist_out)
		applogo = tiapp.generate_infoplist(plist_out, app_id, 'iphone', project_dir, None)

		# Run the compiler to autogenerate .m files
		info("Copying classes, creating generated .m files...")
		compiler = Compiler(project_dir,app_id,app_name,buildMode)
		compiler.compileProject(silent=True)

		#... But we still have to nuke the stuff that gets built that we don't want
		# to bundle.
		ios_build = os.path.join(build_dir,'build')
		if os.path.isdir(ios_build):
			shutil.rmtree(os.path.join(build_dir,'build'))

		# Install applogo/splash/etc.
		info("Copying icons and splash...")
		install_logo(tiapp, applogo, project_dir, template_dir, resources_dir)
		install_defaults(project_dir, template_dir, resources_dir)

		# Get Modules
		detector = ModuleDetector(project_dir)
		missing_modules, modules = detector.find_app_modules(tiapp, 'iphone', buildMode)

		if len(missing_modules) != 0:
			for module in missing_modules:
				info("MISSING MODULE: %s ... Project will not build correctly" % module['id'])
			info("Terminating export: Please fix your modules.")
			sys.exit(1)

		module_search_path, module_asset_dirs = locate_modules(modules, project_dir, resources_dir, info)

		lib_dir = os.path.join(build_dir, 'lib')
		if not os.path.exists(lib_dir):
			os.makedirs(lib_dir)

		if len(module_search_path) > 0:
			info("Copying modules...")
			for module in module_search_path:
				module_name, module_path = module
				info("\t%s..." % module_name)
				shutil.copy(os.path.join(module_path, module_name), lib_dir)
				module[1] = os.path.join(lib_dir, module_name)

			info("Copying module metadata...")
			metadata_dir = os.path.join(build_dir, 'metadata')
			for module in modules:
				module_metadata = os.path.join(module.path,'metadata.json')
				if os.path.exists(module_metadata):
					if not os.path.exists(metadata_dir):
						os.makedirs(metadata_dir)
					target = os.path.join(metadata_dir, "%s.json" % module.manifest.moduleid)
					shutil.copyfile(module_metadata, target)

			# Note: The module link information has to be added to
			# the xcodeproj after it's created.
			# We also have to mod the module_search_path to reference
			# the local 'lib' directory instead of the original
			# module install location
			info("Linking modules...")
			local_modules = []
			for module in module_search_path:
				name = module[0]
				newpath = os.path.join('lib',name)
				local_modules.append([name, newpath])
			link_modules(local_modules, app_name, build_dir, relative=True)

		# Copy libraries
		info("Copying libraries...")
		iphone_dir = os.path.join(sdk_dir, 'iphone')
		for lib in glob.iglob(os.path.join(iphone_dir,'lib*')):
			info("\t%s..." % lib)
			shutil.copy(lib, lib_dir)

		# Process i18n files
		info("Processing i18n...")
		locale_compiler = LocaleCompiler(app_name, project_dir, 'ios', buildMode, resources_dir)
		locale_compiler.compile()

		# Migrate compile scripts
		info("Copying custom Titanium compiler scripts...")
		shutil.copytree(os.path.join(sdk_dir,'common'),titanium_local)
		shutil.copy(os.path.join(sdk_dir,'tiapp.py'),titanium_local)

		iphone_script_dir = os.path.join(titanium_local,'iphone')
		os.mkdir(iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','compiler.py'),iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','tools.py'),iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','run.py'),iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','csspacker.py'),iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','jspacker.py'),iphone_script_dir)
		shutil.copy(os.path.join(sdk_dir,'iphone','titanium_prep'),iphone_script_dir)

		# Add compilation to the build script in project
		info("Modifying pre-compile stage...")
		xcodeproj = os.path.join(build_dir,'%s.xcodeproj' % app_name, 'project.pbxproj')
		contents = codecs.open(xcodeproj,'r',encoding='utf-8').read()

		css_compiler = os.path.join('titanium','css','csscompiler.py')
		ti_compiler = os.path.join('titanium','iphone','compiler.py')
		script = """%s . ios Resources
%s . export-build $TARGETED_DEVICE_FAMILY $SDKROOT %s""" % (css_compiler, ti_compiler, version)
		contents = fix_xcode_script(contents,"Pre-Compile",script)

		# write our new project
		f = codecs.open(xcodeproj,'w',encoding='utf-8')
		f.write(contents)
		f.close()

		info("Finished! Share and Enjoy.")

if __name__ == "__main__":
		main(sys.argv)
		sys.exit(0)
		<|endoftext|>