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

import os,sys
from optparse import OptionParser
import time
import fastq
import preprocesser
from multiprocessing import Process, Queue
import copy
from util import *

def parseCommand():
    usage = "Automatic Filtering, Trimming, and Error Removing for Illumina fastq data(Illumina 1.8 or newer format, see http://support.illumina.com/help/SequencingAnalysisWorkflow/Content/Vault/Informatics/Sequencing_Analysis/CASAVA/swSEQ_mCA_FASTQFiles.htm)\n\nFull command:\n%prog [-d input_dir][-1 read1_file] [-2 read1_file] [-7 index1_file] [-5 index2_file] [-g good_output_folder] [-b bad_output_folder] [-f trim_front] [-t trim_tail] [-q qualified_quality_phred] [-l unqualified_base_limit] [-p poly_size_limit] [-a allow_mismatch_in_poly] [-n n_base_limit] [--debubble=on/off] [--debubble_dir=xxx] [--draw=on/off] [--read1_flag=_R1_] [--read2_flag=_R2_] [--index1_flag=_I1_] [--index2_flag=_I2_] \n\nSimplest usage:\ncd to the folder containing your fastq data, run <python after.py>"
    version = "%prog 0.1.0"
    parser = OptionParser(usage = usage, version = version)
    parser.add_option("-1", "--read1_file", dest = "read1_file",
        help = "file name of read1, required. If input_dir is specified, then this arg is ignored.")
    parser.add_option("-2", "--read2_file", dest = "read2_file", default = None,
        help = "file name of read2, if paired. If input_dir is specified, then this arg is ignored.")
    parser.add_option("-7", "--index1_file", dest = "index1_file", default = None,
        help = "file name of 7' index. If input_dir is specified, then this arg is ignored.")
    parser.add_option("-5", "--index2_file", dest = "index2_file", default = None,
        help = "file name of 5' index. If input_dir is specified, then this arg is ignored.")
    parser.add_option("-d", "--input_dir", dest = "input_dir", default = None,
        help = "the input dir to process automatically. If read1_file are input_dir are not specified, then current dir (.) is specified to input_dir")
    parser.add_option("-g", "--good_output_folder", dest = "good_output_folder", default = "good",
        help = "the folder to store good reads, by default it is the same folder contains read1")
    parser.add_option("-b", "--bad_output_folder", dest = "bad_output_folder", default = "bad",
        help = "the folder to store bad reads, by default it is same as good_output_folder")
    parser.add_option("", "--read1_flag", dest = "read1_flag", default = "_R1_",
        help = "specify the name flag of read1, default is _R1_, which means a file with name *_R1_* is read1 file")
    parser.add_option("", "--read2_flag", dest = "read2_flag", default = "_R2_",
        help = "specify the name flag of read2, default is _R2_, which means a file with name *_R2_* is read2 file")
    parser.add_option("", "--index1_flag", dest = "index1_flag", default = "_I1_",
        help = "specify the name flag of index1, default is _I1_, which means a file with name *_I1_* is index2 file")
    parser.add_option("", "--index2_flag", dest = "index2_flag", default = "_I2_",
        help = "specify the name flag of index2, default is _I2_, which means a file with name *_I2_* is index2 file")
    parser.add_option("-f", "--trim_front", dest = "trim_front", default = -1, type = "int",
        help = "number of bases to be trimmed in the head of read. -1 means auto detect")
    parser.add_option("-t", "--trim_tail", dest = "trim_tail", default = -1, type = "int",
        help = "number of bases to be trimmed in the tail of read. -1 means auto detect")
    parser.add_option("", "--trim_pair_same", dest = "trim_pair_same", default = "true",
        help = "use same trimming configuration for read1 and read2 to keep their sequence length identical, default is true")
    parser.add_option("-q", "--qualified_quality_phred", dest = "qualified_quality_phred", default = 14, type = "int",
        help = "the quality value that a base is qualifyed. Default 14 means phred base quality >=Q14 is qualified.")
    parser.add_option("-u", "--unqualified_base_limit", dest = "unqualified_base_limit", default = 10, type = "int",
        help = "if exists more than unqualified_base_limit bases that quality is lower than qualified quality, then this read/pair is bad. Default is 10")
    parser.add_option("-p", "--poly_size_limit", dest = "poly_size_limit", default = 35, type = "int",
        help = "if exists one polyX(polyG means GGGGGGGGG...), and its length is >= poly_size_limit, then this read/pair is bad. Default is 35")
    parser.add_option("-a", "--allow_mismatch_in_poly", dest = "allow_mismatch_in_poly", default = 5, type = "int",
        help = "the count of allowed mismatches when evaluating poly_X. Default 5 means disallow any mismatches")
    parser.add_option("-n", "--n_base_limit", dest = "n_base_limit", default = 5, type = "int",
        help = "if exists more than maxn bases have N, then this read/pair is bad. Default is 5")
    parser.add_option("-s", "--seq_len_req", dest = "seq_len_req", default = 35, type = "int",
        help = "if the trimmed read is shorter than seq_len_req, then this read/pair is bad. Default is 35")
    parser.add_option("", "--debubble", dest = "debubble", default = "off",
        help = "specify whether apply debubble algorithm to remove the reads in the bubbles. Default is off")
    parser.add_option("", "--debubble_dir", dest = "debubble_dir", default = "debubble",
        help = "specify the folder to store output of debubble algorithm, default is debubble")
    parser.add_option("", "--draw", dest = "draw", default = "on",
        help = "specify whether draw the pictures or not, when use debubble or QC. Default is on")
    parser.add_option("", "--barcode", dest = "barcode", default = "on",
        help = "specify whether deal with barcode sequencing files, default is on, which means all files with barcode_flag in filename will be treated as barcode sequencing files")
    parser.add_option("", "--barcode_length", dest = "barcode_length", default = 12, type="int",
        help = "specify the designed length of barcode")
    parser.add_option("", "--barcode_flag", dest = "barcode_flag", default = "barcode",
        help = "specify the name flag of a barcoded file, default is barcode, which means a file with name *barcode* is a barcoded file")
    parser.add_option("", "--barcode_verify", dest = "barcode_verify", default = "CAGTA",
        help = "specify the verify sequence of a barcode which is adjunct to the barcode")
    parser.add_option("", "--store_overlap", dest = "store_overlap", default = "off",
        help = "specify whether store only overlapped bases of the good reads")
    parser.add_option("", "--overlap_output_folder", dest = "overlap_output_folder", default = "overlap",
        help = "the folder to store only overlapped bases of the good reads")
    return parser.parse_args()

def processDir(folder, options):

    fqext = (".fq", ".fastq", "fq.gz", ".fastq.gz")
    read1name = options.read1_flag
    read2name = options.read2_flag
    index1name = options.index1_flag
    index2name = options.index2_flag
    barcodeflag = options.barcode_flag

    #is not a dir
    if not os.path.isdir(folder):
        return

    options_list = []

    files = os.listdir(folder)
    for f in files:
        path = os.path.join(folder, f)
        if os.path.isdir(path):
            continue

        isfq = False
        for ext in fqext:
            if f.endswith(ext):isfq = True
        if isfq == False:
            continue

        # here we skip those files with name starting with Undetermined
        # because these files are usually with unknown barcode and have no need to be processed
        if f.startswith("Undetermined"):
            continue

        #find read1 file
        if read1name in f:
            print(f)
            opt = copy.copy(options)
            read1 = path
            opt.read1_file = read1
            read2 = read1.replace(read1name, read2name)
            if os.path.exists(read2):opt.read2_file = read2
            index1 = read1.replace(read1name, index1name)
            if os.path.exists(index1):opt.index1_file = index1
            index2 = read1.replace(read1name, index2name)
            if os.path.exists(index2):opt.index2_file = index2
            if barcodeflag in f and parseBool(options.barcode):
                opt.barcode = True
                #for barcode sequencing we don't trim it at front
                opt.trim_front = 0
                opt.trim_front2 = 0
            else:
                opt.barcode = False
            options_list.append(opt)

    if len(options_list) == 0:
        print("no read files to run with, do you call the program correctly?")
        print("see -h for help")
        return

    #create a job for each options
    jobs = [Process(target = processOptions, args = (o, )) for o in options_list]
    for job in jobs: job.start()
    #wait for them to complete
    for job in jobs: job.join()

def processOptions(options):
    filter = preprocesser.seqFilter(options)
    filter.run()

def runDebubble(options):
    #lazy import debubble here because debubble uses PIL, which is not supported by pypy
    #we can run with pypy with debubble off
    import debubble
    if os.path.exists(os.path.join(options.debubble_dir, "circles.csv")):
        return
    print("runDebubble")
    debubble.debubbleDir(options.input_dir, 20, options.debubble_dir, options.draw)

def main():
    time1 = time.time()
    if sys.version_info.major >2:
        print('python3 is not supported yet, please use python2')
        sys.exit(1)

    (options, args) = parseCommand()
    options.debubble = parseBool(options.debubble)
    options.trim_pair_same = parseBool(options.trim_pair_same)
    options.draw = parseBool(options.draw)
    options.store_overlap = parseBool(options.store_overlap)
    options.trim_front2 = options.trim_front
    options.trim_tail2 = options.trim_tail

    if options.input_dir == None and options.read1_file == None:
        print('specify current dir as input dir')
        options.input_dir="."

    if options.input_dir != None:
        if options.debubble:
            runDebubble(options)
        processDir(options.input_dir, options)
    else:
        processOptions(options)

    time2 = time.time()
    print('Time used: ' + str(time2-time1))

if __name__  == "__main__":
    main()
<commit_msg>improve read1/read2/index1/index2 file flags<commit_after> #!/usr/bin/env python

import os,sys
from optparse import OptionParser
import time
import fastq
import preprocesser
from multiprocessing import Process, Queue
import copy
from util import *

def parseCommand():
    usage = "Automatic Filtering, Trimming, and Error Removing for Illumina fastq data(Illumina 1.8 or newer format, see http://support.illumina.com/help/SequencingAnalysisWorkflow/Content/Vault/Informatics/Sequencing_Analysis/CASAVA/swSEQ_mCA_FASTQFiles.htm)\n\nFull command:\n%prog [-d input_dir][-1 read1_file] [-2 read1_file] [-7 index1_file] [-5 index2_file] [-g good_output_folder] [-b bad_output_folder] [-f trim_front] [-t trim_tail] [-q qualified_quality_phred] [-l unqualified_base_limit] [-p poly_size_limit] [-a allow_mismatch_in_poly] [-n n_base_limit] [--debubble=on/off] [--debubble_dir=xxx] [--draw=on/off] [--read1_flag=_R1_] [--read2_flag=_R2_] [--index1_flag=_I1_] [--index2_flag=_I2_] \n\nSimplest usage:\ncd to the folder containing your fastq data, run <python after.py>"
    version = "%prog 0.1.0"
    parser = OptionParser(usage = usage, version = version)
    parser.add_option("-1", "--read1_file", dest = "read1_file",
        help = "file name of read1, required. If input_dir is specified, then this arg is ignored.")
    parser.add_option("-2", "--read2_file", dest = "read2_file", default = None,
        help = "file name of read2, if paired. If input_dir is specified, then this arg is ignored.")
    parser.add_option("-7", "--index1_file", dest = "index1_file", default = None,
        help = "file name of 7' index. If input_dir is specified, then this arg is ignored.")
    parser.add_option("-5", "--index2_file", dest = "index2_file", default = None,
        help = "file name of 5' index. If input_dir is specified, then this arg is ignored.")
    parser.add_option("-d", "--input_dir", dest = "input_dir", default = None,
        help = "the input dir to process automatically. If read1_file are input_dir are not specified, then current dir (.) is specified to input_dir")
    parser.add_option("-g", "--good_output_folder", dest = "good_output_folder", default = "good",
        help = "the folder to store good reads, by default it is the same folder contains read1")
    parser.add_option("-b", "--bad_output_folder", dest = "bad_output_folder", default = "bad",
        help = "the folder to store bad reads, by default it is same as good_output_folder")
    parser.add_option("", "--read1_flag", dest = "read1_flag", default = "R1",
        help = "specify the name flag of read1, default is R1, which means a file with name *R1* is read1 file")
    parser.add_option("", "--read2_flag", dest = "read2_flag", default = "R2",
        help = "specify the name flag of read2, default is R2, which means a file with name *R2* is read2 file")
    parser.add_option("", "--index1_flag", dest = "index1_flag", default = "I1",
        help = "specify the name flag of index1, default is I1, which means a file with name *I1* is index2 file")
    parser.add_option("", "--index2_flag", dest = "index2_flag", default = "I2",
        help = "specify the name flag of index2, default is I2, which means a file with name *I2* is index2 file")
    parser.add_option("-f", "--trim_front", dest = "trim_front", default = -1, type = "int",
        help = "number of bases to be trimmed in the head of read. -1 means auto detect")
    parser.add_option("-t", "--trim_tail", dest = "trim_tail", default = -1, type = "int",
        help = "number of bases to be trimmed in the tail of read. -1 means auto detect")
    parser.add_option("", "--trim_pair_same", dest = "trim_pair_same", default = "true",
        help = "use same trimming configuration for read1 and read2 to keep their sequence length identical, default is true")
    parser.add_option("-q", "--qualified_quality_phred", dest = "qualified_quality_phred", default = 14, type = "int",
        help = "the quality value that a base is qualifyed. Default 14 means phred base quality >=Q14 is qualified.")
    parser.add_option("-u", "--unqualified_base_limit", dest = "unqualified_base_limit", default = 10, type = "int",
        help = "if exists more than unqualified_base_limit bases that quality is lower than qualified quality, then this read/pair is bad. Default is 10")
    parser.add_option("-p", "--poly_size_limit", dest = "poly_size_limit", default = 35, type = "int",
        help = "if exists one polyX(polyG means GGGGGGGGG...), and its length is >= poly_size_limit, then this read/pair is bad. Default is 35")
    parser.add_option("-a", "--allow_mismatch_in_poly", dest = "allow_mismatch_in_poly", default = 5, type = "int",
        help = "the count of allowed mismatches when evaluating poly_X. Default 5 means disallow any mismatches")
    parser.add_option("-n", "--n_base_limit", dest = "n_base_limit", default = 5, type = "int",
        help = "if exists more than maxn bases have N, then this read/pair is bad. Default is 5")
    parser.add_option("-s", "--seq_len_req", dest = "seq_len_req", default = 35, type = "int",
        help = "if the trimmed read is shorter than seq_len_req, then this read/pair is bad. Default is 35")
    parser.add_option("", "--debubble", dest = "debubble", default = "off",
        help = "specify whether apply debubble algorithm to remove the reads in the bubbles. Default is off")
    parser.add_option("", "--debubble_dir", dest = "debubble_dir", default = "debubble",
        help = "specify the folder to store output of debubble algorithm, default is debubble")
    parser.add_option("", "--draw", dest = "draw", default = "on",
        help = "specify whether draw the pictures or not, when use debubble or QC. Default is on")
    parser.add_option("", "--barcode", dest = "barcode", default = "on",
        help = "specify whether deal with barcode sequencing files, default is on, which means all files with barcode_flag in filename will be treated as barcode sequencing files")
    parser.add_option("", "--barcode_length", dest = "barcode_length", default = 12, type="int",
        help = "specify the designed length of barcode")
    parser.add_option("", "--barcode_flag", dest = "barcode_flag", default = "barcode",
        help = "specify the name flag of a barcoded file, default is barcode, which means a file with name *barcode* is a barcoded file")
    parser.add_option("", "--barcode_verify", dest = "barcode_verify", default = "CAGTA",
        help = "specify the verify sequence of a barcode which is adjunct to the barcode")
    parser.add_option("", "--store_overlap", dest = "store_overlap", default = "off",
        help = "specify whether store only overlapped bases of the good reads")
    parser.add_option("", "--overlap_output_folder", dest = "overlap_output_folder", default = "overlap",
        help = "the folder to store only overlapped bases of the good reads")
    return parser.parse_args()

def matchFlag(filename, flag):
    if flag.endswith('.') or flag.endswith('_') or flag.endswith('-'):
        return flag in filename
    else:
        return (flag+"." in filename) or (flag+"_" in filename) or (flag+"-" in filename)

def processDir(folder, options):

    fqext = (".fq", ".fastq", "fq.gz", ".fastq.gz")
    read1name = options.read1_flag
    read2name = options.read2_flag
    index1name = options.index1_flag
    index2name = options.index2_flag
    barcodeflag = options.barcode_flag

    #is not a dir
    if not os.path.isdir(folder):
        return

    options_list = []

    files = os.listdir(folder)
    for f in files:
        path = os.path.join(folder, f)
        if os.path.isdir(path):
            continue

        isfq = False
        for ext in fqext:
            if f.endswith(ext):isfq = True
        if isfq == False:
            continue

        # here we skip those files with name starting with Undetermined
        # because these files are usually with unknown barcode and have no need to be processed
        if f.startswith("Undetermined"):
            continue

        #find read1 file
        if matchFlag(f, read1name):
            print(f)
            opt = copy.copy(options)
            read1 = path
            opt.read1_file = read1
            read2 = read1.replace(read1name, read2name)
            if os.path.exists(read2):opt.read2_file = read2
            index1 = read1.replace(read1name, index1name)
            if os.path.exists(index1):opt.index1_file = index1
            index2 = read1.replace(read1name, index2name)
            if os.path.exists(index2):opt.index2_file = index2
            if barcodeflag in f and parseBool(options.barcode):
                opt.barcode = True
                #for barcode sequencing we don't trim it at front
                opt.trim_front = 0
                opt.trim_front2 = 0
            else:
                opt.barcode = False
            options_list.append(opt)

    if len(options_list) == 0:
        print("no read files to run with, do you call the program correctly?")
        print("see -h for help")
        return

    #create a job for each options
    jobs = [Process(target = processOptions, args = (o, )) for o in options_list]
    for job in jobs: job.start()
    #wait for them to complete
    for job in jobs: job.join()

def processOptions(options):
    filter = preprocesser.seqFilter(options)
    filter.run()

def runDebubble(options):
    #lazy import debubble here because debubble uses PIL, which is not supported by pypy
    #we can run with pypy with debubble off
    import debubble
    if os.path.exists(os.path.join(options.debubble_dir, "circles.csv")):
        return
    print("runDebubble")
    debubble.debubbleDir(options.input_dir, 20, options.debubble_dir, options.draw)

def main():
    time1 = time.time()
    if sys.version_info.major >2:
        print('python3 is not supported yet, please use python2')
        sys.exit(1)

    (options, args) = parseCommand()
    options.debubble = parseBool(options.debubble)
    options.trim_pair_same = parseBool(options.trim_pair_same)
    options.draw = parseBool(options.draw)
    options.store_overlap = parseBool(options.store_overlap)
    options.trim_front2 = options.trim_front
    options.trim_tail2 = options.trim_tail

    if options.input_dir == None and options.read1_file == None:
        print('specify current dir as input dir')
        options.input_dir="."

    if options.input_dir != None:
        if options.debubble:
            runDebubble(options)
        processDir(options.input_dir, options)
    else:
        processOptions(options)

    time2 = time.time()
    print('Time used: ' + str(time2-time1))

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


<commit_before>""" See COPYING for license information """

class Error(StandardError):
    """
    Base class for all errors and exceptions
    """
    pass

class ResponseError(Error):
    """
    A General response error
    """
    def __init__(self, status, reason):
        self.status = status
        self.reason = reason
        Error.__init__(self)

    def __str__(self):
        return '%d - %s' % (self.status, self.reason)

    def __repr__(self):
        return '%d - %s' % (self.status, self.reason)

class HostListError(ResponseError):
    """
    Raised when host list generation fails.
    """
    pass

class ConfigFileError(ResponseError):
    """
    Raised when host list generation fails.
    """
    pass

class ConfigSyncError(ResponseError):
    """
    Raised when issue syncing configs to system's root
    """
    pass

class TemplateFileError(ResponseError):
    """
    Raised when something goes wrong
    manipulating a template file.
    """
    pass

class DiskSetupError(ResponseError):
    """
    Raised when something goes wrong
    setting up the disks.
    """
    pass
<commit_msg>fixing some of the comments and adding new class<commit_after>""" See COPYING for license information """

class Error(StandardError):
    """
    Base class for all errors and exceptions
    """
    pass

class ResponseError(Error):
    """
    A General response error
    """
    def __init__(self, status, reason):
        self.status = status
        self.reason = reason
        Error.__init__(self)

    def __str__(self):
        return '%d - %s' % (self.status, self.reason)

    def __repr__(self):
        return '%d - %s' % (self.status, self.reason)

class HostListError(ResponseError):
    """
    Raised when host list generation fails.
    """
    pass

class ConfigFileError(ResponseError):
    """
    Raised when host list generation fails.
    """
    pass

class ConfigSyncError(ResponseError):
    """
    Raised when issue syncing configs to system's root
    """
    pass

class TemplateFileError(ResponseError):
    """
    Raised when something goes wrong manipulating a template file.
    """
    pass

class DiskSetupError(ResponseError):
    """
    Raised when something goes wrong setting up the disks.
    """
    pass

class UploadTemplatesError(ResponseError):
    """
    Raised when something goes wrong uploading templates
    """
    pass
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
General documentation architecture:

Home
Index

- Getting started
    Getting started with the sequential model
    Getting started with the functional api
    Examples
    FAQ
    Installation guide

- Models
    About Keras models
        explain when one should use Sequential or functional API
        explain compilation step
        explain weight saving, weight loading
        explain serialization, deserialization
    Sequential
    Model (functional API)

- Layers
    About Keras layers
        explain common layer functions: get_weights, set_weights, get_config
        explain input_shape
        explain usage on non-Keras tensors
    Core layers
    Convolutional
    Recurrent
    Embeddings
    Normalization
    Advanced activations
    Noise

- Preprocessing
    Image preprocessing
    Text preprocessing
    Sequence preprocessing

Objectives
Metrics
Optimizers
Activations
Callbacks
Datasets
Backend
Initializations
Regularizers
Constraints
Visualization
Scikit-learn API

'''
from __future__ import print_function
from __future__ import unicode_literals

import re
import inspect
import os
import shutil
import sys
if sys.version[0] == '2':
    reload(sys)
    sys.setdefaultencoding('utf8')

from keras.layers import convolutional
from keras.layers import pooling
from keras.layers import local
from keras.layers import recurrent
from keras.layers import core
from keras.layers import noise
from keras.layers import normalization
from keras.layers import advanced_activations
from keras.layers import embeddings
from keras.layers import wrappers
from keras import optimizers
from keras import callbacks
from keras import models
from keras.engine import topology
from keras import objectives
from keras import metrics
from keras import backend
from keras import constraints
from keras import activations
from keras import regularizers
from keras.utils import io_utils

EXCLUDE = {
    'Optimizer',
    'Wrapper',
    'get_session',
    'set_session',
    'CallbackList',
}

PAGES = [
    {
        'page': 'models/sequential.md',
        'functions': [
            models.Sequential.compile,
            models.Sequential.fit,
            models.Sequential.evaluate,
            models.Sequential.predict,
            models.Sequential.predict_classes,
            models.Sequential.predict_proba,
            models.Sequential.train_on_batch,
            models.Sequential.test_on_batch,
            models.Sequential.predict_on_batch,
            models.Sequential.fit_generator,
            models.Sequential.evaluate_generator,
            models.Sequential.predict_generator,
        ],
    },
    {
        'page': 'models/model.md',
        'functions': [
            models.Model.compile,
            models.Model.fit,
            models.Model.evaluate,
            models.Model.predict,
            models.Model.train_on_batch,
            models.Model.test_on_batch,
            models.Model.predict_on_batch,
            models.Model.fit_generator,
            models.Model.evaluate_generator,
            models.Model.predict_generator,
            models.Model.get_layer,
        ]
    },
    {
        'page': 'layers/core.md',
        'classes': [
            core.Dense,
            core.Activation,
            core.Dropout,
            core.SpatialDropout2D,
            core.SpatialDropout3D,
            core.Flatten,
            core.Reshape,
            core.Permute,
            core.RepeatVector,
            topology.Merge,
            core.Lambda,
            core.ActivityRegularization,
            core.Masking,
            core.Highway,
            core.MaxoutDense,
            core.TimeDistributedDense,
        ],
    },
    {
        'page': 'layers/convolutional.md',
        'classes': [
            convolutional.Convolution1D,
            convolutional.AtrousConvolution1D,
            convolutional.Convolution2D,
            convolutional.AtrousConvolution2D,
            convolutional.SeparableConvolution2D,
            convolutional.Deconvolution2D,
            convolutional.Convolution3D,
            convolutional.Cropping1D,
            convolutional.Cropping2D,
            convolutional.Cropping3D,
            convolutional.UpSampling1D,
            convolutional.UpSampling2D,
            convolutional.UpSampling3D,
            convolutional.ZeroPadding1D,
            convolutional.ZeroPadding2D,
            convolutional.ZeroPadding3D,
        ],
    },
    {
        'page': 'layers/pooling.md',
        'classes': [
            pooling.MaxPooling1D,
            pooling.MaxPooling2D,
            pooling.MaxPooling3D,
            pooling.AveragePooling1D,
            pooling.AveragePooling2D,
            pooling.AveragePooling3D,
            pooling.GlobalMaxPooling1D,
            pooling.GlobalAveragePooling1D,
            pooling.GlobalMaxPooling2D,
            pooling.GlobalAveragePooling2D,
        ],
    },
    {
        'page': 'layers/local.md',
        'classes': [
            local.LocallyConnected1D,
            local.LocallyConnected2D,
        ],
    },
    {
        'page': 'layers/recurrent.md',
        'classes': [
            recurrent.Recurrent,
            recurrent.SimpleRNN,
            recurrent.GRU,
            recurrent.LSTM,
        ],
    },
    {
        'page': 'layers/embeddings.md',
        'classes': [
            embeddings.Embedding,
        ],
    },
    {
        'page': 'layers/normalization.md',
        'classes': [
            normalization.BatchNormalization,
        ],
    },
    {
        'page': 'layers/advanced-activations.md',
        'all_module_classes': [advanced_activations],
    },
    {
        'page': 'layers/noise.md',
        'all_module_classes': [noise],
    },
    {
        'page': 'layers/wrappers.md',
        'all_module_classes': [wrappers],
    },
    {
        'page': 'metrics.md',
        'all_module_functions': [metrics],
    },
    {
        'page': 'optimizers.md',
        'all_module_classes': [optimizers],
    },
    {
        'page': 'callbacks.md',
        'all_module_classes': [callbacks],
    },
    {
        'page': 'backend.md',
        'all_module_functions': [backend],
    },
    {
        'page': 'io_utils.md',
        'classes': [
            io_utils.HDF5Matrix
        ],
    },
]

ROOT = 'http://keras.io/'

def get_earliest_class_that_defined_member(member, cls):
    ancestors = get_classes_ancestors([cls])
    result = None
    for ancestor in ancestors:
        if member in dir(ancestor):
            result = ancestor
    if not result:
        return cls
    return result

def get_classes_ancestors(classes):
    ancestors = []
    for cls in classes:
        ancestors += cls.__bases__
    filtered_ancestors = []
    for ancestor in ancestors:
        if ancestor.__name__ in ['object']:
            continue
        filtered_ancestors.append(ancestor)
    if filtered_ancestors:
        return filtered_ancestors + get_classes_ancestors(filtered_ancestors)
    else:
        return filtered_ancestors

def get_function_signature(function, method=True):
    signature = inspect.getargspec(function)
    defaults = signature.defaults
    if method:
        args = signature.args[1:]
    else:
        args = signature.args
    if defaults:
        kwargs = zip(args[-len(defaults):], defaults)
        args = args[:-len(defaults)]
    else:
        kwargs = []
    st = '%s.%s(' % (function.__module__, function.__name__)
    for a in args:
        st += str(a) + ', '
    for a, v in kwargs:
        if type(v) == str:
            v = '\'' + v + '\''
        st += str(a) + '=' + str(v) + ', '
    if kwargs or args:
        return st[:-2] + ')'
    else:
        return st + ')'

def get_class_signature(cls):
    try:
        class_signature = get_function_signature(cls.__init__)
        class_signature = class_signature.replace('__init__', cls.__name__)
    except:
        # in case the class inherits from object and does not
        # define __init__
        class_signature = cls.__module__ + '.' + cls.__name__ + '()'
    return class_signature

def class_to_docs_link(cls):
    module_name = cls.__module__
    assert module_name[:6] == 'keras.'
    module_name = module_name[6:]
    link = ROOT + module_name.replace('.', '/') + '#' + cls.__name__.lower()
    return link

def class_to_source_link(cls):
    module_name = cls.__module__
    assert module_name[:6] == 'keras.'
    path = module_name.replace('.', '/')
    path += '.py'
    line = inspect.getsourcelines(cls)[-1]
    link = 'https://github.com/fchollet/keras/blob/master/' + path + '#L' + str(line)
    return '[[source]](' + link + ')'

def code_snippet(snippet):
    result = '```python\n'
    result += snippet + '\n'
    result += '```\n'
    return result

def process_class_docstring(docstring):
    docstring = re.sub(r'\n    # (.*)\n',
                       r'\n    __\1__\n\n',
                       docstring)

    docstring = re.sub(r'    ([^\s\\]+):(.*)\n',
                       r'    - __\1__:\2\n',
                       docstring)

    docstring = docstring.replace('    ' * 5, '\t\t')
    docstring = docstring.replace('    ' * 3, '\t')
    docstring = docstring.replace('    ', '')
    return docstring

def process_function_docstring(docstring):
    docstring = re.sub(r'\n    # (.*)\n',
                       r'\n    __\1__\n\n',
                       docstring)
    docstring = re.sub(r'\n        # (.*)\n',
                       r'\n        __\1__\n\n',
                       docstring)

    docstring = re.sub(r'    ([^\s\\]+):(.*)\n',
                       r'    - __\1__:\2\n',
                       docstring)

    docstring = docstring.replace('    ' * 6, '\t\t')
    docstring = docstring.replace('    ' * 4, '\t')
    docstring = docstring.replace('    ', '')
    return docstring

print('Cleaning up existing sources directory.')
if os.path.exists('sources'):
    shutil.rmtree('sources')

print('Populating sources directory with templates.')
for subdir, dirs, fnames in os.walk('templates'):
    for fname in fnames:
        new_subdir = subdir.replace('templates', 'sources')
        if not os.path.exists(new_subdir):
            os.makedirs(new_subdir)
        if fname[-3:] == '.md':
            fpath = os.path.join(subdir, fname)
            new_fpath = fpath.replace('templates', 'sources')
            shutil.copy(fpath, new_fpath)

print('Starting autogeneration.')
for page_data in PAGES:
    blocks = []
    classes = page_data.get('classes', [])
    for module in page_data.get('all_module_classes', []):
        module_classes = []
        for name in dir(module):
            if name[0] == '_' or name in EXCLUDE:
                continue
            module_member = getattr(module, name)
            if inspect.isclass(module_member):
                cls = module_member
                if cls.__module__ == module.__name__:
                    if cls not in module_classes:
                        module_classes.append(cls)
        module_classes.sort(key=lambda x: id(x))
        classes += module_classes

    for cls in classes:
        subblocks = []
        signature = get_class_signature(cls)
        subblocks.append('' + class_to_source_link(cls) + '')
        subblocks.append('### ' + cls.__name__ + '\n')
        subblocks.append(code_snippet(signature))
        docstring = cls.__doc__
        if docstring:
            subblocks.append(process_class_docstring(docstring))
        blocks.append('\n'.join(subblocks))

    functions = page_data.get('functions', [])
    for module in page_data.get('all_module_functions', []):
        module_functions = []
        for name in dir(module):
            if name[0] == '_' or name in EXCLUDE:
                continue
            module_member = getattr(module, name)
            if inspect.isfunction(module_member):
                function = module_member
                if module.__name__ in function.__module__:
                    if function not in module_functions:
                        module_functions.append(function)
        module_functions.sort(key=lambda x: id(x))
        functions += module_functions

    for function in functions:
        subblocks = []
        signature = get_function_signature(function, method=False)
        signature = signature.replace(function.__module__ + '.', '')
        subblocks.append('### ' + function.__name__ + '\n')
        subblocks.append(code_snippet(signature))
        docstring = function.__doc__
        if docstring:
            subblocks.append(process_function_docstring(docstring))
            blocks.append('\n\n'.join(subblocks))

    mkdown = '\n----\n\n'.join(blocks)
    # save module page.
    # Either insert content into existing page,
    # or create page otherwise
    page_name = page_data['page']
    path = os.path.join('sources', page_name)
    if os.path.exists(path):
        template = open(path).read()
        assert '{{autogenerated}}' in template, ('Template found for ' + path +
                                                 ' but missing {{autogenerated}} tag.')
        mkdown = template.replace('{{autogenerated}}', mkdown)
        print('...inserting autogenerated content into template:', path)
    else:
        print('...creating new page with autogenerated content:', path)
    subdir = os.path.dirname(path)
    if not os.path.exists(subdir):
        os.makedirs(subdir)
    open(path, 'w').write(mkdown)
<commit_msg>Add more util docs (#4154)<commit_after># -*- coding: utf-8 -*-
'''
General documentation architecture:

Home
Index

- Getting started
    Getting started with the sequential model
    Getting started with the functional api
    Examples
    FAQ
    Installation guide

- Models
    About Keras models
        explain when one should use Sequential or functional API
        explain compilation step
        explain weight saving, weight loading
        explain serialization, deserialization
    Sequential
    Model (functional API)

- Layers
    About Keras layers
        explain common layer functions: get_weights, set_weights, get_config
        explain input_shape
        explain usage on non-Keras tensors
    Core layers
    Convolutional
    Recurrent
    Embeddings
    Normalization
    Advanced activations
    Noise

- Preprocessing
    Image preprocessing
    Text preprocessing
    Sequence preprocessing

Objectives
Metrics
Optimizers
Activations
Callbacks
Datasets
Backend
Initializations
Regularizers
Constraints
Visualization
Scikit-learn API

'''
from __future__ import print_function
from __future__ import unicode_literals

import re
import inspect
import os
import shutil
import sys
if sys.version[0] == '2':
    reload(sys)
    sys.setdefaultencoding('utf8')

from keras.layers import convolutional
from keras.layers import pooling
from keras.layers import local
from keras.layers import recurrent
from keras.layers import core
from keras.layers import noise
from keras.layers import normalization
from keras.layers import advanced_activations
from keras.layers import embeddings
from keras.layers import wrappers
from keras import optimizers
from keras import callbacks
from keras import models
from keras.engine import topology
from keras import objectives
from keras import metrics
from keras import backend
from keras import constraints
from keras import activations
from keras import regularizers
from keras.utils import data_utils
from keras.utils import io_utils
from keras.utils import layer_utils
from keras.utils import np_utils

EXCLUDE = {
    'Optimizer',
    'Wrapper',
    'get_session',
    'set_session',
    'CallbackList',
}

PAGES = [
    {
        'page': 'models/sequential.md',
        'functions': [
            models.Sequential.compile,
            models.Sequential.fit,
            models.Sequential.evaluate,
            models.Sequential.predict,
            models.Sequential.predict_classes,
            models.Sequential.predict_proba,
            models.Sequential.train_on_batch,
            models.Sequential.test_on_batch,
            models.Sequential.predict_on_batch,
            models.Sequential.fit_generator,
            models.Sequential.evaluate_generator,
            models.Sequential.predict_generator,
        ],
    },
    {
        'page': 'models/model.md',
        'functions': [
            models.Model.compile,
            models.Model.fit,
            models.Model.evaluate,
            models.Model.predict,
            models.Model.train_on_batch,
            models.Model.test_on_batch,
            models.Model.predict_on_batch,
            models.Model.fit_generator,
            models.Model.evaluate_generator,
            models.Model.predict_generator,
            models.Model.get_layer,
        ]
    },
    {
        'page': 'layers/core.md',
        'classes': [
            core.Dense,
            core.Activation,
            core.Dropout,
            core.SpatialDropout2D,
            core.SpatialDropout3D,
            core.Flatten,
            core.Reshape,
            core.Permute,
            core.RepeatVector,
            topology.Merge,
            core.Lambda,
            core.ActivityRegularization,
            core.Masking,
            core.Highway,
            core.MaxoutDense,
            core.TimeDistributedDense,
        ],
    },
    {
        'page': 'layers/convolutional.md',
        'classes': [
            convolutional.Convolution1D,
            convolutional.AtrousConvolution1D,
            convolutional.Convolution2D,
            convolutional.AtrousConvolution2D,
            convolutional.SeparableConvolution2D,
            convolutional.Deconvolution2D,
            convolutional.Convolution3D,
            convolutional.Cropping1D,
            convolutional.Cropping2D,
            convolutional.Cropping3D,
            convolutional.UpSampling1D,
            convolutional.UpSampling2D,
            convolutional.UpSampling3D,
            convolutional.ZeroPadding1D,
            convolutional.ZeroPadding2D,
            convolutional.ZeroPadding3D,
        ],
    },
    {
        'page': 'layers/pooling.md',
        'classes': [
            pooling.MaxPooling1D,
            pooling.MaxPooling2D,
            pooling.MaxPooling3D,
            pooling.AveragePooling1D,
            pooling.AveragePooling2D,
            pooling.AveragePooling3D,
            pooling.GlobalMaxPooling1D,
            pooling.GlobalAveragePooling1D,
            pooling.GlobalMaxPooling2D,
            pooling.GlobalAveragePooling2D,
        ],
    },
    {
        'page': 'layers/local.md',
        'classes': [
            local.LocallyConnected1D,
            local.LocallyConnected2D,
        ],
    },
    {
        'page': 'layers/recurrent.md',
        'classes': [
            recurrent.Recurrent,
            recurrent.SimpleRNN,
            recurrent.GRU,
            recurrent.LSTM,
        ],
    },
    {
        'page': 'layers/embeddings.md',
        'classes': [
            embeddings.Embedding,
        ],
    },
    {
        'page': 'layers/normalization.md',
        'classes': [
            normalization.BatchNormalization,
        ],
    },
    {
        'page': 'layers/advanced-activations.md',
        'all_module_classes': [advanced_activations],
    },
    {
        'page': 'layers/noise.md',
        'all_module_classes': [noise],
    },
    {
        'page': 'layers/wrappers.md',
        'all_module_classes': [wrappers],
    },
    {
        'page': 'metrics.md',
        'all_module_functions': [metrics],
    },
    {
        'page': 'optimizers.md',
        'all_module_classes': [optimizers],
    },
    {
        'page': 'callbacks.md',
        'all_module_classes': [callbacks],
    },
    {
        'page': 'backend.md',
        'all_module_functions': [backend],
    },
    {
        'page': 'utils/data_utils.md',
        'functions': [
            data_utils.get_file,
        ]
    },
    {
        'page': 'utils/io_utils.md',
        'classes': [
            io_utils.HDF5Matrix
        ],
    },
    {
        'page': 'utils/layer_utils.md',
        'functions': [
            layer_utils.layer_from_config,
        ]
    },
    {
        'page': 'utils/np_utils.md',
        'all_module_functions': [np_utils]
    },
]

ROOT = 'http://keras.io/'

def get_earliest_class_that_defined_member(member, cls):
    ancestors = get_classes_ancestors([cls])
    result = None
    for ancestor in ancestors:
        if member in dir(ancestor):
            result = ancestor
    if not result:
        return cls
    return result

def get_classes_ancestors(classes):
    ancestors = []
    for cls in classes:
        ancestors += cls.__bases__
    filtered_ancestors = []
    for ancestor in ancestors:
        if ancestor.__name__ in ['object']:
            continue
        filtered_ancestors.append(ancestor)
    if filtered_ancestors:
        return filtered_ancestors + get_classes_ancestors(filtered_ancestors)
    else:
        return filtered_ancestors

def get_function_signature(function, method=True):
    signature = inspect.getargspec(function)
    defaults = signature.defaults
    if method:
        args = signature.args[1:]
    else:
        args = signature.args
    if defaults:
        kwargs = zip(args[-len(defaults):], defaults)
        args = args[:-len(defaults)]
    else:
        kwargs = []
    st = '%s.%s(' % (function.__module__, function.__name__)
    for a in args:
        st += str(a) + ', '
    for a, v in kwargs:
        if type(v) == str:
            v = '\'' + v + '\''
        st += str(a) + '=' + str(v) + ', '
    if kwargs or args:
        return st[:-2] + ')'
    else:
        return st + ')'

def get_class_signature(cls):
    try:
        class_signature = get_function_signature(cls.__init__)
        class_signature = class_signature.replace('__init__', cls.__name__)
    except:
        # in case the class inherits from object and does not
        # define __init__
        class_signature = cls.__module__ + '.' + cls.__name__ + '()'
    return class_signature

def class_to_docs_link(cls):
    module_name = cls.__module__
    assert module_name[:6] == 'keras.'
    module_name = module_name[6:]
    link = ROOT + module_name.replace('.', '/') + '#' + cls.__name__.lower()
    return link

def class_to_source_link(cls):
    module_name = cls.__module__
    assert module_name[:6] == 'keras.'
    path = module_name.replace('.', '/')
    path += '.py'
    line = inspect.getsourcelines(cls)[-1]
    link = 'https://github.com/fchollet/keras/blob/master/' + path + '#L' + str(line)
    return '[[source]](' + link + ')'

def code_snippet(snippet):
    result = '```python\n'
    result += snippet + '\n'
    result += '```\n'
    return result

def process_class_docstring(docstring):
    docstring = re.sub(r'\n    # (.*)\n',
                       r'\n    __\1__\n\n',
                       docstring)

    docstring = re.sub(r'    ([^\s\\]+):(.*)\n',
                       r'    - __\1__:\2\n',
                       docstring)

    docstring = docstring.replace('    ' * 5, '\t\t')
    docstring = docstring.replace('    ' * 3, '\t')
    docstring = docstring.replace('    ', '')
    return docstring

def process_function_docstring(docstring):
    docstring = re.sub(r'\n    # (.*)\n',
                       r'\n    __\1__\n\n',
                       docstring)
    docstring = re.sub(r'\n        # (.*)\n',
                       r'\n        __\1__\n\n',
                       docstring)

    docstring = re.sub(r'    ([^\s\\]+):(.*)\n',
                       r'    - __\1__:\2\n',
                       docstring)

    docstring = docstring.replace('    ' * 6, '\t\t')
    docstring = docstring.replace('    ' * 4, '\t')
    docstring = docstring.replace('    ', '')
    return docstring

print('Cleaning up existing sources directory.')
if os.path.exists('sources'):
    shutil.rmtree('sources')

print('Populating sources directory with templates.')
for subdir, dirs, fnames in os.walk('templates'):
    for fname in fnames:
        new_subdir = subdir.replace('templates', 'sources')
        if not os.path.exists(new_subdir):
            os.makedirs(new_subdir)
        if fname[-3:] == '.md':
            fpath = os.path.join(subdir, fname)
            new_fpath = fpath.replace('templates', 'sources')
            shutil.copy(fpath, new_fpath)

print('Starting autogeneration.')
for page_data in PAGES:
    blocks = []
    classes = page_data.get('classes', [])
    for module in page_data.get('all_module_classes', []):
        module_classes = []
        for name in dir(module):
            if name[0] == '_' or name in EXCLUDE:
                continue
            module_member = getattr(module, name)
            if inspect.isclass(module_member):
                cls = module_member
                if cls.__module__ == module.__name__:
                    if cls not in module_classes:
                        module_classes.append(cls)
        module_classes.sort(key=lambda x: id(x))
        classes += module_classes

    for cls in classes:
        subblocks = []
        signature = get_class_signature(cls)
        subblocks.append('' + class_to_source_link(cls) + '')
        subblocks.append('### ' + cls.__name__ + '\n')
        subblocks.append(code_snippet(signature))
        docstring = cls.__doc__
        if docstring:
            subblocks.append(process_class_docstring(docstring))
        blocks.append('\n'.join(subblocks))

    functions = page_data.get('functions', [])
    for module in page_data.get('all_module_functions', []):
        module_functions = []
        for name in dir(module):
            if name[0] == '_' or name in EXCLUDE:
                continue
            module_member = getattr(module, name)
            if inspect.isfunction(module_member):
                function = module_member
                if module.__name__ in function.__module__:
                    if function not in module_functions:
                        module_functions.append(function)
        module_functions.sort(key=lambda x: id(x))
        functions += module_functions

    for function in functions:
        subblocks = []
        signature = get_function_signature(function, method=False)
        signature = signature.replace(function.__module__ + '.', '')
        subblocks.append('### ' + function.__name__ + '\n')
        subblocks.append(code_snippet(signature))
        docstring = function.__doc__
        if docstring:
            subblocks.append(process_function_docstring(docstring))
            blocks.append('\n\n'.join(subblocks))

    mkdown = '\n----\n\n'.join(blocks)
    # save module page.
    # Either insert content into existing page,
    # or create page otherwise
    page_name = page_data['page']
    path = os.path.join('sources', page_name)
    if os.path.exists(path):
        template = open(path).read()
        assert '{{autogenerated}}' in template, ('Template found for ' + path +
                                                 ' but missing {{autogenerated}} tag.')
        mkdown = template.replace('{{autogenerated}}', mkdown)
        print('...inserting autogenerated content into template:', path)
    else:
        print('...creating new page with autogenerated content:', path)
    subdir = os.path.dirname(path)
    if not os.path.exists(subdir):
        os.makedirs(subdir)
    open(path, 'w').write(mkdown)
<|endoftext|>


<commit_before># emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""Handling connectivity: combines FreeSurfer surfaces with subcortical volumes."""
import os
from glob import glob
import json

import nibabel as nb
from nibabel import cifti2 as ci
import numpy as np
from nilearn.image import resample_to_img

from nipype.utils.filemanip import split_filename
from nipype.interfaces.base import (
    BaseInterfaceInputSpec, TraitedSpec, File, traits,
    SimpleInterface, Directory
)
from templateflow.api import get as get_template

# CITFI structures with corresponding FS labels
CIFTI_STRUCT_WITH_LABELS = {
    # SURFACES
    'CIFTI_STRUCTURE_CORTEX_LEFT': None,
    'CIFTI_STRUCTURE_CORTEX_RIGHT': None,

    # SUBCORTICAL
    'CIFTI_STRUCTURE_ACCUMBENS_LEFT': [26],
    'CIFTI_STRUCTURE_ACCUMBENS_RIGHT': [58],
    'CIFTI_STRUCTURE_AMYGDALA_LEFT': [18],
    'CIFTI_STRUCTURE_AMYGDALA_RIGHT': [54],
    'CIFTI_STRUCTURE_BRAIN_STEM': [16],
    'CIFTI_STRUCTURE_CAUDATE_LEFT': [11],
    'CIFTI_STRUCTURE_CAUDATE_RIGHT': [50],
    'CIFTI_STRUCTURE_CEREBELLUM_LEFT': [
        6,  # DKT31
        8,  # HCP MNI152
    ],
    'CIFTI_STRUCTURE_CEREBELLUM_RIGHT': [
        45,  # DKT31
        47,  # HCP MNI152
    ],
    'CIFTI_STRUCTURE_DIENCEPHALON_VENTRAL_LEFT': [28],
    'CIFTI_STRUCTURE_DIENCEPHALON_VENTRAL_RIGHT': [60],
    'CIFTI_STRUCTURE_HIPPOCAMPUS_LEFT': [17],
    'CIFTI_STRUCTURE_HIPPOCAMPUS_RIGHT': [53],
    'CIFTI_STRUCTURE_PALLIDUM_LEFT': [13],
    'CIFTI_STRUCTURE_PALLIDUM_RIGHT': [52],
    'CIFTI_STRUCTURE_PUTAMEN_LEFT': [12],
    'CIFTI_STRUCTURE_PUTAMEN_RIGHT': [51],
    'CIFTI_STRUCTURE_THALAMUS_LEFT': [10],
    'CIFTI_STRUCTURE_THALAMUS_RIGHT': [49],
}

class _GenerateCiftiInputSpec(BaseInterfaceInputSpec):
    bold_file = File(mandatory=True, exists=True, desc="input BOLD file")
    volume_target = traits.Enum("MNI152NLin6Asym", "MNI152NLin2009cAsym", usedefault=True,
                                desc="CIFTI volumetric output space")
    surface_target = traits.Enum("fsLR", "fsaverage5", "fsaverage6", usedefault=True,
                                 desc="CIFTI surface target space")
    density = traits.Enum('32k', '59k', usedefault=True,
                          help='surface hemisphere vertices')
    TR = traits.Float(mandatory=True, desc="Repetition time")
    surface_bolds = traits.List(File(exists=True), mandatory=True,
                                desc="list of surface BOLD GIFTI files"
                                     " (length 2 with order [L,R])")
    subjects_dir = Directory(mandatory=True, desc="FreeSurfer SUBJECTS_DIR")

class _GenerateCiftiOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc="generated CIFTI file")
    variant = traits.Str(desc="combination of target spaces label")
    variant_key = File(exists=True, desc='file storing variant space information')

class GenerateCifti(SimpleInterface):
    """
    Generate CIFTI image from BOLD file in target spaces.

    Currently supports ``fsLR``, ``fsaverage5``, or ``fsaverage6`` for template surfaces and
    ``MNI152NLin6Asym`` or ``MNI152NLin2009cAsym`` as template volumes.

    """

    input_spec = _GenerateCiftiInputSpec
    output_spec = _GenerateCiftiOutputSpec

    def _run_interface(self, runtime):
        targets = self._define_variant()
        annotation_files, label_file = self._fetch_data()
        self._results["out_file"] = self._create_cifti_image(
            self.inputs.bold_file,
            label_file,
            self.inputs.surface_bolds,
            targets,
            self.inputs.TR,
            annotation_files=annotation_files)
        return runtime

    def _define_variant(self):
        """Assign arbitrary label to combination of CIFTI spaces."""
        space = None
        variants = {
            'fsLR': ['fsLR', 'MNI152NLin6Asym'],
            'space1': ['fsaverage5', 'MNI152NLin2009cAsym'],
            'space2': ['fsaverage6', 'MNI152NLin2009cAsym'],
        }
        for sp, targets in variants.items():
            if all(
                target in targets for target in
                [self.inputs.surface_target, self.inputs.volume_target]
            ):
                space = sp
        if space is None:
            raise NotImplementedError

        variant_key = os.path.abspath('dtseries_variant.json')
        with open(variant_key, 'w') as fp:
            json.dump({space: variants[space]}, fp)
        self._results['variant_key'] = variant_key
        self._results['variant'] = space
        return variants[space]

    def _fetch_data(self):
        """Converts inputspec to files"""
        if (
            self.inputs.surface_target not in ("fsaverage5", "fsaverage6", "fsLR") or
            self.inputs.volume_target not in ("MNI152NLin2009cAsym", "MNI152NLin6Asym")
        ):
            raise NotImplementedError(
                "Target space (surface: {0}, volume: {1}) is not supported".format(
                    self.inputs.surface_target, self.inputs.volume_target
                )
            )

        tpl_kwargs = {'suffix': 'dseg'}
        annotation_files = None
        if self.inputs.volume_target == "MNI152NLin2009cAsym":
            tpl_kwargs.update({
                'resolution': '2',
                'desc': 'DKT31',
            })
            annotation_files = sorted(
                glob(os.path.join(self.inputs.subjects_dir,
                                  self.inputs.surface_target,
                                  'label',
                                  '*h.aparc.annot'))
            )

        elif self.inputs.volume_target == 'MNI152NLin6Asym':
            res = '2' if self.inputs.density == '32k' else '5'

            tpl_kwargs.update({
                'atlas': 'fsLR',
                'resolution': res,
            })
            annotation_files = [
                str(f) for f in get_template(
                    'fsLR', density=self.inputs.density, desc='nomedialwall', suffix='dparc'
                )
            ]

        if len(annotation_files) != 2:
            raise IOError("Invalid number of surface annotation files")

        label_file = str(get_template(self.inputs.volume_target, **tpl_kwargs))
        return annotation_files, label_file

    @staticmethod
    def _create_cifti_image(bold_file, label_file, bold_surfs,
                            targets, tr, annotation_files):
        """
        Generate CIFTI image in target space.

        Parameters
            bold_file : str
                BOLD volumetric timeseries
            label_file : str
                Subcortical label file
            bold_surfs : list
                BOLD surface timeseries [L,R]
            targets : list
                Surface and volumetric output spaces
            tr : float
                BOLD repetition time
            annotation_files : list
                Surface label files used to remove medial wall

        Returns
            out :
                BOLD data saved as CIFTI dtseries
        """
        bold_img = nb.load(bold_file)
        label_img = nb.load(label_file)
        if label_img.shape != bold_img.shape[:3]:
            bold_img = resample_to_img(bold_img, label_img)

        bold_data = bold_img.get_fdata(dtype='float32')
        timepoints = bold_img.shape[3]
        label_data = np.asanyarray(label_img.dataobj).astype('int16')

        # Create brain models
        idx_offset = 0
        brainmodels = []
        bm_ts = np.empty((timepoints, 0))

        for structure, labels in CIFTI_STRUCT_WITH_LABELS.items():
            if labels is None:  # surface model
                model_type = "CIFTI_MODEL_TYPE_SURFACE"
                # use the corresponding annotation
                hemi = structure.split('_')[-1]
                # currently only supports L/R cortex
                surf = nb.load(bold_surfs[hemi == "RIGHT"])
                surf_verts = len(surf.darrays[0].data)
                if annotation_files[0].endswith('.annot'):
                    annot = nb.freesurfer.read_annot(annotation_files[hemi == "RIGHT"])
                    # remove medial wall
                    medial = np.nonzero(annot[0] != annot[2].index(b'unknown'))[0]
                else:
                    annot = nb.load(annotation_files[hemi == "RIGHT"])
                    medial = np.nonzero(annot.darrays[0].data)[0]
                # extract values across volumes
                ts = np.array([tsarr.data[medial] for tsarr in surf.darrays])

                vert_idx = ci.Cifti2VertexIndices(medial)
                bm = ci.Cifti2BrainModel(index_offset=idx_offset,
                                         index_count=len(vert_idx),
                                         model_type=model_type,
                                         brain_structure=structure,
                                         vertex_indices=vert_idx,
                                         n_surface_vertices=surf_verts)
                bm_ts = np.column_stack((bm_ts, ts))
                idx_offset += len(vert_idx)
                brainmodels.append(bm)
            else:
                model_type = "CIFTI_MODEL_TYPE_VOXELS"
                vox = []
                ts = None
                for label in labels:
                    ijk = np.nonzero(label_data == label)
                    ts = (bold_data[ijk] if ts is None
                          else np.concatenate((ts, bold_data[ijk])))
                    vox += [[ijk[0][ix], ijk[1][ix], ijk[2][ix]]
                            for ix, row in enumerate(ts)]

                bm_ts = np.column_stack((bm_ts, ts.T))

                vox = ci.Cifti2VoxelIndicesIJK(vox)
                bm = ci.Cifti2BrainModel(index_offset=idx_offset,
                                         index_count=len(vox),
                                         model_type=model_type,
                                         brain_structure=structure,
                                         voxel_indices_ijk=vox)
                idx_offset += len(vox)
                brainmodels.append(bm)

        # add volume information
        brainmodels.append(
            ci.Cifti2Volume(
                bold_img.shape[:3],
                ci.Cifti2TransformationMatrixVoxelIndicesIJKtoXYZ(-3, bold_img.affine)
            )
        )

        # generate Matrix information
        series_map = ci.Cifti2MatrixIndicesMap(
            (0, ),
            'CIFTI_INDEX_TYPE_SERIES',
            number_of_series_points=timepoints,
            series_exponent=0,
            series_start=0.,
            series_step=tr,
            series_unit='SECOND'
        )
        geometry_map = ci.Cifti2MatrixIndicesMap(
            (1, ),
            'CIFTI_INDEX_TYPE_BRAIN_MODELS',
            maps=brainmodels
        )
        # provide some metadata to CIFTI matrix
        meta = {
            "target_surface": targets[0],
            "target_volume": targets[1],
        }
        # generate and save CIFTI image
        matrix = ci.Cifti2Matrix()
        matrix.append(series_map)
        matrix.append(geometry_map)
        matrix.metadata = ci.Cifti2MetaData(meta)
        hdr = ci.Cifti2Header(matrix)
        img = ci.Cifti2Image(bm_ts, hdr)
        img.nifti_header.set_intent('NIFTI_INTENT_CONNECTIVITY_DENSE_SERIES')

        _, out_base, _ = split_filename(bold_file)
        out_file = "{}.dtseries.nii".format(out_base)
        ci.save(img, out_file)
        return os.path.join(os.getcwd(), out_file)

class _CiftiNameSourceInputSpec(BaseInterfaceInputSpec):
    variant = traits.Str(mandatory=True,
                         desc=('unique label of spaces used in combination to'
                               ' generate CIFTI file'))

class _CiftiNameSourceOutputSpec(TraitedSpec):
    out_name = traits.Str(desc='(partial) filename formatted according to template')

class CiftiNameSource(SimpleInterface):
    """Construct new filename based on unique label of spaces used to generate a CIFTI file."""

    input_spec = _CiftiNameSourceInputSpec
    output_spec = _CiftiNameSourceOutputSpec

    def _run_interface(self, runtime):
        suffix = 'bold.dtseries'
        if 'hcp' in self.inputs.variant:
            suffix = 'space-hcp_bold.dtseries'
        self._results['out_name'] = suffix
        return runtime
<commit_msg>fix: atlas key<commit_after># emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""Handling connectivity: combines FreeSurfer surfaces with subcortical volumes."""
import os
from glob import glob
import json

import nibabel as nb
from nibabel import cifti2 as ci
import numpy as np
from nilearn.image import resample_to_img

from nipype.utils.filemanip import split_filename
from nipype.interfaces.base import (
    BaseInterfaceInputSpec, TraitedSpec, File, traits,
    SimpleInterface, Directory
)
from templateflow.api import get as get_template

# CITFI structures with corresponding FS labels
CIFTI_STRUCT_WITH_LABELS = {
    # SURFACES
    'CIFTI_STRUCTURE_CORTEX_LEFT': None,
    'CIFTI_STRUCTURE_CORTEX_RIGHT': None,

    # SUBCORTICAL
    'CIFTI_STRUCTURE_ACCUMBENS_LEFT': [26],
    'CIFTI_STRUCTURE_ACCUMBENS_RIGHT': [58],
    'CIFTI_STRUCTURE_AMYGDALA_LEFT': [18],
    'CIFTI_STRUCTURE_AMYGDALA_RIGHT': [54],
    'CIFTI_STRUCTURE_BRAIN_STEM': [16],
    'CIFTI_STRUCTURE_CAUDATE_LEFT': [11],
    'CIFTI_STRUCTURE_CAUDATE_RIGHT': [50],
    'CIFTI_STRUCTURE_CEREBELLUM_LEFT': [
        6,  # DKT31
        8,  # HCP MNI152
    ],
    'CIFTI_STRUCTURE_CEREBELLUM_RIGHT': [
        45,  # DKT31
        47,  # HCP MNI152
    ],
    'CIFTI_STRUCTURE_DIENCEPHALON_VENTRAL_LEFT': [28],
    'CIFTI_STRUCTURE_DIENCEPHALON_VENTRAL_RIGHT': [60],
    'CIFTI_STRUCTURE_HIPPOCAMPUS_LEFT': [17],
    'CIFTI_STRUCTURE_HIPPOCAMPUS_RIGHT': [53],
    'CIFTI_STRUCTURE_PALLIDUM_LEFT': [13],
    'CIFTI_STRUCTURE_PALLIDUM_RIGHT': [52],
    'CIFTI_STRUCTURE_PUTAMEN_LEFT': [12],
    'CIFTI_STRUCTURE_PUTAMEN_RIGHT': [51],
    'CIFTI_STRUCTURE_THALAMUS_LEFT': [10],
    'CIFTI_STRUCTURE_THALAMUS_RIGHT': [49],
}

class _GenerateCiftiInputSpec(BaseInterfaceInputSpec):
    bold_file = File(mandatory=True, exists=True, desc="input BOLD file")
    volume_target = traits.Enum("MNI152NLin6Asym", "MNI152NLin2009cAsym", usedefault=True,
                                desc="CIFTI volumetric output space")
    surface_target = traits.Enum("fsLR", "fsaverage5", "fsaverage6", usedefault=True,
                                 desc="CIFTI surface target space")
    density = traits.Enum('32k', '59k', usedefault=True,
                          help='surface hemisphere vertices')
    TR = traits.Float(mandatory=True, desc="Repetition time")
    surface_bolds = traits.List(File(exists=True), mandatory=True,
                                desc="list of surface BOLD GIFTI files"
                                     " (length 2 with order [L,R])")
    subjects_dir = Directory(mandatory=True, desc="FreeSurfer SUBJECTS_DIR")

class _GenerateCiftiOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc="generated CIFTI file")
    variant = traits.Str(desc="combination of target spaces label")
    variant_key = File(exists=True, desc='file storing variant space information')

class GenerateCifti(SimpleInterface):
    """
    Generate CIFTI image from BOLD file in target spaces.

    Currently supports ``fsLR``, ``fsaverage5``, or ``fsaverage6`` for template surfaces and
    ``MNI152NLin6Asym`` or ``MNI152NLin2009cAsym`` as template volumes.

    """

    input_spec = _GenerateCiftiInputSpec
    output_spec = _GenerateCiftiOutputSpec

    def _run_interface(self, runtime):
        targets = self._define_variant()
        annotation_files, label_file = self._fetch_data()
        self._results["out_file"] = self._create_cifti_image(
            self.inputs.bold_file,
            label_file,
            self.inputs.surface_bolds,
            targets,
            self.inputs.TR,
            annotation_files=annotation_files)
        return runtime

    def _define_variant(self):
        """Assign arbitrary label to combination of CIFTI spaces."""
        space = None
        variants = {
            'fsLR': ['fsLR', 'MNI152NLin6Asym'],
            'space1': ['fsaverage5', 'MNI152NLin2009cAsym'],
            'space2': ['fsaverage6', 'MNI152NLin2009cAsym'],
        }
        for sp, targets in variants.items():
            if all(
                target in targets for target in
                [self.inputs.surface_target, self.inputs.volume_target]
            ):
                space = sp
        if space is None:
            raise NotImplementedError

        variant_key = os.path.abspath('dtseries_variant.json')
        with open(variant_key, 'w') as fp:
            json.dump({space: variants[space]}, fp)
        self._results['variant_key'] = variant_key
        self._results['variant'] = space
        return variants[space]

    def _fetch_data(self):
        """Converts inputspec to files"""
        if (
            self.inputs.surface_target not in ("fsaverage5", "fsaverage6", "fsLR") or
            self.inputs.volume_target not in ("MNI152NLin2009cAsym", "MNI152NLin6Asym")
        ):
            raise NotImplementedError(
                "Target space (surface: {0}, volume: {1}) is not supported".format(
                    self.inputs.surface_target, self.inputs.volume_target
                )
            )

        tpl_kwargs = {'suffix': 'dseg'}
        annotation_files = None
        if self.inputs.volume_target == "MNI152NLin2009cAsym":
            tpl_kwargs.update({
                'resolution': '2',
                'desc': 'DKT31',
            })
            annotation_files = sorted(
                glob(os.path.join(self.inputs.subjects_dir,
                                  self.inputs.surface_target,
                                  'label',
                                  '*h.aparc.annot'))
            )

        elif self.inputs.volume_target == 'MNI152NLin6Asym':
            res = '2' if self.inputs.density == '32k' else '5'

            tpl_kwargs.update({
                'atlas': 'HCP',
                'resolution': res,
            })
            annotation_files = [
                str(f) for f in get_template(
                    'fsLR', density=self.inputs.density, desc='nomedialwall', suffix='dparc'
                )
            ]

        if len(annotation_files) != 2:
            raise IOError("Invalid number of surface annotation files")

        label_file = str(get_template(self.inputs.volume_target, **tpl_kwargs))
        return annotation_files, label_file

    @staticmethod
    def _create_cifti_image(bold_file, label_file, bold_surfs,
                            targets, tr, annotation_files):
        """
        Generate CIFTI image in target space.

        Parameters
            bold_file : str
                BOLD volumetric timeseries
            label_file : str
                Subcortical label file
            bold_surfs : list
                BOLD surface timeseries [L,R]
            targets : list
                Surface and volumetric output spaces
            tr : float
                BOLD repetition time
            annotation_files : list
                Surface label files used to remove medial wall

        Returns
            out :
                BOLD data saved as CIFTI dtseries
        """
        bold_img = nb.load(bold_file)
        label_img = nb.load(label_file)
        if label_img.shape != bold_img.shape[:3]:
            bold_img = resample_to_img(bold_img, label_img)

        bold_data = bold_img.get_fdata(dtype='float32')
        timepoints = bold_img.shape[3]
        label_data = np.asanyarray(label_img.dataobj).astype('int16')

        # Create brain models
        idx_offset = 0
        brainmodels = []
        bm_ts = np.empty((timepoints, 0))

        for structure, labels in CIFTI_STRUCT_WITH_LABELS.items():
            if labels is None:  # surface model
                model_type = "CIFTI_MODEL_TYPE_SURFACE"
                # use the corresponding annotation
                hemi = structure.split('_')[-1]
                # currently only supports L/R cortex
                surf = nb.load(bold_surfs[hemi == "RIGHT"])
                surf_verts = len(surf.darrays[0].data)
                if annotation_files[0].endswith('.annot'):
                    annot = nb.freesurfer.read_annot(annotation_files[hemi == "RIGHT"])
                    # remove medial wall
                    medial = np.nonzero(annot[0] != annot[2].index(b'unknown'))[0]
                else:
                    annot = nb.load(annotation_files[hemi == "RIGHT"])
                    medial = np.nonzero(annot.darrays[0].data)[0]
                # extract values across volumes
                ts = np.array([tsarr.data[medial] for tsarr in surf.darrays])

                vert_idx = ci.Cifti2VertexIndices(medial)
                bm = ci.Cifti2BrainModel(index_offset=idx_offset,
                                         index_count=len(vert_idx),
                                         model_type=model_type,
                                         brain_structure=structure,
                                         vertex_indices=vert_idx,
                                         n_surface_vertices=surf_verts)
                bm_ts = np.column_stack((bm_ts, ts))
                idx_offset += len(vert_idx)
                brainmodels.append(bm)
            else:
                model_type = "CIFTI_MODEL_TYPE_VOXELS"
                vox = []
                ts = None
                for label in labels:
                    ijk = np.nonzero(label_data == label)
                    ts = (bold_data[ijk] if ts is None
                          else np.concatenate((ts, bold_data[ijk])))
                    vox += [[ijk[0][ix], ijk[1][ix], ijk[2][ix]]
                            for ix, row in enumerate(ts)]

                bm_ts = np.column_stack((bm_ts, ts.T))

                vox = ci.Cifti2VoxelIndicesIJK(vox)
                bm = ci.Cifti2BrainModel(index_offset=idx_offset,
                                         index_count=len(vox),
                                         model_type=model_type,
                                         brain_structure=structure,
                                         voxel_indices_ijk=vox)
                idx_offset += len(vox)
                brainmodels.append(bm)

        # add volume information
        brainmodels.append(
            ci.Cifti2Volume(
                bold_img.shape[:3],
                ci.Cifti2TransformationMatrixVoxelIndicesIJKtoXYZ(-3, bold_img.affine)
            )
        )

        # generate Matrix information
        series_map = ci.Cifti2MatrixIndicesMap(
            (0, ),
            'CIFTI_INDEX_TYPE_SERIES',
            number_of_series_points=timepoints,
            series_exponent=0,
            series_start=0.,
            series_step=tr,
            series_unit='SECOND'
        )
        geometry_map = ci.Cifti2MatrixIndicesMap(
            (1, ),
            'CIFTI_INDEX_TYPE_BRAIN_MODELS',
            maps=brainmodels
        )
        # provide some metadata to CIFTI matrix
        meta = {
            "target_surface": targets[0],
            "target_volume": targets[1],
        }
        # generate and save CIFTI image
        matrix = ci.Cifti2Matrix()
        matrix.append(series_map)
        matrix.append(geometry_map)
        matrix.metadata = ci.Cifti2MetaData(meta)
        hdr = ci.Cifti2Header(matrix)
        img = ci.Cifti2Image(bm_ts, hdr)
        img.nifti_header.set_intent('NIFTI_INTENT_CONNECTIVITY_DENSE_SERIES')

        _, out_base, _ = split_filename(bold_file)
        out_file = "{}.dtseries.nii".format(out_base)
        ci.save(img, out_file)
        return os.path.join(os.getcwd(), out_file)

class _CiftiNameSourceInputSpec(BaseInterfaceInputSpec):
    variant = traits.Str(mandatory=True,
                         desc=('unique label of spaces used in combination to'
                               ' generate CIFTI file'))

class _CiftiNameSourceOutputSpec(TraitedSpec):
    out_name = traits.Str(desc='(partial) filename formatted according to template')

class CiftiNameSource(SimpleInterface):
    """Construct new filename based on unique label of spaces used to generate a CIFTI file."""

    input_spec = _CiftiNameSourceInputSpec
    output_spec = _CiftiNameSourceOutputSpec

    def _run_interface(self, runtime):
        suffix = 'bold.dtseries'
        if 'hcp' in self.inputs.variant:
            suffix = 'space-hcp_bold.dtseries'
        self._results['out_name'] = suffix
        return runtime
<|endoftext|>


<commit_before>from __future__ import unicode_literals

from rest_framework import status, response

from nodeconductor.core import models
from nodeconductor.core.exceptions import IncorrectStateException

class UpdateOnlyStableMixin(object):
    """
    Allow modification of entities in stable state only.
    """

    def initial(self, request, *args, **kwargs):
        acceptable_states = {
            'update': models.SynchronizationStates.STABLE_STATES,
            'partial_update': models.SynchronizationStates.STABLE_STATES,
            'destroy': models.SynchronizationStates.STABLE_STATES | {models.SynchronizationStates.NEW},
        }
        if self.action in acceptable_states.keys():
            obj = self.get_object()
            if obj and isinstance(obj, models.SynchronizableMixin):
                if obj.state not in acceptable_states[self.action]:
                    raise IncorrectStateException(
                        'Modification allowed in stable states only.')

        return super(UpdateOnlyStableMixin, self).initial(request, *args, **kwargs)

class UserContextMixin(object):
    """ Pass current user to serializer context """

    def get_serializer_context(self):
        context = super(UserContextMixin, self).get_serializer_context()
        context['user'] = self.request.user
        return context

class StateMixin(object):
    """ Raise exception if object is not in correct state for action """

    acceptable_states = {}

    def initial(self, request, *args, **kwargs):
        States = models.StateMixin.States
        acceptable_states = {
            'update': [States.OK],
            'partial_update': [States.OK],
            'destroy': [States.OK, States.ERRED],
        }
        acceptable_states.update(self.acceptable_states)
        acceptable_state = acceptable_states.get(self.action)
        if acceptable_state:
            obj = self.get_object()
            if obj.state not in acceptable_state:
                raise IncorrectStateException('Modification allowed in stable states only.')

        return super(StateMixin, self).initial(request, *args, **kwargs)

class RuntimeStateMixin(object):
    acceptable_states = {}

    def initial(self, request, *args, **kwargs):
        if self.action:
            self.check_operation(request, self.get_object(), self.action)
        return super(RuntimeStateMixin, self).initial(request, *args, **kwargs)

    def check_operation(self, request, obj, action):
        acceptable_state = self.acceptable_states.get(action)
        if acceptable_state:
            if obj.state != models.StateMixin.States.OK or obj.runtime_state != acceptable_state:
                raise IncorrectStateException(
                    'Performing %s operation is not allowed for resource in its current state.' % action)

class AsyncExecutor(object):
    async_executor = True

class CreateExecutorMixin(AsyncExecutor):
    create_executor = NotImplemented

    def perform_create(self, serializer):
        instance = serializer.save()
        self.create_executor.execute(instance, async=self.async_executor)
        instance.refresh_from_db()

class UpdateExecutorMixin(AsyncExecutor):
    update_executor = NotImplemented

    def perform_update(self, serializer):
        instance = self.get_object()
        # Save all instance fields before update.
        # To avoid additional DB queries - store foreign keys as ids.
        # Warning! M2M fields will be ignored.
        before_update_fields = {f: getattr(instance, f.attname) for f in instance._meta.fields}
        super(UpdateExecutorMixin, self).perform_update(serializer)
        instance.refresh_from_db()
        updated_fields = {f.name for f, v in before_update_fields.items() if v != getattr(instance, f.attname)}
        self.update_executor.execute(instance, async=self.async_executor, updated_fields=updated_fields)
        instance.refresh_from_db()

class DeleteExecutorMixin(AsyncExecutor):
    delete_executor = NotImplemented

    def destroy(self, request, *args, **kwargs):
        instance = self.get_object()
        self.delete_executor.execute(
            instance, async=self.async_executor, force=instance.state == models.StateMixin.States.ERRED)
        return response.Response(
            {'detail': 'Deletion was scheduled'}, status=status.HTTP_202_ACCEPTED)

class EagerLoadMixin(object):
    """ Reduce number of requests to DB.

        Serializer should implement static method "eager_load", that selects
        objects that are necessary for serialization.
    """

    def get_queryset(self):
        queryset = super(EagerLoadMixin, self).get_queryset()
        serializer_class = self.get_serializer_class()
        if self.action in ('list', 'retrieve') and hasattr(serializer_class, 'eager_load'):
            queryset = serializer_class.eager_load(queryset)
        return queryset
<commit_msg>Fix typo (NC-1490)<commit_after>from __future__ import unicode_literals

from rest_framework import status, response

from nodeconductor.core import models
from nodeconductor.core.exceptions import IncorrectStateException

class UpdateOnlyStableMixin(object):
    """
    Allow modification of entities in stable state only.
    """

    def initial(self, request, *args, **kwargs):
        acceptable_states = {
            'update': models.SynchronizationStates.STABLE_STATES,
            'partial_update': models.SynchronizationStates.STABLE_STATES,
            'destroy': models.SynchronizationStates.STABLE_STATES | {models.SynchronizationStates.NEW},
        }
        if self.action in acceptable_states.keys():
            obj = self.get_object()
            if obj and isinstance(obj, models.SynchronizableMixin):
                if obj.state not in acceptable_states[self.action]:
                    raise IncorrectStateException(
                        'Modification allowed in stable states only.')

        return super(UpdateOnlyStableMixin, self).initial(request, *args, **kwargs)

class UserContextMixin(object):
    """ Pass current user to serializer context """

    def get_serializer_context(self):
        context = super(UserContextMixin, self).get_serializer_context()
        context['user'] = self.request.user
        return context

class StateMixin(object):
    """ Raise exception if object is not in correct state for action """

    acceptable_states = {}

    def initial(self, request, *args, **kwargs):
        States = models.StateMixin.States
        acceptable_states = {
            'update': [States.OK],
            'partial_update': [States.OK],
            'destroy': [States.OK, States.ERRED],
        }
        acceptable_states.update(self.acceptable_states)
        acceptable_state = acceptable_states.get(self.action)
        if acceptable_state:
            obj = self.get_object()
            if obj.state not in acceptable_state:
                raise IncorrectStateException('Modification allowed in stable states only.')

        return super(StateMixin, self).initial(request, *args, **kwargs)

class RuntimeStateMixin(object):
    acceptable_states = {}

    def initial(self, request, *args, **kwargs):
        if self.action in self.acceptable_states:
            self.check_operation(request, self.get_object(), self.action)
        return super(RuntimeStateMixin, self).initial(request, *args, **kwargs)

    def check_operation(self, request, obj, action):
        acceptable_state = self.acceptable_states.get(action)
        if acceptable_state:
            if obj.state != models.StateMixin.States.OK or obj.runtime_state != acceptable_state:
                raise IncorrectStateException(
                    'Performing %s operation is not allowed for resource in its current state.' % action)

class AsyncExecutor(object):
    async_executor = True

class CreateExecutorMixin(AsyncExecutor):
    create_executor = NotImplemented

    def perform_create(self, serializer):
        instance = serializer.save()
        self.create_executor.execute(instance, async=self.async_executor)
        instance.refresh_from_db()

class UpdateExecutorMixin(AsyncExecutor):
    update_executor = NotImplemented

    def perform_update(self, serializer):
        instance = self.get_object()
        # Save all instance fields before update.
        # To avoid additional DB queries - store foreign keys as ids.
        # Warning! M2M fields will be ignored.
        before_update_fields = {f: getattr(instance, f.attname) for f in instance._meta.fields}
        super(UpdateExecutorMixin, self).perform_update(serializer)
        instance.refresh_from_db()
        updated_fields = {f.name for f, v in before_update_fields.items() if v != getattr(instance, f.attname)}
        self.update_executor.execute(instance, async=self.async_executor, updated_fields=updated_fields)
        instance.refresh_from_db()

class DeleteExecutorMixin(AsyncExecutor):
    delete_executor = NotImplemented

    def destroy(self, request, *args, **kwargs):
        instance = self.get_object()
        self.delete_executor.execute(
            instance, async=self.async_executor, force=instance.state == models.StateMixin.States.ERRED)
        return response.Response(
            {'detail': 'Deletion was scheduled'}, status=status.HTTP_202_ACCEPTED)

class EagerLoadMixin(object):
    """ Reduce number of requests to DB.

        Serializer should implement static method "eager_load", that selects
        objects that are necessary for serialization.
    """

    def get_queryset(self):
        queryset = super(EagerLoadMixin, self).get_queryset()
        serializer_class = self.get_serializer_class()
        if self.action in ('list', 'retrieve') and hasattr(serializer_class, 'eager_load'):
            queryset = serializer_class.eager_load(queryset)
        return queryset
<|endoftext|>