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

from __future__ import print_function

import hashlib
import os
import sys
from difflib import SequenceMatcher, unified_diff
from functools import wraps
from io import open
from itertools import groupby

from junit_xml import TestSuite, TestCase
from pony.orm import Database, Required, db_session, Set, Optional, select, \
    PrimaryKey, RowNotFound, ERDiagramError, OperationalError

from mutmut import BAD_TIMEOUT, OK_SUSPICIOUS, BAD_SURVIVED, UNTESTED, \
    OK_KILLED, MutationID, Context, mutate

try:
    from itertools import zip_longest
except ImportError:  # pragma: no cover (python2)
    from itertools import izip_longest as zip_longest

if sys.version_info < (3, 0):   # pragma: no cover (python 2 specific)
    # noinspection PyUnresolvedReferences
    text_type = unicode  # noqa: F821
    # This little hack is needed to get the click tester working on python 2.7
    orig_print = print

    def print(x='', **kwargs):
        x = x.decode("utf-8")
        orig_print(x.encode("utf-8"), **kwargs)

else:
    text_type = str

db = Database()

current_db_version = 2

NO_TESTS_FOUND = 'NO TESTS FOUND'

class MiscData(db.Entity):
    key = PrimaryKey(text_type, auto=True)
    value = Optional(text_type, autostrip=False)

class SourceFile(db.Entity):
    filename = Required(text_type, autostrip=False)
    lines = Set('Line')

class Line(db.Entity):
    sourcefile = Required(SourceFile)
    line = Optional(text_type, autostrip=False)
    line_number = Required(int)
    mutants = Set('Mutant')

class Mutant(db.Entity):
    line = Required(Line)
    index = Required(int)
    tested_against_hash = Optional(text_type, autostrip=False)
    status = Required(text_type, autostrip=False)  # really an enum of mutant_statuses

def init_db(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        if db.provider is None:
            cache_filename = os.path.join(os.getcwd(), '.mutmut-cache')
            db.bind(provider='sqlite', filename=cache_filename, create_db=True)

            try:
                db.generate_mapping(create_tables=True)
            except OperationalError:
                pass

            if os.path.exists(cache_filename):
                # If the existing cache file is out of data, delete it and start over
                with db_session:
                    try:
                        v = MiscData.get(key='version')
                        if v is None:
                            existing_db_version = 1
                        else:
                            existing_db_version = int(v.value)
                    except (RowNotFound, ERDiagramError, OperationalError):
                        existing_db_version = 1

                if existing_db_version != current_db_version:
                    print('mutmut cache is out of date, clearing it...')
                    db.drop_all_tables(with_all_data=True)
                    db.schema = None  # Pony otherwise thinks we've already created the tables
                    db.generate_mapping(create_tables=True)

            with db_session:
                v = get_or_create(MiscData, key='version')
                v.value = str(current_db_version)

        return f(*args, **kwargs)
    return wrapper

def hash_of(filename):
    with open(filename, 'rb') as f:
        m = hashlib.sha256()
        m.update(f.read())
        return m.hexdigest()

def hash_of_tests(tests_dirs):
    m = hashlib.sha256()
    found_something = False
    for tests_dir in tests_dirs:
        for root, dirs, files in os.walk(tests_dir):
            for filename in files:
                with open(os.path.join(root, filename), 'rb') as f:
                    m.update(f.read())
                    found_something = True
    if not found_something:
        return NO_TESTS_FOUND
    return m.hexdigest()

def get_apply_line(mutant):
    apply_line = 'mutmut apply %s' % mutant.id
    return apply_line

@init_db
@db_session
def print_result_cache():
    print('To apply a mutant on disk:')
    print('    mutmut apply <id>')
    print('')
    print('To show a mutant:')
    print('    mutmut show <id>')
    print('')

    def print_stuff(title, mutant_query):
        mutant_list = list(mutant_query)
        if mutant_list:
            print('')
            print("{} ({})".format(title, len(mutant_list)))
            for filename, mutants in groupby(mutant_list, key=lambda x: x.line.sourcefile.filename):
                mutants = list(mutants)
                print('')
                print("---- {} ({}) ----".format(filename, len(mutants)))
                print('')
                print(', '.join([str(x.id) for x in mutants]))

    print_stuff('Timed out ⏰', select(x for x in Mutant if x.status == BAD_TIMEOUT))
    print_stuff('Suspicious 🤔', select(x for x in Mutant if x.status == OK_SUSPICIOUS))
    print_stuff('Survived 🙁', select(x for x in Mutant if x.status == BAD_SURVIVED))
    print_stuff('Untested', select(x for x in Mutant if x.status == UNTESTED))

def get_unified_diff(argument, dict_synonyms):
    filename, mutation_id = filename_and_mutation_id_from_pk(argument)
    with open(filename) as f:
        source = f.read()
    context = Context(
        source=source,
        filename=filename,
        mutation_id=mutation_id,
        dict_synonyms=dict_synonyms,
    )
    mutated_source, number_of_mutations_performed = mutate(context)
    if not number_of_mutations_performed:
        return ""

    output = ""
    for line in unified_diff(source.split('\n'), mutated_source.split('\n'), fromfile=filename, tofile=filename, lineterm=''):
        output += line + "\n"
    return output

@init_db
@db_session
def print_result_cache_junitxml(dict_synonyms, suspicious_policy, untested_policy):
    test_cases = []
    mutant_list = list(select(x for x in Mutant))
    for filename, mutants in groupby(mutant_list, key=lambda x: x.line.sourcefile.filename):
        for mutant in mutants:
            tc = TestCase("Mutant #{}".format(mutant.id), file=filename, line=mutant.line.line_number, stdout=mutant.line.line)
            if mutant.status == BAD_SURVIVED:
                tc.add_failure_info(message=mutant.status, output=get_unified_diff(mutant.id, dict_synonyms))
            if mutant.status == BAD_TIMEOUT:
                tc.add_error_info(message=mutant.status, error_type="timeout", output=get_unified_diff(mutant.id, dict_synonyms))
            if mutant.status == OK_SUSPICIOUS:
                if suspicious_policy != 'ignore':
                    func = getattr(tc, 'add_{}_info'.format(suspicious_policy))
                    func(message=mutant.status, output=get_unified_diff(mutant.id, dict_synonyms))
            if mutant.status == UNTESTED:
                if untested_policy != 'ignore':
                    func = getattr(tc, 'add_{}_info'.format(untested_policy))
                    func(message=mutant.status, output=get_unified_diff(mutant.id, dict_synonyms))

            test_cases.append(tc)

    ts = TestSuite("mutmut", test_cases)
    print(TestSuite.to_xml_string([ts]))

def get_or_create(model, defaults=None, **params):
    if defaults is None:
        defaults = {}
    obj = model.get(**params)
    if obj is None:
        params = params.copy()
        for k, v in defaults.items():
            if k not in params:
                params[k] = v
        return model(**params)
    else:
        return obj

def sequence_ops(a, b):
    sequence_matcher = SequenceMatcher(a=a, b=b)

    for tag, i1, i2, j1, j2 in sequence_matcher.get_opcodes():
        a_sub_sequence = a[i1:i2]
        b_sub_sequence = b[j1:j2]
        for x in zip_longest(a_sub_sequence, range(i1, i2), b_sub_sequence, range(j1, j2)):
            yield (tag,) + x

@init_db
@db_session
def update_line_numbers(filename):
    sourcefile = get_or_create(SourceFile, filename=filename)

    cached_line_objects = list(sourcefile.lines.order_by(Line.line_number))

    cached_lines = [x.line for x in cached_line_objects]

    with open(filename) as f:
        existing_lines = [x.strip('\n') for x in f.readlines()]

    if not cached_lines:
        for i, line in enumerate(existing_lines):
            Line(sourcefile=sourcefile, line=line, line_number=i)
        return

    for command, a, a_index, b, b_index in sequence_ops(cached_lines, existing_lines):
        if command == 'equal':
            if a_index != b_index:
                cached_obj = cached_line_objects[a_index]
                assert cached_obj.line == existing_lines[b_index]
                cached_obj.line_number = b_index

        elif command == 'delete':
            cached_line_objects[a_index].delete()

        elif command == 'insert':
            if b is not None:
                Line(sourcefile=sourcefile, line=b, line_number=b_index)

        elif command == 'replace':
            if a_index is not None:
                cached_line_objects[a_index].delete()
            if b is not None:
                Line(sourcefile=sourcefile, line=b, line_number=b_index)

        else:
            assert False, 'unknown opcode from SequenceMatcher: %s' % command

@init_db
@db_session
def register_mutants(mutations_by_file):
    for filename, mutation_ids in mutations_by_file.items():
        sourcefile = get_or_create(SourceFile, filename=filename)

        for mutation_id in mutation_ids:
            line = Line.get(sourcefile=sourcefile, line=mutation_id.line, line_number=mutation_id.line_number)
            assert line is not None
            get_or_create(Mutant, line=line, index=mutation_id.index, defaults=dict(status=UNTESTED))

@init_db
@db_session
def update_mutant_status(file_to_mutate, mutation_id, status, tests_hash):
    sourcefile = SourceFile.get(filename=file_to_mutate)
    line = Line.get(sourcefile=sourcefile, line=mutation_id.line, line_number=mutation_id.line_number)
    mutant = Mutant.get(line=line, index=mutation_id.index)
    mutant.status = status
    mutant.tested_against_hash = tests_hash

@init_db
@db_session
def cached_mutation_status(filename, mutation_id, hash_of_tests):
    sourcefile = SourceFile.get(filename=filename)
    line = Line.get(sourcefile=sourcefile, line=mutation_id.line, line_number=mutation_id.line_number)
    mutant = Mutant.get(line=line, index=mutation_id.index)

    if mutant.status == OK_KILLED:
        # We assume that if a mutant was killed, a change to the test suite will mean it's still killed
        return OK_KILLED

    if mutant.tested_against_hash != hash_of_tests or mutant.tested_against_hash == NO_TESTS_FOUND or hash_of_tests == NO_TESTS_FOUND:
        return UNTESTED

    return mutant.status

@init_db
@db_session
def mutation_id_from_pk(pk):
    mutant = Mutant.get(id=pk)
    return MutationID(line=mutant.line.line, index=mutant.index, line_number=mutant.line.line_number)

@init_db
@db_session
def filename_and_mutation_id_from_pk(pk):
    mutant = Mutant.get(id=pk)
    return mutant.line.sourcefile.filename, mutation_id_from_pk(pk)

@init_db
@db_session
def cached_test_time():
    d = MiscData.get(key='baseline_time_elapsed')
    return float(d.value) if d else None

@init_db
@db_session
def set_cached_test_time(baseline_time_elapsed):
    get_or_create(MiscData, key='baseline_time_elapsed').value = str(baseline_time_elapsed)
<commit_msg>Fixed a bug with grouping the result output<commit_after># -*- coding: utf-8 -*-

from __future__ import print_function

import hashlib
import os
import sys
from difflib import SequenceMatcher, unified_diff
from functools import wraps
from io import open
from itertools import groupby

from junit_xml import TestSuite, TestCase
from pony.orm import Database, Required, db_session, Set, Optional, select, \
    PrimaryKey, RowNotFound, ERDiagramError, OperationalError

from mutmut import BAD_TIMEOUT, OK_SUSPICIOUS, BAD_SURVIVED, UNTESTED, \
    OK_KILLED, MutationID, Context, mutate

try:
    from itertools import zip_longest
except ImportError:  # pragma: no cover (python2)
    from itertools import izip_longest as zip_longest

if sys.version_info < (3, 0):   # pragma: no cover (python 2 specific)
    # noinspection PyUnresolvedReferences
    text_type = unicode  # noqa: F821
    # This little hack is needed to get the click tester working on python 2.7
    orig_print = print

    def print(x='', **kwargs):
        x = x.decode("utf-8")
        orig_print(x.encode("utf-8"), **kwargs)

else:
    text_type = str

db = Database()

current_db_version = 2

NO_TESTS_FOUND = 'NO TESTS FOUND'

class MiscData(db.Entity):
    key = PrimaryKey(text_type, auto=True)
    value = Optional(text_type, autostrip=False)

class SourceFile(db.Entity):
    filename = Required(text_type, autostrip=False)
    lines = Set('Line')

class Line(db.Entity):
    sourcefile = Required(SourceFile)
    line = Optional(text_type, autostrip=False)
    line_number = Required(int)
    mutants = Set('Mutant')

class Mutant(db.Entity):
    line = Required(Line)
    index = Required(int)
    tested_against_hash = Optional(text_type, autostrip=False)
    status = Required(text_type, autostrip=False)  # really an enum of mutant_statuses

def init_db(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        if db.provider is None:
            cache_filename = os.path.join(os.getcwd(), '.mutmut-cache')
            db.bind(provider='sqlite', filename=cache_filename, create_db=True)

            try:
                db.generate_mapping(create_tables=True)
            except OperationalError:
                pass

            if os.path.exists(cache_filename):
                # If the existing cache file is out of data, delete it and start over
                with db_session:
                    try:
                        v = MiscData.get(key='version')
                        if v is None:
                            existing_db_version = 1
                        else:
                            existing_db_version = int(v.value)
                    except (RowNotFound, ERDiagramError, OperationalError):
                        existing_db_version = 1

                if existing_db_version != current_db_version:
                    print('mutmut cache is out of date, clearing it...')
                    db.drop_all_tables(with_all_data=True)
                    db.schema = None  # Pony otherwise thinks we've already created the tables
                    db.generate_mapping(create_tables=True)

            with db_session:
                v = get_or_create(MiscData, key='version')
                v.value = str(current_db_version)

        return f(*args, **kwargs)
    return wrapper

def hash_of(filename):
    with open(filename, 'rb') as f:
        m = hashlib.sha256()
        m.update(f.read())
        return m.hexdigest()

def hash_of_tests(tests_dirs):
    m = hashlib.sha256()
    found_something = False
    for tests_dir in tests_dirs:
        for root, dirs, files in os.walk(tests_dir):
            for filename in files:
                with open(os.path.join(root, filename), 'rb') as f:
                    m.update(f.read())
                    found_something = True
    if not found_something:
        return NO_TESTS_FOUND
    return m.hexdigest()

def get_apply_line(mutant):
    apply_line = 'mutmut apply %s' % mutant.id
    return apply_line

@init_db
@db_session
def print_result_cache():
    print('To apply a mutant on disk:')
    print('    mutmut apply <id>')
    print('')
    print('To show a mutant:')
    print('    mutmut show <id>')
    print('')

    def print_stuff(title, mutant_query):
        mutant_list = list(sorted(mutant_query, key=lambda x: x.line.sourcefile.filename))
        if mutant_list:
            print('')
            print("{} ({})".format(title, len(mutant_list)))
            for filename, mutants in groupby(mutant_list, key=lambda x: x.line.sourcefile.filename):
                mutants = list(mutants)
                print('')
                print("---- {} ({}) ----".format(filename, len(mutants)))
                print('')
                print(', '.join([str(x.id) for x in mutants]))

    print_stuff('Timed out ⏰', select(x for x in Mutant if x.status == BAD_TIMEOUT))
    print_stuff('Suspicious 🤔', select(x for x in Mutant if x.status == OK_SUSPICIOUS))
    print_stuff('Survived 🙁', select(x for x in Mutant if x.status == BAD_SURVIVED))
    print_stuff('Untested', select(x for x in Mutant if x.status == UNTESTED))

def get_unified_diff(argument, dict_synonyms):
    filename, mutation_id = filename_and_mutation_id_from_pk(argument)
    with open(filename) as f:
        source = f.read()
    context = Context(
        source=source,
        filename=filename,
        mutation_id=mutation_id,
        dict_synonyms=dict_synonyms,
    )
    mutated_source, number_of_mutations_performed = mutate(context)
    if not number_of_mutations_performed:
        return ""

    output = ""
    for line in unified_diff(source.split('\n'), mutated_source.split('\n'), fromfile=filename, tofile=filename, lineterm=''):
        output += line + "\n"
    return output

@init_db
@db_session
def print_result_cache_junitxml(dict_synonyms, suspicious_policy, untested_policy):
    test_cases = []
    mutant_list = list(select(x for x in Mutant))
    for filename, mutants in groupby(mutant_list, key=lambda x: x.line.sourcefile.filename):
        for mutant in mutants:
            tc = TestCase("Mutant #{}".format(mutant.id), file=filename, line=mutant.line.line_number, stdout=mutant.line.line)
            if mutant.status == BAD_SURVIVED:
                tc.add_failure_info(message=mutant.status, output=get_unified_diff(mutant.id, dict_synonyms))
            if mutant.status == BAD_TIMEOUT:
                tc.add_error_info(message=mutant.status, error_type="timeout", output=get_unified_diff(mutant.id, dict_synonyms))
            if mutant.status == OK_SUSPICIOUS:
                if suspicious_policy != 'ignore':
                    func = getattr(tc, 'add_{}_info'.format(suspicious_policy))
                    func(message=mutant.status, output=get_unified_diff(mutant.id, dict_synonyms))
            if mutant.status == UNTESTED:
                if untested_policy != 'ignore':
                    func = getattr(tc, 'add_{}_info'.format(untested_policy))
                    func(message=mutant.status, output=get_unified_diff(mutant.id, dict_synonyms))

            test_cases.append(tc)

    ts = TestSuite("mutmut", test_cases)
    print(TestSuite.to_xml_string([ts]))

def get_or_create(model, defaults=None, **params):
    if defaults is None:
        defaults = {}
    obj = model.get(**params)
    if obj is None:
        params = params.copy()
        for k, v in defaults.items():
            if k not in params:
                params[k] = v
        return model(**params)
    else:
        return obj

def sequence_ops(a, b):
    sequence_matcher = SequenceMatcher(a=a, b=b)

    for tag, i1, i2, j1, j2 in sequence_matcher.get_opcodes():
        a_sub_sequence = a[i1:i2]
        b_sub_sequence = b[j1:j2]
        for x in zip_longest(a_sub_sequence, range(i1, i2), b_sub_sequence, range(j1, j2)):
            yield (tag,) + x

@init_db
@db_session
def update_line_numbers(filename):
    sourcefile = get_or_create(SourceFile, filename=filename)

    cached_line_objects = list(sourcefile.lines.order_by(Line.line_number))

    cached_lines = [x.line for x in cached_line_objects]

    with open(filename) as f:
        existing_lines = [x.strip('\n') for x in f.readlines()]

    if not cached_lines:
        for i, line in enumerate(existing_lines):
            Line(sourcefile=sourcefile, line=line, line_number=i)
        return

    for command, a, a_index, b, b_index in sequence_ops(cached_lines, existing_lines):
        if command == 'equal':
            if a_index != b_index:
                cached_obj = cached_line_objects[a_index]
                assert cached_obj.line == existing_lines[b_index]
                cached_obj.line_number = b_index

        elif command == 'delete':
            cached_line_objects[a_index].delete()

        elif command == 'insert':
            if b is not None:
                Line(sourcefile=sourcefile, line=b, line_number=b_index)

        elif command == 'replace':
            if a_index is not None:
                cached_line_objects[a_index].delete()
            if b is not None:
                Line(sourcefile=sourcefile, line=b, line_number=b_index)

        else:
            assert False, 'unknown opcode from SequenceMatcher: %s' % command

@init_db
@db_session
def register_mutants(mutations_by_file):
    for filename, mutation_ids in mutations_by_file.items():
        sourcefile = get_or_create(SourceFile, filename=filename)

        for mutation_id in mutation_ids:
            line = Line.get(sourcefile=sourcefile, line=mutation_id.line, line_number=mutation_id.line_number)
            assert line is not None
            get_or_create(Mutant, line=line, index=mutation_id.index, defaults=dict(status=UNTESTED))

@init_db
@db_session
def update_mutant_status(file_to_mutate, mutation_id, status, tests_hash):
    sourcefile = SourceFile.get(filename=file_to_mutate)
    line = Line.get(sourcefile=sourcefile, line=mutation_id.line, line_number=mutation_id.line_number)
    mutant = Mutant.get(line=line, index=mutation_id.index)
    mutant.status = status
    mutant.tested_against_hash = tests_hash

@init_db
@db_session
def cached_mutation_status(filename, mutation_id, hash_of_tests):
    sourcefile = SourceFile.get(filename=filename)
    line = Line.get(sourcefile=sourcefile, line=mutation_id.line, line_number=mutation_id.line_number)
    mutant = Mutant.get(line=line, index=mutation_id.index)

    if mutant.status == OK_KILLED:
        # We assume that if a mutant was killed, a change to the test suite will mean it's still killed
        return OK_KILLED

    if mutant.tested_against_hash != hash_of_tests or mutant.tested_against_hash == NO_TESTS_FOUND or hash_of_tests == NO_TESTS_FOUND:
        return UNTESTED

    return mutant.status

@init_db
@db_session
def mutation_id_from_pk(pk):
    mutant = Mutant.get(id=pk)
    return MutationID(line=mutant.line.line, index=mutant.index, line_number=mutant.line.line_number)

@init_db
@db_session
def filename_and_mutation_id_from_pk(pk):
    mutant = Mutant.get(id=pk)
    return mutant.line.sourcefile.filename, mutation_id_from_pk(pk)

@init_db
@db_session
def cached_test_time():
    d = MiscData.get(key='baseline_time_elapsed')
    return float(d.value) if d else None

@init_db
@db_session
def set_cached_test_time(baseline_time_elapsed):
    get_or_create(MiscData, key='baseline_time_elapsed').value = str(baseline_time_elapsed)
<|endoftext|>


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

import codecs
from datetime import datetime
import json
import time
import urllib
import subprocess

from flask import Markup, g, render_template, request
from slimit import minify
from smartypants import smartypants

import app_config
import copytext

class BetterJSONEncoder(json.JSONEncoder):
    """
    A JSON encoder that intelligently handles datetimes.
    """
    def default(self, obj):
        if isinstance(obj, datetime):
            encoded_object = obj.isoformat()
        else:
            encoded_object = json.JSONEncoder.default(self, obj)

        return encoded_object

class Includer(object):
    """
    Base class for Javascript and CSS psuedo-template-tags.

    See `make_context` for an explanation of `asset_depth`.
    """
    def __init__(self, asset_depth=0):
        self.includes = []
        self.tag_string = None
        self.asset_depth = asset_depth

    def push(self, path):
        self.includes.append(path)

        return ''

    def _compress(self):
        raise NotImplementedError()

    def _relativize_path(self, path):
        relative_path = path
        depth = len(request.path.split('/')) - (2 + self.asset_depth)

        while depth > 0:
            relative_path = '../%s' % relative_path
            depth -= 1

        return relative_path

    def render(self, path):
        if getattr(g, 'compile_includes', False):
            if path in g.compiled_includes:
                timestamp_path = g.compiled_includes[path]
            else:
                # Add a querystring to the rendered filename to prevent caching
                timestamp_path = '%s?%i' % (path, int(time.time()))

                out_path = 'www/%s' % path

                if path not in g.compiled_includes:
                    print 'Rendering %s' % out_path

                    with codecs.open(out_path, 'w', encoding='utf-8') as f:
                        f.write(self._compress())

                # See "fab render"
                g.compiled_includes[path] = timestamp_path

            markup = Markup(self.tag_string % self._relativize_path(timestamp_path))
        else:
            response = ','.join(self.includes)

            response = '\n'.join([
                self.tag_string % self._relativize_path(src) for src in self.includes
            ])

            markup = Markup(response)

        del self.includes[:]

        return markup

class JavascriptIncluder(Includer):
    """
    Psuedo-template tag that handles collecting Javascript and serving appropriate clean or compressed versions.
    """
    def __init__(self, *args, **kwargs):
        Includer.__init__(self, *args, **kwargs)

        self.tag_string = '<script type="text/javascript" src="%s"></script>'

    def _compress(self):
        output = []
        src_paths = []

        for src in self.includes:
            src_paths.append('www/%s' % src)

            with codecs.open('www/%s' % src, encoding='utf-8') as f:
                print '- compressing %s' % src
                output.append(minify(f.read()))

        context = make_context()
        context['paths'] = src_paths

        header = render_template('_js_header.js', **context)
        output.insert(0, header)

        return '\n'.join(output)

class CSSIncluder(Includer):
    """
    Psuedo-template tag that handles collecting CSS and serving appropriate clean or compressed versions.
    """
    def __init__(self, *args, **kwargs):
        Includer.__init__(self, *args, **kwargs)

        self.tag_string = '<link rel="stylesheet" type="text/css" href="%s" />'

    def _compress(self):
        output = []

        src_paths = []

        for src in self.includes:

            src_paths.append('%s' % src)

            try:
                compressed_src = subprocess.check_output(["node_modules/less/bin/lessc", "-x", src])
                output.append(compressed_src)
            except:
                print 'It looks like "lessc" isn\'t installed. Try running: "npm install"'
                raise

        context = make_context()
        context['paths'] = src_paths

        header = render_template('_css_header.css', **context)
        output.insert(0, header)

        return '\n'.join(output)

def flatten_app_config():
    """
    Returns a copy of app_config containing only
    configuration variables.
    """
    config = {}

    # Only all-caps [constant] vars get included
    for k, v in app_config.__dict__.items():
        if k.upper() == k:
            config[k] = v

    return config

def make_context(asset_depth=0):
    """
    Create a base-context for rendering views.
    Includes app_config and JS/CSS includers.

    `asset_depth` indicates how far into the url hierarchy
    the assets are hosted. If 0, then they are at the root.
    If 1 then at /foo/, etc.
    """
    context = flatten_app_config()

    try:
        context['COPY'] = copytext.Copy(app_config.COPY_PATH)
    except copytext.CopyException:
        pass

    context['JS'] = JavascriptIncluder(asset_depth=asset_depth)
    context['CSS'] = CSSIncluder(asset_depth=asset_depth)

    return context

def urlencode_filter(s):
    """
    Filter to urlencode strings.
    """
    if type(s) == 'Markup':
        s = s.unescape()

    # Evaulate COPY elements
    if type(s) is not unicode:
        s = unicode(s)

    s = s.encode('utf8')
    s = urllib.quote_plus(s)

    return Markup(s)

def smarty_filter(s):
    """
    Filter to smartypants strings.
    """
    if type(s) == 'Markup':
        s = s.unescape()

    # Evaulate COPY elements
    if type(s) is not unicode:
        s = unicode(s)

    s = s.encode('utf-8')
    s = smartypants(s)

    return Markup(s)

<commit_msg>if smartypants fails, print offending string to the console. closes #522<commit_after>#!/usr/bin/env python

import codecs
from datetime import datetime
import json
import time
import urllib
import subprocess

from flask import Markup, g, render_template, request
from slimit import minify
from smartypants import smartypants

import app_config
import copytext

class BetterJSONEncoder(json.JSONEncoder):
    """
    A JSON encoder that intelligently handles datetimes.
    """
    def default(self, obj):
        if isinstance(obj, datetime):
            encoded_object = obj.isoformat()
        else:
            encoded_object = json.JSONEncoder.default(self, obj)

        return encoded_object

class Includer(object):
    """
    Base class for Javascript and CSS psuedo-template-tags.

    See `make_context` for an explanation of `asset_depth`.
    """
    def __init__(self, asset_depth=0):
        self.includes = []
        self.tag_string = None
        self.asset_depth = asset_depth

    def push(self, path):
        self.includes.append(path)

        return ''

    def _compress(self):
        raise NotImplementedError()

    def _relativize_path(self, path):
        relative_path = path
        depth = len(request.path.split('/')) - (2 + self.asset_depth)

        while depth > 0:
            relative_path = '../%s' % relative_path
            depth -= 1

        return relative_path

    def render(self, path):
        if getattr(g, 'compile_includes', False):
            if path in g.compiled_includes:
                timestamp_path = g.compiled_includes[path]
            else:
                # Add a querystring to the rendered filename to prevent caching
                timestamp_path = '%s?%i' % (path, int(time.time()))

                out_path = 'www/%s' % path

                if path not in g.compiled_includes:
                    print 'Rendering %s' % out_path

                    with codecs.open(out_path, 'w', encoding='utf-8') as f:
                        f.write(self._compress())

                # See "fab render"
                g.compiled_includes[path] = timestamp_path

            markup = Markup(self.tag_string % self._relativize_path(timestamp_path))
        else:
            response = ','.join(self.includes)

            response = '\n'.join([
                self.tag_string % self._relativize_path(src) for src in self.includes
            ])

            markup = Markup(response)

        del self.includes[:]

        return markup

class JavascriptIncluder(Includer):
    """
    Psuedo-template tag that handles collecting Javascript and serving appropriate clean or compressed versions.
    """
    def __init__(self, *args, **kwargs):
        Includer.__init__(self, *args, **kwargs)

        self.tag_string = '<script type="text/javascript" src="%s"></script>'

    def _compress(self):
        output = []
        src_paths = []

        for src in self.includes:
            src_paths.append('www/%s' % src)

            with codecs.open('www/%s' % src, encoding='utf-8') as f:
                print '- compressing %s' % src
                output.append(minify(f.read()))

        context = make_context()
        context['paths'] = src_paths

        header = render_template('_js_header.js', **context)
        output.insert(0, header)

        return '\n'.join(output)

class CSSIncluder(Includer):
    """
    Psuedo-template tag that handles collecting CSS and serving appropriate clean or compressed versions.
    """
    def __init__(self, *args, **kwargs):
        Includer.__init__(self, *args, **kwargs)

        self.tag_string = '<link rel="stylesheet" type="text/css" href="%s" />'

    def _compress(self):
        output = []

        src_paths = []

        for src in self.includes:

            src_paths.append('%s' % src)

            try:
                compressed_src = subprocess.check_output(["node_modules/less/bin/lessc", "-x", src])
                output.append(compressed_src)
            except:
                print 'It looks like "lessc" isn\'t installed. Try running: "npm install"'
                raise

        context = make_context()
        context['paths'] = src_paths

        header = render_template('_css_header.css', **context)
        output.insert(0, header)

        return '\n'.join(output)

def flatten_app_config():
    """
    Returns a copy of app_config containing only
    configuration variables.
    """
    config = {}

    # Only all-caps [constant] vars get included
    for k, v in app_config.__dict__.items():
        if k.upper() == k:
            config[k] = v

    return config

def make_context(asset_depth=0):
    """
    Create a base-context for rendering views.
    Includes app_config and JS/CSS includers.

    `asset_depth` indicates how far into the url hierarchy
    the assets are hosted. If 0, then they are at the root.
    If 1 then at /foo/, etc.
    """
    context = flatten_app_config()

    try:
        context['COPY'] = copytext.Copy(app_config.COPY_PATH)
    except copytext.CopyException:
        pass

    context['JS'] = JavascriptIncluder(asset_depth=asset_depth)
    context['CSS'] = CSSIncluder(asset_depth=asset_depth)

    return context

def urlencode_filter(s):
    """
    Filter to urlencode strings.
    """
    if type(s) == 'Markup':
        s = s.unescape()

    # Evaulate COPY elements
    if type(s) is not unicode:
        s = unicode(s)

    s = s.encode('utf8')
    s = urllib.quote_plus(s)

    return Markup(s)

def smarty_filter(s):
    """
    Filter to smartypants strings.
    """
    if type(s) == 'Markup':
        s = s.unescape()

    # Evaulate COPY elements
    if type(s) is not unicode:
        s = unicode(s)

    s = s.encode('utf-8')
    s = smartypants(s)

    try:
        return Markup(s)
    except:
        print s
        return Markup(s)

<|endoftext|>


<commit_before>
"""
This package provides the database models using PonyORM.

Note that we do not rely on the auto increment feature as the previous
SQLAlchemy implementation did not set AUTO INCREMENT on the ID fields,
as PonyORM would.
"""

from flask import url_for
from flux import app, config, utils

import datetime
import hashlib
import os
import pony.orm as orm
import shutil
import uuid

db = orm.Database(**config.database)
session = orm.db_session
commit = orm.commit
rollback = orm.rollback
select = orm.select
desc = orm.desc

class User(db.Entity):
  _table_ = 'users'

  id = orm.PrimaryKey(int)
  name = orm.Required(str, unique=True)
  passhash = orm.Required(str)
  can_manage = orm.Required(bool)
  can_download_artifacts = orm.Required(bool)
  can_view_buildlogs = orm.Required(bool)
  login_tokens = orm.Set('LoginToken')

  def __init__(self, **kwargs):
    if 'id' not in kwargs:
      kwargs['id'] = (orm.max(x.id for x in User) or 0) + 1
    super().__init__(**kwargs)

  def set_password(self, password):
    self.passhash = utils.hash_pw(password)

  @classmethod
  def get_by_login_details(cls, user_name, password):
    passhash = utils.hash_pw(password)
    return orm.select(x for x in cls if x.name == user_name and
                        x.passhash == passhash).first()

  @classmethod
  def get_root_user(cls):
    return orm.select(x for x in cls if x.name == config.root_user).first()

  @classmethod
  def create_or_update_root(cls):
    root = cls.get_root_user()
    if root:
      # Make sure the root has all privileges.
      root.can_manage = True
      root.can_download_artifacts = True
      root.can_view_buildlogs = True
      root.set_password(config.root_password)
      root.name = config.root_user
    else:
      # Create a new root user.
      app.logger.info('Creating new root user: {!r}'.format(config.root_user))
      root = cls(
        name=config.root_user,
        passhash=utils.hash_pw(config.root_password),
        can_manage=True,
        can_download_artifacts=True,
        can_view_buildlogs=True)
    return root

  def url(self):
    return url_for('edit_user', user_id=self.id)

class LoginToken(db.Entity):
  """
  A login token represents the credentials that we can savely store in the
  browser's session and it will not reveal any information about the users
  password. For additional security, a login token is bound to the IP that
  the user logged in from an has an expiration date.

  The expiration duration can be set with the `login_token_duration`
  configuration value. Setting this option to #None will prevent tokens
  from expiring.
  """

  _table_ = 'logintokens'

  id = orm.PrimaryKey(int)
  ip = orm.Required(str)
  user = orm.Required(User)
  token = orm.Required(str, unique=True)
  created = orm.Required(datetime.datetime)

  @classmethod
  def create(cls, ip, user):
    " Create a new login token assigned to the specified IP and user. "

    id = (orm.max(x.id for x in cls) or 0) + 1
    created = datetime.datetime.now()
    token = str(uuid.uuid4()).replace('-', '')
    token += hashlib.md5((token + str(created)).encode()).hexdigest()
    return cls(id=id, ip=ip, user=user, token=token, created=created)

  def expired(self):
    " Returns #True if the token is expired, #False otherwise. "

    if config.login_token_duration is None:
      return False
    now = datetime.datetime.now()
    return (self.created + config.login_token_duration) < now

class Repository(db.Entity):
  """
  Represents a repository for which push events are being accepted. The Git
  server specified at the `clone_url` must accept the Flux server's public
  key.
  """

  _table_ = 'repos'

  id = orm.PrimaryKey(int)
  name = orm.Required(str)
  secret = orm.Required(str)
  clone_url = orm.Required(str)
  build_count = orm.Required(int, default=0)
  builds = orm.Set('Build')
  ref_whitelist = orm.Optional(str)  # newline separated list of accepted Git refs

  def __init__(self, **kwargs):
    if 'id' not in kwargs:
      kwargs['id'] = (orm.max(x.id for x in Repository) or 0) + 1
    super().__init__(**kwargs)

  def url(self, **kwargs):
    return url_for('view_repo', path=self.name, **kwargs)

  def check_accept_ref(self, ref):
    whitelist = list(filter(bool, self.ref_whitelist.split('\n')))
    if not whitelist or ref in whitelist:
      return True
    return False

  def validate_ref_whitelist(self, value, oldvalue, initiator):
    return '\n'.join(filter(bool, (x.strip() for x in value.split('\n'))))

  def most_recent_build(self):
    return self.builds.select().order_by(desc(Build.date_started)).first()

class Build(db.Entity):
  """
  Represents a build that is generated on a push to a repository. The build is
  initially queued and then processed when a slot is available. The build
  directory is generated from the configured root directory and the build
  #uuid. The log file has the exact same path with the `.log` suffix appended.

  After the build is complete (whether successful or errornous), the build
  directory is zipped and the original directory is removed.
  """

  _table_ = 'builds'

  Status_Queued = 'queued'
  Status_Building = 'building'
  Status_Error = 'error'
  Status_Success = 'success'
  Status_Stopped = 'stopped'
  Status = [Status_Queued, Status_Building, Status_Error, Status_Success, Status_Stopped]

  Data_BuildDir = 'build_dir'
  Data_OverrideDir = 'override_dir'
  Data_Artifact = 'artifact'
  Data_Log = 'log'

  class CanNotDelete(Exception):
    pass

  id = orm.PrimaryKey(int)
  repo = orm.Required(Repository, column='repo_id')
  ref = orm.Required(str)
  commit_sha = orm.Required(str)
  num = orm.Required(int)
  status = orm.Required(str)  # One of the Status strings
  date_queued = orm.Required(datetime.datetime, default=datetime.datetime.now)
  date_started = orm.Optional(datetime.datetime)
  date_finished = orm.Optional(datetime.datetime)

  def __init__(self, **kwargs):
    # Backwards compatibility for when SQLAlchemy was used, Auto Increment
    # was not enabled there.
    if 'id' not in kwargs:
      kwargs['id'] = (orm.max(x.id for x in Build) or 0) + 1
    super(Build, self).__init__(**kwargs)

  def url(self, data=None, **kwargs):
    path = self.repo.name + '/' + str(self.num)
    if not data:
      return url_for('view_build', path=path, **kwargs)
    elif data in (self.Data_Artifact, self.Data_Log):
      return url_for('download', build_id=self.id, data=data, **kwargs)
    else:
      raise ValueError('invalid mode: {!r}'.format(mode))

  def path(self, data=Data_BuildDir):
    base = os.path.join(config.build_dir, self.repo.name.replace('/', os.sep), str(self.num))
    if data == self.Data_BuildDir:
      return base
    elif data == self.Data_Artifact:
      return base + '.zip'
    elif data == self.Data_Log:
      return base + '.log'
    elif data == self.Data_OverrideDir:
      return os.path.join(config.override_dir, self.repo.name.replace('/', os.sep))
    else:
      raise ValueError('invalid value for "data": {!r}'.format(data))

  def exists(self, data):
    return os.path.exists(self.path(data))

  def log_contents(self):
    path = self.path(self.Data_Log)
    if os.path.isfile(path):
      with open(path, 'r') as fp:
        return fp.read()
    return None

  def check_download_permission(self, data, user):
    if data == self.Data_Artifact:
      return user.can_download_artifacts and (
        self.status == self.Status_Success or user.can_view_buildlogs)
    elif data == self.Data_Log:
      return user.can_view_buildlogs
    else:
      raise ValueError('invalid value for data: {!r}'.format(data))

  def delete_build(self):
    if self.status == self.Status_Building:
      raise self.CanNotDelete('can not delete build in progress')
    try:
      os.remove(self.path(self.Data_Artifact))
    except OSError as exc:
      app.logger.exception(exc)
    try:
      os.remove(self.path(self.Data_Log))
    except OSError as exc:
      app.logger.exception(exc)

  # db.Entity Overrides

  def before_delete(self):
    self.delete_build()

def get_target_for(path):
  """
  Given an URL path, returns either a #Repository or #Build that the path
  identifies. #None will be retunred if the path points to an unknown
  repository or build.

  Examples:

      /User/repo    => Repository(User/repo)
      /User/repo/1  => Build(1, Repository(User/repo))
  """

  parts = path.split('/')
  if len(parts) not in (2, 3):
    return None
  repo_name = parts[0] + '/' + parts[1]
  repo = Repository.get(name=repo_name)
  if not repo:
    return None
  if len(parts) == 3:
    try: num = int(parts[2])
    except ValueError: return None
    return Build.get(repo=repo, num=num)
  return repo

db.generate_mapping(create_tables=True)
<commit_msg>Fixed secret column in Repository model to be optional.<commit_after>
"""
This package provides the database models using PonyORM.

Note that we do not rely on the auto increment feature as the previous
SQLAlchemy implementation did not set AUTO INCREMENT on the ID fields,
as PonyORM would.
"""

from flask import url_for
from flux import app, config, utils

import datetime
import hashlib
import os
import pony.orm as orm
import shutil
import uuid

db = orm.Database(**config.database)
session = orm.db_session
commit = orm.commit
rollback = orm.rollback
select = orm.select
desc = orm.desc

class User(db.Entity):
  _table_ = 'users'

  id = orm.PrimaryKey(int)
  name = orm.Required(str, unique=True)
  passhash = orm.Required(str)
  can_manage = orm.Required(bool)
  can_download_artifacts = orm.Required(bool)
  can_view_buildlogs = orm.Required(bool)
  login_tokens = orm.Set('LoginToken')

  def __init__(self, **kwargs):
    if 'id' not in kwargs:
      kwargs['id'] = (orm.max(x.id for x in User) or 0) + 1
    super().__init__(**kwargs)

  def set_password(self, password):
    self.passhash = utils.hash_pw(password)

  @classmethod
  def get_by_login_details(cls, user_name, password):
    passhash = utils.hash_pw(password)
    return orm.select(x for x in cls if x.name == user_name and
                        x.passhash == passhash).first()

  @classmethod
  def get_root_user(cls):
    return orm.select(x for x in cls if x.name == config.root_user).first()

  @classmethod
  def create_or_update_root(cls):
    root = cls.get_root_user()
    if root:
      # Make sure the root has all privileges.
      root.can_manage = True
      root.can_download_artifacts = True
      root.can_view_buildlogs = True
      root.set_password(config.root_password)
      root.name = config.root_user
    else:
      # Create a new root user.
      app.logger.info('Creating new root user: {!r}'.format(config.root_user))
      root = cls(
        name=config.root_user,
        passhash=utils.hash_pw(config.root_password),
        can_manage=True,
        can_download_artifacts=True,
        can_view_buildlogs=True)
    return root

  def url(self):
    return url_for('edit_user', user_id=self.id)

class LoginToken(db.Entity):
  """
  A login token represents the credentials that we can savely store in the
  browser's session and it will not reveal any information about the users
  password. For additional security, a login token is bound to the IP that
  the user logged in from an has an expiration date.

  The expiration duration can be set with the `login_token_duration`
  configuration value. Setting this option to #None will prevent tokens
  from expiring.
  """

  _table_ = 'logintokens'

  id = orm.PrimaryKey(int)
  ip = orm.Required(str)
  user = orm.Required(User)
  token = orm.Required(str, unique=True)
  created = orm.Required(datetime.datetime)

  @classmethod
  def create(cls, ip, user):
    " Create a new login token assigned to the specified IP and user. "

    id = (orm.max(x.id for x in cls) or 0) + 1
    created = datetime.datetime.now()
    token = str(uuid.uuid4()).replace('-', '')
    token += hashlib.md5((token + str(created)).encode()).hexdigest()
    return cls(id=id, ip=ip, user=user, token=token, created=created)

  def expired(self):
    " Returns #True if the token is expired, #False otherwise. "

    if config.login_token_duration is None:
      return False
    now = datetime.datetime.now()
    return (self.created + config.login_token_duration) < now

class Repository(db.Entity):
  """
  Represents a repository for which push events are being accepted. The Git
  server specified at the `clone_url` must accept the Flux server's public
  key.
  """

  _table_ = 'repos'

  id = orm.PrimaryKey(int)
  name = orm.Required(str)
  secret = orm.Optional(str)
  clone_url = orm.Required(str)
  build_count = orm.Required(int, default=0)
  builds = orm.Set('Build')
  ref_whitelist = orm.Optional(str)  # newline separated list of accepted Git refs

  def __init__(self, **kwargs):
    if 'id' not in kwargs:
      kwargs['id'] = (orm.max(x.id for x in Repository) or 0) + 1
    super().__init__(**kwargs)

  def url(self, **kwargs):
    return url_for('view_repo', path=self.name, **kwargs)

  def check_accept_ref(self, ref):
    whitelist = list(filter(bool, self.ref_whitelist.split('\n')))
    if not whitelist or ref in whitelist:
      return True
    return False

  def validate_ref_whitelist(self, value, oldvalue, initiator):
    return '\n'.join(filter(bool, (x.strip() for x in value.split('\n'))))

  def most_recent_build(self):
    return self.builds.select().order_by(desc(Build.date_started)).first()

class Build(db.Entity):
  """
  Represents a build that is generated on a push to a repository. The build is
  initially queued and then processed when a slot is available. The build
  directory is generated from the configured root directory and the build
  #uuid. The log file has the exact same path with the `.log` suffix appended.

  After the build is complete (whether successful or errornous), the build
  directory is zipped and the original directory is removed.
  """

  _table_ = 'builds'

  Status_Queued = 'queued'
  Status_Building = 'building'
  Status_Error = 'error'
  Status_Success = 'success'
  Status_Stopped = 'stopped'
  Status = [Status_Queued, Status_Building, Status_Error, Status_Success, Status_Stopped]

  Data_BuildDir = 'build_dir'
  Data_OverrideDir = 'override_dir'
  Data_Artifact = 'artifact'
  Data_Log = 'log'

  class CanNotDelete(Exception):
    pass

  id = orm.PrimaryKey(int)
  repo = orm.Required(Repository, column='repo_id')
  ref = orm.Required(str)
  commit_sha = orm.Required(str)
  num = orm.Required(int)
  status = orm.Required(str)  # One of the Status strings
  date_queued = orm.Required(datetime.datetime, default=datetime.datetime.now)
  date_started = orm.Optional(datetime.datetime)
  date_finished = orm.Optional(datetime.datetime)

  def __init__(self, **kwargs):
    # Backwards compatibility for when SQLAlchemy was used, Auto Increment
    # was not enabled there.
    if 'id' not in kwargs:
      kwargs['id'] = (orm.max(x.id for x in Build) or 0) + 1
    super(Build, self).__init__(**kwargs)

  def url(self, data=None, **kwargs):
    path = self.repo.name + '/' + str(self.num)
    if not data:
      return url_for('view_build', path=path, **kwargs)
    elif data in (self.Data_Artifact, self.Data_Log):
      return url_for('download', build_id=self.id, data=data, **kwargs)
    else:
      raise ValueError('invalid mode: {!r}'.format(mode))

  def path(self, data=Data_BuildDir):
    base = os.path.join(config.build_dir, self.repo.name.replace('/', os.sep), str(self.num))
    if data == self.Data_BuildDir:
      return base
    elif data == self.Data_Artifact:
      return base + '.zip'
    elif data == self.Data_Log:
      return base + '.log'
    elif data == self.Data_OverrideDir:
      return os.path.join(config.override_dir, self.repo.name.replace('/', os.sep))
    else:
      raise ValueError('invalid value for "data": {!r}'.format(data))

  def exists(self, data):
    return os.path.exists(self.path(data))

  def log_contents(self):
    path = self.path(self.Data_Log)
    if os.path.isfile(path):
      with open(path, 'r') as fp:
        return fp.read()
    return None

  def check_download_permission(self, data, user):
    if data == self.Data_Artifact:
      return user.can_download_artifacts and (
        self.status == self.Status_Success or user.can_view_buildlogs)
    elif data == self.Data_Log:
      return user.can_view_buildlogs
    else:
      raise ValueError('invalid value for data: {!r}'.format(data))

  def delete_build(self):
    if self.status == self.Status_Building:
      raise self.CanNotDelete('can not delete build in progress')
    try:
      os.remove(self.path(self.Data_Artifact))
    except OSError as exc:
      app.logger.exception(exc)
    try:
      os.remove(self.path(self.Data_Log))
    except OSError as exc:
      app.logger.exception(exc)

  # db.Entity Overrides

  def before_delete(self):
    self.delete_build()

def get_target_for(path):
  """
  Given an URL path, returns either a #Repository or #Build that the path
  identifies. #None will be retunred if the path points to an unknown
  repository or build.

  Examples:

      /User/repo    => Repository(User/repo)
      /User/repo/1  => Build(1, Repository(User/repo))
  """

  parts = path.split('/')
  if len(parts) not in (2, 3):
    return None
  repo_name = parts[0] + '/' + parts[1]
  repo = Repository.get(name=repo_name)
  if not repo:
    return None
  if len(parts) == 3:
    try: num = int(parts[2])
    except ValueError: return None
    return Build.get(repo=repo, num=num)
  return repo

db.generate_mapping(create_tables=True)
<|endoftext|>


<commit_before>#!/bin/env python

import sys,pysam,argparse

def main(args):
    targetbam = pysam.Samfile(args.targetbam, 'rb')
    donorbam  = pysam.Samfile(args.donorbam, 'rb')
    outputbam = pysam.Samfile(args.outputbam, 'wb', template=targetbam)

    # load reads from donorbam into dict
    sys.stderr.write("loading donor reads into dictionary...\n")
    nr = 0
    rdict = {}
    for read in donorbam.fetch(until_eof=True):
        pairname = 'F' # read is first in pair
        if read.is_read2:
            pairname = 'S' # read is second in pair
        if not read.is_paired:
            pairname = 'U' # read is unpaired

        extqname = ','.join((read.qname,pairname))
        rdict[extqname] = read
        nr += 1

    sys.stderr.write("loaded " + str(nr) + " reads, reading " + args.targetbam + "...\n")
    for read in targetbam.fetch(until_eof=True):
        pairname = 'F' # read is first in pair
        if read.is_read2:
            pairname = 'S' # read is second in pair
        if not read.is_paired:
            pairname = 'U' # read is unpaired

        extqname = ','.join((read.qname,pairname))
        if extqname in rdict:
            outputbam.write(rdict[extqname])
        else:
            outputbam.write(read)

    targetbam.close()
    donorbam.close()
    outputbam.close()

if __name__=='__main__':
    parser = argparse.ArgumentParser(description='replaces aligned reads in bamfile1 with aligned reads from bamfile2')
    parser.add_argument('-b', '--bam', dest='targetbam', required=True,
                        help='original .bam')
    parser.add_argument('-r', '--replacebam', dest='donorbam', required=True,
                        help='.bam with reads to replace original bam')
    parser.add_argument('-o', '--outputbam', dest='outputbam', required=True)
    args = parser.parse_args()
    main(args)
<commit_msg>added all option to replacereads<commit_after>#!/bin/env python

import sys,pysam,argparse

def main(args):
    targetbam = pysam.Samfile(args.targetbam, 'rb')
    donorbam  = pysam.Samfile(args.donorbam, 'rb')
    outputbam = pysam.Samfile(args.outputbam, 'wb', template=targetbam)

    # load reads from donorbam into dict
    sys.stderr.write("loading donor reads into dictionary...\n")
    nr = 0
    rdict = {}
    for read in donorbam.fetch(until_eof=True):
        pairname = 'F' # read is first in pair
        if read.is_read2:
            pairname = 'S' # read is second in pair
        if not read.is_paired:
            pairname = 'U' # read is unpaired

        extqname = ','.join((read.qname,pairname))
        rdict[extqname] = read
        nr += 1

    sys.stderr.write("loaded " + str(nr) + " reads, reading " + args.targetbam + "...\n")

    used = {}
    for read in targetbam.fetch(until_eof=True):
        pairname = 'F' # read is first in pair
        if read.is_read2:
            pairname = 'S' # read is second in pair
        if not read.is_paired:
            pairname = 'U' # read is unpaired

        extqname = ','.join((read.qname,pairname))
        if extqname in rdict:
            outputbam.write(rdict[extqname])
            used[extqname] = True
        else:
            outputbam.write(read)

        # dump the unused reads from the donor if requested with --all
        if args.all:
            for extqname in rdict.keys():
                if extqname not in used:
                    outputbam.write(rdict[extqname])

    targetbam.close()
    donorbam.close()
    outputbam.close()

if __name__=='__main__':
    parser = argparse.ArgumentParser(description='replaces aligned reads in bamfile1 with aligned reads from bamfile2')
    parser.add_argument('-b', '--bam', dest='targetbam', required=True,
                        help='original .bam')
    parser.add_argument('-r', '--replacebam', dest='donorbam', required=True,
                        help='.bam with reads to replace original bam')
    parser.add_argument('-o', '--outputbam', dest='outputbam', required=True)
    parser.add_argument('--all', action='store_true', default=False, help="append reads that don't match target .bam")
    args = parser.parse_args()
    main(args)
<|endoftext|>


<commit_before>def usage(base_parser, subparsers): #pylint: disable=unused-variable
  parser = subparsers.add_parser('fsl', parents=[base_parser])
  parser.set_author('Robert E. Smith (robert.smith@florey.edu.au)')
  parser.set_synopsis('Use FSL commands to generate the 5TT image based on a T1-weighted image')
  parser.add_citation('', 'Smith, S. M. Fast robust automated brain extraction. Human Brain Mapping, 2002, 17, 143-155', True)
  parser.add_citation('', 'Zhang, Y.; Brady, M. & Smith, S. Segmentation of brain MR images through a hidden Markov random field model and the expectation-maximization algorithm. IEEE Transactions on Medical Imaging, 2001, 20, 45-57', True)
  parser.add_citation('', 'Patenaude, B.; Smith, S. M.; Kennedy, D. N. & Jenkinson, M. A Bayesian model of shape and appearance for subcortical brain segmentation. NeuroImage, 2011, 56, 907-922', True)
  parser.add_citation('', 'Smith, S. M.; Jenkinson, M.; Woolrich, M. W.; Beckmann, C. F.; Behrens, T. E.; Johansen-Berg, H.; Bannister, P. R.; De Luca, M.; Drobnjak, I.; Flitney, D. E.; Niazy, R. K.; Saunders, J.; Vickers, J.; Zhang, Y.; De Stefano, N.; Brady, J. M. & Matthews, P. M. Advances in functional and structural MR image analysis and implementation as FSL. NeuroImage, 2004, 23, S208-S219', True)
  parser.add_argument('input',  help='The input T1-weighted image')
  parser.add_argument('output', help='The output 5TT image')
  options = parser.add_argument_group('Options specific to the \'fsl\' algorithm')
  options.add_argument('-t2', metavar='<T2 image>', help='Provide a T2-weighted image in addition to the default T1-weighted image; this will be used as a second input to FSL FAST')
  options.add_argument('-mask', help='Manually provide a brain mask, rather than deriving one in the script')
  options.add_argument('-premasked', action='store_true', help='Indicate that brain masking has already been applied to the input image')
  parser.flag_mutually_exclusive_options( [ 'mask', 'premasked' ] )

def check_output_paths(): #pylint: disable=unused-variable
  from mrtrix3 import app
  app.check_output_path(app.ARGS.output)

def get_inputs(): #pylint: disable=unused-variable
  from mrtrix3 import app, image, MRtrixError, path, run
  image.check_3d_nonunity(path.from_user(app.ARGS.input, False))
  run.command('mrconvert ' + path.from_user(app.ARGS.input) + ' ' + path.to_scratch('input.mif'))
  if app.ARGS.mask:
    run.command('mrconvert ' + path.from_user(app.ARGS.mask) + ' ' + path.to_scratch('mask.mif') + ' -datatype bit -strides -1,+2,+3')
  if app.ARGS.t2:
    if not image.match(path.from_user(app.ARGS.input, False), path.from_user(app.ARGS.t2, False)):
      raise MRtrixError('Provided T2 image does not match input T1 image')
    run.command('mrconvert ' + path.from_user(app.ARGS.t2) + ' ' + path.to_scratch('T2.nii') + ' -strides -1,+2,+3')

def execute(): #pylint: disable=unused-variable
  import math, os
  from distutils import find_executable
  from mrtrix3 import app, fsl, image, is_windows, MRtrixError, path, run

  if is_windows():
    raise MRtrixError('\'fsl\' algorithm of 5ttgen script cannot be run on Windows: FSL not available on Windows')

  fsl_path = os.environ.get('FSLDIR', '')
  if not fsl_path:
    raise MRtrixError('Environment variable FSLDIR is not set; please run appropriate FSL configuration script')

  bet_cmd = fsl.exe_name('bet')
  fast_cmd = fsl.exe_name('fast')
  first_cmd = fsl.exe_name('run_first_all')
  ssroi_cmd = fsl.exe_name('standard_space_roi')

  first_atlas_path = os.path.join(fsl_path, 'data', 'first', 'models_336_bin')
  if not os.path.isdir(first_atlas_path):
    raise MRtrixError('Atlases required for FSL\'s FIRST program not installed; please install fsl-first-data using your relevant package manager')

  fsl_suffix = fsl.suffix()

  sgm_structures = [ 'L_Accu', 'R_Accu', 'L_Caud', 'R_Caud', 'L_Pall', 'R_Pall', 'L_Puta', 'R_Puta', 'L_Thal', 'R_Thal' ]
  if app.ARGS.sgm_amyg_hipp:
    sgm_structures.extend([ 'L_Amyg', 'R_Amyg', 'L_Hipp', 'R_Hipp' ])

  t1_spacing = image.Header('input.mif').spacing()
  upsample_for_first = False
  # If voxel size is 1.25mm or larger, make a guess that the user has erroneously re-gridded their data
  if math.pow(t1_spacing[0] * t1_spacing[1] * t1_spacing[2], 1.0/3.0) > 1.225:
    app.warn('Voxel size larger than expected for T1-weighted images (' + str(t1_spacing) + '); '
             'note that ACT does not require re-gridding of T1 image to DWI space, and indeed '
             'retaining the original higher resolution of the T1 image is preferable')
    upsample_for_first = True

  run.command('mrconvert input.mif T1.nii -strides -1,+2,+3')

  fast_t1_input = 'T1.nii'
  fast_t2_input = ''

  # Decide whether or not we're going to do any brain masking
  if os.path.exists('mask.mif'):

    fast_t1_input = 'T1_masked' + fsl_suffix

    # Check to see if the mask matches the T1 image
    if image.match('T1.nii', 'mask.mif'):
      run.command('mrcalc T1.nii mask.mif -mult ' + fast_t1_input)
      mask_path = 'mask.mif'
    else:
      app.warn('Mask image does not match input image - re-gridding')
      run.command('mrtransform mask.mif mask_regrid.mif -template T1.nii -datatype bit')
      run.command('mrcalc T1.nii mask_regrid.mif -mult ' + fast_t1_input)
      mask_path = 'mask_regrid.mif'

    if os.path.exists('T2.nii'):
      fast_t2_input = 'T2_masked' + fsl_suffix
      run.command('mrcalc T2.nii ' + mask_path + ' -mult ' + fast_t2_input)

  elif app.ARGS.premasked:

    fast_t1_input = 'T1.nii'
    if os.path.exists('T2.nii'):
      fast_t2_input = 'T2.nii'

  else:

    # Use FSL command standard_space_roi to do an initial masking of the image before BET
    # Also reduce the FoV of the image
    # Using MNI 1mm dilated brain mask rather than the -b option in standard_space_roi (which uses the 2mm mask); the latter looks 'buggy' to me... Unfortunately even with the 1mm 'dilated' mask, it can still cut into some brain areas, hence the explicit dilation
    mni_mask_path = os.path.join(fsl_path, 'data', 'standard', 'MNI152_T1_1mm_brain_mask_dil.nii.gz')
    mni_mask_dilation = 0
    if os.path.exists (mni_mask_path):
      mni_mask_dilation = 4
    else:
      mni_mask_path = os.path.join(fsl_path, 'data', 'standard', 'MNI152_T1_2mm_brain_mask_dil.nii.gz')
      if os.path.exists (mni_mask_path):
        mni_mask_dilation = 2
    try:
      if mni_mask_dilation:
        run.command('maskfilter ' + mni_mask_path + ' dilate mni_mask.nii -npass ' + str(mni_mask_dilation))
        if app.ARGS.nocrop:
          ssroi_roi_option = ' -roiNONE'
        else:
          ssroi_roi_option = ' -roiFOV'
        run.command(ssroi_cmd + ' T1.nii T1_preBET' + fsl_suffix + ' -maskMASK mni_mask.nii' + ssroi_roi_option)
      else:
        run.command(ssroi_cmd + ' T1.nii T1_preBET' + fsl_suffix + ' -b')
    except run.MRtrixCmdError:
      pass
    try:
      pre_bet_image = fsl.find_image('T1_preBET')
    except MRtrixError:
      app.warn('FSL script \'standard_space_roi\' did not complete successfully' + \
               ('' if find_executable('dc') else ' (possibly due to program \'dc\' not being installed') + '; ' + \
               'attempting to continue by providing un-cropped image to BET')
      pre_bet_image = 'T1.nii'

    # BET
    run.command(bet_cmd + ' ' + pre_bet_image + ' T1_BET' + fsl_suffix + ' -f 0.15 -R')
    fast_t1_input = fsl.find_image('T1_BET' + fsl_suffix)

    if os.path.exists('T2.nii'):
      if app.ARGS.nocrop:
        fast_t2_input = 'T2.nii'
      else:
        # Just a reduction of FoV, no sub-voxel interpolation going on
        run.command('mrtransform T2.nii T2_cropped.nii -template ' + fast_t1_input + ' -interp nearest')
        fast_t2_input = 'T2_cropped.nii'

  # Finish branching based on brain masking

  # FAST
  if fast_t2_input:
    run.command(fast_cmd + ' -S 2 ' + fast_t2_input + ' ' + fast_t1_input)
  else:
    run.command(fast_cmd + ' ' + fast_t1_input)

  # FIRST
  first_input = 'T1.nii'
  if upsample_for_first:
    app.warn('Generating 1mm isotropic T1 image for FIRST in hope of preventing failure, since input image is of lower resolution')
    run.command('mrresize T1.nii T1_1mm.nii -voxel 1.0 -interp sinc')
    first_input = 'T1_1mm.nii'
  first_brain_extracted_option = ''
  if app.ARGS.premasked:
    first_brain_extracted_option = ' -b'
  first_debug_option = ''
  if not app.DO_CLEANUP:
    first_debug_option = ' -d'
  first_verbosity_option = ''
  if app.VERBOSITY == 3:
    first_verbosity_option = ' -v'
  run.command(first_cmd + ' -m none -s ' + ','.join(sgm_structures) + ' -i ' + first_input + ' -o first' + first_brain_extracted_option + first_debug_option + first_verbosity_option)
  fsl.check_first('first', sgm_structures)

  # Convert FIRST meshes to partial volume images
  pve_image_list = [ ]
  progress = app.ProgressBar('Generating partial volume images for SGM structures', len(sgm_structures))
  for struct in sgm_structures:
    pve_image_path = 'mesh2voxel_' + struct + '.mif'
    vtk_in_path = 'first-' + struct + '_first.vtk'
    vtk_temp_path = struct + '.vtk'
    run.command('meshconvert ' + vtk_in_path + ' ' + vtk_temp_path + ' -transform first2real ' + first_input)
    run.command('mesh2voxel ' + vtk_temp_path + ' ' + fast_t1_input + ' ' + pve_image_path)
    pve_image_list.append(pve_image_path)
    progress.increment()
  progress.done()
  run.command(['mrmath', pve_image_list, 'sum', '-', '|', \
               'mrcalc', '-', '1.0', '-min', 'all_sgms.mif'])

  # Combine the tissue images into the 5TT format within the script itself
  fast_output_prefix = fast_t1_input.split('.')[0]
  fast_csf_output = fsl.find_image(fast_output_prefix + '_pve_0')
  fast_gm_output = fsl.find_image(fast_output_prefix + '_pve_1')
  fast_wm_output = fsl.find_image(fast_output_prefix + '_pve_2')
  # Step 1: Run LCC on the WM image
  run.command('mrthreshold ' + fast_wm_output + ' - -abs 0.001 | maskfilter - connect - -connectivity | mrcalc 1 - 1 -gt -sub remove_unconnected_wm_mask.mif -datatype bit')
  # Step 2: Generate the images in the same fashion as the old 5ttgen binary used to:
  #   - Preserve CSF as-is
  #   - Preserve SGM, unless it results in a sum of volume fractions greater than 1, in which case clamp
  #   - Multiply the FAST volume fractions of GM and CSF, so that the sum of CSF, SGM, CGM and WM is 1.0
  run.command('mrcalc ' + fast_csf_output + ' remove_unconnected_wm_mask.mif -mult csf.mif')
  run.command('mrcalc 1.0 csf.mif -sub all_sgms.mif -min sgm.mif')
  run.command('mrcalc 1.0 csf.mif sgm.mif -add -sub ' + fast_gm_output + ' ' + fast_wm_output + ' -add -div multiplier.mif')
  run.command('mrcalc multiplier.mif -finite multiplier.mif 0.0 -if multiplier_noNAN.mif')
  run.command('mrcalc ' + fast_gm_output + ' multiplier_noNAN.mif -mult remove_unconnected_wm_mask.mif -mult cgm.mif')
  run.command('mrcalc ' + fast_wm_output + ' multiplier_noNAN.mif -mult remove_unconnected_wm_mask.mif -mult wm.mif')
  run.command('mrcalc 0 wm.mif -min path.mif')
  run.command('mrcat cgm.mif sgm.mif wm.mif csf.mif path.mif - -axis 3 | mrconvert - combined_precrop.mif -strides +2,+3,+4,+1')

  # Use mrcrop to reduce file size (improves caching of image data during tracking)
  if app.ARGS.nocrop:
    run.function(os.rename, 'combined_precrop.mif', 'result.mif')
  else:
    run.command('mrmath combined_precrop.mif sum - -axis 3 | mrthreshold - - -abs 0.5 | mrcrop combined_precrop.mif result.mif -mask -')

  run.command('mrconvert result.mif ' + path.from_user(app.ARGS.output, True) + app.mrconvert_output_option(path.from_user(app.ARGS.input, True)))
<commit_msg>5ttgen fsl: Fix import error<commit_after>def usage(base_parser, subparsers): #pylint: disable=unused-variable
  parser = subparsers.add_parser('fsl', parents=[base_parser])
  parser.set_author('Robert E. Smith (robert.smith@florey.edu.au)')
  parser.set_synopsis('Use FSL commands to generate the 5TT image based on a T1-weighted image')
  parser.add_citation('', 'Smith, S. M. Fast robust automated brain extraction. Human Brain Mapping, 2002, 17, 143-155', True)
  parser.add_citation('', 'Zhang, Y.; Brady, M. & Smith, S. Segmentation of brain MR images through a hidden Markov random field model and the expectation-maximization algorithm. IEEE Transactions on Medical Imaging, 2001, 20, 45-57', True)
  parser.add_citation('', 'Patenaude, B.; Smith, S. M.; Kennedy, D. N. & Jenkinson, M. A Bayesian model of shape and appearance for subcortical brain segmentation. NeuroImage, 2011, 56, 907-922', True)
  parser.add_citation('', 'Smith, S. M.; Jenkinson, M.; Woolrich, M. W.; Beckmann, C. F.; Behrens, T. E.; Johansen-Berg, H.; Bannister, P. R.; De Luca, M.; Drobnjak, I.; Flitney, D. E.; Niazy, R. K.; Saunders, J.; Vickers, J.; Zhang, Y.; De Stefano, N.; Brady, J. M. & Matthews, P. M. Advances in functional and structural MR image analysis and implementation as FSL. NeuroImage, 2004, 23, S208-S219', True)
  parser.add_argument('input',  help='The input T1-weighted image')
  parser.add_argument('output', help='The output 5TT image')
  options = parser.add_argument_group('Options specific to the \'fsl\' algorithm')
  options.add_argument('-t2', metavar='<T2 image>', help='Provide a T2-weighted image in addition to the default T1-weighted image; this will be used as a second input to FSL FAST')
  options.add_argument('-mask', help='Manually provide a brain mask, rather than deriving one in the script')
  options.add_argument('-premasked', action='store_true', help='Indicate that brain masking has already been applied to the input image')
  parser.flag_mutually_exclusive_options( [ 'mask', 'premasked' ] )

def check_output_paths(): #pylint: disable=unused-variable
  from mrtrix3 import app
  app.check_output_path(app.ARGS.output)

def get_inputs(): #pylint: disable=unused-variable
  from mrtrix3 import app, image, MRtrixError, path, run
  image.check_3d_nonunity(path.from_user(app.ARGS.input, False))
  run.command('mrconvert ' + path.from_user(app.ARGS.input) + ' ' + path.to_scratch('input.mif'))
  if app.ARGS.mask:
    run.command('mrconvert ' + path.from_user(app.ARGS.mask) + ' ' + path.to_scratch('mask.mif') + ' -datatype bit -strides -1,+2,+3')
  if app.ARGS.t2:
    if not image.match(path.from_user(app.ARGS.input, False), path.from_user(app.ARGS.t2, False)):
      raise MRtrixError('Provided T2 image does not match input T1 image')
    run.command('mrconvert ' + path.from_user(app.ARGS.t2) + ' ' + path.to_scratch('T2.nii') + ' -strides -1,+2,+3')

def execute(): #pylint: disable=unused-variable
  import math, os
  from distutils.spawn import find_executable
  from mrtrix3 import app, fsl, image, is_windows, MRtrixError, path, run

  if is_windows():
    raise MRtrixError('\'fsl\' algorithm of 5ttgen script cannot be run on Windows: FSL not available on Windows')

  fsl_path = os.environ.get('FSLDIR', '')
  if not fsl_path:
    raise MRtrixError('Environment variable FSLDIR is not set; please run appropriate FSL configuration script')

  bet_cmd = fsl.exe_name('bet')
  fast_cmd = fsl.exe_name('fast')
  first_cmd = fsl.exe_name('run_first_all')
  ssroi_cmd = fsl.exe_name('standard_space_roi')

  first_atlas_path = os.path.join(fsl_path, 'data', 'first', 'models_336_bin')
  if not os.path.isdir(first_atlas_path):
    raise MRtrixError('Atlases required for FSL\'s FIRST program not installed; please install fsl-first-data using your relevant package manager')

  fsl_suffix = fsl.suffix()

  sgm_structures = [ 'L_Accu', 'R_Accu', 'L_Caud', 'R_Caud', 'L_Pall', 'R_Pall', 'L_Puta', 'R_Puta', 'L_Thal', 'R_Thal' ]
  if app.ARGS.sgm_amyg_hipp:
    sgm_structures.extend([ 'L_Amyg', 'R_Amyg', 'L_Hipp', 'R_Hipp' ])

  t1_spacing = image.Header('input.mif').spacing()
  upsample_for_first = False
  # If voxel size is 1.25mm or larger, make a guess that the user has erroneously re-gridded their data
  if math.pow(t1_spacing[0] * t1_spacing[1] * t1_spacing[2], 1.0/3.0) > 1.225:
    app.warn('Voxel size larger than expected for T1-weighted images (' + str(t1_spacing) + '); '
             'note that ACT does not require re-gridding of T1 image to DWI space, and indeed '
             'retaining the original higher resolution of the T1 image is preferable')
    upsample_for_first = True

  run.command('mrconvert input.mif T1.nii -strides -1,+2,+3')

  fast_t1_input = 'T1.nii'
  fast_t2_input = ''

  # Decide whether or not we're going to do any brain masking
  if os.path.exists('mask.mif'):

    fast_t1_input = 'T1_masked' + fsl_suffix

    # Check to see if the mask matches the T1 image
    if image.match('T1.nii', 'mask.mif'):
      run.command('mrcalc T1.nii mask.mif -mult ' + fast_t1_input)
      mask_path = 'mask.mif'
    else:
      app.warn('Mask image does not match input image - re-gridding')
      run.command('mrtransform mask.mif mask_regrid.mif -template T1.nii -datatype bit')
      run.command('mrcalc T1.nii mask_regrid.mif -mult ' + fast_t1_input)
      mask_path = 'mask_regrid.mif'

    if os.path.exists('T2.nii'):
      fast_t2_input = 'T2_masked' + fsl_suffix
      run.command('mrcalc T2.nii ' + mask_path + ' -mult ' + fast_t2_input)

  elif app.ARGS.premasked:

    fast_t1_input = 'T1.nii'
    if os.path.exists('T2.nii'):
      fast_t2_input = 'T2.nii'

  else:

    # Use FSL command standard_space_roi to do an initial masking of the image before BET
    # Also reduce the FoV of the image
    # Using MNI 1mm dilated brain mask rather than the -b option in standard_space_roi (which uses the 2mm mask); the latter looks 'buggy' to me... Unfortunately even with the 1mm 'dilated' mask, it can still cut into some brain areas, hence the explicit dilation
    mni_mask_path = os.path.join(fsl_path, 'data', 'standard', 'MNI152_T1_1mm_brain_mask_dil.nii.gz')
    mni_mask_dilation = 0
    if os.path.exists (mni_mask_path):
      mni_mask_dilation = 4
    else:
      mni_mask_path = os.path.join(fsl_path, 'data', 'standard', 'MNI152_T1_2mm_brain_mask_dil.nii.gz')
      if os.path.exists (mni_mask_path):
        mni_mask_dilation = 2
    try:
      if mni_mask_dilation:
        run.command('maskfilter ' + mni_mask_path + ' dilate mni_mask.nii -npass ' + str(mni_mask_dilation))
        if app.ARGS.nocrop:
          ssroi_roi_option = ' -roiNONE'
        else:
          ssroi_roi_option = ' -roiFOV'
        run.command(ssroi_cmd + ' T1.nii T1_preBET' + fsl_suffix + ' -maskMASK mni_mask.nii' + ssroi_roi_option)
      else:
        run.command(ssroi_cmd + ' T1.nii T1_preBET' + fsl_suffix + ' -b')
    except run.MRtrixCmdError:
      pass
    try:
      pre_bet_image = fsl.find_image('T1_preBET')
    except MRtrixError:
      app.warn('FSL script \'standard_space_roi\' did not complete successfully' + \
               ('' if find_executable('dc') else ' (possibly due to program \'dc\' not being installed') + '; ' + \
               'attempting to continue by providing un-cropped image to BET')
      pre_bet_image = 'T1.nii'

    # BET
    run.command(bet_cmd + ' ' + pre_bet_image + ' T1_BET' + fsl_suffix + ' -f 0.15 -R')
    fast_t1_input = fsl.find_image('T1_BET' + fsl_suffix)

    if os.path.exists('T2.nii'):
      if app.ARGS.nocrop:
        fast_t2_input = 'T2.nii'
      else:
        # Just a reduction of FoV, no sub-voxel interpolation going on
        run.command('mrtransform T2.nii T2_cropped.nii -template ' + fast_t1_input + ' -interp nearest')
        fast_t2_input = 'T2_cropped.nii'

  # Finish branching based on brain masking

  # FAST
  if fast_t2_input:
    run.command(fast_cmd + ' -S 2 ' + fast_t2_input + ' ' + fast_t1_input)
  else:
    run.command(fast_cmd + ' ' + fast_t1_input)

  # FIRST
  first_input = 'T1.nii'
  if upsample_for_first:
    app.warn('Generating 1mm isotropic T1 image for FIRST in hope of preventing failure, since input image is of lower resolution')
    run.command('mrresize T1.nii T1_1mm.nii -voxel 1.0 -interp sinc')
    first_input = 'T1_1mm.nii'
  first_brain_extracted_option = ''
  if app.ARGS.premasked:
    first_brain_extracted_option = ' -b'
  first_debug_option = ''
  if not app.DO_CLEANUP:
    first_debug_option = ' -d'
  first_verbosity_option = ''
  if app.VERBOSITY == 3:
    first_verbosity_option = ' -v'
  run.command(first_cmd + ' -m none -s ' + ','.join(sgm_structures) + ' -i ' + first_input + ' -o first' + first_brain_extracted_option + first_debug_option + first_verbosity_option)
  fsl.check_first('first', sgm_structures)

  # Convert FIRST meshes to partial volume images
  pve_image_list = [ ]
  progress = app.ProgressBar('Generating partial volume images for SGM structures', len(sgm_structures))
  for struct in sgm_structures:
    pve_image_path = 'mesh2voxel_' + struct + '.mif'
    vtk_in_path = 'first-' + struct + '_first.vtk'
    vtk_temp_path = struct + '.vtk'
    run.command('meshconvert ' + vtk_in_path + ' ' + vtk_temp_path + ' -transform first2real ' + first_input)
    run.command('mesh2voxel ' + vtk_temp_path + ' ' + fast_t1_input + ' ' + pve_image_path)
    pve_image_list.append(pve_image_path)
    progress.increment()
  progress.done()
  run.command(['mrmath', pve_image_list, 'sum', '-', '|', \
               'mrcalc', '-', '1.0', '-min', 'all_sgms.mif'])

  # Combine the tissue images into the 5TT format within the script itself
  fast_output_prefix = fast_t1_input.split('.')[0]
  fast_csf_output = fsl.find_image(fast_output_prefix + '_pve_0')
  fast_gm_output = fsl.find_image(fast_output_prefix + '_pve_1')
  fast_wm_output = fsl.find_image(fast_output_prefix + '_pve_2')
  # Step 1: Run LCC on the WM image
  run.command('mrthreshold ' + fast_wm_output + ' - -abs 0.001 | maskfilter - connect - -connectivity | mrcalc 1 - 1 -gt -sub remove_unconnected_wm_mask.mif -datatype bit')
  # Step 2: Generate the images in the same fashion as the old 5ttgen binary used to:
  #   - Preserve CSF as-is
  #   - Preserve SGM, unless it results in a sum of volume fractions greater than 1, in which case clamp
  #   - Multiply the FAST volume fractions of GM and CSF, so that the sum of CSF, SGM, CGM and WM is 1.0
  run.command('mrcalc ' + fast_csf_output + ' remove_unconnected_wm_mask.mif -mult csf.mif')
  run.command('mrcalc 1.0 csf.mif -sub all_sgms.mif -min sgm.mif')
  run.command('mrcalc 1.0 csf.mif sgm.mif -add -sub ' + fast_gm_output + ' ' + fast_wm_output + ' -add -div multiplier.mif')
  run.command('mrcalc multiplier.mif -finite multiplier.mif 0.0 -if multiplier_noNAN.mif')
  run.command('mrcalc ' + fast_gm_output + ' multiplier_noNAN.mif -mult remove_unconnected_wm_mask.mif -mult cgm.mif')
  run.command('mrcalc ' + fast_wm_output + ' multiplier_noNAN.mif -mult remove_unconnected_wm_mask.mif -mult wm.mif')
  run.command('mrcalc 0 wm.mif -min path.mif')
  run.command('mrcat cgm.mif sgm.mif wm.mif csf.mif path.mif - -axis 3 | mrconvert - combined_precrop.mif -strides +2,+3,+4,+1')

  # Use mrcrop to reduce file size (improves caching of image data during tracking)
  if app.ARGS.nocrop:
    run.function(os.rename, 'combined_precrop.mif', 'result.mif')
  else:
    run.command('mrmath combined_precrop.mif sum - -axis 3 | mrthreshold - - -abs 0.5 | mrcrop combined_precrop.mif result.mif -mask -')

  run.command('mrconvert result.mif ' + path.from_user(app.ARGS.output, True) + app.mrconvert_output_option(path.from_user(app.ARGS.input, True)))
<|endoftext|>