<commit_before># -*- encoding: utf-8 -*-
##############################################################################
#
#    Copyright (C) 2015 Compassion CH (http://www.compassion.ch)
#    Releasing children from poverty in Jesus' name
#    @author: Stephane Eicher <eicher31@hotmail.com>
#
#    The licence is in the file __openerp__.py
#
##############################################################################
import sys
import base64
import logging

from io import BytesIO
from smb.SMBConnection import SMBConnection

from . import translate_connector

from openerp import models, api, fields, _
from openerp.tools.config import config
from openerp.exceptions import Warning
from openerp.addons.sbc_compassion.models.correspondence_page import \
    BOX_SEPARATOR

logger = logging.getLogger(__name__)

class S2BGenerator(models.Model):
    _inherit = 'correspondence.s2b.generator'

    selection_domain = fields.Char(
        default="[('correspondant_id.category_id', '=', 23),"
                " ('state', '=', 'active')]"
    )

class Correspondence(models.Model):
    """ This class intercepts a letter before it is sent to GMC.
        Letters are pushed to local translation platform if needed.
        """

    _inherit = 'correspondence'

    ##########################################################################
    #                              ORM METHODS                               #
    ##########################################################################
    @api.model
    def create(self, vals):
        """ Create a message for sending the CommKit after be translated on
             the local translate plaform.
        """
        if vals.get('direction') == "Beneficiary To Supporter":
            correspondence = super(Correspondence, self).create(vals)
        else:
            sponsorship = self.env['recurring.contract'].browse(
                vals['sponsorship_id'])

            original_lang = self.env['res.lang.compassion'].browse(
                vals.get('original_language_id'))

            # TODO Remove this fix when HAITI case is resolved
            # For now, we switch French to Creole for avoiding translation
            if 'HA' in sponsorship.child_id.local_id:
                french = self.env.ref(
                    'child_compassion.lang_compassion_french')
                creole = self.env.ref(
                    'child_compassion.lang_compassion_haitien_creole')
                if original_lang == french:
                    vals['original_language_id'] = creole.id

            if original_lang.translatable and original_lang not in sponsorship\
                    .child_id.project_id.field_office_id.spoken_language_ids:
                correspondence = super(Correspondence, self.with_context(
                    no_comm_kit=True)).create(vals)
                correspondence.send_local_translate()
            else:
                correspondence = super(Correspondence, self).create(vals)

        return correspondence

    ##########################################################################
    #                             PUBLIC METHODS                             #
    ##########################################################################
    @api.multi
    def process_letter(self):
        """ Called when B2S letter is Published. Check if translation is
         needed and upload to translation platform. """
        letters_to_send = self.env[self._name]
        for letter in self:
            # if letter.original_language_id in \
            #         letter.supporter_languages_ids or \
            if (letter.beneficiary_language_ids &
                    letter.supporter_languages_ids) or \
                    letter.has_valid_language:
                if super(Correspondence, letter).process_letter():
                    letters_to_send += letter
            else:
                letter.download_attach_letter_image()
                letter.send_local_translate()

        letters_to_send.send_communication()
        return True

    @api.one
    def send_local_translate(self):
        """
        Sends the letter to the local translation platform.
        :return: None
        """
        child = self.sponsorship_id.child_id

        # Specify the src and dst language
        src_lang_id, dst_lang_id = self._get_translation_langs()

        # File name
        sponsor = self.sponsorship_id.partner_id
        # TODO : replace by global_id once all ids are fetched
        file_name = "_".join(
            (child.local_id, sponsor.ref, str(self.id))) + '.pdf'

        # Send letter to local translate platform
        tc = translate_connector.TranslateConnect()
        text_id = tc.upsert_text(
            self, file_name, tc.get_lang_id(src_lang_id),
            tc.get_lang_id(dst_lang_id))
        translation_id = tc.upsert_translation(text_id, self)
        tc.upsert_translation_status(translation_id)

        # Transfer file on the NAS
        self._transfer_file_on_nas(file_name)
        self.state = 'Global Partner translation queue'

    @api.one
    def remove_local_translate(self):
        """
        Remove a letter from local translation platform and change state of
        letter in Odoo
        :return: None
        """
        tc = translate_connector.TranslateConnect()
        tc.remove_translation_with_odoo_id(self.id)
        if self.direction == 'Supporter To Beneficiary':
            self.state = 'Received in the system'
            self.create_commkit()
        else:
            self.state = 'Published to Global Partner'

    @api.one
    def update_translation(self, translate_lang, translate_text, translator):
        """
        Puts the translated text into the correspondence.
        :param translate_lang: code_iso of the language of the translation
        :param translate_text: text of the translation
        :return: None
        """
        translate_lang_id = self.env['res.lang.compassion'].search(
            [('code_iso', '=', translate_lang)]).id
        translator_partner = self.env['res.partner'].search([
            ('ref', '=', translator)])

        if self.direction == 'Supporter To Beneficiary':
            state = 'Received in the system'
            target_text = 'original_text'
            language_field = 'original_language_id'
            # Remove #BOX# in the text, as supporter letters don't have boxes
            translate_text = translate_text.replace(BOX_SEPARATOR, '\n')
            # TODO Remove this fix when HAITI case is resolved
            # For now we switch French to Creole
            if 'HA' in self.child_id.local_id:
                french = self.env.ref(
                    'child_compassion.lang_compassion_french')
                creole = self.env.ref(
                    'child_compassion.lang_compassion_haitien_creole')
                if translate_lang_id == french.id:
                    translate_lang_id = creole.id
        else:
            state = 'Published to Global Partner'
            target_text = 'translated_text'
            language_field = 'translation_language_id'
            # TODO Remove this when new translation tool is working
            # Workaround to avoid overlapping translation : everything goes
            # in L6 template
            self.b2s_layout_id = self.env.ref('sbc_compassion.b2s_l6')

        # Check that layout L4 translation gets on second page
        if self.b2s_layout_id == self.env.ref('sbc_compassion.b2s_l4') and \
                not translate_text.startswith('#PAGE#'):
            translate_text = '#PAGE#' + translate_text
        self.write({
            target_text: translate_text.replace('\r', ''),
            'state': state,
            language_field: translate_lang_id,
            'translator_id': translator_partner.id})

        # Send to GMC
        if self.direction == 'Supporter To Beneficiary':
            self.create_commkit()
        else:
            # Recompose the letter image and process letter
            self.letter_image.unlink()
            if super(Correspondence, self).process_letter():
                self.send_communication()

    ##########################################################################
    #                             PRIVATE METHODS                            #
    ##########################################################################
    def _get_translation_langs(self):
        """
        Finds the source_language et destination_language suited for
        translation of the given letter.

        S2B:
            - src_lang is the original language of the letter
            - dst_lang is the lang of the child if translatable, else
              english

        B2S:
            - src_lang is the original language if translatable, else
              english
            - dst_lang is the main language of the sponsor
        :return: src_lang, dst_lang
        :rtype: res.lang.compassion, res.lang.compassion
        """
        self.ensure_one()
        src_lang_id = False
        dst_lang_id = False
        if self.direction == 'Supporter To Beneficiary':
            # Check that the letter is not yet sent to GMC
            if self.kit_identifier:
                raise Warning(_("Letter already sent to GMC cannot be "
                                "translated! [%s]") % self.kit_identifier)

            src_lang_id = self.original_language_id
            child_langs = self.beneficiary_language_ids.filtered(
                'translatable')
            if child_langs:
                dst_lang_id = child_langs[-1]
            else:
                dst_lang_id = self.env.ref(
                    'child_compassion.lang_compassion_english')

        elif self.direction == 'Beneficiary To Supporter':
            if self.original_language_id and \
                    self.original_language_id.translatable:
                src_lang_id = self.original_language_id
            else:
                src_lang_id = self.env.ref(
                    'child_compassion.lang_compassion_english')
            dst_lang_id = self.supporter_languages_ids.filtered(
                lambda lang: lang.lang_id and lang.lang_id.code ==
                self.correspondant_id.lang)

        return src_lang_id, dst_lang_id

    def _transfer_file_on_nas(self, file_name):
        """
        Puts the letter file on the NAS folder for the translation platform.
        :return: None
        """
        self.ensure_one()
        # Retrieve configuration
        smb_user = config.get('smb_user')
        smb_pass = config.get('smb_pwd')
        smb_ip = config.get('smb_ip')
        smb_port = int(config.get('smb_port', 0))
        if not (smb_user and smb_pass and smb_ip and smb_port):
            raise Exception('No config SMB in file .conf')

        # Copy file in the imported letter folder
        smb_conn = SMBConnection(smb_user, smb_pass, 'openerp', 'nas')
        if smb_conn.connect(smb_ip, smb_port):
            file_ = BytesIO(base64.b64decode(
                self.letter_image.with_context(
                    bin_size=False).datas))
            nas_share_name = self.env.ref(
                'sbc_switzerland.nas_share_name').value

            nas_letters_store_path = self.env.ref(
                'sbc_switzerland.nas_letters_store_path').value + file_name
            smb_conn.storeFile(nas_share_name,
                               nas_letters_store_path, file_)

            logger.info('File {} store on NAS with success'
                        .format(self.letter_image.name))
        else:
            raise Warning(_('Connection to NAS failed'))

    # CRON Methods
    ##############
    @api.model
    def check_local_translation_done(self):
        reload(sys)
        sys.setdefaultencoding('UTF8')
        tc = translate_connector.TranslateConnect()
        letters_to_update = tc.get_translated_letters()

        for letter in letters_to_update:
            correspondence = self.browse(letter["letter_odoo_id"])
            logger.info(".....CHECK TRANSLATION FOR LETTER {}".format(
                correspondence.id))
            if not correspondence.exists():
                logger.warning(("The correspondence id {} doesn't exist in the"
                                "Odoo DB. Remove it manually on MySQL DB. \
                                'todo_id' is set to 5 => 'Pas sur Odoo'")
                               .format(correspondence.id))
                tc.update_translation_to_not_in_odoo(letter["id"])
                continue

            correspondence.update_translation(letter["target_lang"],
                                              letter["text"],
                                              letter["translator"])
            # tc.remove_letter(letter["text_id"])
            # update: don't remove letter but set todo id to 'Traité'
            tc.update_translation_to_treated(letter["id"])
            self.env.cr.commit()
        return True
<commit_msg>Inverse L3 pages<commit_after># -*- encoding: utf-8 -*-
##############################################################################
#
#    Copyright (C) 2015 Compassion CH (http://www.compassion.ch)
#    Releasing children from poverty in Jesus' name
#    @author: Stephane Eicher <eicher31@hotmail.com>
#
#    The licence is in the file __openerp__.py
#
##############################################################################
import sys
import base64
import logging

from io import BytesIO

from pyPdf.pdf import PdfFileReader, PdfFileWriter
from smb.SMBConnection import SMBConnection

from . import translate_connector

from openerp import models, api, fields, _
from openerp.tools.config import config
from openerp.exceptions import Warning
from openerp.addons.sbc_compassion.models.correspondence_page import \
    BOX_SEPARATOR

logger = logging.getLogger(__name__)

class S2BGenerator(models.Model):
    _inherit = 'correspondence.s2b.generator'

    selection_domain = fields.Char(
        default="[('correspondant_id.category_id', '=', 23),"
                " ('state', '=', 'active')]"
    )

class Correspondence(models.Model):
    """ This class intercepts a letter before it is sent to GMC.
        Letters are pushed to local translation platform if needed.
        """

    _inherit = 'correspondence'

    ##########################################################################
    #                              ORM METHODS                               #
    ##########################################################################
    @api.model
    def create(self, vals):
        """ Create a message for sending the CommKit after be translated on
             the local translate plaform.
        """
        if vals.get('direction') == "Beneficiary To Supporter":
            correspondence = super(Correspondence, self).create(vals)
        else:
            sponsorship = self.env['recurring.contract'].browse(
                vals['sponsorship_id'])

            original_lang = self.env['res.lang.compassion'].browse(
                vals.get('original_language_id'))

            # TODO Remove this fix when HAITI case is resolved
            # For now, we switch French to Creole for avoiding translation
            if 'HA' in sponsorship.child_id.local_id:
                french = self.env.ref(
                    'child_compassion.lang_compassion_french')
                creole = self.env.ref(
                    'child_compassion.lang_compassion_haitien_creole')
                if original_lang == french:
                    vals['original_language_id'] = creole.id

            if original_lang.translatable and original_lang not in sponsorship\
                    .child_id.project_id.field_office_id.spoken_language_ids:
                correspondence = super(Correspondence, self.with_context(
                    no_comm_kit=True)).create(vals)
                correspondence.send_local_translate()
            else:
                correspondence = super(Correspondence, self).create(vals)

        # Swap pages for L3 layouts as we scan in wrong order
        if correspondence.template_id.layout == 'CH-A-3S01-1' and \
                correspondence.source != 'compassion':
            input_pdf = PdfFileReader(BytesIO(base64.b64decode(
                correspondence.letter_image.datas)))
            output_pdf = PdfFileWriter()
            nb_pages = input_pdf.numPages
            if nb_pages >= 2:
                output_pdf.addPage(input_pdf.getPage(1))
                output_pdf.addPage(input_pdf.getPage(0))
                if nb_pages > 2:
                    for i in range(2, nb_pages):
                        output_pdf.addPage(input_pdf.getPage(i))
                letter_data = BytesIO()
                output_pdf.write(letter_data)
                letter_data.seek(0)
                correspondence.write({
                    'letter_image': base64.b64encode(letter_data.read())
                })

        return correspondence

    ##########################################################################
    #                             PUBLIC METHODS                             #
    ##########################################################################
    @api.multi
    def process_letter(self):
        """ Called when B2S letter is Published. Check if translation is
         needed and upload to translation platform. """
        letters_to_send = self.env[self._name]
        for letter in self:
            # if letter.original_language_id in \
            #         letter.supporter_languages_ids or \
            if (letter.beneficiary_language_ids &
                    letter.supporter_languages_ids) or \
                    letter.has_valid_language:
                if super(Correspondence, letter).process_letter():
                    letters_to_send += letter
            else:
                letter.download_attach_letter_image()
                letter.send_local_translate()

        letters_to_send.send_communication()
        return True

    @api.one
    def send_local_translate(self):
        """
        Sends the letter to the local translation platform.
        :return: None
        """
        child = self.sponsorship_id.child_id

        # Specify the src and dst language
        src_lang_id, dst_lang_id = self._get_translation_langs()

        # File name
        sponsor = self.sponsorship_id.partner_id
        # TODO : replace by global_id once all ids are fetched
        file_name = "_".join(
            (child.local_id, sponsor.ref, str(self.id))) + '.pdf'

        # Send letter to local translate platform
        tc = translate_connector.TranslateConnect()
        text_id = tc.upsert_text(
            self, file_name, tc.get_lang_id(src_lang_id),
            tc.get_lang_id(dst_lang_id))
        translation_id = tc.upsert_translation(text_id, self)
        tc.upsert_translation_status(translation_id)

        # Transfer file on the NAS
        self._transfer_file_on_nas(file_name)
        self.state = 'Global Partner translation queue'

    @api.one
    def remove_local_translate(self):
        """
        Remove a letter from local translation platform and change state of
        letter in Odoo
        :return: None
        """
        tc = translate_connector.TranslateConnect()
        tc.remove_translation_with_odoo_id(self.id)
        if self.direction == 'Supporter To Beneficiary':
            self.state = 'Received in the system'
            self.create_commkit()
        else:
            self.state = 'Published to Global Partner'

    @api.one
    def update_translation(self, translate_lang, translate_text, translator):
        """
        Puts the translated text into the correspondence.
        :param translate_lang: code_iso of the language of the translation
        :param translate_text: text of the translation
        :return: None
        """
        translate_lang_id = self.env['res.lang.compassion'].search(
            [('code_iso', '=', translate_lang)]).id
        translator_partner = self.env['res.partner'].search([
            ('ref', '=', translator)])

        if self.direction == 'Supporter To Beneficiary':
            state = 'Received in the system'
            target_text = 'original_text'
            language_field = 'original_language_id'
            # Remove #BOX# in the text, as supporter letters don't have boxes
            translate_text = translate_text.replace(BOX_SEPARATOR, '\n')
            # TODO Remove this fix when HAITI case is resolved
            # For now we switch French to Creole
            if 'HA' in self.child_id.local_id:
                french = self.env.ref(
                    'child_compassion.lang_compassion_french')
                creole = self.env.ref(
                    'child_compassion.lang_compassion_haitien_creole')
                if translate_lang_id == french.id:
                    translate_lang_id = creole.id
        else:
            state = 'Published to Global Partner'
            target_text = 'translated_text'
            language_field = 'translation_language_id'
            # TODO Remove this when new translation tool is working
            # Workaround to avoid overlapping translation : everything goes
            # in L6 template
            self.b2s_layout_id = self.env.ref('sbc_compassion.b2s_l6')

        # Check that layout L4 translation gets on second page
        if self.b2s_layout_id == self.env.ref('sbc_compassion.b2s_l4') and \
                not translate_text.startswith('#PAGE#'):
            translate_text = '#PAGE#' + translate_text
        self.write({
            target_text: translate_text.replace('\r', ''),
            'state': state,
            language_field: translate_lang_id,
            'translator_id': translator_partner.id})

        # Send to GMC
        if self.direction == 'Supporter To Beneficiary':
            self.create_commkit()
        else:
            # Recompose the letter image and process letter
            self.letter_image.unlink()
            if super(Correspondence, self).process_letter():
                self.send_communication()

    ##########################################################################
    #                             PRIVATE METHODS                            #
    ##########################################################################
    def _get_translation_langs(self):
        """
        Finds the source_language et destination_language suited for
        translation of the given letter.

        S2B:
            - src_lang is the original language of the letter
            - dst_lang is the lang of the child if translatable, else
              english

        B2S:
            - src_lang is the original language if translatable, else
              english
            - dst_lang is the main language of the sponsor
        :return: src_lang, dst_lang
        :rtype: res.lang.compassion, res.lang.compassion
        """
        self.ensure_one()
        src_lang_id = False
        dst_lang_id = False
        if self.direction == 'Supporter To Beneficiary':
            # Check that the letter is not yet sent to GMC
            if self.kit_identifier:
                raise Warning(_("Letter already sent to GMC cannot be "
                                "translated! [%s]") % self.kit_identifier)

            src_lang_id = self.original_language_id
            child_langs = self.beneficiary_language_ids.filtered(
                'translatable')
            if child_langs:
                dst_lang_id = child_langs[-1]
            else:
                dst_lang_id = self.env.ref(
                    'child_compassion.lang_compassion_english')

        elif self.direction == 'Beneficiary To Supporter':
            if self.original_language_id and \
                    self.original_language_id.translatable:
                src_lang_id = self.original_language_id
            else:
                src_lang_id = self.env.ref(
                    'child_compassion.lang_compassion_english')
            dst_lang_id = self.supporter_languages_ids.filtered(
                lambda lang: lang.lang_id and lang.lang_id.code ==
                self.correspondant_id.lang)

        return src_lang_id, dst_lang_id

    def _transfer_file_on_nas(self, file_name):
        """
        Puts the letter file on the NAS folder for the translation platform.
        :return: None
        """
        self.ensure_one()
        # Retrieve configuration
        smb_user = config.get('smb_user')
        smb_pass = config.get('smb_pwd')
        smb_ip = config.get('smb_ip')
        smb_port = int(config.get('smb_port', 0))
        if not (smb_user and smb_pass and smb_ip and smb_port):
            raise Exception('No config SMB in file .conf')

        # Copy file in the imported letter folder
        smb_conn = SMBConnection(smb_user, smb_pass, 'openerp', 'nas')
        if smb_conn.connect(smb_ip, smb_port):
            file_ = BytesIO(base64.b64decode(
                self.letter_image.with_context(
                    bin_size=False).datas))
            nas_share_name = self.env.ref(
                'sbc_switzerland.nas_share_name').value

            nas_letters_store_path = self.env.ref(
                'sbc_switzerland.nas_letters_store_path').value + file_name
            smb_conn.storeFile(nas_share_name,
                               nas_letters_store_path, file_)

            logger.info('File {} store on NAS with success'
                        .format(self.letter_image.name))
        else:
            raise Warning(_('Connection to NAS failed'))

    # CRON Methods
    ##############
    @api.model
    def check_local_translation_done(self):
        reload(sys)
        sys.setdefaultencoding('UTF8')
        tc = translate_connector.TranslateConnect()
        letters_to_update = tc.get_translated_letters()

        for letter in letters_to_update:
            correspondence = self.browse(letter["letter_odoo_id"])
            logger.info(".....CHECK TRANSLATION FOR LETTER {}".format(
                correspondence.id))
            if not correspondence.exists():
                logger.warning(("The correspondence id {} doesn't exist in the"
                                "Odoo DB. Remove it manually on MySQL DB. \
                                'todo_id' is set to 5 => 'Pas sur Odoo'")
                               .format(correspondence.id))
                tc.update_translation_to_not_in_odoo(letter["id"])
                continue

            correspondence.update_translation(letter["target_lang"],
                                              letter["text"],
                                              letter["translator"])
            # tc.remove_letter(letter["text_id"])
            # update: don't remove letter but set todo id to 'Traité'
            tc.update_translation_to_treated(letter["id"])
            self.env.cr.commit()
        return True
<|endoftext|>


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

# Script to re-package unzipped bags as standard transfers, utilizing checksums from bag manifest.
# Assumes bags are structured as either bag/data/(content) or bag/data/objects/(content).
# Enables use of scripts to add metadata to SIP without failing transfer at bag validation.

from __future__ import print_function, unicode_literals
import os
import shutil
import sys

def main(transfer_path):
    transfer_path = os.path.abspath(transfer_path)

    # check if transfer is an unzipped bag
    if not os.path.isfile(os.path.join(transfer_path, 'bag-info.txt')):
        return 1

    # move files in data up one level if 'objects' folder already exists
    data_path = os.path.join(transfer_path, 'data')
    if os.path.isdir(os.path.join(data_path, 'objects')):
        data_contents = os.listdir(data_path)
        data_contents = [os.path.join(data_path, filename) for filename in data_contents]
        for f in data_contents:
            shutil.move(f, transfer_path)
    # otherwise, rename data to objects
    else:
        os.rename(data_path, os.path.join(transfer_path, 'objects'))

    # create metadata and subdoc folders if don't already exist
    metadata_dir = os.path.join(transfer_path, 'metadata')
    subdoc_dir = os.path.join(metadata_dir, 'submissionDocumentation')
    if not os.path.isdir(metadata_dir):
        os.mkdir(metadata_dir)
    if not os.path.isdir(subdoc_dir):
        os.mkdir(subdoc_dir)

    # write manifest checksums to checksum file
    with open(os.path.join(transfer_path, 'manifest-md5.txt'), 'r') as old_file:
        with open (os.path.join(metadata_dir, 'checksum.md5'), 'w') as new_file:
                for line in old_file:
                    if "data/objects/" in line:
                        new_line = line.replace("data/objects/", "../objects/")
                    else:
                        new_line = line.replace("data/", "../objects/")
                    new_file.write(new_line)

    # move bag files to submissionDocumentation
    for bagfile in 'bag-info.txt', 'bagit.txt', 'manifest-md5.txt', 'tagmanifest-md5.txt':
        shutil.copy2(os.path.join(transfer_path, bagfile), os.path.join(subdoc_dir, bagfile))
        os.remove(os.path.join(transfer_path, bagfile))

    return 0

if __name__ == '__main__':
    transfer_path = sys.argv[1]
    sys.exit(main(transfer_path))
<commit_msg>Replace bag file copy/delete with shutil.move<commit_after>#!/usr/bin/env python

# Script to re-package unzipped bags as standard transfers, utilizing checksums from bag manifest.
# Assumes bags are structured as either bag/data/(content) or bag/data/objects/(content).
# Enables use of scripts to add metadata to SIP without failing transfer at bag validation.

from __future__ import print_function, unicode_literals
import os
import shutil
import sys

def main(transfer_path):
    transfer_path = os.path.abspath(transfer_path)

    # check if transfer is an unzipped bag
    if not os.path.isfile(os.path.join(transfer_path, 'bag-info.txt')):
        return 1

    # move files in data up one level if 'objects' folder already exists
    data_path = os.path.join(transfer_path, 'data')
    if os.path.isdir(os.path.join(data_path, 'objects')):
        data_contents = os.listdir(data_path)
        data_contents = [os.path.join(data_path, filename) for filename in data_contents]
        for f in data_contents:
            shutil.move(f, transfer_path)
    # otherwise, rename data to objects
    else:
        os.rename(data_path, os.path.join(transfer_path, 'objects'))

    # create metadata and subdoc folders if don't already exist
    metadata_dir = os.path.join(transfer_path, 'metadata')
    subdoc_dir = os.path.join(metadata_dir, 'submissionDocumentation')
    if not os.path.isdir(metadata_dir):
        os.mkdir(metadata_dir)
    if not os.path.isdir(subdoc_dir):
        os.mkdir(subdoc_dir)

    # write manifest checksums to checksum file
    with open(os.path.join(transfer_path, 'manifest-md5.txt'), 'r') as old_file:
        with open (os.path.join(metadata_dir, 'checksum.md5'), 'w') as new_file:
                for line in old_file:
                    if "data/objects/" in line:
                        new_line = line.replace("data/objects/", "../objects/")
                    else:
                        new_line = line.replace("data/", "../objects/")
                    new_file.write(new_line)

    # move bag files to submissionDocumentation
    for bagfile in 'bag-info.txt', 'bagit.txt', 'manifest-md5.txt', 'tagmanifest-md5.txt':
        shutil.move(os.path.join(transfer_path, bagfile), subdoc_dir)

    return 0

if __name__ == '__main__':
    transfer_path = sys.argv[1]
    sys.exit(main(transfer_path))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright (c) 2016-2017 Ircam
# Copyright (c) 2016-2017 Guillaume Pellerin
# Copyright (c) 2016-2017 Emilie Zawadzki

# This file is part of mezzanine-organization.

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

__version__ = 1.0
<commit_msg>goto 1.7<commit_after># -*- coding: utf-8 -*-
#
# Copyright (c) 2016-2017 Ircam
# Copyright (c) 2016-2017 Guillaume Pellerin
# Copyright (c) 2016-2017 Emilie Zawadzki

# This file is part of mezzanine-organization.

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

__version__ = 1.7
<|endoftext|>


<commit_before>"""Estimate contamination using VerifyBamID2.
"""
import os
import shutil
import subprocess
import sys

from bcbio import utils
from bcbio.log import logger
from bcbio.bam import ref
from bcbio.distributed.transaction import file_transaction
from bcbio.heterogeneity import chromhacks
from bcbio.pipeline import datadict as dd
from bcbio.provenance import do
from bcbio.variation import samtools

def run(bam_file, data, out_dir):
    out_base = os.path.join(utils.safe_makedir(out_dir),
                            "%s-verifybamid" % (dd.get_sample_name(data)))
    out_file = out_base + ".selfSM"
    failed_file = out_base + ".failed"
    exts = [".out"]
    out = {}
    if not utils.file_exists(out_file) and not utils.file_exists(failed_file):
        with file_transaction(data, out_base) as tx_out_base:
            cmd = ["verifybamid2", "1000g.phase3", "100k", "b38" if dd.get_genome_build(data) == "hg38" else "b37",
                   "--Reference", dd.get_ref_file(data), "--Output", tx_out_base, "--DisableSanityCheck"]
            cmd += _get_input_args(bam_file, data, out_base)
            try:
                do.run(cmd, "VerifyBamID contamination checks")
            except subprocess.CalledProcessError, msg:
                def allowed_errors(l):
                    return (l.find("Insufficient Available markers") >= 0 or
                            l.find("No reads found in any of the regions") >= 0)
                if any([allowed_errors(l) for l in str(msg).split("\n")]):
                    logger.info("Skipping VerifyBamID, not enough overlapping markers found: %s" %
                                dd.get_sample_name(data))
                    with open(failed_file, "w") as out_handle:
                        out_handle.write(str(msg))
                else:
                    logger.warning(str(msg))
                    raise
            else:
                # Fix any sample name problems, for pileups
                shutil.move(tx_out_base + ".selfSM", tx_out_base + ".selfSM.orig")
                with open(tx_out_base + ".selfSM.orig") as in_handle:
                    with open(tx_out_base + ".selfSM", "w") as out_handle:
                        sample_name = None
                        for line in in_handle:
                            if line.startswith("DefaultSampleName"):
                                line = line.replace("DefaultSampleName", dd.get_sample_name(data))
                            # work around bug in finding SM from BAM RG at end of line
                            if len(line.strip().split("\t")) == 1:
                                sample_name = line.strip()
                                line = None
                            elif sample_name:
                                parts = line.split("\t")
                                parts[0] = sample_name
                                line = "\t".join(parts)
                                sample_name = None
                            if line:
                                out_handle.write(line)
                for e in exts + [".selfSM"]:
                    if os.path.exists(tx_out_base + e):
                        shutil.copy(tx_out_base + e, out_base + e)
    if utils.file_exists(out_file):
        out["base"] = out_file
        out["secondary"] = [out_base + e for e in exts if os.path.exists(out_base + e)]
    return out

def _get_input_args(bam_file, data, out_base):
    """Retrieve input args, depending on genome build.

    VerifyBamID2 only handles GRCh37 (1, 2, 3) not hg19, so need to generate
    a pileup for hg19 and fix chromosome naming.
    """
    if dd.get_genome_build(data) in ["hg19"]:
        out_file = "%s-mpileup.txt" % out_base
        if not utils.file_exists(out_file):
            with file_transaction(data, out_file) as tx_out_file:
                mpileup_cl = samtools.prep_mpileup([bam_file], dd.get_ref_file(data), data["config"], want_bcf=False,
                                                   target_regions=_get_autosomal_bed(data, tx_out_file))
                cl = ("{mpileup_cl} | sed 's/^chr//' > {tx_out_file}")
                do.run(cl.format(**locals()), "Create pileup from BAM input")
        return ["--PileupFile", out_file]
    else:
        return ["--BamFile", bam_file]

def _get_autosomal_bed(data, base_file):
    out_file = "%s-stdchroms.bed" % utils.splitext_plus(base_file)[0]
    with open(out_file, "w") as out_handle:
        for r in ref.file_contigs(dd.get_ref_file(data)):
            if chromhacks.is_autosomal(r.name):
                out_handle.write("%s\t0\t%s\n" % (r.name, r.size))
    return out_file

<commit_msg>VerifyBamID: re-enable probe sanity checking and filtering<commit_after>"""Estimate contamination using VerifyBamID2.
"""
import os
import shutil
import subprocess
import sys

from bcbio import utils
from bcbio.log import logger
from bcbio.bam import ref
from bcbio.distributed.transaction import file_transaction
from bcbio.heterogeneity import chromhacks
from bcbio.pipeline import datadict as dd
from bcbio.provenance import do
from bcbio.variation import samtools

def run(bam_file, data, out_dir):
    out_base = os.path.join(utils.safe_makedir(out_dir),
                            "%s-verifybamid" % (dd.get_sample_name(data)))
    out_file = out_base + ".selfSM"
    failed_file = out_base + ".failed"
    exts = [".out"]
    out = {}
    if not utils.file_exists(out_file) and not utils.file_exists(failed_file):
        with file_transaction(data, out_base) as tx_out_base:
            cmd = ["verifybamid2", "1000g.phase3", "100k", "b38" if dd.get_genome_build(data) == "hg38" else "b37",
                   "--Reference", dd.get_ref_file(data), "--Output", tx_out_base]
            cmd += _get_input_args(bam_file, data, out_base)
            try:
                do.run(cmd, "VerifyBamID contamination checks")
            except subprocess.CalledProcessError, msg:
                def allowed_errors(l):
                    return (l.find("Insufficient Available markers") >= 0 or
                            l.find("No reads found in any of the regions") >= 0)
                if any([allowed_errors(l) for l in str(msg).split("\n")]):
                    logger.info("Skipping VerifyBamID, not enough overlapping markers found: %s" %
                                dd.get_sample_name(data))
                    with open(failed_file, "w") as out_handle:
                        out_handle.write(str(msg))
                else:
                    logger.warning(str(msg))
                    raise
            else:
                # Fix any sample name problems, for pileups
                shutil.move(tx_out_base + ".selfSM", tx_out_base + ".selfSM.orig")
                with open(tx_out_base + ".selfSM.orig") as in_handle:
                    with open(tx_out_base + ".selfSM", "w") as out_handle:
                        sample_name = None
                        for line in in_handle:
                            if line.startswith("DefaultSampleName"):
                                line = line.replace("DefaultSampleName", dd.get_sample_name(data))
                            # work around bug in finding SM from BAM RG at end of line
                            if len(line.strip().split("\t")) == 1:
                                sample_name = line.strip()
                                line = None
                            elif sample_name:
                                parts = line.split("\t")
                                parts[0] = sample_name
                                line = "\t".join(parts)
                                sample_name = None
                            if line:
                                out_handle.write(line)
                for e in exts + [".selfSM"]:
                    if os.path.exists(tx_out_base + e):
                        shutil.copy(tx_out_base + e, out_base + e)
    if utils.file_exists(out_file):
        out["base"] = out_file
        out["secondary"] = [out_base + e for e in exts if os.path.exists(out_base + e)]
    return out

def _get_input_args(bam_file, data, out_base):
    """Retrieve input args, depending on genome build.

    VerifyBamID2 only handles GRCh37 (1, 2, 3) not hg19, so need to generate
    a pileup for hg19 and fix chromosome naming.
    """
    if dd.get_genome_build(data) in ["hg19"]:
        out_file = "%s-mpileup.txt" % out_base
        if not utils.file_exists(out_file):
            with file_transaction(data, out_file) as tx_out_file:
                mpileup_cl = samtools.prep_mpileup([bam_file], dd.get_ref_file(data), data["config"], want_bcf=False,
                                                   target_regions=_get_autosomal_bed(data, tx_out_file))
                cl = ("{mpileup_cl} | sed 's/^chr//' > {tx_out_file}")
                do.run(cl.format(**locals()), "Create pileup from BAM input")
        return ["--PileupFile", out_file]
    else:
        return ["--BamFile", bam_file]

def _get_autosomal_bed(data, base_file):
    out_file = "%s-stdchroms.bed" % utils.splitext_plus(base_file)[0]
    with open(out_file, "w") as out_handle:
        for r in ref.file_contigs(dd.get_ref_file(data)):
            if chromhacks.is_autosomal(r.name):
                out_handle.write("%s\t0\t%s\n" % (r.name, r.size))
    return out_file

<|endoftext|>


<commit_before>"""Assess transcript abundance in RNA-seq experiments using Cufflinks.

http://cufflinks.cbcb.umd.edu/manual.html
"""
import os
import tempfile

from bcbio.utils import get_in, file_exists, safe_makedir
from bcbio.distributed.transaction import file_transaction
from bcbio.pipeline import config_utils
from bcbio.provenance import do
from bcbio.rnaseq import gtf, annotate_gtf
import pandas as pd

def run(align_file, ref_file, data):
    config = data["config"]
    cmd = _get_general_options(align_file, config)
    cmd.extend(_get_no_assembly_options(ref_file, data))
    out_dir = _get_output_dir(align_file, data)
    tracking_file = os.path.join(out_dir, "genes.fpkm_tracking")
    fpkm_file = os.path.join(out_dir, data['rgnames']['sample']) + ".fpkm"
    tracking_file_isoform = os.path.join(out_dir, "isoforms.fpkm_tracking")
    fpkm_file_isoform = os.path.join(out_dir, data['rgnames']['sample']) + ".isoform.fpkm"
    if not file_exists(fpkm_file):
        with file_transaction(data, out_dir) as tmp_out_dir:
            cmd.extend(["--output-dir", tmp_out_dir])
            cmd.extend([align_file])
            cmd = map(str, cmd)
            do.run(cmd, "Cufflinks on %s." % (align_file))
        fpkm_file = gene_tracking_to_fpkm(tracking_file, fpkm_file)
        fpkm_file_isoform = gene_tracking_to_fpkm(tracking_file_isoform, fpkm_file_isoform)
    return out_dir, fpkm_file, fpkm_file_isoform

def gene_tracking_to_fpkm(tracking_file, out_file):
    """
    take a gene-level tracking file from Cufflinks and output a two column
    table with the first column as IDs and the second column as FPKM for the
    sample. combines FPKM from the same genes into one FPKM value to fix
    this bug: http://seqanswers.com/forums/showthread.php?t=5224&page=2
    """
    if file_exists(out_file):
        return out_file
    df = pd.io.parsers.read_table(tracking_file, sep="\t", header=0)
    df = df[['tracking_id', 'FPKM']]
    df = df.groupby(['tracking_id']).sum()
    df.to_csv(out_file, sep="\t", header=False, index_label=False)
    return out_file

def _get_general_options(align_file, config):
    options = []
    cufflinks = config_utils.get_program("cufflinks", config)
    options.extend([cufflinks])
    options.extend(["--num-threads", config["algorithm"].get("num_cores", 1)])
    options.extend(["--quiet"])
    options.extend(["--no-update-check"])
    options.extend(["--max-bundle-frags", 2000000])
    options.extend(_get_stranded_flag(config))
    return options

def _get_no_assembly_options(ref_file, data):
    options = []
    options.extend(["--frag-bias-correct", ref_file])
    options.extend(["--multi-read-correct"])
    options.extend(["--upper-quartile-norm"])
    gtf_file = data["genome_resources"]["rnaseq"].get("transcripts", "")
    if gtf_file:
        options.extend(["--GTF", gtf_file])
    mask_file = data["genome_resources"]["rnaseq"].get("transcripts_mask", "")
    if mask_file:
        options.extend(["--mask-file", mask_file])

    return options

def _get_stranded_flag(config):
    strand_flag = {"unstranded": "fr-unstranded",
                   "firststrand": "fr-firststrand",
                   "secondstrand": "fr-secondstrand"}
    stranded = get_in(config, ("algorithm", "strandedness"), "unstranded").lower()
    assert stranded in strand_flag, ("%s is not a valid strandedness value. "
                                     "Valid values are 'firststrand', "
                                     "'secondstrand' and 'unstranded" % (stranded))
    flag = strand_flag[stranded]
    return ["--library-type", flag]

def _get_output_dir(align_file, data, sample_dir=True):
    config = data["config"]
    name = data["rgnames"]["sample"] if sample_dir else ""
    return os.path.join(get_in(data, ("dirs", "work")), "cufflinks", name)

def assemble(bam_file, ref_file, num_cores, out_dir, data):
    safe_makedir(out_dir)
    out_file = os.path.join(out_dir, "assembly", "transcripts.gtf")
    if file_exists(out_file):
        return out_file
    with file_transaction(data, out_dir) as tmp_out_dir:
        cmd = ("cufflinks --output-dir {tmp_out_dir} --num-threads {num_cores} "
               "--frag-bias-correct {ref_file} "
               "--multi-read-correct --upper-quartile-norm {bam_file}")
        cmd = cmd.format(**locals())
        do.run(cmd, "Assembling transcripts with Cufflinks using %s." % bam_file)
    return out_file

def clean_assembly(gtf_file, clean=None, dirty=None):
    """
    clean the likely garbage transcripts from the GTF file including:
    1. any novel single-exon transcripts
    2. any features with an unknown strand
    """
    base, ext = os.path.splitext(gtf_file)
    db = gtf.get_gtf_db(gtf_file, in_memory=True)
    clean = clean if clean else base + ".clean" + ext
    dirty = dirty if dirty else base + ".dirty" + ext
    if file_exists(clean):
        return clean, dirty
    with open(clean, "w") as clean_handle, open(dirty, "w") as dirty_handle:
        for gene in db.features_of_type('gene'):
            for transcript in db.children(gene, level=1):
                if is_likely_noise(db, transcript):
                    write_transcript(db, dirty_handle, transcript)
                else:
                    write_transcript(db, clean_handle, transcript)
    return clean, dirty

def write_transcript(db, handle, transcript):
    for feature in db.children(transcript):
        handle.write(str(feature) + "\n")

def is_likely_noise(db, transcript):
    if is_novel_single_exon(db, transcript):
        return True
    if strand_unknown(db, transcript):
        return True

def strand_unknown(db, transcript):
    """
    for unstranded data with novel transcripts single exon genes
    will have no strand information. single exon novel genes are also
    a source of noise in the Cufflinks assembly so this removes them
    """
    features = list(db.children(transcript))
    strand = features[0].strand
    if strand == ".":
        return True
    else:
        return False

def is_novel_single_exon(db, transcript):
    features = list(db.children(transcript))
    exons = [x for x in features if x.featuretype == "exon"]
    class_code = features[0].attributes.get("class_code", None)[0]
    if len(exons) == 1 and class_code == "u":
        return True
    return False

def fix_cufflinks_attributes(ref_gtf, merged_gtf, data, out_file=None):
    """
    replace the cufflinks gene_id and transcript_id with the
    gene_id and transcript_id from ref_gtf, where available

    """
    ref_db = gtf.get_gtf_db(ref_gtf, in_memory=True)
    merged_db = gtf.get_gtf_db(merged_gtf, in_memory=True)
    base, ext = os.path.splitext(merged_gtf)
    fixed = out_file if out_file else base + ".clean.fixed" + ext
    if file_exists(fixed):
        return fixed

    ref_tid_to_gid = {}
    for gene in ref_db.features_of_type('gene'):
        for transcript in ref_db.children(gene, level=1):
            ref_tid_to_gid[transcript.id] = gene.id

    ctid_to_cgid = {}
    ctid_to_oid = {}
    for gene in merged_db.features_of_type('gene'):
        for transcript in merged_db.children(gene, level=1):
            ctid_to_cgid[transcript.id] = gene.id
            feature = list(merged_db.children(transcript))[0]
            oid = feature.attributes.get("oId", [None])[0]
            if oid:
                ctid_to_oid[transcript.id] = oid
    cgid_to_gid = {}
    for ctid, oid in ctid_to_oid.items():
        cgid = ctid_to_cgid.get(ctid, None)
        oid = ctid_to_oid.get(ctid, None)
        gid = ref_tid_to_gid.get(oid, None) if oid else None
        if cgid and gid:
            cgid_to_gid[cgid] = gid

    with file_transaction(data, fixed) as tmp_fixed_file:
        with open(tmp_fixed_file, "w") as out_handle:
            for gene in merged_db.features_of_type('gene'):
                for transcript in merged_db.children(gene, level=1):
                    for feature in merged_db.children(transcript):
                        cgid = feature.attributes.get("gene_id", [None])[0]
                        gid = cgid_to_gid.get(cgid, None)
                        ctid = None
                        if gid:
                            feature.attributes["gene_id"][0] = gid
                            ctid = feature.attributes.get("transcript_id",
                                                          [None])[0]
                        tid = ctid_to_oid.get(ctid, None)
                        if tid:
                            feature.attributes["transcript_id"][0] = tid
                        if "nearest_ref" in feature.attributes:
                            del feature.attributes["nearest_ref"]
                        if "oId" in feature.attributes:
                            del feature.attributes["oId"]
                        out_handle.write(str(feature) + "\n")
    return fixed

def merge(assembled_gtfs, ref_file, gtf_file, num_cores, data):
    """
    run cuffmerge on a set of assembled GTF files
    """
    assembled_file = tempfile.NamedTemporaryFile(delete=False).name
    with open(assembled_file, "w") as temp_handle:
        for assembled in assembled_gtfs:
            temp_handle.write(assembled + "\n")
    out_dir = os.path.join("assembly", "cuffmerge")
    merged_file = os.path.join(out_dir, "merged.gtf")
    out_file = os.path.join(out_dir, "assembled.gtf")
    if file_exists(out_file):
        return out_file
    with file_transaction(data, out_dir) as tmp_out_dir:
        cmd = ("cuffmerge -o {tmp_out_dir} --ref-gtf {gtf_file} "
               "--num-threads {num_cores} --ref-sequence {ref_file} "
               "{assembled_file}")
        cmd = cmd.format(**locals())
        do.run(cmd, "Merging transcript assemblies with reference.")
    clean, _ = clean_assembly(merged_file)
    fixed = fix_cufflinks_attributes(gtf_file, clean, data)
    classified = annotate_gtf.annotate_novel_coding(fixed, gtf_file, ref_file)
    filtered = annotate_gtf.cleanup_transcripts(classified, gtf_file, ref_file)
    os.rename(filtered, out_file)
    return out_file
<commit_msg>Move files to sample specific directory during assembly.<commit_after>"""Assess transcript abundance in RNA-seq experiments using Cufflinks.

http://cufflinks.cbcb.umd.edu/manual.html
"""
import os
import tempfile

from bcbio.utils import get_in, file_exists, safe_makedir
from bcbio.distributed.transaction import file_transaction
from bcbio.pipeline import config_utils
from bcbio.provenance import do
from bcbio.rnaseq import gtf, annotate_gtf
import pandas as pd

def run(align_file, ref_file, data):
    config = data["config"]
    cmd = _get_general_options(align_file, config)
    cmd.extend(_get_no_assembly_options(ref_file, data))
    out_dir = _get_output_dir(align_file, data)
    tracking_file = os.path.join(out_dir, "genes.fpkm_tracking")
    fpkm_file = os.path.join(out_dir, data['rgnames']['sample']) + ".fpkm"
    tracking_file_isoform = os.path.join(out_dir, "isoforms.fpkm_tracking")
    fpkm_file_isoform = os.path.join(out_dir, data['rgnames']['sample']) + ".isoform.fpkm"
    if not file_exists(fpkm_file):
        with file_transaction(data, out_dir) as tmp_out_dir:
            cmd.extend(["--output-dir", tmp_out_dir])
            cmd.extend([align_file])
            cmd = map(str, cmd)
            do.run(cmd, "Cufflinks on %s." % (align_file))
        fpkm_file = gene_tracking_to_fpkm(tracking_file, fpkm_file)
        fpkm_file_isoform = gene_tracking_to_fpkm(tracking_file_isoform, fpkm_file_isoform)
    return out_dir, fpkm_file, fpkm_file_isoform

def gene_tracking_to_fpkm(tracking_file, out_file):
    """
    take a gene-level tracking file from Cufflinks and output a two column
    table with the first column as IDs and the second column as FPKM for the
    sample. combines FPKM from the same genes into one FPKM value to fix
    this bug: http://seqanswers.com/forums/showthread.php?t=5224&page=2
    """
    if file_exists(out_file):
        return out_file
    df = pd.io.parsers.read_table(tracking_file, sep="\t", header=0)
    df = df[['tracking_id', 'FPKM']]
    df = df.groupby(['tracking_id']).sum()
    df.to_csv(out_file, sep="\t", header=False, index_label=False)
    return out_file

def _get_general_options(align_file, config):
    options = []
    cufflinks = config_utils.get_program("cufflinks", config)
    options.extend([cufflinks])
    options.extend(["--num-threads", config["algorithm"].get("num_cores", 1)])
    options.extend(["--quiet"])
    options.extend(["--no-update-check"])
    options.extend(["--max-bundle-frags", 2000000])
    options.extend(_get_stranded_flag(config))
    return options

def _get_no_assembly_options(ref_file, data):
    options = []
    options.extend(["--frag-bias-correct", ref_file])
    options.extend(["--multi-read-correct"])
    options.extend(["--upper-quartile-norm"])
    gtf_file = data["genome_resources"]["rnaseq"].get("transcripts", "")
    if gtf_file:
        options.extend(["--GTF", gtf_file])
    mask_file = data["genome_resources"]["rnaseq"].get("transcripts_mask", "")
    if mask_file:
        options.extend(["--mask-file", mask_file])

    return options

def _get_stranded_flag(config):
    strand_flag = {"unstranded": "fr-unstranded",
                   "firststrand": "fr-firststrand",
                   "secondstrand": "fr-secondstrand"}
    stranded = get_in(config, ("algorithm", "strandedness"), "unstranded").lower()
    assert stranded in strand_flag, ("%s is not a valid strandedness value. "
                                     "Valid values are 'firststrand', "
                                     "'secondstrand' and 'unstranded" % (stranded))
    flag = strand_flag[stranded]
    return ["--library-type", flag]

def _get_output_dir(align_file, data, sample_dir=True):
    config = data["config"]
    name = data["rgnames"]["sample"] if sample_dir else ""
    return os.path.join(get_in(data, ("dirs", "work")), "cufflinks", name)

def assemble(bam_file, ref_file, num_cores, out_dir, data):
    out_dir = os.path.join(out_dir, data["rgnames"]["sample"])
    safe_makedir(out_dir)
    out_file = os.path.join(out_dir, data["rgnames"]["sample"], "transcripts.gtf")
    if file_exists(out_file):
        return out_file
    with file_transaction(data, out_dir) as tmp_out_dir:
        cmd = ("cufflinks --output-dir {tmp_out_dir} --num-threads {num_cores} "
               "--frag-bias-correct {ref_file} "
               "--multi-read-correct --upper-quartile-norm {bam_file}")
        cmd = cmd.format(**locals())
        do.run(cmd, "Assembling transcripts with Cufflinks using %s." % bam_file)
    return out_file

def clean_assembly(gtf_file, clean=None, dirty=None):
    """
    clean the likely garbage transcripts from the GTF file including:
    1. any novel single-exon transcripts
    2. any features with an unknown strand
    """
    base, ext = os.path.splitext(gtf_file)
    db = gtf.get_gtf_db(gtf_file, in_memory=True)
    clean = clean if clean else base + ".clean" + ext
    dirty = dirty if dirty else base + ".dirty" + ext
    if file_exists(clean):
        return clean, dirty
    with open(clean, "w") as clean_handle, open(dirty, "w") as dirty_handle:
        for gene in db.features_of_type('gene'):
            for transcript in db.children(gene, level=1):
                if is_likely_noise(db, transcript):
                    write_transcript(db, dirty_handle, transcript)
                else:
                    write_transcript(db, clean_handle, transcript)
    return clean, dirty

def write_transcript(db, handle, transcript):
    for feature in db.children(transcript):
        handle.write(str(feature) + "\n")

def is_likely_noise(db, transcript):
    if is_novel_single_exon(db, transcript):
        return True
    if strand_unknown(db, transcript):
        return True

def strand_unknown(db, transcript):
    """
    for unstranded data with novel transcripts single exon genes
    will have no strand information. single exon novel genes are also
    a source of noise in the Cufflinks assembly so this removes them
    """
    features = list(db.children(transcript))
    strand = features[0].strand
    if strand == ".":
        return True
    else:
        return False

def is_novel_single_exon(db, transcript):
    features = list(db.children(transcript))
    exons = [x for x in features if x.featuretype == "exon"]
    class_code = features[0].attributes.get("class_code", None)[0]
    if len(exons) == 1 and class_code == "u":
        return True
    return False

def fix_cufflinks_attributes(ref_gtf, merged_gtf, data, out_file=None):
    """
    replace the cufflinks gene_id and transcript_id with the
    gene_id and transcript_id from ref_gtf, where available

    """
    ref_db = gtf.get_gtf_db(ref_gtf, in_memory=True)
    merged_db = gtf.get_gtf_db(merged_gtf, in_memory=True)
    base, ext = os.path.splitext(merged_gtf)
    fixed = out_file if out_file else base + ".clean.fixed" + ext
    if file_exists(fixed):
        return fixed

    ref_tid_to_gid = {}
    for gene in ref_db.features_of_type('gene'):
        for transcript in ref_db.children(gene, level=1):
            ref_tid_to_gid[transcript.id] = gene.id

    ctid_to_cgid = {}
    ctid_to_oid = {}
    for gene in merged_db.features_of_type('gene'):
        for transcript in merged_db.children(gene, level=1):
            ctid_to_cgid[transcript.id] = gene.id
            feature = list(merged_db.children(transcript))[0]
            oid = feature.attributes.get("oId", [None])[0]
            if oid:
                ctid_to_oid[transcript.id] = oid
    cgid_to_gid = {}
    for ctid, oid in ctid_to_oid.items():
        cgid = ctid_to_cgid.get(ctid, None)
        oid = ctid_to_oid.get(ctid, None)
        gid = ref_tid_to_gid.get(oid, None) if oid else None
        if cgid and gid:
            cgid_to_gid[cgid] = gid

    with file_transaction(data, fixed) as tmp_fixed_file:
        with open(tmp_fixed_file, "w") as out_handle:
            for gene in merged_db.features_of_type('gene'):
                for transcript in merged_db.children(gene, level=1):
                    for feature in merged_db.children(transcript):
                        cgid = feature.attributes.get("gene_id", [None])[0]
                        gid = cgid_to_gid.get(cgid, None)
                        ctid = None
                        if gid:
                            feature.attributes["gene_id"][0] = gid
                            ctid = feature.attributes.get("transcript_id",
                                                          [None])[0]
                        tid = ctid_to_oid.get(ctid, None)
                        if tid:
                            feature.attributes["transcript_id"][0] = tid
                        if "nearest_ref" in feature.attributes:
                            del feature.attributes["nearest_ref"]
                        if "oId" in feature.attributes:
                            del feature.attributes["oId"]
                        out_handle.write(str(feature) + "\n")
    return fixed

def merge(assembled_gtfs, ref_file, gtf_file, num_cores, data):
    """
    run cuffmerge on a set of assembled GTF files
    """
    assembled_file = tempfile.NamedTemporaryFile(delete=False).name
    with open(assembled_file, "w") as temp_handle:
        for assembled in assembled_gtfs:
            temp_handle.write(assembled + "\n")
    out_dir = os.path.join("assembly", "cuffmerge")
    merged_file = os.path.join(out_dir, "merged.gtf")
    out_file = os.path.join(out_dir, "assembled.gtf")
    if file_exists(out_file):
        return out_file
    with file_transaction(data, out_dir) as tmp_out_dir:
        cmd = ("cuffmerge -o {tmp_out_dir} --ref-gtf {gtf_file} "
               "--num-threads {num_cores} --ref-sequence {ref_file} "
               "{assembled_file}")
        cmd = cmd.format(**locals())
        do.run(cmd, "Merging transcript assemblies with reference.")
    clean, _ = clean_assembly(merged_file)
    fixed = fix_cufflinks_attributes(gtf_file, clean, data)
    classified = annotate_gtf.annotate_novel_coding(fixed, gtf_file, ref_file)
    filtered = annotate_gtf.cleanup_transcripts(classified, gtf_file, ref_file)
    os.rename(filtered, out_file)
    return out_file
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright 2016, Major Hayden
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Analyzes the differences between two OpenStack-Ansible commits."""
import argparse
import json
import logging
import os
import sys

from git import Repo
import jinja2
import requests
import yaml

def create_parser():
    """Setup argument Parsing."""
    description = """OpenStack-Ansible Release Diff Generator
----------------------------------------

Finds changes in OpenStack projects and OpenStack-Ansible roles between two
commits in OpenStack-Ansible.

"""

    parser = argparse.ArgumentParser(
        usage='%(prog)s',
        description=description,
        epilog='Licensed "Apache 2.0"',
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        'old_commit',
        action='store',
        nargs=1,
        help="Git SHA of the older commit",
    )
    parser.add_argument(
        'new_commit',
        action='store',
        nargs=1,
        help="Git SHA of the newer commit",
    )
    parser.add_argument(
        '--debug',
        action='store_true',
        default=False,
        help="Enable debug output",
    )
    parser.add_argument(
        '-d', '--directory',
        action='store',
        default="~/.osa-differ",
        help="Git repo storage directory (default: ~/.osa-differ)",
    )
    parser.add_argument(
        '-u', '--update',
        action='store_true',
        default=False,
        help="Fetch latest changes to repo",
    )
    display_opts = parser.add_argument_group("Limit scope")
    display_opts.add_argument(
        "--skip-projects",
        action="store_true",
        help="Skip checking for changes in OpenStack projects"
    )
    display_opts.add_argument(
        "--skip-roles",
        action="store_true",
        help="Skip checking for changes in OpenStack-Ansible roles"
    )
    output_desc = ("Note: Output is always printed to stdout unless "
                   "--quiet is provided.")
    output_opts = parser.add_argument_group('Output options', output_desc)
    output_opts.add_argument(
        '--quiet',
        action='store_true',
        default=False,
        help="Do not output to stdout",
    )
    output_opts.add_argument(
        '--gist',
        action='store_true',
        default=False,
        help="Output into a GitHub Gist",
    )
    output_opts.add_argument(
        '--file',
        metavar="FILENAME",
        action='store',
        help="Output to a file",
    )
    return parser

def get_commits(repo_dir, old_commit, new_commit, hide_merges=True):
    """Find all commits between two commit SHAs."""
    repo = Repo(repo_dir)
    commits = repo.iter_commits(rev="{0}..{1}".format(old_commit, new_commit))
    if hide_merges:
        return [x for x in commits if not x.summary.startswith("Merge ")]
    else:
        return list(commits)

def get_commit_url(repo_url):
    """Determine URL to view commits for repo."""
    if "github.com" in repo_url:
        return repo_url
    if "git.openstack.org" in repo_url:
        uri = '/'.join(repo_url.split('/')[-2:])
        return "https://github.com/{0}".format(uri)

    # If it didn't match these conditions, just return it.
    return repo_url

def get_logger(debug=False):
    """Set up the logger."""
    log_format = "%(asctime)s - %(levelname)s - %(message)s"
    logging.basicConfig(level=logging.WARNING, format=log_format)
    logger = logging.getLogger(__name__)

    if debug:
        logger.setLevel(logging.DEBUG)

    return logger

def get_projects(osa_repo_dir, yaml_files, commit):
    """Get all projects from multiple YAML files."""
    yaml_parsed = []
    for yaml_file in yaml_files:
        # Check out the correct commit SHA from the repository
        repo = Repo(osa_repo_dir)
        repo.head.reference = repo.commit(commit)
        repo.head.reset(index=True, working_tree=True)

        filename = "{0}/{1}".format(osa_repo_dir, yaml_file)
        with open(filename, 'r') as f:
            yaml_parsed.append(yaml.load(f))

    merged_dicts = {k: v for d in yaml_parsed for k, v in d.items()}

    return merged_dicts

def get_roles(osa_repo_dir, commit):
    """Read OSA role information at a particular commit."""
    repo = Repo(osa_repo_dir)
    repo.head.reference = repo.commit(commit)
    repo.head.reset(index=True, working_tree=True)

    filename = "{0}/ansible-role-requirements.yml".format(osa_repo_dir)
    with open(filename, 'r') as f:
        roles_yaml = yaml.load(f)

    return roles_yaml

def parse_arguments():
    """Parse arguments."""
    parser = create_parser()
    args = parser.parse_args()
    return args

def post_gist(report_data, old_sha, new_sha):
    """Post the report to a GitHub Gist and return the URL of the gist."""
    payload = {
        "description": ("Changes in OpenStack-Ansible between "
                        "{0} and {1}".format(old_sha, new_sha)),
        "public": True,
        "files": {
            "osa-diff-{0}-{1}.rst".format(old_sha, new_sha): {
                "content": report_data
            }
        }
    }
    url = "https://api.github.com/gists"
    r = requests.post(url, data=json.dumps(payload))
    response = r.json()
    return response['html_url']

def prepare_storage_dir(storage_directory):
    """Prepare the storage directory."""
    storage_directory = os.path.expanduser(storage_directory)
    if not os.path.exists(storage_directory):
        os.mkdir(storage_directory)

    return storage_directory

def render_template(template_file, template_vars):
    """Render a jinja template."""
    # Load our Jinja templates
    template_dir = "{0}/templates".format(
        os.path.dirname(os.path.abspath(__file__))
    )
    jinja_env = jinja2.Environment(
        loader=jinja2.FileSystemLoader(template_dir),
        trim_blocks=True
    )
    rendered = jinja_env.get_template(template_file).render(template_vars)

    return rendered

def repo_clone(repo_dir, repo_url):
    """Clone repository to this host."""
    repo = Repo.clone_from(repo_url, repo_dir)
    return repo

def repo_pull(repo_dir, repo_url, fetch=False):
    """Reset repository and optionally update it."""
    # Make sure the repository is reset to the master branch.
    repo = Repo(repo_dir)
    repo.heads.master.checkout()
    repo.head.reset(index=True, working_tree=True)

    # Only get the latest updates if requested.
    if fetch:
        repo.remotes['origin'].pull()

    return repo

def update_repo(repo_dir, repo_url, fetch=False):
    """Clone the repo if it doesn't exist already, otherwise update it."""
    repo_exists = os.path.exists(repo_dir)
    if repo_exists:
        repo = repo_pull(repo_dir, repo_url, fetch)
    else:
        repo = repo_clone(repo_dir, repo_url)

    return repo

def validate_commits(repo_dir, commits):
    """Test if a commit is valid for the repository."""
    repo = Repo(repo_dir)
    for commit in commits:
        try:
            commit = repo.commit(commit)
        except:
            raise Exception("Commit {0} could not be found".format(commit))

    return True

def validate_commit_range(repo_dir, old_commit, new_commit):
    """Check if commit range is valid. Flip it if needed."""
    # Are there any commits between the two commits that were provided?
    commits = get_commits(repo_dir, old_commit, new_commit)
    if len(commits) == 0:
        # The user might have gotten their commits out of order. Let's flip
        # the order of the commits and try again.
        commits = get_commits(repo_dir, new_commit, old_commit)
        if len(commits) == 0:
            # Okay, so there really are no commits between the two commits
            # provided by the user. :)
            msg = "The commit range provided is invalid."
            raise Exception(msg)
        else:
            return 'flip'

    return True

def run_osa_differ():
    """The script starts here."""
    # Get our arguments from the command line
    args = parse_arguments()

    # Configure logging
    logger = get_logger(args.debug)

    # Create the storage directory if it doesn't exist already.
    try:
        storage_directory = prepare_storage_dir(args.directory)
    except OSError:
        print("ERROR: Couldn't create the storage directory {0}. "
              "Please create it manually.".format(args.directory))
        sys.exit(1)

    # Prepare the main OpenStack-Ansible repository
    osa_old_commit = args.old_commit[0]
    osa_new_commit = args.new_commit[0]
    osa_repo_url = "https://git.openstack.org/openstack/openstack-ansible"
    osa_repo_dir = "{0}/openstack-ansible".format(storage_directory)
    update_repo(osa_repo_dir, osa_repo_url)

    # Are these commits valid?
    validate_commits(osa_repo_dir, [osa_old_commit, osa_new_commit])

    # Do we have a valid commit range?
    validate_commit_range(osa_repo_dir, osa_old_commit, osa_new_commit)

    # Get the commits in the range
    commits = get_commits(osa_repo_dir, osa_old_commit, osa_new_commit)

    # Start off our report with a header and our OpenStack-Ansible commits.
    template_vars = {
        'args': args,
        'repo': 'openstack-ansible',
        'commits': commits,
        'commit_base_url': get_commit_url(osa_repo_url),
        'old_sha': osa_old_commit,
        'new_sha': osa_new_commit
    }
    report = render_template('offline-header.j2', template_vars)

    # Get the list of OpenStack roles from the newer and older commits.
    role_yaml = get_roles(osa_repo_dir, osa_old_commit)
    role_yaml_latest = get_roles(osa_repo_dir, osa_new_commit)

    # Don't loop through the roles if the user asked us not to do so.
    if args.skip_roles:
        role_yaml_latest = []
    else:
        report += ("OpenStack-Ansible Roles\n"
                   "-----------------------")

    for role in role_yaml_latest:
        # Prepare our repo directory and clone the repo if needed. Only pull
        # if the user requests it.
        repo_dir = "{0}/{1}".format(storage_directory, role['name'])
        repo_url = role['src']
        update_repo(repo_dir, repo_url, args.update)

        # Get the commit SHAs for this role from the older and newer
        # OpenStack-Ansible commits.
        role_old_sha = next(x['version'] for x in role_yaml
                            if x['name'] == role['name'])
        role_new_sha = role['version']

        # Loop through the commits and render our template.
        commits = get_commits(repo_dir, role_old_sha, role_new_sha)
        template_vars = {
            'repo': role['name'],
            'commits': commits,
            'commit_base_url': get_commit_url(repo_url),
            'old_sha': role_old_sha,
            'new_sha': role_new_sha
        }
        rst = render_template('offline-repo-changes.j2', template_vars)
        report += rst

    # Get the list of OpenStack projects from newer commit and older commit.
    yaml_files = [
        'playbooks/defaults/repo_packages/openstack_services.yml',
        'playbooks/defaults/repo_packages/openstack_other.yml'
    ]
    project_yaml = get_projects(osa_repo_dir, yaml_files, osa_old_commit)
    project_yaml_latest = get_projects(osa_repo_dir, yaml_files,
                                       osa_new_commit)

    # Narrow down a list of projects from the latest YAML we retrieved.
    projects = sorted([x[:-9] for x in project_yaml_latest.keys()
                      if x.endswith('git_repo')])

    # Don't loop through the projects if the user asked us not to do so.
    if args.skip_projects:
        project_yaml_latest = []
    else:
        report += ("OpenStack Projects\n"
                   "------------------")

    # Loop through each project and find changes. If we've been asked to update
    # the repository, let's do that too.
    logger.info("Looping through projects to check for changes")
    for project in projects:
        # Prepare our repo directory and clone the repo if needed. Only pull
        # if the user requests it.
        repo_dir = "{0}/{1}".format(storage_directory, project)
        repo_url = project_yaml_latest['{0}_git_repo'.format(project)]
        update_repo(repo_dir, repo_url, args.update)

        # Get the commit SHAs for this project from the older and newer
        # OpenStack-Ansible commits.
        key_for_sha = "{0}_git_install_branch".format(project)
        project_new_sha = project_yaml_latest[key_for_sha]

        # There's a chance that this project was recently added and it might
        # not have existed in the YAML files in the previous OSA release. If
        # that happens, let's just skip it.
        try:
            project_old_sha = project_yaml[key_for_sha]
        except KeyError:
            continue

        # Loop through the commits and render our template.
        commits = get_commits(repo_dir, project_old_sha, project_new_sha)
        template_vars = {
            'repo': project,
            'commits': commits,
            'commit_base_url': get_commit_url(repo_url),
            'old_sha': project_old_sha,
            'new_sha': project_new_sha
        }
        rst = render_template('offline-repo-changes.j2', template_vars)
        report += rst

    # Print the report to stdout unless the user specified --quiet.
    if not args.quiet:
        print(report)

    if args.gist:
        gist_url = post_gist(report, osa_old_commit, osa_new_commit)
        print("Report posted to GitHub Gist: {0}".format(gist_url))

    if args.file is not None:
        print("Report written to file: {0}".format(args.file))
        with open(args.file, 'w') as f:
            f.write(report)

if __name__ == "__main__":
    run_osa_differ()
<commit_msg>Handle newly added roles properly<commit_after>#!/usr/bin/env python
# Copyright 2016, Major Hayden
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Analyzes the differences between two OpenStack-Ansible commits."""
import argparse
import json
import logging
import os
import sys

from git import Repo
import jinja2
import requests
import yaml

def create_parser():
    """Setup argument Parsing."""
    description = """OpenStack-Ansible Release Diff Generator
----------------------------------------

Finds changes in OpenStack projects and OpenStack-Ansible roles between two
commits in OpenStack-Ansible.

"""

    parser = argparse.ArgumentParser(
        usage='%(prog)s',
        description=description,
        epilog='Licensed "Apache 2.0"',
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        'old_commit',
        action='store',
        nargs=1,
        help="Git SHA of the older commit",
    )
    parser.add_argument(
        'new_commit',
        action='store',
        nargs=1,
        help="Git SHA of the newer commit",
    )
    parser.add_argument(
        '--debug',
        action='store_true',
        default=False,
        help="Enable debug output",
    )
    parser.add_argument(
        '-d', '--directory',
        action='store',
        default="~/.osa-differ",
        help="Git repo storage directory (default: ~/.osa-differ)",
    )
    parser.add_argument(
        '-u', '--update',
        action='store_true',
        default=False,
        help="Fetch latest changes to repo",
    )
    display_opts = parser.add_argument_group("Limit scope")
    display_opts.add_argument(
        "--skip-projects",
        action="store_true",
        help="Skip checking for changes in OpenStack projects"
    )
    display_opts.add_argument(
        "--skip-roles",
        action="store_true",
        help="Skip checking for changes in OpenStack-Ansible roles"
    )
    output_desc = ("Note: Output is always printed to stdout unless "
                   "--quiet is provided.")
    output_opts = parser.add_argument_group('Output options', output_desc)
    output_opts.add_argument(
        '--quiet',
        action='store_true',
        default=False,
        help="Do not output to stdout",
    )
    output_opts.add_argument(
        '--gist',
        action='store_true',
        default=False,
        help="Output into a GitHub Gist",
    )
    output_opts.add_argument(
        '--file',
        metavar="FILENAME",
        action='store',
        help="Output to a file",
    )
    return parser

def get_commits(repo_dir, old_commit, new_commit, hide_merges=True):
    """Find all commits between two commit SHAs."""
    repo = Repo(repo_dir)
    commits = repo.iter_commits(rev="{0}..{1}".format(old_commit, new_commit))
    if hide_merges:
        return [x for x in commits if not x.summary.startswith("Merge ")]
    else:
        return list(commits)

def get_commit_url(repo_url):
    """Determine URL to view commits for repo."""
    if "github.com" in repo_url:
        return repo_url
    if "git.openstack.org" in repo_url:
        uri = '/'.join(repo_url.split('/')[-2:])
        return "https://github.com/{0}".format(uri)

    # If it didn't match these conditions, just return it.
    return repo_url

def get_logger(debug=False):
    """Set up the logger."""
    log_format = "%(asctime)s - %(levelname)s - %(message)s"
    logging.basicConfig(level=logging.WARNING, format=log_format)
    logger = logging.getLogger(__name__)

    if debug:
        logger.setLevel(logging.DEBUG)

    return logger

def get_projects(osa_repo_dir, yaml_files, commit):
    """Get all projects from multiple YAML files."""
    yaml_parsed = []
    for yaml_file in yaml_files:
        # Check out the correct commit SHA from the repository
        repo = Repo(osa_repo_dir)
        repo.head.reference = repo.commit(commit)
        repo.head.reset(index=True, working_tree=True)

        filename = "{0}/{1}".format(osa_repo_dir, yaml_file)
        with open(filename, 'r') as f:
            yaml_parsed.append(yaml.load(f))

    merged_dicts = {k: v for d in yaml_parsed for k, v in d.items()}

    return merged_dicts

def get_roles(osa_repo_dir, commit):
    """Read OSA role information at a particular commit."""
    repo = Repo(osa_repo_dir)
    repo.head.reference = repo.commit(commit)
    repo.head.reset(index=True, working_tree=True)

    filename = "{0}/ansible-role-requirements.yml".format(osa_repo_dir)
    with open(filename, 'r') as f:
        roles_yaml = yaml.load(f)

    return roles_yaml

def parse_arguments():
    """Parse arguments."""
    parser = create_parser()
    args = parser.parse_args()
    return args

def post_gist(report_data, old_sha, new_sha):
    """Post the report to a GitHub Gist and return the URL of the gist."""
    payload = {
        "description": ("Changes in OpenStack-Ansible between "
                        "{0} and {1}".format(old_sha, new_sha)),
        "public": True,
        "files": {
            "osa-diff-{0}-{1}.rst".format(old_sha, new_sha): {
                "content": report_data
            }
        }
    }
    url = "https://api.github.com/gists"
    r = requests.post(url, data=json.dumps(payload))
    response = r.json()
    return response['html_url']

def prepare_storage_dir(storage_directory):
    """Prepare the storage directory."""
    storage_directory = os.path.expanduser(storage_directory)
    if not os.path.exists(storage_directory):
        os.mkdir(storage_directory)

    return storage_directory

def render_template(template_file, template_vars):
    """Render a jinja template."""
    # Load our Jinja templates
    template_dir = "{0}/templates".format(
        os.path.dirname(os.path.abspath(__file__))
    )
    jinja_env = jinja2.Environment(
        loader=jinja2.FileSystemLoader(template_dir),
        trim_blocks=True
    )
    rendered = jinja_env.get_template(template_file).render(template_vars)

    return rendered

def repo_clone(repo_dir, repo_url):
    """Clone repository to this host."""
    repo = Repo.clone_from(repo_url, repo_dir)
    return repo

def repo_pull(repo_dir, repo_url, fetch=False):
    """Reset repository and optionally update it."""
    # Make sure the repository is reset to the master branch.
    repo = Repo(repo_dir)
    repo.heads.master.checkout()
    repo.head.reset(index=True, working_tree=True)

    # Only get the latest updates if requested.
    if fetch:
        repo.remotes['origin'].pull()

    return repo

def update_repo(repo_dir, repo_url, fetch=False):
    """Clone the repo if it doesn't exist already, otherwise update it."""
    repo_exists = os.path.exists(repo_dir)
    if repo_exists:
        repo = repo_pull(repo_dir, repo_url, fetch)
    else:
        repo = repo_clone(repo_dir, repo_url)

    return repo

def validate_commits(repo_dir, commits):
    """Test if a commit is valid for the repository."""
    repo = Repo(repo_dir)
    for commit in commits:
        try:
            commit = repo.commit(commit)
        except:
            raise Exception("Commit {0} could not be found".format(commit))

    return True

def validate_commit_range(repo_dir, old_commit, new_commit):
    """Check if commit range is valid. Flip it if needed."""
    # Are there any commits between the two commits that were provided?
    commits = get_commits(repo_dir, old_commit, new_commit)
    if len(commits) == 0:
        # The user might have gotten their commits out of order. Let's flip
        # the order of the commits and try again.
        commits = get_commits(repo_dir, new_commit, old_commit)
        if len(commits) == 0:
            # Okay, so there really are no commits between the two commits
            # provided by the user. :)
            msg = "The commit range provided is invalid."
            raise Exception(msg)
        else:
            return 'flip'

    return True

def run_osa_differ():
    """The script starts here."""
    # Get our arguments from the command line
    args = parse_arguments()

    # Configure logging
    logger = get_logger(args.debug)

    # Create the storage directory if it doesn't exist already.
    try:
        storage_directory = prepare_storage_dir(args.directory)
    except OSError:
        print("ERROR: Couldn't create the storage directory {0}. "
              "Please create it manually.".format(args.directory))
        sys.exit(1)

    # Prepare the main OpenStack-Ansible repository
    osa_old_commit = args.old_commit[0]
    osa_new_commit = args.new_commit[0]
    osa_repo_url = "https://git.openstack.org/openstack/openstack-ansible"
    osa_repo_dir = "{0}/openstack-ansible".format(storage_directory)
    update_repo(osa_repo_dir, osa_repo_url)

    # Are these commits valid?
    validate_commits(osa_repo_dir, [osa_old_commit, osa_new_commit])

    # Do we have a valid commit range?
    validate_commit_range(osa_repo_dir, osa_old_commit, osa_new_commit)

    # Get the commits in the range
    commits = get_commits(osa_repo_dir, osa_old_commit, osa_new_commit)

    # Start off our report with a header and our OpenStack-Ansible commits.
    template_vars = {
        'args': args,
        'repo': 'openstack-ansible',
        'commits': commits,
        'commit_base_url': get_commit_url(osa_repo_url),
        'old_sha': osa_old_commit,
        'new_sha': osa_new_commit
    }
    report = render_template('offline-header.j2', template_vars)

    # Get the list of OpenStack roles from the newer and older commits.
    role_yaml = get_roles(osa_repo_dir, osa_old_commit)
    role_yaml_latest = get_roles(osa_repo_dir, osa_new_commit)

    # Don't loop through the roles if the user asked us not to do so.
    if args.skip_roles:
        role_yaml_latest = []
    else:
        report += ("OpenStack-Ansible Roles\n"
                   "-----------------------")

    for role in role_yaml_latest:
        # Prepare our repo directory and clone the repo if needed. Only pull
        # if the user requests it.
        repo_dir = "{0}/{1}".format(storage_directory, role['name'])
        repo_url = role['src']
        update_repo(repo_dir, repo_url, args.update)

        # Get the commit SHAs for this role from the older and newer
        # OpenStack-Ansible commits. If we can't get the older commit SHA, it
        # could be because the role is newly added and there wasn't an older
        # revision. In that case, skip it.
        role_new_sha = role['version']
        try:
            role_old_sha = next(x['version'] for x in role_yaml
                                if x['name'] == role['name'])
        except:
            continue

        # Loop through the commits and render our template.
        commits = get_commits(repo_dir, role_old_sha, role_new_sha)
        template_vars = {
            'repo': role['name'],
            'commits': commits,
            'commit_base_url': get_commit_url(repo_url),
            'old_sha': role_old_sha,
            'new_sha': role_new_sha
        }
        rst = render_template('offline-repo-changes.j2', template_vars)
        report += rst

    # Get the list of OpenStack projects from newer commit and older commit.
    yaml_files = [
        'playbooks/defaults/repo_packages/openstack_services.yml',
        'playbooks/defaults/repo_packages/openstack_other.yml'
    ]
    project_yaml = get_projects(osa_repo_dir, yaml_files, osa_old_commit)
    project_yaml_latest = get_projects(osa_repo_dir, yaml_files,
                                       osa_new_commit)

    # Narrow down a list of projects from the latest YAML we retrieved.
    projects = sorted([x[:-9] for x in project_yaml_latest.keys()
                      if x.endswith('git_repo')])

    # Don't loop through the projects if the user asked us not to do so.
    if args.skip_projects:
        project_yaml_latest = []
    else:
        report += ("OpenStack Projects\n"
                   "------------------")

    # Loop through each project and find changes. If we've been asked to update
    # the repository, let's do that too.
    logger.info("Looping through projects to check for changes")
    for project in projects:
        # Prepare our repo directory and clone the repo if needed. Only pull
        # if the user requests it.
        repo_dir = "{0}/{1}".format(storage_directory, project)
        repo_url = project_yaml_latest['{0}_git_repo'.format(project)]
        update_repo(repo_dir, repo_url, args.update)

        # Get the commit SHAs for this project from the older and newer
        # OpenStack-Ansible commits.
        key_for_sha = "{0}_git_install_branch".format(project)
        project_new_sha = project_yaml_latest[key_for_sha]

        # There's a chance that this project was recently added and it might
        # not have existed in the YAML files in the previous OSA release. If
        # that happens, let's just skip it.
        try:
            project_old_sha = project_yaml[key_for_sha]
        except KeyError:
            continue

        # Loop through the commits and render our template.
        commits = get_commits(repo_dir, project_old_sha, project_new_sha)
        template_vars = {
            'repo': project,
            'commits': commits,
            'commit_base_url': get_commit_url(repo_url),
            'old_sha': project_old_sha,
            'new_sha': project_new_sha
        }
        rst = render_template('offline-repo-changes.j2', template_vars)
        report += rst

    # Print the report to stdout unless the user specified --quiet.
    if not args.quiet:
        print(report)

    if args.gist:
        gist_url = post_gist(report, osa_old_commit, osa_new_commit)
        print("Report posted to GitHub Gist: {0}".format(gist_url))

    if args.file is not None:
        print("Report written to file: {0}".format(args.file))
        with open(args.file, 'w') as f:
            f.write(report)

if __name__ == "__main__":
    run_osa_differ()
<|endoftext|>