<commit_before>import logging
from typing import Union

import flask
from flask import Blueprint
from flask import Response
from flask import flash
from flask import redirect
from flask import render_template
from flask import request
from flask import url_for
from flask_login import login_required, current_user
from pyswagger.contrib.client import flask
from sqlalchemy import asc
from sqlalchemy import or_
from sqlalchemy.orm import joinedload

from waitlist import db
from waitlist.blueprints.settings import add_menu_entry
from waitlist.permissions import perm_manager
from waitlist.permissions.manager import StaticPermissions, StaticRoles
from waitlist.signal.signals import send_account_created, send_roles_changed, send_account_status_change
from waitlist.storage.database import Account, Character, Role, linked_chars, APICacheCharacterInfo
from waitlist.utility.settings import sget_resident_mail, sget_tbadge_mail, sget_other_mail, sget_other_topic, \
    sget_tbadge_topic, sget_resident_topic
from waitlist.utility.utils import get_random_token

from waitlist.utility import outgate

bp = Blueprint('accounts', __name__)
logger = logging.getLogger(__name__)

perm_manager.define_permission('include_in_accountlist')
perm_manager.define_permission('accounts_edit')
perm_manager.define_permission('settings_access')
perm_manager.define_permission('admin')
perm_manager.define_permission('accounts_download_list')

@bp.route("/", methods=["GET", "POST"])
@login_required
@perm_manager.require('accounts_edit')
def accounts():
    if request.method == "POST":
        acc_name = request.form['account_name']

        acc_roles = request.form.getlist('account_roles')

        char_name = request.form['default_char_name']
        char_name = char_name.strip()

        note = request.form['change_note'].strip()

        char_info = outgate.character.get_info_by_name(char_name)
        if char_info is None:
            flash(f"A Character named {char_name} does not exist!")
        else:
            char_id = char_info.id
            acc = Account()
            acc.username = acc_name

            acc.login_token = get_random_token(16)

            if len(acc_roles) > 0:
                db_roles = db.session.query(Role).filter(or_(Role.name == name for name in acc_roles)).all()
                for role in db_roles:
                    acc.roles.append(role)

            db.session.add(acc)

            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            acc.characters.append(character)

            db.session.flush()

            acc.current_char = char_id

            db.session.commit()
            send_account_created(accounts, acc.id, current_user.id, acc_roles, 'Creating account. ' + note)

    roles = db.session.query(Role).order_by(Role.name).all()
    accs = db.session.query(Account).order_by(asc(Account.disabled)).order_by(Account.username).all()
    mails = {
        'resident': [sget_resident_mail(), sget_resident_topic()],
        'tbadge': [sget_tbadge_mail(), sget_tbadge_topic()],
        'other': [sget_other_mail(), sget_other_topic()]
    }

    return render_template("settings/accounts.html", roles=roles, accounts=accs, mails=mails)

@bp.route("/edit", methods=["POST"])
@login_required
@perm_manager.require('accounts_edit')
def account_edit():
    acc_id = int(request.form['account_id'])
    acc_name = request.form['account_name']

    note = request.form['change_note'].strip()

    acc_roles = request.form.getlist('account_roles')

    char_name = request.form['default_char_name']
    char_name = char_name.strip()
    if char_name == "":
        char_name = None

    acc = db.session.query(Account).filter(Account.id == acc_id).first()
    if acc is None:
        return flask.abort(400)

    if acc.username != acc_name:
        acc.username = acc_name

    # if there are roles, add new ones, remove the ones that aren't there
    if len(acc_roles) > 0:
        roles_new = {}
        for r in acc_roles:
            roles_new[r] = True

        # db_roles = session.query(Role).filter(or_(Role.name == name for name in acc_roles)).all()
        roles_to_remove = []
        for role in acc.roles:
            if role.name in roles_new:
                del roles_new[role.name]  # remove because it is already in the db
            else:
                # remove the roles because it not submitted anymore
                # only remove admin if current user is an admin
                if role.name == StaticRoles.ADMIN and not perm_manager.get_permission(StaticPermissions.ADMIN).can():
                    continue
                roles_to_remove.append(role)  # mark for removal

        for role in roles_to_remove:
            acc.roles.remove(role)

        # if it is not an admin remove admin role from new roles
        if not perm_manager.get_permission('admin').can():
            if 'admin' in roles_new:
                del roles_new['admin']

        # add remaining roles
        if len(roles_new) > 0:
            new_db_roles = db.session.query(Role).filter(or_(Role.name == name for name in roles_new))
            for role in new_db_roles:
                acc.roles.append(role)
        if len(roles_new) > 0 or len(roles_to_remove) > 0:
            send_roles_changed(account_edit, acc.id, current_user.id, [x for x in roles_new],
                               [x.name for x in roles_to_remove], note)
    else:
        # make sure all roles are removed
        roles_to_remove = []
        for role in acc.roles:
            # only remove admin if current user is an admin
            if role.name == StaticRoles.ADMIN and not perm_manager.get_permission(StaticPermissions.ADMIN).can():
                continue
            roles_to_remove.append(role)

        if len(roles_to_remove) > 0:
            for role in roles_to_remove:
                acc.roles.remove(role)
            db.session.flush()
            send_roles_changed(account_edit, acc.id, current_user.id, [x.name for x in roles_to_remove], [], note)

    if char_name is not None:
        char_info = outgate.character.get_info_by_name(char_name)
        if char_info is None:
            flash(f"Character with name {char_name} could not be found!")
        else:
            char_id = char_info.id
            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                # lets make sure we have the correct name (case)
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            # check if character is linked to this account
            link = db.session.query(linked_chars) \
                .filter((linked_chars.c.id == acc_id) & (linked_chars.c.char_id == char_id)).first()
            if link is None:
                acc.characters.append(character)

            db.session.flush()
            acc.current_char = char_id

    db.session.commit()
    return redirect(url_for('.accounts'), code=303)

@bp.route("/self_edit", methods=["POST"])
@login_required
@perm_manager.require('settings_access')
def account_self_edit():
    acc_id = current_user.id

    char_name = request.form['default_char_name']
    char_name = char_name.strip()
    if char_name == "":
        char_name = None

    acc = db.session.query(Account).filter(Account.id == acc_id).first()
    if acc is None:
        return flask.abort(400)

    if char_name is not None:
        char_info = outgate.character.get_info_by_name(char_name)

        if char_info is None:
            flash("Character with name {char_name} could not be found!")
        else:
            char_id = char_info.id
            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                # lets make sure we have the correct name (case)
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            # check if character is linked to this account
            link = db.session.query(linked_chars) \
                .filter((linked_chars.c.id == acc_id) & (linked_chars.c.char_id == char_id)).first()
            if link is None:
                acc.characters.append(character)

            db.session.flush()
            if acc.current_char != char_id:
                # remove all the access tokens
                if acc.ssoToken is not None:
                    db.session.delete(acc.ssoToken)
                acc.current_char = char_id

    db.session.commit()
    return redirect(url_for('.account_self'), code=303)

@bp.route("/self", methods=["GET"])
@login_required
@perm_manager.require('settings_access')
def account_self():
    acc = db.session.query(Account).filter(Account.id == current_user.id).first()
    return render_template("settings/self.html", account=acc)

@bp.route("/api/account/disabled", methods=['POST'])
@login_required
@perm_manager.require('settings_access')
def account_disabled():
    accid: int = int(request.form['id'])
    acc: Account = db.session.query(Account).filter(Account.id == accid).first()
    status: Union(str, bool) = request.form['disabled']
    send_account_status_change(account_disabled, acc.id, current_user.id, status)
    logger.info("%s sets account %s to %s", current_user.username, acc.username, status)
    if status == 'false':
        status = False
    else:
        status = True

    acc.disabled = status
    db.session.commit()
    return "OK"

@bp.route("/api/account/<int:acc_id>", methods=["DELETE"])
@login_required
@perm_manager.require('admin')
def api_account_delete(acc_id: int) -> Response:
    db.session.query(Account).filter(Account.id == acc_id).delete()
    db.session.commit()
    return flask.jsonify(status="OK")

@bp.route('/accounts/downloadlist/cvs')
@login_required
@perm_manager.require('accounts_download_list')
def accounts_download_csv() -> Response:
    def iter_accs(data):
        for account in data:
            for ci, char in enumerate(account.characters):
                if ci > 0:
                    yield ", " + char.eve_name
                else:
                    yield char.eve_name
            yield '\n'

    permission = perm_manager.get_permission('include_in_accountlist')
    role_check = not Account.disabled
    for role_need in permission.needs:
        if role_need.method != 'role':
            continue
        role_check &= Role.name == role_need.value

    # noinspection PyPep8
    accs = db.session.query(Account).options(joinedload('characters')).join(Account.roles).filter(
        role_check ).order_by(
        Account.username).all()

    response = Response(iter_accs(accs), mimetype='text/csv')
    response.headers['Content-Disposition'] = 'attachment; filename=accounts.csv'
    return response

add_menu_entry('accounts.accounts', 'Accounts', perm_manager.get_permission('accounts_edit').can)
add_menu_entry('accounts.account_self', 'Own Settings', lambda: True)
<commit_msg>fixed removed roles added to profile as being grated in some circumstances<commit_after>import logging
from typing import Union

import flask
from flask import Blueprint
from flask import Response
from flask import flash
from flask import redirect
from flask import render_template
from flask import request
from flask import url_for
from flask_login import login_required, current_user
from pyswagger.contrib.client import flask
from sqlalchemy import asc
from sqlalchemy import or_
from sqlalchemy.orm import joinedload

from waitlist import db
from waitlist.blueprints.settings import add_menu_entry
from waitlist.permissions import perm_manager
from waitlist.permissions.manager import StaticPermissions, StaticRoles
from waitlist.signal.signals import send_account_created, send_roles_changed, send_account_status_change
from waitlist.storage.database import Account, Character, Role, linked_chars, APICacheCharacterInfo
from waitlist.utility.settings import sget_resident_mail, sget_tbadge_mail, sget_other_mail, sget_other_topic, \
    sget_tbadge_topic, sget_resident_topic
from waitlist.utility.utils import get_random_token

from waitlist.utility import outgate

bp = Blueprint('accounts', __name__)
logger = logging.getLogger(__name__)

perm_manager.define_permission('include_in_accountlist')
perm_manager.define_permission('accounts_edit')
perm_manager.define_permission('settings_access')
perm_manager.define_permission('admin')
perm_manager.define_permission('accounts_download_list')

@bp.route("/", methods=["GET", "POST"])
@login_required
@perm_manager.require('accounts_edit')
def accounts():
    if request.method == "POST":
        acc_name = request.form['account_name']

        acc_roles = request.form.getlist('account_roles')

        char_name = request.form['default_char_name']
        char_name = char_name.strip()

        note = request.form['change_note'].strip()

        char_info = outgate.character.get_info_by_name(char_name)
        if char_info is None:
            flash(f"A Character named {char_name} does not exist!")
        else:
            char_id = char_info.id
            acc = Account()
            acc.username = acc_name

            acc.login_token = get_random_token(16)

            if len(acc_roles) > 0:
                db_roles = db.session.query(Role).filter(or_(Role.name == name for name in acc_roles)).all()
                for role in db_roles:
                    acc.roles.append(role)

            db.session.add(acc)

            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            acc.characters.append(character)

            db.session.flush()

            acc.current_char = char_id

            db.session.commit()
            send_account_created(accounts, acc.id, current_user.id, acc_roles, 'Creating account. ' + note)

    roles = db.session.query(Role).order_by(Role.name).all()
    accs = db.session.query(Account).order_by(asc(Account.disabled)).order_by(Account.username).all()
    mails = {
        'resident': [sget_resident_mail(), sget_resident_topic()],
        'tbadge': [sget_tbadge_mail(), sget_tbadge_topic()],
        'other': [sget_other_mail(), sget_other_topic()]
    }

    return render_template("settings/accounts.html", roles=roles, accounts=accs, mails=mails)

@bp.route("/edit", methods=["POST"])
@login_required
@perm_manager.require('accounts_edit')
def account_edit():
    acc_id = int(request.form['account_id'])
    acc_name = request.form['account_name']

    note = request.form['change_note'].strip()

    acc_roles = request.form.getlist('account_roles')

    char_name = request.form['default_char_name']
    char_name = char_name.strip()
    if char_name == "":
        char_name = None

    acc = db.session.query(Account).filter(Account.id == acc_id).first()
    if acc is None:
        return flask.abort(400)

    if acc.username != acc_name:
        acc.username = acc_name

    # if there are roles, add new ones, remove the ones that aren't there
    if len(acc_roles) > 0:
        roles_new = {}
        for r in acc_roles:
            roles_new[r] = True

        # db_roles = session.query(Role).filter(or_(Role.name == name for name in acc_roles)).all()
        roles_to_remove = []
        for role in acc.roles:
            if role.name in roles_new:
                del roles_new[role.name]  # remove because it is already in the db
            else:
                # remove the roles because it not submitted anymore
                # only remove admin if current user is an admin
                if role.name == StaticRoles.ADMIN and not perm_manager.get_permission(StaticPermissions.ADMIN).can():
                    continue
                roles_to_remove.append(role)  # mark for removal

        for role in roles_to_remove:
            acc.roles.remove(role)

        # if it is not an admin remove admin role from new roles
        if not perm_manager.get_permission('admin').can():
            if 'admin' in roles_new:
                del roles_new['admin']

        # add remaining roles
        if len(roles_new) > 0:
            new_db_roles = db.session.query(Role).filter(or_(Role.name == name for name in roles_new))
            for role in new_db_roles:
                acc.roles.append(role)
        if len(roles_new) > 0 or len(roles_to_remove) > 0:
            send_roles_changed(account_edit, acc.id, current_user.id, [x for x in roles_new],
                               [x.name for x in roles_to_remove], note)
    else:
        # make sure all roles are removed
        roles_to_remove = []
        for role in acc.roles:
            # only remove admin if current user is an admin
            if role.name == StaticRoles.ADMIN and not perm_manager.get_permission(StaticPermissions.ADMIN).can():
                continue
            roles_to_remove.append(role)

        if len(roles_to_remove) > 0:
            for role in roles_to_remove:
                acc.roles.remove(role)
            db.session.flush()
            send_roles_changed(account_edit, acc.id, current_user.id, [], [x.name for x in roles_to_remove], note)

    if char_name is not None:
        char_info = outgate.character.get_info_by_name(char_name)
        if char_info is None:
            flash(f"Character with name {char_name} could not be found!")
        else:
            char_id = char_info.id
            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                # lets make sure we have the correct name (case)
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            # check if character is linked to this account
            link = db.session.query(linked_chars) \
                .filter((linked_chars.c.id == acc_id) & (linked_chars.c.char_id == char_id)).first()
            if link is None:
                acc.characters.append(character)

            db.session.flush()
            acc.current_char = char_id

    db.session.commit()
    return redirect(url_for('.accounts'), code=303)

@bp.route("/self_edit", methods=["POST"])
@login_required
@perm_manager.require('settings_access')
def account_self_edit():
    acc_id = current_user.id

    char_name = request.form['default_char_name']
    char_name = char_name.strip()
    if char_name == "":
        char_name = None

    acc = db.session.query(Account).filter(Account.id == acc_id).first()
    if acc is None:
        return flask.abort(400)

    if char_name is not None:
        char_info = outgate.character.get_info_by_name(char_name)

        if char_info is None:
            flash("Character with name {char_name} could not be found!")
        else:
            char_id = char_info.id
            # find out if there is a character like that in the database
            character = db.session.query(Character).filter(Character.id == char_id).first()

            if character is None:
                # lets make sure we have the correct name (case)
                char_info: APICacheCharacterInfo = outgate.character.get_info(char_id)
                character = Character()
                character.eve_name = char_info.characterName
                character.id = char_id

            # check if character is linked to this account
            link = db.session.query(linked_chars) \
                .filter((linked_chars.c.id == acc_id) & (linked_chars.c.char_id == char_id)).first()
            if link is None:
                acc.characters.append(character)

            db.session.flush()
            if acc.current_char != char_id:
                # remove all the access tokens
                if acc.ssoToken is not None:
                    db.session.delete(acc.ssoToken)
                acc.current_char = char_id

    db.session.commit()
    return redirect(url_for('.account_self'), code=303)

@bp.route("/self", methods=["GET"])
@login_required
@perm_manager.require('settings_access')
def account_self():
    acc = db.session.query(Account).filter(Account.id == current_user.id).first()
    return render_template("settings/self.html", account=acc)

@bp.route("/api/account/disabled", methods=['POST'])
@login_required
@perm_manager.require('settings_access')
def account_disabled():
    accid: int = int(request.form['id'])
    acc: Account = db.session.query(Account).filter(Account.id == accid).first()
    status: Union(str, bool) = request.form['disabled']
    send_account_status_change(account_disabled, acc.id, current_user.id, status)
    logger.info("%s sets account %s to %s", current_user.username, acc.username, status)
    if status == 'false':
        status = False
    else:
        status = True

    acc.disabled = status
    db.session.commit()
    return "OK"

@bp.route("/api/account/<int:acc_id>", methods=["DELETE"])
@login_required
@perm_manager.require('admin')
def api_account_delete(acc_id: int) -> Response:
    db.session.query(Account).filter(Account.id == acc_id).delete()
    db.session.commit()
    return flask.jsonify(status="OK")

@bp.route('/accounts/downloadlist/cvs')
@login_required
@perm_manager.require('accounts_download_list')
def accounts_download_csv() -> Response:
    def iter_accs(data):
        for account in data:
            for ci, char in enumerate(account.characters):
                if ci > 0:
                    yield ", " + char.eve_name
                else:
                    yield char.eve_name
            yield '\n'

    permission = perm_manager.get_permission('include_in_accountlist')
    role_check = not Account.disabled
    for role_need in permission.needs:
        if role_need.method != 'role':
            continue
        role_check &= Role.name == role_need.value

    # noinspection PyPep8
    accs = db.session.query(Account).options(joinedload('characters')).join(Account.roles).filter(
        role_check ).order_by(
        Account.username).all()

    response = Response(iter_accs(accs), mimetype='text/csv')
    response.headers['Content-Disposition'] = 'attachment; filename=accounts.csv'
    return response

add_menu_entry('accounts.accounts', 'Accounts', perm_manager.get_permission('accounts_edit').can)
add_menu_entry('accounts.account_self', 'Own Settings', lambda: True)
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""Prepare combined VCF files of COSMIC resource for cancer variant calling.

This prepares the specified version and copies the updates into your bcbio installation
in the correct locations in the variation directory.

Because of licensing restrictions, download from COSMIC requires registration:

https://cancer.sanger.ac.uk/cosmic/download

Usage:

  export COSMIC_USER="your@registered.email.edu"
  export COSMIC_PASS="cosmic_password"
  bcbio_python prepare_cosmic.py <cosmic_version> </path/to/bcbio>

References:
http://gatkforums.broadinstitute.org/discussion/2226/cosmic-and-dbsnp-files-for-mutect
"""
import gzip
import logging
import os
import re
import requests
import subprocess
import sys
import tempfile
import shutil
from argparse import ArgumentParser

from bcbio import utils
from bcbio.variation import vcfutils

logging.basicConfig(format='%(asctime)s [%(levelname).1s] %(message)s', level=logging.INFO)

def main(cosmic_version, bcbio_genome_dir, overwrite=False, clean=False):
    work_dir = utils.safe_makedir(os.path.join(os.getcwd(), "cosmic-prep"))
    os.chdir(work_dir)

    for genome_build, bcbio_build, add_chr in [("GRCh37", "GRCh37", False), ("GRCh38", "hg38", True)]:
        bcbio_base = os.path.join(bcbio_genome_dir, "genomes", "Hsapiens", bcbio_build)
        installed_file = os.path.join(bcbio_base, "variation", f"cosmic-v{cosmic_version}.vcf.gz")
        installed_link = os.path.join(bcbio_base, "variation", "cosmic.vcf.gz")
        logging.info(f"Beginning COSMIC v{cosmic_version} prep for {genome_build}.")
        if not os.path.exists(bcbio_base):
            continue
        if os.path.exists(installed_file):
            if not overwrite:
                logging.info(f"{installed_file} exists, please use the --overwrite flag to overwrite the existing files if you want to reinstall.")
                continue
            else:
                logging.info(f"{installed_file} exists, removing.")
                remove_installed(installed_file, installed_link)
        bcbio_ref = os.path.join(bcbio_base, "seq", f"{bcbio_build}.fa")
        cosmic_vcf_files = get_cosmic_vcf_files(genome_build, cosmic_version, clean)
        sorted_inputs = []
        for fname in cosmic_vcf_files:
            sorted_inputs.append(sort_to_ref(fname, bcbio_ref, add_chr=add_chr))
        out_dir = utils.safe_makedir(os.path.join(f"v{cosmic_version}", "bcbio_ready", bcbio_build))
        out_file = os.path.join(out_dir, "cosmic.vcf.gz")
        ready_cosmic = combine_cosmic(sorted_inputs, bcbio_ref, out_file)
        variation_dir = utils.safe_makedir(os.path.join(bcbio_base, "variation"))
        utils.copy_plus(ready_cosmic, installed_file)
        logging.info(f"Created COSMIC v{cosmic_version} resource in {installed_file}.")
        logging.info(f"Linking {installed_file} as {installed_link}.")
        make_links(installed_file, installed_link)
        logging.info(f"Finished COSMIC v{cosmic_version} prep for {genome_build}.")
        # prepare hg19 from the GRCh37 file
        if bcbio_build == "GRCh37":
            genome_build = "hg19"
            logging.info(f"Prepping COSMIC v{cosmic_version} for {genome_build} from the GRCh37 preparation.")
            bcbio_base = os.path.join(bcbio_genome_dir, "genomes", "Hsapiens", genome_build)
            if not os.path.exists(bcbio_base):
                continue
            if os.path.exists(installed_file):
                installed_file = os.path.join(bcbio_base, "variation", f"cosmic-v{cosmic_version}.vcf.gz")
                installed_link = os.path.join(bcbio_base, "variation", "cosmic.vcf.gz")
                if not overwrite:
                    logging.info(f"{installed_file} exists, please use the --overwrite flag to overwrite the existing files if you want to reinstall.")
                    continue
                else:
                    logging.info(f"{installed_file} exists, removing.")
                    remove_installed(installed_file, installed_link)
            out_dir = utils.safe_makedir(os.path.join(f"v{cosmic_version}", "bcbio_ready", genome_build))
            out_file = os.path.join(out_dir, f"cosmic-v{cosmic_version}.vcf.gz")
            logging.info(f"Translating GRCh37 chromosome names to hg19 chromosome names.")
            hg19_cosmic = map_coords_to_ucsc(ready_cosmic, bcbio_ref, out_file)
            variation_dir = utils.safe_makedir(os.path.join(bcbio_base, "variation"))
            utils.copy_plus(hg19_cosmic, installed_file)
            logging.info(f"Created COSMIC v{cosmic_version} resource in {installed_file}.")
            logging.info(f"Linking {installed_file} as {installed_link}.")
            make_links(installed_file, installed_link)
            logging.info(f"Finished COSMIC v{cosmic_version} prep for {genome_build}.")

def remove_installed(installed_file, installed_link):
    logging.info(f"Removing {installed_file}.")
    if os.path.exists(installed_file):
        os.remove(installed_file)
    installed_index = installed_file + ".tbi"
    logging.info(f"Removing {installed_index}.")
    if os.path.exists(installed_index):
        os.remove(installed_index)
    logging.info(f"Removing {installed_link}.")
    if os.path.lexists(installed_link):
        os.remove(installed_link)
    installed_index = installed_link + ".tbi"
    logging.info(f"Removing {installed_index}.")
    if os.path.lexists(installed_index):
        os.remove(installed_index)

def make_links(installed_file, installed_link):
    os.symlink(os.path.basename(installed_file), installed_link)
    os.symlink(os.path.basename(installed_file + ".tbi"), installed_link + ".tbi")

def map_coords_to_ucsc(grc_cosmic, ref_file, out_file):
    hg19_ref_file = ref_file.replace("GRCh37", "hg19")
    if not os.path.exists(out_file):
        contig_cl = vcfutils.add_contig_to_header_cl(hg19_ref_file, out_file)
        cmd = ("zcat {grc_cosmic} | "
               r'sed "s/^\([0-9]\+\)\t/chr\1\t/g" | sed "s/^MT/chrM/g" | sed "s/^X/chrX/g" | sed "s/^Y/chrY/g" '
               "| {contig_cl} "
               "| bgzip -c > {out_file}")
        subprocess.check_call(cmd.format(**locals()), shell=True)
    if os.path.exists("%s-header.txt" % utils.splitext_plus(out_file)[0]):
        os.remove("%s-header.txt" % utils.splitext_plus(out_file)[0])
    return vcfutils.bgzip_and_index(out_file, {})

def _rename_to_ucsc(line):
    chrom, rest = line.split("\t", 1)
    if chrom == "MT":
        new_chrom = "chrM"
    else:
        new_chrom = "chr%s" % chrom
    return "%s\t%s" % (new_chrom, rest)

def combine_cosmic(fnames, ref_file, out_file):
    logging.info(f"Combining COSMIC files to {out_file}.")
    if not os.path.exists(out_file):
        cmd = ["picard", "MergeVcfs", "O=%s" % out_file, "D=%s" % ref_file.replace(".fa", ".dict")] + \
              ["I=%s" % x for x in fnames] + \
              ["USE_JDK_DEFLATER=true", "USE_JDK_INFLATER=true", "CREATE_INDEX=false"]
        subprocess.check_call(cmd)
    return vcfutils.bgzip_and_index(out_file, {})

def sort_to_ref(fname, ref_file, add_chr):
    """Match reference genome ordering.
    """
    logging.info(f"Sorting {fname} to match the order of {ref_file}.")
    out_file = "%s-prep.vcf.gz" % (fname.replace(".vcf.gz", ""))
    if not os.path.exists(out_file):
        if add_chr:
            fix_chrom = r'| sed "s/^\([0-9]\+\)\t/chr\1\t/g" | sed "s/^MT/chrM/g" | sed "s/^X/chrX/g" | sed "s/^Y/chrY/g" '
        else:
            fix_chrom = ''
        contig_cl = vcfutils.add_contig_to_header_cl(ref_file, out_file)
        cmd = ("gunzip -c {fname} {fix_chrom} | "
               "bcftools norm --check-ref s --do-not-normalize -f {ref_file} |"
               "bcftools view -e 'SNP=1' |"
               "gsort /dev/stdin {ref_file}.fai | {contig_cl} | "
               "bgzip -c > {out_file}")
        subprocess.check_call(cmd.format(**locals()), shell=True)
    logging.info(f"bgzipping and indexing {out_file}.")
    return vcfutils.bgzip_and_index(out_file, {})

def get_cosmic_vcf_files(genome_build, cosmic_version, clean):
    """Retrieve using new authentication based download approach.

    GRCh38/cosmic/v85/VCF/CosmicCodingMuts.vcf.gz
    GRCh38/cosmic/v85/VCF/CosmicNonCodingVariants.vcf.gz
    """
    vdir = os.path.join("v%s" % cosmic_version, genome_build)
    if os.path.exists(vdir):
        if not clean:
            logging.info(f"{vdir} files exist, please use the --clean flag to overwrite the existing files if you want to reinstall.")
        else:
            logging.info(f"{vdir} exists, removing.")
            remove_cosmic_directory(vdir)
    logging.info("Downloading COSMIC VCF files.")
    url = "https://cancer.sanger.ac.uk/cosmic/file_download/"
    out_dir = utils.safe_makedir(os.path.join("v%s" % cosmic_version, genome_build))
    fnames = []
    for ctype in ["CosmicCodingMuts", "CosmicNonCodingVariants"]:
        filename = os.path.join(out_dir, "%s.vcf.gz" % ctype)
        if not os.path.exists(filename):
            filepath = "%s/cosmic/v%s/VCF/%s.vcf.gz" % (genome_build, cosmic_version, ctype)
            logging.info("Downloading %s" % (url + filepath))
            try:
                r = requests.get(url + filepath, auth=(os.environ["COSMIC_USER"], os.environ["COSMIC_PASS"]))
            except KeyError as e:
                print("KeyError: {} not found. Be sure to export your COSMIC_USER and COSMIC_PASS before running in order to download the files".format(e))
                raise e
            download_url = r.json()["url"]
            r = requests.get(download_url)
            with open(filename, "wb") as f:
                f.write(r.content)
        fnames.append(filename)
    return fnames

def remove_cosmic_directory(installed_directory):
    logging.info(f"Removing {installed_directory}.")
    shutil.rmtree(installed_directory)

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("cosmic_version", help="COSMIC version to install.", default="89")
    parser.add_argument("bcbio_directory", help="Path to bcbio installation. Should contain the 'genomes' directory.")
    parser.add_argument("--overwrite", action="store_true", default=False)
    parser.add_argument("--clean", action="store_true", default=False)
    args = parser.parse_args()
    main(args.cosmic_version, args.bcbio_directory, args.overwrite, args.clean)
<commit_msg>Add help messages for the --clean and --overwrite options.<commit_after>#!/usr/bin/env python
"""Prepare combined VCF files of COSMIC resource for cancer variant calling.

This prepares the specified version and copies the updates into your bcbio installation
in the correct locations in the variation directory.

Because of licensing restrictions, download from COSMIC requires registration:

https://cancer.sanger.ac.uk/cosmic/download

Usage:

  export COSMIC_USER="your@registered.email.edu"
  export COSMIC_PASS="cosmic_password"
  bcbio_python prepare_cosmic.py <cosmic_version> </path/to/bcbio>

References:
http://gatkforums.broadinstitute.org/discussion/2226/cosmic-and-dbsnp-files-for-mutect
"""
import gzip
import logging
import os
import re
import requests
import subprocess
import sys
import tempfile
import shutil
from argparse import ArgumentParser

from bcbio import utils
from bcbio.variation import vcfutils

logging.basicConfig(format='%(asctime)s [%(levelname).1s] %(message)s', level=logging.INFO)

def main(cosmic_version, bcbio_genome_dir, overwrite=False, clean=False):
    work_dir = utils.safe_makedir(os.path.join(os.getcwd(), "cosmic-prep"))
    os.chdir(work_dir)

    for genome_build, bcbio_build, add_chr in [("GRCh37", "GRCh37", False), ("GRCh38", "hg38", True)]:
        bcbio_base = os.path.join(bcbio_genome_dir, "genomes", "Hsapiens", bcbio_build)
        installed_file = os.path.join(bcbio_base, "variation", f"cosmic-v{cosmic_version}.vcf.gz")
        installed_link = os.path.join(bcbio_base, "variation", "cosmic.vcf.gz")
        logging.info(f"Beginning COSMIC v{cosmic_version} prep for {genome_build}.")
        if not os.path.exists(bcbio_base):
            continue
        if os.path.exists(installed_file):
            if not overwrite:
                logging.info(f"{installed_file} exists, please use the --overwrite flag to overwrite the existing files if you want to reinstall.")
                continue
            else:
                logging.info(f"{installed_file} exists, removing.")
                remove_installed(installed_file, installed_link)
        bcbio_ref = os.path.join(bcbio_base, "seq", f"{bcbio_build}.fa")
        cosmic_vcf_files = get_cosmic_vcf_files(genome_build, cosmic_version, clean)
        sorted_inputs = []
        for fname in cosmic_vcf_files:
            sorted_inputs.append(sort_to_ref(fname, bcbio_ref, add_chr=add_chr))
        out_dir = utils.safe_makedir(os.path.join(f"v{cosmic_version}", "bcbio_ready", bcbio_build))
        out_file = os.path.join(out_dir, "cosmic.vcf.gz")
        ready_cosmic = combine_cosmic(sorted_inputs, bcbio_ref, out_file)
        variation_dir = utils.safe_makedir(os.path.join(bcbio_base, "variation"))
        utils.copy_plus(ready_cosmic, installed_file)
        logging.info(f"Created COSMIC v{cosmic_version} resource in {installed_file}.")
        logging.info(f"Linking {installed_file} as {installed_link}.")
        make_links(installed_file, installed_link)
        logging.info(f"Finished COSMIC v{cosmic_version} prep for {genome_build}.")
        # prepare hg19 from the GRCh37 file
        if bcbio_build == "GRCh37":
            genome_build = "hg19"
            logging.info(f"Prepping COSMIC v{cosmic_version} for {genome_build} from the GRCh37 preparation.")
            bcbio_base = os.path.join(bcbio_genome_dir, "genomes", "Hsapiens", genome_build)
            if not os.path.exists(bcbio_base):
                continue
            if os.path.exists(installed_file):
                installed_file = os.path.join(bcbio_base, "variation", f"cosmic-v{cosmic_version}.vcf.gz")
                installed_link = os.path.join(bcbio_base, "variation", "cosmic.vcf.gz")
                if not overwrite:
                    logging.info(f"{installed_file} exists, please use the --overwrite flag to overwrite the existing files if you want to reinstall.")
                    continue
                else:
                    logging.info(f"{installed_file} exists, removing.")
                    remove_installed(installed_file, installed_link)
            out_dir = utils.safe_makedir(os.path.join(f"v{cosmic_version}", "bcbio_ready", genome_build))
            out_file = os.path.join(out_dir, f"cosmic-v{cosmic_version}.vcf.gz")
            logging.info(f"Translating GRCh37 chromosome names to hg19 chromosome names.")
            hg19_cosmic = map_coords_to_ucsc(ready_cosmic, bcbio_ref, out_file)
            variation_dir = utils.safe_makedir(os.path.join(bcbio_base, "variation"))
            utils.copy_plus(hg19_cosmic, installed_file)
            logging.info(f"Created COSMIC v{cosmic_version} resource in {installed_file}.")
            logging.info(f"Linking {installed_file} as {installed_link}.")
            make_links(installed_file, installed_link)
            logging.info(f"Finished COSMIC v{cosmic_version} prep for {genome_build}.")

def remove_installed(installed_file, installed_link):
    logging.info(f"Removing {installed_file}.")
    if os.path.exists(installed_file):
        os.remove(installed_file)
    installed_index = installed_file + ".tbi"
    logging.info(f"Removing {installed_index}.")
    if os.path.exists(installed_index):
        os.remove(installed_index)
    logging.info(f"Removing {installed_link}.")
    if os.path.lexists(installed_link):
        os.remove(installed_link)
    installed_index = installed_link + ".tbi"
    logging.info(f"Removing {installed_index}.")
    if os.path.lexists(installed_index):
        os.remove(installed_index)

def make_links(installed_file, installed_link):
    os.symlink(os.path.basename(installed_file), installed_link)
    os.symlink(os.path.basename(installed_file + ".tbi"), installed_link + ".tbi")

def map_coords_to_ucsc(grc_cosmic, ref_file, out_file):
    hg19_ref_file = ref_file.replace("GRCh37", "hg19")
    if not os.path.exists(out_file):
        contig_cl = vcfutils.add_contig_to_header_cl(hg19_ref_file, out_file)
        cmd = ("zcat {grc_cosmic} | "
               r'sed "s/^\([0-9]\+\)\t/chr\1\t/g" | sed "s/^MT/chrM/g" | sed "s/^X/chrX/g" | sed "s/^Y/chrY/g" '
               "| {contig_cl} "
               "| bgzip -c > {out_file}")
        subprocess.check_call(cmd.format(**locals()), shell=True)
    if os.path.exists("%s-header.txt" % utils.splitext_plus(out_file)[0]):
        os.remove("%s-header.txt" % utils.splitext_plus(out_file)[0])
    return vcfutils.bgzip_and_index(out_file, {})

def _rename_to_ucsc(line):
    chrom, rest = line.split("\t", 1)
    if chrom == "MT":
        new_chrom = "chrM"
    else:
        new_chrom = "chr%s" % chrom
    return "%s\t%s" % (new_chrom, rest)

def combine_cosmic(fnames, ref_file, out_file):
    logging.info(f"Combining COSMIC files to {out_file}.")
    if not os.path.exists(out_file):
        cmd = ["picard", "MergeVcfs", "O=%s" % out_file, "D=%s" % ref_file.replace(".fa", ".dict")] + \
              ["I=%s" % x for x in fnames] + \
              ["USE_JDK_DEFLATER=true", "USE_JDK_INFLATER=true", "CREATE_INDEX=false"]
        subprocess.check_call(cmd)
    return vcfutils.bgzip_and_index(out_file, {})

def sort_to_ref(fname, ref_file, add_chr):
    """Match reference genome ordering.
    """
    logging.info(f"Sorting {fname} to match the order of {ref_file}.")
    out_file = "%s-prep.vcf.gz" % (fname.replace(".vcf.gz", ""))
    if not os.path.exists(out_file):
        if add_chr:
            fix_chrom = r'| sed "s/^\([0-9]\+\)\t/chr\1\t/g" | sed "s/^MT/chrM/g" | sed "s/^X/chrX/g" | sed "s/^Y/chrY/g" '
        else:
            fix_chrom = ''
        contig_cl = vcfutils.add_contig_to_header_cl(ref_file, out_file)
        cmd = ("gunzip -c {fname} {fix_chrom} | "
               "bcftools norm --check-ref s --do-not-normalize -f {ref_file} |"
               "bcftools view -e 'SNP=1' |"
               "gsort /dev/stdin {ref_file}.fai | {contig_cl} | "
               "bgzip -c > {out_file}")
        subprocess.check_call(cmd.format(**locals()), shell=True)
    logging.info(f"bgzipping and indexing {out_file}.")
    return vcfutils.bgzip_and_index(out_file, {})

def get_cosmic_vcf_files(genome_build, cosmic_version, clean):
    """Retrieve using new authentication based download approach.

    GRCh38/cosmic/v85/VCF/CosmicCodingMuts.vcf.gz
    GRCh38/cosmic/v85/VCF/CosmicNonCodingVariants.vcf.gz
    """
    vdir = os.path.join("v%s" % cosmic_version, genome_build)
    if os.path.exists(vdir):
        if not clean:
            logging.info(f"{vdir} files exist, please use the --clean flag to overwrite the existing files if you want to reinstall.")
        else:
            logging.info(f"{vdir} exists, removing.")
            remove_cosmic_directory(vdir)
    logging.info("Downloading COSMIC VCF files.")
    url = "https://cancer.sanger.ac.uk/cosmic/file_download/"
    out_dir = utils.safe_makedir(os.path.join("v%s" % cosmic_version, genome_build))
    fnames = []
    for ctype in ["CosmicCodingMuts", "CosmicNonCodingVariants"]:
        filename = os.path.join(out_dir, "%s.vcf.gz" % ctype)
        if not os.path.exists(filename):
            filepath = "%s/cosmic/v%s/VCF/%s.vcf.gz" % (genome_build, cosmic_version, ctype)
            logging.info("Downloading %s" % (url + filepath))
            try:
                r = requests.get(url + filepath, auth=(os.environ["COSMIC_USER"], os.environ["COSMIC_PASS"]))
            except KeyError as e:
                print("KeyError: {} not found. Be sure to export your COSMIC_USER and COSMIC_PASS before running in order to download the files".format(e))
                raise e
            download_url = r.json()["url"]
            r = requests.get(download_url)
            with open(filename, "wb") as f:
                f.write(r.content)
        fnames.append(filename)
    return fnames

def remove_cosmic_directory(installed_directory):
    logging.info(f"Removing {installed_directory}.")
    shutil.rmtree(installed_directory)

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("cosmic_version", help="COSMIC version to install.", default="89")
    parser.add_argument("bcbio_directory", help="Path to bcbio installation. Should contain the 'genomes' directory.")
    parser.add_argument("--overwrite", action="store_true", default=False, help="Overwrite existing cosmic.vcf.gz files.")
    parser.add_argument("--clean", action="store_true", default=False, help="Remove pre-downloaded files, if available.")
    args = parser.parse_args()
    main(args.cosmic_version, args.bcbio_directory, args.overwrite, args.clean)
<|endoftext|>


<commit_before>#!/usr/local/bin/python3.4
# encoding: utf-8
'''
pyverse.bin.gcode -- gets a normalized version of the code.

pyverse.bin.gcode -- Normalized version of code is all white spaces are translated to ' ' space.
                     All start of line comments are removed.
                     Adds the DELIMITER query at start and end of code.
                     In future version it will also add the databse name to the object name, if missing.

It defines classes_and_methods

@author:     Itay Moav

@copyright:  2008 All rights reserved.

@license:    license

@contact:    user_email
@deffield    updated: Updated
'''

import sys
import os
path, filename = os.path.split(__file__)
sys.path.insert(0, path+'/..')
import config
import app.commands

def main():
    '''Command line options.'''

    try:
        # Setup argument parser
        parser = app.ArgumentParser(description=config.program_license, formatter_class=app.RawDescriptionHelpFormatter)
        parser.add_argument("--version",action="version",version=config.program_version_message)
        parser.add_argument("--all", dest="handle_all", action="store_true", help="Specifying this flag will fetch the entire project")
        parser.add_argument("-v", "--verbose", dest="verbosity", action="store_true", help=config.help_common_language['verbosity'])

        parser.add_argument("-s","--stored_proc", dest="stored_proc", action="store",nargs='?', default=False, const='All',                             \
                                help="fetch code for all stored procedures, or the folder/*.sql specified. Root folder is the database name.")

        parser.add_argument("-w","--views", dest="views", action="store",nargs='?', default=False, const='All',                                         \
                                help="fetch code for all views, or the folder/*.sql specified. Root folder is the database name.")

        parser.add_argument("-t","--triggers", dest="triggers", action="store",nargs='?',  default=False, const='All',                                  \
                                help="fetch code for all triggers, or the folder/*.sql specified. Root folder is the database name.")

        parser.add_argument("-f","--functions", dest="functions", action="store",nargs='?',  default=False, const='All',                                \
                                help="fetch code for all functions, or the folder/*.sql specified. Root folder is the database name.")

        parser.add_argument("-c","--scripts", dest="scripts", action="store",nargs='?', default=False, const='All',                                     \
                                help="fetch code for all scripts, or the folder/*.sql specified. Root folder is the database name.")

        parser.add_argument("--db", dest="code_source", action="store",nargs='?', default='assets', const='db',                                         \
                                help="If flag specified, will bring the code from the DB.")

        parser.add_argument("--original", dest="dont_clean_code", action="store",nargs='?', default=False, const=True,                                  \
                                help="If flag specified, will bring the code as is from Assets folder.")

        parser.add_argument("-a", "--assets", dest="assets_path", action="store", nargs='?',  help=config.help_common_language['assets'])

        CodeCleaner = app.commands.CleanCodeObj(parser)
        CodeCleaner.run()

    except Exception as e:
        if config.DEBUG:
            raise(e)
        indent = len(config.program_name) * " "
        sys.stderr.write(config.program_name + ": " + repr(e) + "\n")
        sys.stderr.write(indent + "  for help use --help")
        return 2

#++++++++++++++++++++++++++++++++++++ MAIN ENTRY POINT ++++++++++++++++++++++++++++++++++
if len(sys.argv) == 1: # no params given, do --help
    sys.argv.append("-h")
sys.exit(main())
<commit_msg>same include path commit<commit_after>#!/usr/local/bin/python3.4
# encoding: utf-8
'''
pyverse.bin.gcode -- gets a normalized version of the code.

pyverse.bin.gcode -- Normalized version of code is all white spaces are translated to ' ' space.
                     All start of line comments are removed.
                     Adds the DELIMITER query at start and end of code.
                     In future version it will also add the databse name to the object name, if missing.

It defines classes_and_methods

@author:     Itay Moav

@copyright:  2008 All rights reserved.

@license:    license

@contact:    user_email
@deffield    updated: Updated
'''

import sys
import os

sys.path.insert(0, os.path.dirname(os.path.realpath(__file__)) + '/..')

import config
import app.commands

def main():
    '''Command line options.'''

    try:
        # Setup argument parser
        parser = app.ArgumentParser(description=config.program_license, formatter_class=app.RawDescriptionHelpFormatter)
        parser.add_argument("--version",action="version",version=config.program_version_message)
        parser.add_argument("--all", dest="handle_all", action="store_true", help="Specifying this flag will fetch the entire project")
        parser.add_argument("-v", "--verbose", dest="verbosity", action="store_true", help=config.help_common_language['verbosity'])

        parser.add_argument("-s","--stored_proc", dest="stored_proc", action="store",nargs='?', default=False, const='All',                             \
                                help="fetch code for all stored procedures, or the folder/*.sql specified. Root folder is the database name.")

        parser.add_argument("-w","--views", dest="views", action="store",nargs='?', default=False, const='All',                                         \
                                help="fetch code for all views, or the folder/*.sql specified. Root folder is the database name.")

        parser.add_argument("-t","--triggers", dest="triggers", action="store",nargs='?',  default=False, const='All',                                  \
                                help="fetch code for all triggers, or the folder/*.sql specified. Root folder is the database name.")

        parser.add_argument("-f","--functions", dest="functions", action="store",nargs='?',  default=False, const='All',                                \
                                help="fetch code for all functions, or the folder/*.sql specified. Root folder is the database name.")

        parser.add_argument("-c","--scripts", dest="scripts", action="store",nargs='?', default=False, const='All',                                     \
                                help="fetch code for all scripts, or the folder/*.sql specified. Root folder is the database name.")

        parser.add_argument("--db", dest="code_source", action="store",nargs='?', default='assets', const='db',                                         \
                                help="If flag specified, will bring the code from the DB.")

        parser.add_argument("--original", dest="dont_clean_code", action="store",nargs='?', default=False, const=True,                                  \
                                help="If flag specified, will bring the code as is from Assets folder.")

        parser.add_argument("-a", "--assets", dest="assets_path", action="store", nargs='?',  help=config.help_common_language['assets'])

        CodeCleaner = app.commands.CleanCodeObj(parser)
        CodeCleaner.run()

    except Exception as e:
        if config.DEBUG:
            raise(e)
        indent = len(config.program_name) * " "
        sys.stderr.write(config.program_name + ": " + repr(e) + "\n")
        sys.stderr.write(indent + "  for help use --help")
        return 2

#++++++++++++++++++++++++++++++++++++ MAIN ENTRY POINT ++++++++++++++++++++++++++++++++++
if len(sys.argv) == 1: # no params given, do --help
    sys.argv.append("-h")
sys.exit(main())
<|endoftext|>


<commit_before>import json
import base64
import asyncio

import furl

from waterbutler.core import streams
from waterbutler.core import provider
from waterbutler.core import exceptions
from waterbutler.core import utils

from waterbutler.providers.github import settings
from waterbutler.providers.github.metadata import GitHubRevision
from waterbutler.providers.github.metadata import GitHubFileContentMetadata
from waterbutler.providers.github.metadata import GitHubFolderContentMetadata
from waterbutler.providers.github.metadata import GitHubFileTreeMetadata
from waterbutler.providers.github.metadata import GitHubFolderTreeMetadata

GIT_EMPTY_SHA = '4b825dc642cb6eb9a060e54bf8d69288fbee4904'

class GitHubPath(utils.WaterButlerPath):

    def __init__(self, path, prefix=False, suffix=False):
        super().__init__(path, prefix=prefix, suffix=suffix)

class GitHubProvider(provider.BaseProvider):

    BASE_URL = settings.BASE_URL

    def __init__(self, auth, credentials, settings):
        super().__init__(auth, credentials, settings)
        self.name = self.auth['name']
        self.email = self.auth['email']
        self.token = self.credentials['token']
        self.owner = self.settings['owner']
        self.repo = self.settings['repo']

    @property
    def default_headers(self):
        return {'Authorization': 'token {}'.format(self.token)}

    @property
    def committer(self):
        return {
            'name': self.name,
            'email': self.email,
        }

    def build_repo_url(self, *segments, **query):
        segments = ('repos', self.owner, self.repo) + segments
        return self.build_url(*segments, **query)

    @asyncio.coroutine
    def download(self, path, ref=None, **kwargs):
        if self._is_sha(ref):
            resp = yield from self.make_request(
                'GET',
                self.build_repo_url('git', 'blobs', ref),
                headers={'Accept': 'application/vnd.github.VERSION.raw'},
                expects=(200, ),
                throws=exceptions.DownloadError,
            )
        else:
            path = GitHubPath(path)
            url = furl.furl(self.build_repo_url('contents', path.path))
            if ref:
                url.args.update({'ref': ref})
            resp = yield from self.make_request(
                'GET',
                url.url,
                headers={'Accept': 'application/vnd.github.VERSION.raw'},
                expects=(200, ),
                throws=exceptions.DownloadError,
            )

        return streams.ResponseStreamReader(resp)

    @asyncio.coroutine
    def upload(self, stream, path, message=None, branch=None, **kwargs):
        path = GitHubPath(path)

        content = yield from stream.read()
        encoded = base64.b64encode(content)
        message = message or 'File uploaded on behalf of WaterButler'

        data = {
            'path': path.path,
            'message': message,
            'content': encoded.decode('utf-8'),
            'committer': self.committer,
        }
        if branch is not None:
            data['branch'] = branch
        # Check whether file already exists in tree; if it does, add the SHA
        # to the payload. This changes the call from a create to an update.
        tree = yield from self.metadata(str(path.parent), ref=branch)
        existing = next(
            (each for each in tree if each['path'] == str(path)),
            None
        )
        if existing:
            data['sha'] = existing['extra']['fileSha']

        resp = yield from self.make_request(
            'PUT',
            self.build_repo_url('contents', path.path),
            data=json.dumps(data),
            expects=(200, 201),
            throws=exceptions.UploadError,
        )
        data = yield from resp.json()
        return GitHubFileContentMetadata(data['content'], commit=data['commit']).serialized(), (not existing)

    @asyncio.coroutine
    def delete(self, path, sha=None, message=None, branch=None, **kwargs):
        path = GitHubPath(path)

        if path.is_dir:
            yield from self._delete_folder(path, message, branch, **kwargs)
        else:
            yield from self._delete_file(path, sha, message, branch, **kwargs)

    @asyncio.coroutine
    def metadata(self, path, ref=None, recursive=False, **kwargs):
        """Get Metadata about the requested file or folder
        :param str path: The path to a file or folder
        :param str ref: A branch or a commit SHA
        :rtype dict:
        :rtype list:
        """
        path = GitHubPath(path)
        if path.is_dir:
            return (yield from self._metadata_folder(path, ref=ref, recursive=recursive, **kwargs))
        else:
            return (yield from self._metadata_file(path, ref=ref, **kwargs))

    @asyncio.coroutine
    def revisions(self, path, sha=None, **kwargs):
        path = GitHubPath(path)

        resp = yield from self.make_request(
            'GET',
            self.build_repo_url('commits', path=path.path, sha=sha),
            expects=(200, ),
            throws=exceptions.RevisionsError
        )

        return [
            GitHubRevision(item).serialized()
            for item in (yield from resp.json())
        ]

    @asyncio.coroutine
    def _delete_file(self, path, sha=None, message=None, branch=None, **kwargs):
        if not sha:
            raise exceptions.MetadataError('A sha is required for deleting')

        message = message or 'File deleted on behalf of WaterButler'
        data = {
            'message': message,
            'sha': sha,
            'committer': self.committer,
        }
        if branch:
            data['branch'] = branch
        yield from self.make_request(
            'DELETE',
            self.build_repo_url('contents', path.path),
            headers={'Content-Type': 'application/json'},
            data=json.dumps(data),
            expects=(200, ),
            throws=exceptions.DeleteError,
        )

    @asyncio.coroutine
    def _delete_folder(self, path, message=None, branch=None, **kwargs):
        if not branch:
            repo = yield from self._fetch_repo()
            branch = repo['default_branch']

        branch_data = yield from self._fetch_branch(branch)
        old_commit_sha = branch_data['commit']['sha']
        old_commit_tree_sha = branch_data['commit']['commit']['tree']['sha']

        # e.g. 'level1', 'level2', or ''
        tree_paths = path.parts
        trees = [{
            'target': tree_paths[0],
            'tree': [
                {
                    'path': item['path'],
                    'mode': item['mode'],
                    'type': item['type'],
                    'sha': item['sha'],
                }
                for item in (yield from self._fetch_tree(old_commit_tree_sha))['tree']
            ]
        }]
        for idx, tree_path in enumerate(tree_paths[:-1]):
            try:
                tree_sha = next(x for x in trees[-1]['tree'] if x['path'] == tree_path)['sha']
            except StopIteration:
                raise exceptions.MetadataError(
                    'Could not delete folder \'{0}\''.format(path),
                    code=404,
                )
            trees.append({
                'target': tree_paths[idx + 1],
                'tree': [
                    {
                        'path': item['path'],
                        'mode': item['mode'],
                        'type': item['type'],
                        'sha': item['sha'],
                    }
                    for item in (yield from self._fetch_tree(tree_sha))['tree']
                ]
            })

        # The last tree's structure is rewritten w/o the target folder, all others
        # in the hierarchy are simply updated to reflect this change.
        tree = trees.pop()
        if tree['target'] == '':
            # Git Empty SHA
            tree_sha = GIT_EMPTY_SHA
        else:
            # Delete the folder from the tree
            for item in tree['tree']:
                if item['path'] == tree['target']:
                    tree['tree'].remove(item)
                    break
            tree_data = yield from self._create_tree({'tree': tree['tree']})
            tree_sha = tree_data['sha']
            # Update parent tree(s)
            for tree in reversed(trees):
                for item in tree['tree']:
                    if item['path'] == tree['target']:
                        item['sha'] = tree_sha
                        break
                tree_data = yield from self._create_tree({'tree': tree['tree']})
                tree_sha = tree_data['sha']

        # Create a new commit which references our top most tree change.
        message = message or 'Folder deleted on behalf of WaterButler'
        commit_resp = yield from self.make_request(
            'POST',
            self.build_repo_url('git', 'commits'),
            headers={'Content-Type': 'application/json'},
            data=json.dumps({
                'message': message,
                'tree': tree_sha,
                'parents': [
                    old_commit_sha,
                ],
            }),
            expects=(201, ),
            throws=exceptions.DeleteError,
        )
        commit_data = yield from commit_resp.json()
        commit_sha = commit_data['sha']

        # Update repository reference, point to the newly created commit.
        ref_resp = yield from self.make_request(
            'PATCH',
            self.build_repo_url('git', 'refs', 'heads', branch),
            headers={'Content-Type': 'application/json'},
            data=json.dumps({
                'sha': commit_sha,
            }),
            expects=(200, ),
            throws=exceptions.DeleteError,
        )
        yield from ref_resp.json()

    @asyncio.coroutine
    def _fetch_branch(self, branch):
        resp = yield from self.make_request(
            'GET',
            self.build_repo_url('branches', branch)
        )
        return (yield from resp.json())

    @asyncio.coroutine
    def _fetch_contents(self, path, ref=None):
        url = furl.furl(self.build_repo_url('contents', path.path))
        if ref:
            url.args.update({'ref': ref})
        resp = yield from self.make_request(
            'GET',
            url.url,
            expects=(200, ),
            throws=exceptions.MetadataError
        )
        return (yield from resp.json())

    @asyncio.coroutine
    def _fetch_repo(self):
        resp = yield from self.make_request(
            'GET',
            self.build_repo_url(),
            expects=(200, ),
            throws=exceptions.MetadataError
        )
        return (yield from resp.json())

    @asyncio.coroutine
    def _fetch_tree(self, sha, recursive=False):
        url = furl.furl(self.build_repo_url('git', 'trees', sha))
        if recursive:
            url.args.update({'recursive': 1})
        resp = yield from self.make_request(
            'GET',
            url.url,
            expects=(200, ),
            throws=exceptions.MetadataError
        )
        return (yield from resp.json())

    @asyncio.coroutine
    def _create_tree(self, tree):
        resp = yield from self.make_request(
            'POST',
            self.build_repo_url('git', 'trees'),
            headers={'Content-Type': 'application/json'},
            data=json.dumps(tree),
            expects=(201, ),
            throws=exceptions.DeleteError,
        )
        return (yield from resp.json())

    def _is_sha(self, ref):
        # sha1 is always 40 characters in length
        try:
            if len(ref) != 40:
                return False
            # sha1 is always base 16 (hex)
            int(ref, 16)
        except (TypeError, ValueError, ):
            return False
        return True

    @asyncio.coroutine
    def _metadata_folder(self, path, ref=None, recursive=False, **kwargs):
        parent_path = path.parent

        # if we have a sha or recursive lookup specified we'll need to perform
        # the operation using the git/trees api which requires a sha.
        if self._is_sha(ref) or recursive:
            if self._is_sha(ref):
                tree_sha = ref
            elif parent_path.is_root:
                if not ref:
                    repo = yield from self._fetch_repo()
                    ref = repo['default_branch']
                branch_data = yield from self._fetch_branch(ref)
                tree_sha = branch_data['commit']['commit']['tree']['sha']
            else:
                data = yield from self._fetch_contents(parent_path, ref=ref)
                try:
                    tree_sha = next(x for x in data if x['path'] == path.path)['sha']
                except StopIteration:
                    raise exceptions.MetadataError(
                        'Could not find folder \'{0}\''.format(path),
                        code=404,
                    )

            data = yield from self._fetch_tree(tree_sha, recursive=recursive)

            ret = []
            for item in data['tree']:
                if item['type'] == 'tree':
                    ret.append(GitHubFolderTreeMetadata(item, folder=path.path).serialized())
                else:
                    ret.append(GitHubFileTreeMetadata(item, folder=path.path).serialized())
            return ret
        else:
            data = yield from self._fetch_contents(path, ref=ref)

            ret = []
            for item in data:
                if item['type'] == 'dir':
                    ret.append(GitHubFolderContentMetadata(item).serialized())
                else:
                    ret.append(GitHubFileContentMetadata(item).serialized())
            return ret

    @asyncio.coroutine
    def _metadata_file(self, path, ref=None, **kwargs):
        data = yield from self._fetch_contents(path, ref)

        if isinstance(data, list):
            raise exceptions.MetadataError(
                'Could not retrieve file \'{0}\''.format(path),
                code=404,
            )

        return GitHubFileContentMetadata(data).serialized()
<commit_msg>patch for github, content was exhausting server memory<commit_after>import json
import base64
import asyncio

import furl

from waterbutler.core import streams
from waterbutler.core import provider
from waterbutler.core import exceptions
from waterbutler.core import utils

from waterbutler.providers.github import settings
from waterbutler.providers.github.metadata import GitHubRevision
from waterbutler.providers.github.metadata import GitHubFileContentMetadata
from waterbutler.providers.github.metadata import GitHubFolderContentMetadata
from waterbutler.providers.github.metadata import GitHubFileTreeMetadata
from waterbutler.providers.github.metadata import GitHubFolderTreeMetadata

GIT_EMPTY_SHA = '4b825dc642cb6eb9a060e54bf8d69288fbee4904'

class GitHubPath(utils.WaterButlerPath):

    def __init__(self, path, prefix=False, suffix=False):
        super().__init__(path, prefix=prefix, suffix=suffix)

class GitHubProvider(provider.BaseProvider):

    BASE_URL = settings.BASE_URL

    def __init__(self, auth, credentials, settings):
        super().__init__(auth, credentials, settings)
        self.name = self.auth['name']
        self.email = self.auth['email']
        self.token = self.credentials['token']
        self.owner = self.settings['owner']
        self.repo = self.settings['repo']

    @property
    def default_headers(self):
        return {'Authorization': 'token {}'.format(self.token)}

    @property
    def committer(self):
        return {
            'name': self.name,
            'email': self.email,
        }

    def build_repo_url(self, *segments, **query):
        segments = ('repos', self.owner, self.repo) + segments
        return self.build_url(*segments, **query)

    @asyncio.coroutine
    def download(self, path, ref=None, **kwargs):
        if self._is_sha(ref):
            resp = yield from self.make_request(
                'GET',
                self.build_repo_url('git', 'blobs', ref),
                headers={'Accept': 'application/vnd.github.VERSION.raw'},
                expects=(200, ),
                throws=exceptions.DownloadError,
            )
        else:
            path = GitHubPath(path)
            url = furl.furl(self.build_repo_url('contents', path.path))
            if ref:
                url.args.update({'ref': ref})
            resp = yield from self.make_request(
                'GET',
                url.url,
                headers={'Accept': 'application/vnd.github.VERSION.raw'},
                expects=(200, ),
                throws=exceptions.DownloadError,
            )

        return streams.ResponseStreamReader(resp)

    @asyncio.coroutine
    def upload(self, stream, path, message=None, branch=None, **kwargs):
        path = GitHubPath(path)

        content = yield from stream.read()
        content = base64.b64encode(content)
        content = content.decode('utf-8')
        message = message or 'File uploaded on behalf of WaterButler'

        data = {
            'path': path.path,
            'message': message,
            'content': content,
            'committer': self.committer,
        }
        if branch is not None:
            data['branch'] = branch
        # Check whether file already exists in tree; if it does, add the SHA
        # to the payload. This changes the call from a create to an update.
        tree = yield from self.metadata(str(path.parent), ref=branch)
        existing = next(
            (each for each in tree if each['path'] == str(path)),
            None
        )
        if existing:
            data['sha'] = existing['extra']['fileSha']

        resp = yield from self.make_request(
            'PUT',
            self.build_repo_url('contents', path.path),
            data=json.dumps(data),
            expects=(200, 201),
            throws=exceptions.UploadError,
        )
        data = yield from resp.json()
        return GitHubFileContentMetadata(data['content'], commit=data['commit']).serialized(), (not existing)

    @asyncio.coroutine
    def delete(self, path, sha=None, message=None, branch=None, **kwargs):
        path = GitHubPath(path)

        if path.is_dir:
            yield from self._delete_folder(path, message, branch, **kwargs)
        else:
            yield from self._delete_file(path, sha, message, branch, **kwargs)

    @asyncio.coroutine
    def metadata(self, path, ref=None, recursive=False, **kwargs):
        """Get Metadata about the requested file or folder
        :param str path: The path to a file or folder
        :param str ref: A branch or a commit SHA
        :rtype dict:
        :rtype list:
        """
        path = GitHubPath(path)
        if path.is_dir:
            return (yield from self._metadata_folder(path, ref=ref, recursive=recursive, **kwargs))
        else:
            return (yield from self._metadata_file(path, ref=ref, **kwargs))

    @asyncio.coroutine
    def revisions(self, path, sha=None, **kwargs):
        path = GitHubPath(path)

        resp = yield from self.make_request(
            'GET',
            self.build_repo_url('commits', path=path.path, sha=sha),
            expects=(200, ),
            throws=exceptions.RevisionsError
        )

        return [
            GitHubRevision(item).serialized()
            for item in (yield from resp.json())
        ]

    @asyncio.coroutine
    def _delete_file(self, path, sha=None, message=None, branch=None, **kwargs):
        if not sha:
            raise exceptions.MetadataError('A sha is required for deleting')

        message = message or 'File deleted on behalf of WaterButler'
        data = {
            'message': message,
            'sha': sha,
            'committer': self.committer,
        }
        if branch:
            data['branch'] = branch
        yield from self.make_request(
            'DELETE',
            self.build_repo_url('contents', path.path),
            headers={'Content-Type': 'application/json'},
            data=json.dumps(data),
            expects=(200, ),
            throws=exceptions.DeleteError,
        )

    @asyncio.coroutine
    def _delete_folder(self, path, message=None, branch=None, **kwargs):
        if not branch:
            repo = yield from self._fetch_repo()
            branch = repo['default_branch']

        branch_data = yield from self._fetch_branch(branch)
        old_commit_sha = branch_data['commit']['sha']
        old_commit_tree_sha = branch_data['commit']['commit']['tree']['sha']

        # e.g. 'level1', 'level2', or ''
        tree_paths = path.parts
        trees = [{
            'target': tree_paths[0],
            'tree': [
                {
                    'path': item['path'],
                    'mode': item['mode'],
                    'type': item['type'],
                    'sha': item['sha'],
                }
                for item in (yield from self._fetch_tree(old_commit_tree_sha))['tree']
            ]
        }]
        for idx, tree_path in enumerate(tree_paths[:-1]):
            try:
                tree_sha = next(x for x in trees[-1]['tree'] if x['path'] == tree_path)['sha']
            except StopIteration:
                raise exceptions.MetadataError(
                    'Could not delete folder \'{0}\''.format(path),
                    code=404,
                )
            trees.append({
                'target': tree_paths[idx + 1],
                'tree': [
                    {
                        'path': item['path'],
                        'mode': item['mode'],
                        'type': item['type'],
                        'sha': item['sha'],
                    }
                    for item in (yield from self._fetch_tree(tree_sha))['tree']
                ]
            })

        # The last tree's structure is rewritten w/o the target folder, all others
        # in the hierarchy are simply updated to reflect this change.
        tree = trees.pop()
        if tree['target'] == '':
            # Git Empty SHA
            tree_sha = GIT_EMPTY_SHA
        else:
            # Delete the folder from the tree
            for item in tree['tree']:
                if item['path'] == tree['target']:
                    tree['tree'].remove(item)
                    break
            tree_data = yield from self._create_tree({'tree': tree['tree']})
            tree_sha = tree_data['sha']
            # Update parent tree(s)
            for tree in reversed(trees):
                for item in tree['tree']:
                    if item['path'] == tree['target']:
                        item['sha'] = tree_sha
                        break
                tree_data = yield from self._create_tree({'tree': tree['tree']})
                tree_sha = tree_data['sha']

        # Create a new commit which references our top most tree change.
        message = message or 'Folder deleted on behalf of WaterButler'
        commit_resp = yield from self.make_request(
            'POST',
            self.build_repo_url('git', 'commits'),
            headers={'Content-Type': 'application/json'},
            data=json.dumps({
                'message': message,
                'tree': tree_sha,
                'parents': [
                    old_commit_sha,
                ],
            }),
            expects=(201, ),
            throws=exceptions.DeleteError,
        )
        commit_data = yield from commit_resp.json()
        commit_sha = commit_data['sha']

        # Update repository reference, point to the newly created commit.
        ref_resp = yield from self.make_request(
            'PATCH',
            self.build_repo_url('git', 'refs', 'heads', branch),
            headers={'Content-Type': 'application/json'},
            data=json.dumps({
                'sha': commit_sha,
            }),
            expects=(200, ),
            throws=exceptions.DeleteError,
        )
        yield from ref_resp.json()

    @asyncio.coroutine
    def _fetch_branch(self, branch):
        resp = yield from self.make_request(
            'GET',
            self.build_repo_url('branches', branch)
        )
        return (yield from resp.json())

    @asyncio.coroutine
    def _fetch_contents(self, path, ref=None):
        url = furl.furl(self.build_repo_url('contents', path.path))
        if ref:
            url.args.update({'ref': ref})
        resp = yield from self.make_request(
            'GET',
            url.url,
            expects=(200, ),
            throws=exceptions.MetadataError
        )
        return (yield from resp.json())

    @asyncio.coroutine
    def _fetch_repo(self):
        resp = yield from self.make_request(
            'GET',
            self.build_repo_url(),
            expects=(200, ),
            throws=exceptions.MetadataError
        )
        return (yield from resp.json())

    @asyncio.coroutine
    def _fetch_tree(self, sha, recursive=False):
        url = furl.furl(self.build_repo_url('git', 'trees', sha))
        if recursive:
            url.args.update({'recursive': 1})
        resp = yield from self.make_request(
            'GET',
            url.url,
            expects=(200, ),
            throws=exceptions.MetadataError
        )
        return (yield from resp.json())

    @asyncio.coroutine
    def _create_tree(self, tree):
        resp = yield from self.make_request(
            'POST',
            self.build_repo_url('git', 'trees'),
            headers={'Content-Type': 'application/json'},
            data=json.dumps(tree),
            expects=(201, ),
            throws=exceptions.DeleteError,
        )
        return (yield from resp.json())

    def _is_sha(self, ref):
        # sha1 is always 40 characters in length
        try:
            if len(ref) != 40:
                return False
            # sha1 is always base 16 (hex)
            int(ref, 16)
        except (TypeError, ValueError, ):
            return False
        return True

    @asyncio.coroutine
    def _metadata_folder(self, path, ref=None, recursive=False, **kwargs):
        parent_path = path.parent

        # if we have a sha or recursive lookup specified we'll need to perform
        # the operation using the git/trees api which requires a sha.
        if self._is_sha(ref) or recursive:
            if self._is_sha(ref):
                tree_sha = ref
            elif parent_path.is_root:
                if not ref:
                    repo = yield from self._fetch_repo()
                    ref = repo['default_branch']
                branch_data = yield from self._fetch_branch(ref)
                tree_sha = branch_data['commit']['commit']['tree']['sha']
            else:
                data = yield from self._fetch_contents(parent_path, ref=ref)
                try:
                    tree_sha = next(x for x in data if x['path'] == path.path)['sha']
                except StopIteration:
                    raise exceptions.MetadataError(
                        'Could not find folder \'{0}\''.format(path),
                        code=404,
                    )

            data = yield from self._fetch_tree(tree_sha, recursive=recursive)

            ret = []
            for item in data['tree']:
                if item['type'] == 'tree':
                    ret.append(GitHubFolderTreeMetadata(item, folder=path.path).serialized())
                else:
                    ret.append(GitHubFileTreeMetadata(item, folder=path.path).serialized())
            return ret
        else:
            data = yield from self._fetch_contents(path, ref=ref)

            ret = []
            for item in data:
                if item['type'] == 'dir':
                    ret.append(GitHubFolderContentMetadata(item).serialized())
                else:
                    ret.append(GitHubFileContentMetadata(item).serialized())
            return ret

    @asyncio.coroutine
    def _metadata_file(self, path, ref=None, **kwargs):
        data = yield from self._fetch_contents(path, ref)

        if isinstance(data, list):
            raise exceptions.MetadataError(
                'Could not retrieve file \'{0}\''.format(path),
                code=404,
            )

        return GitHubFileContentMetadata(data).serialized()
<|endoftext|>


<commit_before>import sys
import matplotlib as mpl
mpl.rcParams['text.usetex'] = False
import matplotlib.pyplot as plt
from desc.twinkles import validate_ic
from lsst.sims.catUtils.utils import ObservationMetaDataGenerator

if __name__ == "__main__":

    cat_name = sys.argv[1]

    # opsim_db = sys.argv[2]

    # visit_num = sys.argv[3]

    val_cat = validate_ic()

    df_gal, df_agn, df_sne = val_cat.load_cat(cat_name, 'Dynamic')

    spr_agn = val_cat.process_sprinkled_agn(df_agn)

    agn_lens_gals = val_cat.process_lenses(spr_agn, df_gal)

    agn_location_test = val_cat.compare_agn_location(spr_agn, agn_lens_gals)

    #spr_sne = val_cat.process_sprinkled_sne(df_sne)

    #sne_lens_gals = val_cat.process_lenses(spr_sne, df_gal)

    #sne_location_test = val_cat.compare_sne_location(spr_sne, sne_lens_gals)

    test_agn_inputs = val_cat.compare_agn_inputs(spr_agn, agn_lens_gals)

    #test_sne_inputs = val_cat.compare_sne_inputs(spr_sne, sne_lens_gals)

    test_agn_lens_mags = val_cat.compare_agn_lens_mags(spr_agn, agn_lens_gals)

    #test_sne_lens_mags = val_cat.compare_sne_lens_mags(spr_sne, sne_lens_gals)

    test_agn_image_mags = val_cat.compare_agn_image_mags(spr_agn, agn_lens_gals,
                                                         59844.22, 'r')
<commit_msg>Finalized test call.<commit_after>import os
import sys
import matplotlib as mpl
mpl.rcParams['text.usetex'] = False
import matplotlib.pyplot as plt
from desc.twinkles import validate_ic

if __name__ == "__main__":

    cat_folder = sys.argv[1]

    visit_num = int(sys.argv[2])

    filter_list = ['u', 'g', 'r', 'i', 'z', 'y']

    print("Loading visit info")
    with open(os.path.join(cat_folder, 'phosim_cat_%i.txt' % visit_num)) as f:
        for line in f:
            line_info = line.split(' ')
            if line_info[0] == 'mjd':
                visit_mjd = float(line_info[1])
            elif line_info[0] == 'filter':
                visit_band = filter_list[int(line_info[1])]
            elif line_info[0] == 'vistime':
                delta_t = float(line_info[1])/2.

    # This converts the way phosim wants the time to opsim time
    visit_mjd -= delta_t/86400.0

    sne_SED_file_dir = 'Dynamic'

    print("Running tests")
    val_cat = validate_ic()

    df_gal, df_agn, df_sne = val_cat.load_cat(cat_folder, visit_num,
                                              sne_SED_file_dir)

    spr_agn = val_cat.process_sprinkled_agn(df_agn)

    agn_lens_gals = val_cat.process_agn_lenses(spr_agn, df_gal)

    agn_location_test = val_cat.compare_agn_location(spr_agn, agn_lens_gals)

    spr_sne = val_cat.process_sprinkled_sne(df_sne)

    sne_lens_gals = val_cat.process_sne_lenses(df_gal)

    sne_location_test = val_cat.compare_sne_location(spr_sne, sne_lens_gals)

    test_agn_inputs = val_cat.compare_agn_inputs(spr_agn, agn_lens_gals)

    test_sne_lens_inputs = val_cat.compare_sne_lens_inputs(sne_lens_gals)

    test_sne_image_inputs = val_cat.compare_sne_image_inputs(spr_sne,
                                                             sne_lens_gals,
                                                             visit_mjd,
                                                             sne_SED_file_dir)

    test_agn_lens_mags = val_cat.compare_agn_lens_mags(spr_agn, agn_lens_gals)

    test_sne_lens_mags = val_cat.compare_sne_lens_mags(sne_lens_gals)

    test_agn_image_mags = val_cat.compare_agn_image_mags(spr_agn, agn_lens_gals,
                                                         visit_mjd, visit_band)

    test_sne_image_mags = val_cat.compare_sne_image_mags(spr_sne, sne_lens_gals,
                                                         visit_mjd, visit_band)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# For debugging, use this command to start neovim:
#
# NVIM_PYTHON_LOG_FILE=nvim.log NVIM_PYTHON_LOG_LEVEL=INFO nvim
#
#
# Please register source before executing any other code, this allow cm_core to
# read basic information about the source without loading the whole module, and
# modules required by this module
from cm import get_src, register_source, get_pos, getLogger

# A completion source with CTRL-X CTRL-N like feature
#
# sort=0 for not using NCM's builtin sorting
register_source(name='cm-keyword-continue',
                priority=6,
                abbreviation='',
                word_pattern=r'\S+',
                sort=0,
                cm_refresh_patterns=[r'\s+$',r'^$'],)

import re
import copy

logger = getLogger(__name__)

class Source:

    def __init__(self,nvim):

        self._nvim = nvim

    def cm_refresh(self,info,ctx,*args):

        force = ctx.get('force',False)

        compiled = re.compile(info['word_pattern'])

        typed = ctx['typed']
        if typed.strip()=='' and not force:
            # At the beginning of the line, need force to trigger the popup,
            # Otherwise this will be annoying.
            return
        try:
            # fetch the previous line for better sorting
            last_line = self._nvim.current.buffer[ctx['lnum']-2]
            typed = last_line + '\n' + typed
        except:
            pass

        typed_words = re.findall(compiled,typed)

        if not typed_words:
            return

        prev_word = ''
        if ctx['base']=='':
            prev_word = typed_words[-1]
            prev_words = typed_words
        else:
            if len(typed_words)<2:
                return
            prev_word = typed_words[-2]
            prev_words = typed_words[0:-1]

        if not isinstance(prev_word,str):
            prev_word = prev_word[0]
            prev_words = [e[0] for e in prev_words]

        reversed_prev_words = list(reversed(prev_words))
        matches = []

        # rank for sorting
        def get_rank(word,span,line,last_line):
            prev = last_line+"\n"+line[0:span[0]]
            words = re.findall(compiled,prev)
            if not words:
                return 0
            if not isinstance(words[0],str):
                words = [e[0] for e in words]
            ret = 0
            reserved_words = list(reversed(words))
            for z in zip(reversed_prev_words,reserved_words):
                if z[0].lower()==z[1].lower():
                    ret += 1
                else:
                    break
            return ret

        for buffer in self._nvim.buffers:

            def word_generator():
                step = 200
                line_cnt = len(buffer)
                for i in range(0,line_cnt,step):
                    lines = buffer[i:i+step]
                    last_line = ''
                    for line in lines:
                        for word in re.finditer(compiled,line):
                            yield word.group(),word.span(),line,last_line
                        last_line = line

            try:
                tmp_prev_word = ''
                for word,span,line,last_line in word_generator():
                    if tmp_prev_word==prev_word:
                        rest_of_line = line[span[1]:]
                        hint = rest_of_line
                        matched = re.compile('\s*(\S+(\s+|$))*').search(rest_of_line,0,50)
                        logger.info('hint: [%s]', hint)
                        if matched:
                            hint = matched.group()
                            logger.info('new hint: [%s]', hint)
                        hint = hint.strip()
                        matches.append(dict(word=word + re.findall(r'\s*',line[span[1]:])[0], menu=hint, _rest_of_line=rest_of_line, _rank=get_rank(word,span,line,last_line)))
                    tmp_prev_word = word
            except Exception as ex:
                logger.exception("Parsing buffer [%s] failed", buffer)

        # sort the result based on total match
        matches.sort(key=lambda e: e['_rank'], reverse=True)

        if matches:
            # add rest_of_line completion for the highest rank
            e = copy.deepcopy(matches[0])
            # e['abbr'] = e['word'] + e['menu'] + '...'
            e['abbr'] = '<the rest> '
            e['menu'] = e['word'] + e['menu'] + '...'
            e['word'] = e['word'].rstrip() + e['_rest_of_line']
            matches.insert(1,e)

        if not force:
            # filter by ranking
            matches = [e for e in matches if e['_rank']>=3 ]

        logger.info('matches %s', matches)
        ret = self._nvim.call('cm#complete', info['name'], ctx, ctx['startcol'], matches)

<commit_msg>refine the display of rest of line. #24<commit_after># -*- coding: utf-8 -*-

# For debugging, use this command to start neovim:
#
# NVIM_PYTHON_LOG_FILE=nvim.log NVIM_PYTHON_LOG_LEVEL=INFO nvim
#
#
# Please register source before executing any other code, this allow cm_core to
# read basic information about the source without loading the whole module, and
# modules required by this module
from cm import get_src, register_source, get_pos, getLogger

# A completion source with CTRL-X CTRL-N like feature
#
# sort=0 for not using NCM's builtin sorting
register_source(name='cm-keyword-continue',
                priority=6,
                abbreviation='',
                word_pattern=r'\S+',
                sort=0,
                cm_refresh_patterns=[r'\s+$',r'^$'],)

import re
import copy

logger = getLogger(__name__)

class Source:

    def __init__(self,nvim):

        self._nvim = nvim

    def cm_refresh(self,info,ctx,*args):

        force = ctx.get('force',False)

        compiled = re.compile(info['word_pattern'])

        typed = ctx['typed']
        if typed.strip()=='' and not force:
            # At the beginning of the line, need force to trigger the popup,
            # Otherwise this will be annoying.
            return
        try:
            # fetch the previous line for better sorting
            last_line = self._nvim.current.buffer[ctx['lnum']-2]
            typed = last_line + '\n' + typed
        except:
            pass

        typed_words = re.findall(compiled,typed)

        if not typed_words:
            return

        prev_word = ''
        if ctx['base']=='':
            prev_word = typed_words[-1]
            prev_words = typed_words
        else:
            if len(typed_words)<2:
                return
            prev_word = typed_words[-2]
            prev_words = typed_words[0:-1]

        if not isinstance(prev_word,str):
            prev_word = prev_word[0]
            prev_words = [e[0] for e in prev_words]

        reversed_prev_words = list(reversed(prev_words))
        matches = []

        # rank for sorting
        def get_rank(word,span,line,last_line):
            prev = last_line+"\n"+line[0:span[0]]
            words = re.findall(compiled,prev)
            if not words:
                return 0
            if not isinstance(words[0],str):
                words = [e[0] for e in words]
            ret = 0
            reserved_words = list(reversed(words))
            for z in zip(reversed_prev_words,reserved_words):
                if z[0].lower()==z[1].lower():
                    ret += 1
                else:
                    break
            return ret

        for buffer in self._nvim.buffers:

            def word_generator():
                step = 200
                line_cnt = len(buffer)
                for i in range(0,line_cnt,step):
                    lines = buffer[i:i+step]
                    last_line = ''
                    for line in lines:
                        for word in re.finditer(compiled,line):
                            yield word.group(),word.span(),line,last_line
                        last_line = line

            try:
                tmp_prev_word = ''
                for word,span,line,last_line in word_generator():
                    if tmp_prev_word==prev_word:
                        rest_of_line = line[span[0]:]
                        hint = rest_of_line
                        matched = re.compile('\s*(\S+(\s+|$))*').search(rest_of_line,0,50)
                        logger.info('hint: [%s]', hint)
                        if matched:
                            hint = matched.group()
                            logger.info('new hint: [%s]', hint)
                        hint = hint.strip()
                        matches.append(dict(word=word + re.findall(r'\s*',line[span[1]:])[0], menu=hint, _rest_of_line=rest_of_line, _rank=get_rank(word,span,line,last_line)))
                    tmp_prev_word = word
            except Exception as ex:
                logger.exception("Parsing buffer [%s] failed", buffer)

        # sort the result based on total match
        matches.sort(key=lambda e: e['_rank'], reverse=True)

        if matches:
            # add rest_of_line completion for the highest rank
            e = copy.deepcopy(matches[0])
            # e['abbr'] = e['word'] + e['menu'] + '...'
            e['abbr'] = 'the rest> '
            e['menu'] = e['word'] + e['menu'] + '...'
            e['word'] = e['_rest_of_line']
            matches.insert(1,e)

        if not force:
            # filter by ranking
            matches = [e for e in matches if e['_rank']>=3 ]

        logger.info('matches %s', matches)
        ret = self._nvim.call('cm#complete', info['name'], ctx, ctx['startcol'], matches)

<|endoftext|>