<commit_before># Brain Tumor Classification
# Script for Training General CNN Models
# Author: Qixun Qu
# Create on: 2017/10/14
# Modify on: 2017/11/22

#     ,,,         ,,,
#   ;"   ';     ;'   ",
#   ;  @.ss$$$$$$s.@  ;
#   `s$$$$$$$$$$$$$$$'
#   $$$$$$$$$$$$$$$$$$
#  $$$$P""Y$$$Y""W$$$$$
#  $$$$  p"$$$"q  $$$$$
#  $$$$  .$$$$$.  $$$$'
#   $$$DaU$$O$$DaU$$$'
#    '$$$$'.^.'$$$$'
#       '&$$$$$&'

'''

Class BTCTrainCNN

-1- Models are defined in class BTCModels.
-2- Hyper-parameters can be set in btc_parameters.py.
-3- Loading tfrecords for training and validating by
    functions in class BTCTFRecords.

'''

from __future__ import print_function

import os
import time
import argparse
import numpy as np
import tensorflow as tf
from btc_settings import *
from btc_train import BTCTrain
from btc_cnn_parameters import cnn_parameters

class BTCTrainCNN(BTCTrain):

    def __init__(self, net, paras, save_path, logs_path):
        '''__INIT__

            Initialization of class BTCTrain to set parameters
            for constructing, training and validating models.

            Inputs:
            -------
            - net: string, the name of the model applied to train
            - paras: dict, parameters for training the model, defined
                     in btc_parameters.py
            - save_path: string, the path of the folder to save models
            - logs_path: string, the path of the folder to save logs

        '''

        super().__init__(paras)

        self.net_name = self.set_net_name(net)
        self.model_path = self.set_dir_path(save_path, self.net_name)
        self.logs_path = self.set_dir_path(logs_path, self.net_name)

        self.network = self._get_network(net)

        return

    def _get_network(self, net):
        '''_GET_NETWORK

            Return network function according to the given net's name.

            Input:
            ------
            - net: string, the net's name

            Output:
            -------
            - the choosen network

        '''

        # Set models by given variable
        if net == CNN:
            network = self.models.cnn
        elif net == FULL_CNN:
            network = self.models.full_cnn
        elif net == RES_CNN:
            network = self.models.res_cnn
        elif net == DENSE_CNN:
            network = self.models.dense_cnn
        else:  # Raise error if model cannot be found
            raise ValueError("Could not found model.")

        return network

    def train(self):
        '''TRAIN

            Train and validate the choosen model.

        '''

        with tf.device("/cpu:0"):
            tra_data, tra_labels, val_data, val_labels = self.load_data()
            x, y_input, is_training, learning_rate = self.inputs()

        with tf.device("/gpu:0"):
            # Obtain logits from the model
            y_output = self.network(x, is_training)

        # Compute loss and accuracy and merge summary
        # The summary can be displayed by TensorBoard
        loss = self.get_softmax_loss(y_input, y_output)
        accuracy = self.get_accuracy(y_input, y_output)
        merged = tf.summary.merge_all()

        train_op = self.create_optimizer(learning_rate, loss)

        # Create a saver to save model while training
        saver = tf.train.Saver()
        sess = tf.InteractiveSession()
        sess.run(self.initialize_variables())
        tra_writer, val_writer = self.create_writers(self.logs_path, sess.graph)

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)

        self.blue_print("\nTraining and Validating model: {}\n".format(self.net_name))

        # Initialize counter
        tra_iters, val_iters, epoch_no = 0, 0, 0
        one_tra_iters, one_val_iters = 0, 0
        best_val_lmean_oss = np.inf

        # Lists to save loss and accuracy of each training step
        tloss_list, taccuracy_list = [], []

        # Initialize the timer to count time of one epoch
        epoch_time = time.time()

        try:
            while not coord.should_stop():
                # Initialize the timer to count time of one training step
                tra_step_time = time.time()

                # Training step
                # Feed the graph, run optimizer and get metrics
                tx, ty = sess.run([tra_data, tra_labels])
                tra_fd = {x: tx, y_input: ty, is_training: True, learning_rate: self.learning_rates[epoch_no]}
                tsummary, tloss, taccuracy, _ = sess.run([merged, loss, accuracy, train_op], feed_dict=tra_fd)

                # Get the time of one training step
                tstime = self.get_time(tra_step_time)

                tra_iters += 1
                one_tra_iters += 1

                # Record metrics of training steps
                tloss_list.append(tloss)
                taccuracy_list.append(taccuracy)
                tra_writer.add_summary(tsummary, tra_iters)
                self.train_metrics.append([tloss, taccuracy])
                self.print_metrics("Train", epoch_no + 1, one_tra_iters, tstime, tloss, taccuracy)

                if tra_iters % self.tepoch_iters[epoch_no] == 0:
                    # Validating step
                    # Lists to save loss and accuracy of each validating step
                    vloss_list, vaccuracy_list = [], []
                    while val_iters < self.vepoch_iters[epoch_no]:
                        # Initialize the timer to count time of one validating step
                        val_step_time = time.time()

                        # Feed the graph, get metrics
                        vx, vy = sess.run([val_data, val_labels])
                        val_fd = {x: vx, y_input: vy, is_training: False}
                        vsummary, vloss, vaccuracy = sess.run([merged, loss, accuracy], feed_dict=val_fd)

                        # Get the time of one validating step
                        vstime = self.get_time(val_step_time)

                        val_iters += 1
                        one_val_iters += 1

                        # Record metrics of validating steps
                        vloss_list.append(vloss)
                        vaccuracy_list.append(vaccuracy)
                        val_writer.add_summary(vsummary, val_iters)
                        self.validate_metrics.append([vloss, vaccuracy])
                        self.print_metrics("Validate", epoch_no + 1, one_val_iters, vstime, vloss, vaccuracy)

                    # Get the time of one epoch
                    self.print_time(epoch_no + 1, self.get_time(epoch_time))
                    epoch_time = time.time()

                    # Compute mean loss and mean accuracy of training steps
                    # in one epoch, and empty lists for next epoch
                    self.print_mean_metrics("Train", epoch_no + 1, tloss_list, taccuracy_list)
                    tloss_list, taccuracy_list = [], []

                    # Compute mean loss and mean accuracy of validating steps in one epoch
                    val_mean_loss = self.print_mean_metrics("Validate", epoch_no + 1, vloss_list, vaccuracy_list)

                    # Save the model with the lowest validating loss
                    if val_mean_loss < best_val_lmean_oss:
                        best_val_lmean_oss = val_mean_loss
                        # Save model after each epoch
                        self.save_model_per_epoch(sess, saver, epoch_no + 1, "best")

                    # Save model after every epoch
                    self.save_model_per_epoch(sess, saver, epoch_no + 1, "last")

                    one_tra_iters = 0
                    one_val_iters = 0
                    epoch_no += 1
                    print()

                    if epoch_no > self.num_epoches:
                        break

        except tf.errors.OutOfRangeError:
            # Stop training
            self.blue_print("Training has stopped.")
            # Save metrics into json files
            self.save_metrics("train_metrics.json", self.train_metrics)
            self.save_metrics("validate_metrics.json", self.validate_metrics)
            self.blue_print("Logs have been saved in: {}\n".format(self.logs_path))
        finally:
            coord.request_stop()

        coord.join(threads)
        sess.close()

        return

if __name__ == "__main__":

    '''

        Example of commandline:
        python btc_train.py --model=cnn
        python btc_train.py --model=full_cnn
        python btc_train.py --model=res_cnn
        python btc_train.py --model=dense_cnn

    '''

    parser = argparse.ArgumentParser()

    help_str = "Select a model in 'cnn', 'full_cnn', 'res_cnn' or 'dense_cnn'."
    parser.add_argument("--model", action="store", default="cnn",
                        dest="model", help=help_str)
    args = parser.parse_args()

    parent_dir = os.path.dirname(os.getcwd())
    save_path = os.path.join(parent_dir, "models")
    logs_path = os.path.join(parent_dir, "logs")

    btc = BTCTrainCNN(args.model, cnn_parameters, save_path, logs_path)
    btc.train()
<commit_msg>Add one new option for model<commit_after># Brain Tumor Classification
# Script for Training General CNN Models
# Author: Qixun Qu
# Create on: 2017/10/14
# Modify on: 2017/11/22

#     ,,,         ,,,
#   ;"   ';     ;'   ",
#   ;  @.ss$$$$$$s.@  ;
#   `s$$$$$$$$$$$$$$$'
#   $$$$$$$$$$$$$$$$$$
#  $$$$P""Y$$$Y""W$$$$$
#  $$$$  p"$$$"q  $$$$$
#  $$$$  .$$$$$.  $$$$'
#   $$$DaU$$O$$DaU$$$'
#    '$$$$'.^.'$$$$'
#       '&$$$$$&'

'''

Class BTCTrainCNN

-1- Models are defined in class BTCModels.
-2- Hyper-parameters can be set in btc_parameters.py.
-3- Loading tfrecords for training and validating by
    functions in class BTCTFRecords.

'''

from __future__ import print_function

import os
import time
import argparse
import numpy as np
import tensorflow as tf
from btc_settings import *
from btc_train import BTCTrain
from btc_cnn_parameters import cnn_parameters

class BTCTrainCNN(BTCTrain):

    def __init__(self, net, paras, save_path, logs_path):
        '''__INIT__

            Initialization of class BTCTrain to set parameters
            for constructing, training and validating models.

            Inputs:
            -------
            - net: string, the name of the model applied to train
            - paras: dict, parameters for training the model, defined
                     in btc_parameters.py
            - save_path: string, the path of the folder to save models
            - logs_path: string, the path of the folder to save logs

        '''

        super().__init__(paras)

        self.net_name = self.set_net_name(net)
        self.model_path = self.set_dir_path(save_path, self.net_name)
        self.logs_path = self.set_dir_path(logs_path, self.net_name)

        self.network = self._get_network(net)

        return

    def _get_network(self, net):
        '''_GET_NETWORK

            Return network function according to the given net's name.

            Input:
            ------
            - net: string, the net's name

            Output:
            -------
            - the choosen network

        '''

        # Set models by given variable
        if net == CNN:
            network = self.models.cnn
        elif net == MULTI_CNN:
            network = self.models.multi_cnn
        elif net == FULL_CNN:
            network = self.models.full_cnn
        elif net == RES_CNN:
            network = self.models.res_cnn
        elif net == DENSE_CNN:
            network = self.models.dense_cnn
        else:  # Raise error if model cannot be found
            raise ValueError("Could not found model.")

        return network

    def train(self):
        '''TRAIN

            Train and validate the choosen model.

        '''

        with tf.device("/cpu:0"):
            tra_data, tra_labels, val_data, val_labels = self.load_data()
            x, y_input, is_training, learning_rate = self.inputs()

        with tf.device("/gpu:0"):
            # Obtain logits from the model
            y_output = self.network(x, is_training)

        # Compute loss and accuracy and merge summary
        # The summary can be displayed by TensorBoard
        loss = self.get_softmax_loss(y_input, y_output)
        accuracy = self.get_accuracy(y_input, y_output)
        merged = tf.summary.merge_all()

        train_op = self.create_optimizer(learning_rate, loss)

        # Create a saver to save model while training
        saver = tf.train.Saver()
        sess = tf.InteractiveSession()
        sess.run(self.initialize_variables())
        tra_writer, val_writer = self.create_writers(self.logs_path, sess.graph)

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)

        self.blue_print("\nTraining and Validating model: {}\n".format(self.net_name))

        # Initialize counter
        tra_iters, val_iters, epoch_no = 0, 0, 0
        one_tra_iters, one_val_iters = 0, 0
        best_val_lmean_oss = np.inf

        # Lists to save loss and accuracy of each training step
        tloss_list, taccuracy_list = [], []

        # Initialize the timer to count time of one epoch
        epoch_time = time.time()

        try:
            while not coord.should_stop():
                # Initialize the timer to count time of one training step
                tra_step_time = time.time()

                # Training step
                # Feed the graph, run optimizer and get metrics
                tx, ty = sess.run([tra_data, tra_labels])
                tra_fd = {x: tx, y_input: ty, is_training: True, learning_rate: self.learning_rates[epoch_no]}
                tsummary, tloss, taccuracy, _ = sess.run([merged, loss, accuracy, train_op], feed_dict=tra_fd)

                # Get the time of one training step
                tstime = self.get_time(tra_step_time)

                tra_iters += 1
                one_tra_iters += 1

                # Record metrics of training steps
                tloss_list.append(tloss)
                taccuracy_list.append(taccuracy)
                tra_writer.add_summary(tsummary, tra_iters)
                self.train_metrics.append([tloss, taccuracy])
                self.print_metrics("Train", epoch_no + 1, one_tra_iters, tstime, tloss, taccuracy)

                if tra_iters % self.tepoch_iters[epoch_no] == 0:
                    # Validating step
                    # Lists to save loss and accuracy of each validating step
                    vloss_list, vaccuracy_list = [], []
                    while val_iters < self.vepoch_iters[epoch_no]:
                        # Initialize the timer to count time of one validating step
                        val_step_time = time.time()

                        # Feed the graph, get metrics
                        vx, vy = sess.run([val_data, val_labels])
                        val_fd = {x: vx, y_input: vy, is_training: False}
                        vsummary, vloss, vaccuracy = sess.run([merged, loss, accuracy], feed_dict=val_fd)

                        # Get the time of one validating step
                        vstime = self.get_time(val_step_time)

                        val_iters += 1
                        one_val_iters += 1

                        # Record metrics of validating steps
                        vloss_list.append(vloss)
                        vaccuracy_list.append(vaccuracy)
                        val_writer.add_summary(vsummary, val_iters)
                        self.validate_metrics.append([vloss, vaccuracy])
                        self.print_metrics("Validate", epoch_no + 1, one_val_iters, vstime, vloss, vaccuracy)

                    # Get the time of one epoch
                    self.print_time(epoch_no + 1, self.get_time(epoch_time))
                    epoch_time = time.time()

                    # Compute mean loss and mean accuracy of training steps
                    # in one epoch, and empty lists for next epoch
                    self.print_mean_metrics("Train", epoch_no + 1, tloss_list, taccuracy_list)
                    tloss_list, taccuracy_list = [], []

                    # Compute mean loss and mean accuracy of validating steps in one epoch
                    val_mean_loss = self.print_mean_metrics("Validate", epoch_no + 1, vloss_list, vaccuracy_list)

                    # Save the model with the lowest validating loss
                    if val_mean_loss < best_val_lmean_oss:
                        best_val_lmean_oss = val_mean_loss
                        # Save model after each epoch
                        self.save_model_per_epoch(sess, saver, epoch_no + 1, "best")

                    # Save model after every epoch
                    self.save_model_per_epoch(sess, saver, epoch_no + 1, "last")

                    one_tra_iters = 0
                    one_val_iters = 0
                    epoch_no += 1
                    print()

                    if epoch_no > self.num_epoches:
                        break

        except tf.errors.OutOfRangeError:
            # Stop training
            self.blue_print("Training has stopped.")
            # Save metrics into json files
            self.save_metrics("train_metrics.json", self.train_metrics)
            self.save_metrics("validate_metrics.json", self.validate_metrics)
            self.blue_print("Logs have been saved in: {}\n".format(self.logs_path))
        finally:
            coord.request_stop()

        coord.join(threads)
        sess.close()

        return

if __name__ == "__main__":

    '''

        Example of commandline:
        python btc_train.py --model=cnn
        python btc_train.py --model=full_cnn
        python btc_train.py --model=res_cnn
        python btc_train.py --model=dense_cnn

    '''

    parser = argparse.ArgumentParser()

    help_str = "Select a model in 'cnn', 'multi_cnn', 'full_cnn', 'res_cnn' or 'dense_cnn'."
    parser.add_argument("--model", action="store", default="cnn",
                        dest="model", help=help_str)
    args = parser.parse_args()

    parent_dir = os.path.dirname(os.getcwd())
    save_path = os.path.join(parent_dir, "models")
    logs_path = os.path.join(parent_dir, "logs")

    btc = BTCTrainCNN(args.model, cnn_parameters, save_path, logs_path)
    btc.train()
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""Module to calculate pn ps."""
from Bio import AlignIO
from Bio.Align import MultipleSeqAlignment
from Bio.Data import CodonTable
from divergence import parse_options
from itertools import product
import logging as log
import re
import sys

def calculate_pnps(genome_ids_a, genome_ids_b, sico_files):
    """Calculate pN, pS, pN/pS & SFS values for all sico_files per clade, and write out statistics per SICO."""
    #For each alignment create separate alignments for clade A & clade B genomes
    alignments = (AlignIO.read(sico_file, 'fasta') for sico_file in sico_files)
    for ali in alignments:
        alignment_a = MultipleSeqAlignment(seqr for seqr in ali if seqr.id.split('|')[0] in genome_ids_a)
        alignment_b = MultipleSeqAlignment(seqr for seqr in ali if seqr.id.split('|')[0] in genome_ids_b)

        #Calculate pn ps for the subaligments of each clade
        _perform_calculations(alignment_a)
        _perform_calculations(alignment_b)

#Using the standard NCBI Bacterial, Archaeal and Plant Plastid Code translation table (11).
BACTERIAL_CODON_TABLE = CodonTable.unambiguous_dna_by_id.get(11)

def _four_fold_degenerate_patterns():
    """Find patterns of 4-fold degenerate codons, wherein all third site substitutions code for the same amino acid."""
    letters = BACTERIAL_CODON_TABLE.nucleotide_alphabet.letters
    #Any combination of letters of length two
    for site12 in [''.join(prod) for prod in product(letters, repeat = 2)]:
        #4-fold when the length of the unique encoded amino acids for all possible third site nucleotides is exactly 1
        if 1 == len(set([BACTERIAL_CODON_TABLE.forward_table.get(site12 + site3) for site3 in letters])):
            #Add regular expression pattern to the set of patterns
            yield '{0}[{1}]'.format(site12, letters)

FOUR_FOLD_DEGENERATE_PATTERNS = set(_four_fold_degenerate_patterns())

def _sfs2str(sfs):
    """Convert Site Frequency Spectrum to human readable textual representation."""
    single = double = triple = ''
    if 1 in sfs:
        single = '{0} singleton'.format(sfs[1]) + ('s' if 1 < sfs[1] else '')
    if 2 in sfs:
        double = '{0} doubleton'.format(sfs[2]) + ('s' if 1 < sfs[2] else '')
    if 3 in sfs:
        triple = '{0} tripleton'.format(sfs[3]) + ('s' if 1 < sfs[3] else '')
    return ', '.join(val for val in [single, double, triple] if val)

def _perform_calculations(alignment):
    """Perform actual calculations on the alignment to determine pN, pS, SFS & the number of ignored cases per SICO."""
    #Print codons for easy debugging
    for seqr in alignment:
        seq = str(seqr.seq)
        log.info('  '.join(seq[idx:idx + 3] for idx in range(0, len(seq), 3)))

    synonymous_sfs = {}
    four_fold_syn_sfs = {}
    non_synonymous_sfs = {}
    mixed_synonymous_polymorphisms = 0
    multiple_site_polymorphisms = 0

    #Calculate range_end here so we can handle alignments that are not multiples of three
    range_end = len(alignment[0]) - len(alignment[0]) % 3
    #Split into codon_alignments
    codon_alignments = [alignment[:, index:index + 3] for index in range(0, range_end, 3)]
    for codon_alignment in codon_alignments:
        #Get string representations of codons for simplicity
        codons = [str(seqr.seq) for seqr in codon_alignment]

        #As per AEW: ignore codons with gaps, and codons with unresolved bases: Basically anything but ACGT
        if 0 < len(''.join(codons).translate(None, 'ACGTactg')):
            continue

        #Stop codons should have already been removed, but lets be sure anyway
        codons_nonstop = [codon for codon in codons if codon not in BACTERIAL_CODON_TABLE.stop_codons]

        #Retrieve translations of codons now that inconclusive & stop-codons have been removed
        translations = [BACTERIAL_CODON_TABLE.forward_table.get(codon) for codon in codons_nonstop]

        #Count unique translations across strains
        translation_usage = dict((aa, translations.count(aa)) for aa in set(translations))

        #Mutations are synonymous when all codons encode the same AA, and there are no gaps
        synonymous = len(translation_usage) == 1 and len(translations) == len(codon_alignment)

        #Retrieve nucleotides per site within the codon
        site1 = [nucl for nucl in codon_alignment[:, 0]]
        site2 = [nucl for nucl in codon_alignment[:, 1]]
        site3 = [nucl for nucl in codon_alignment[:, 2]]

        #Count occurrences of distinct nucleotides across strains
        site1_usage = dict((nucl, site1.count(nucl)) for nucl in set(site1))
        site2_usage = dict((nucl, site2.count(nucl)) for nucl in set(site2))
        site3_usage = dict((nucl, site3.count(nucl)) for nucl in set(site3))

        #Sites are polymorphic if they contain more than one nucleotide
        site1_polymorphic = 1 < len(site1_usage)
        site2_polymorphic = 1 < len(site2_usage)
        site3_polymorphic = 1 < len(site3_usage)
        polymorphisms = site1_polymorphic, site2_polymorphic, site3_polymorphic

        #Continue with next codon if none of the sites is polymorphic
        if not any(polymorphisms):
            continue

        #Determine if only one site is polymorphic by using boolean xor and not all
        single_site_polymorphism = site1_polymorphic ^ site2_polymorphic ^ site3_polymorphic and not all(polymorphisms)

        #Debug print statement
        log.info('1:{0}  2:{1}  3:{2}  Translations:{3}\tSynonymous:{4}\tSingle-site:{5}' \
        .format(site1_usage, site2_usage, site3_usage, translation_usage, synonymous, single_site_polymorphism))

        #Skip multiple site polymorphisms, but do keep a count of how many we encounter
        if not single_site_polymorphism:
            multiple_site_polymorphisms += 1
            continue

        #Determine which site_usage is the single site polymorphism
        polymorph_site_usage = site1_usage if site1_polymorphic else site2_usage if site2_polymorphic else site3_usage

        #Find the 'reference' nucleotide as (one of) the most occurring occupations in this site, so we can -1 later
        psu_values = polymorph_site_usage.values()
        reference_allele_count = max(psu_values)

        #Calculate the local site frequency spectrum, to be added to the gene-wide SFS later
        local_sfs = dict((ntimes, psu_values.count(ntimes)) for ntimes in set(psu_values))
        #Deduct one for the reference_allele_count, which should not count towards the SFS
        local_sfs[reference_allele_count] = local_sfs[reference_allele_count] - 1
        if local_sfs[reference_allele_count] == 0:
            del local_sfs[reference_allele_count]

        if synonymous:
            #If all polymorphisms encode for the same AA, we have multiple synonymous polymorphisms, where:
            #2 nucleotides = 1 polymorphism, 3 nucleotides = 2 polymorphisms, 4 nucleotides = 3 polymorphisms

            #Debug log statement
            log.info('local SFS: {0}'.format(_sfs2str(local_sfs)))

            #Update synonymous SFS by adding values from local SFS
            for maf, count in local_sfs.iteritems():
                prev_occupations = synonymous_sfs.get(maf, 0)
                synonymous_sfs[maf] = prev_occupations + count

            #Codon is four fold degenerate if it matches any pattern in FOUR_FOLD_DEGENERATE_PATTERNS
            if site3_polymorphic:
                codon = site1[0] + site2[0] + site3[0]
                for pattern in FOUR_FOLD_DEGENERATE_PATTERNS:
                    if re.match(pattern, codon):
                        #Update four fold degenerate SFS by adding values from local SFS
                        for maf, count in local_sfs.iteritems():
                            prev_occupations = four_fold_syn_sfs.get(maf, 0)
                            four_fold_syn_sfs[maf] = prev_occupations + count
            continue

        if not synonymous:
            if len(polymorph_site_usage) == len(translation_usage):
                #If all polymorphisms encode for different AA, we have multiple non-synonymous polymorphisms, where:
                #2 nucleotides = 1 polymorphism, 3 nucleotides = 2 polymorphisms, 4 nucleotides = 3 polymorphisms

                #Debug log statement
                log.info('local SFS: {0}'.format(_sfs2str(local_sfs)))

                #Update non synonymous SFS by adding values from local SFS
                for maf, count in local_sfs.iteritems():
                    prev_occupations = non_synonymous_sfs.get(maf, 0)
                    non_synonymous_sfs[maf] = prev_occupations + count
                continue
            else:
                #Some, but not all polymorphisms encode for different AA, making it unclear how this should be scored
                mixed_synonymous_polymorphisms += 1
                continue

    #Use Site Frequency Spectrum to calculate the number of synonymous and non synonymous polymorphisms
    #Note: this requires a complete SFS across synonymous & non_synonymous polymorphisms, be careful when updating code
    synonymous_polymorphisms = sum(synonymous_sfs.values())
    four_fold_syn_polymorphisms = sum(four_fold_syn_sfs.values())
    non_synonymous_polymorphisms = sum(non_synonymous_sfs.values())

    log.info('\n')

    log.info('synonymous_polymorphisms %i', synonymous_polymorphisms)
    log.info('synonymous_sfs:\t%s\n', _sfs2str(synonymous_sfs))

    log.info('four fold synonymous_polymorphisms %i', four_fold_syn_polymorphisms)
    log.info('four fold synonymous_sfs:\t%s\n', _sfs2str(four_fold_syn_sfs))

    log.info('non_synonymous_polymorphisms %i', non_synonymous_polymorphisms)
    log.info('non_synonymous_sfs:\t%s\n', _sfs2str(non_synonymous_sfs))

    log.info('mixed_synonymous_polymorphisms %i', mixed_synonymous_polymorphisms)
    log.info('multiple_site_polymorphisms %i', multiple_site_polymorphisms)

    #1. gene name
    gene_name = 'xyz'

    #2. number of strains (this will typically be the same for all genes)
    number_of_strains = len(alignment)

    #3. number of non-synonymous sites (see below for calculation) = LnP
    #TODOnumber_of_non_syn_sites = ?

    #4. SFS for non-synonymous polymorphsims
    non_synonymous_sfs

    #5. The sum of SFS for non-synonymous polymorphisms = Pn
    pn = sum(non_synonymous_sfs.values())

    #6. number of synonymous sites (see below for calculation) = LsP
    #TODO number_of_non_syn_sites = ?

    #7. SFS for synonymous prolymorphisms
    synonymous_sfs

    #8. The sum of the SFS for synonymous polymorphisms = Ps
    ps = sum(synonymous_sfs.values())

    #9. number of 4-fold synonymous sites = L4

    #10. SFS for 4-fold synonymous polymorphisms

    #11. Sum of the SFS for 4-folds = P4

    #12. number of non-synonymous sites for divergence from PAML (this might be different to 3, because this will come from two randomly chosen sequences) = LnD

    #13. number of non-synonymous substitutions from PAML = Dn

    #14. number of synonymous sites from PAML = LsD

    #15. number of synonymous substitutions from PAML = Ds

    #16. Direction of selection = Dn/(Dn+Ds) - Pn/(Pn+Ps)

    #I would calculate LnP and LsP as follows:
    #
    #LnP = L * LnD/(LnD+LsD)
    #LsP = L * LsD/(LnD+LsD)
    #
    #where L is the length of the sequence from which the polymorphism data is taken (i.e. 3* no of codons used)
    #
    #Finally, we might a line which gives the sum for columns 3 to 15 plus the mean of column 16

def main(args):
    """Main function called when run from command line or as part of pipeline."""
    usage = """
Usage: calculate_pnps.py
--alignment=FILE    alignment file in FASTA format
"""
    options = ['alignment']
    (alignment_file,) = parse_options(usage, options, args)

    alignment = AlignIO.read(alignment_file, 'fasta')
    _perform_calculations(alignment)

if __name__ == '__main__':
    main(sys.argv[1:])
<commit_msg>Also handle x-ton larger than single-, double- & tripleton when printing SFS<commit_after>#!/usr/bin/env python
"""Module to calculate pn ps."""
from Bio import AlignIO
from Bio.Align import MultipleSeqAlignment
from Bio.Data import CodonTable
from divergence import parse_options
from itertools import product
import logging as log
import re
import sys

def calculate_pnps(genome_ids_a, genome_ids_b, sico_files):
    """Calculate pN, pS, pN/pS & SFS values for all sico_files per clade, and write out statistics per SICO."""
    #For each alignment create separate alignments for clade A & clade B genomes
    alignments = (AlignIO.read(sico_file, 'fasta') for sico_file in sico_files)
    for ali in alignments:
        alignment_a = MultipleSeqAlignment(seqr for seqr in ali if seqr.id.split('|')[0] in genome_ids_a)
        alignment_b = MultipleSeqAlignment(seqr for seqr in ali if seqr.id.split('|')[0] in genome_ids_b)

        #Calculate pn ps for the subaligments of each clade
        _perform_calculations(alignment_a)
        _perform_calculations(alignment_b)

#Using the standard NCBI Bacterial, Archaeal and Plant Plastid Code translation table (11).
BACTERIAL_CODON_TABLE = CodonTable.unambiguous_dna_by_id.get(11)

def _four_fold_degenerate_patterns():
    """Find patterns of 4-fold degenerate codons, wherein all third site substitutions code for the same amino acid."""
    letters = BACTERIAL_CODON_TABLE.nucleotide_alphabet.letters
    #Any combination of letters of length two
    for site12 in [''.join(prod) for prod in product(letters, repeat = 2)]:
        #4-fold when the length of the unique encoded amino acids for all possible third site nucleotides is exactly 1
        if 1 == len(set([BACTERIAL_CODON_TABLE.forward_table.get(site12 + site3) for site3 in letters])):
            #Add regular expression pattern to the set of patterns
            yield '{0}[{1}]'.format(site12, letters)

FOUR_FOLD_DEGENERATE_PATTERNS = set(_four_fold_degenerate_patterns())

def _sfs2str(sfs):
    """Convert Site Frequency Spectrum to human readable textual representation."""
    def _generator(sfs):
        """Return successive values for - and names of multipleton(s)."""
        multiple_names = {1:'singleton', 2:'doubleton', 3:'tripleton', 4:'quadrupleton', 5:'quintupleton'}
        for xton in sorted(sfs.keys()):
            name = multiple_names.get(xton, '"{0}"-ton'.format(xton))
            yield '{0} {1}'.format(sfs[xton], name) + ('s' if 1 < sfs[1] else '')
    return ', '.join(_generator(sfs))

def _perform_calculations(alignment):
    """Perform actual calculations on the alignment to determine pN, pS, SFS & the number of ignored cases per SICO."""
    #Print codons for easy debugging
    for seqr in alignment:
        seq = str(seqr.seq)
        log.info('  '.join(seq[idx:idx + 3] for idx in range(0, len(seq), 3)))

    synonymous_sfs = {}
    four_fold_syn_sfs = {}
    non_synonymous_sfs = {}
    mixed_synonymous_polymorphisms = 0
    multiple_site_polymorphisms = 0

    #Calculate range_end here so we can handle alignments that are not multiples of three
    range_end = len(alignment[0]) - len(alignment[0]) % 3
    #Split into codon_alignments
    codon_alignments = [alignment[:, index:index + 3] for index in range(0, range_end, 3)]
    for codon_alignment in codon_alignments:
        #Get string representations of codons for simplicity
        codons = [str(seqr.seq) for seqr in codon_alignment]

        #As per AEW: ignore codons with gaps, and codons with unresolved bases: Basically anything but ACGT
        if 0 < len(''.join(codons).translate(None, 'ACGTactg')):
            continue

        #Stop codons should have already been removed, but lets be sure anyway
        codons_nonstop = [codon for codon in codons if codon not in BACTERIAL_CODON_TABLE.stop_codons]

        #Retrieve translations of codons now that inconclusive & stop-codons have been removed
        translations = [BACTERIAL_CODON_TABLE.forward_table.get(codon) for codon in codons_nonstop]

        #Count unique translations across strains
        translation_usage = dict((aa, translations.count(aa)) for aa in set(translations))

        #Mutations are synonymous when all codons encode the same AA, and there are no gaps
        synonymous = len(translation_usage) == 1 and len(translations) == len(codon_alignment)

        #Retrieve nucleotides per site within the codon
        site1 = [nucl for nucl in codon_alignment[:, 0]]
        site2 = [nucl for nucl in codon_alignment[:, 1]]
        site3 = [nucl for nucl in codon_alignment[:, 2]]

        #Count occurrences of distinct nucleotides across strains
        site1_usage = dict((nucl, site1.count(nucl)) for nucl in set(site1))
        site2_usage = dict((nucl, site2.count(nucl)) for nucl in set(site2))
        site3_usage = dict((nucl, site3.count(nucl)) for nucl in set(site3))

        #Sites are polymorphic if they contain more than one nucleotide
        site1_polymorphic = 1 < len(site1_usage)
        site2_polymorphic = 1 < len(site2_usage)
        site3_polymorphic = 1 < len(site3_usage)
        polymorphisms = site1_polymorphic, site2_polymorphic, site3_polymorphic

        #Continue with next codon if none of the sites is polymorphic
        if not any(polymorphisms):
            continue

        #Determine if only one site is polymorphic by using boolean xor and not all
        single_site_polymorphism = site1_polymorphic ^ site2_polymorphic ^ site3_polymorphic and not all(polymorphisms)

        #Debug print statement
        log.info('1:{0}  2:{1}  3:{2}  Translations:{3}\tSynonymous:{4}\tSingle-site:{5}' \
        .format(site1_usage, site2_usage, site3_usage, translation_usage, synonymous, single_site_polymorphism))

        #Skip multiple site polymorphisms, but do keep a count of how many we encounter
        if not single_site_polymorphism:
            multiple_site_polymorphisms += 1
            continue

        #Determine which site_usage is the single site polymorphism
        polymorph_site_usage = site1_usage if site1_polymorphic else site2_usage if site2_polymorphic else site3_usage

        #Find the 'reference' nucleotide as (one of) the most occurring occupations in this site, so we can -1 later
        psu_values = polymorph_site_usage.values()
        reference_allele_count = max(psu_values)

        #Calculate the local site frequency spectrum, to be added to the gene-wide SFS later
        local_sfs = dict((ntimes, psu_values.count(ntimes)) for ntimes in set(psu_values))
        #Deduct one for the reference_allele_count, which should not count towards the SFS
        local_sfs[reference_allele_count] = local_sfs[reference_allele_count] - 1
        if local_sfs[reference_allele_count] == 0:
            del local_sfs[reference_allele_count]

        if synonymous:
            #If all polymorphisms encode for the same AA, we have multiple synonymous polymorphisms, where:
            #2 nucleotides = 1 polymorphism, 3 nucleotides = 2 polymorphisms, 4 nucleotides = 3 polymorphisms

            #Debug log statement
            log.info('local SFS: {0}'.format(local_sfs))

            #Update synonymous SFS by adding values from local SFS
            for maf, count in local_sfs.iteritems():
                prev_occupations = synonymous_sfs.get(maf, 0)
                synonymous_sfs[maf] = prev_occupations + count

            #Codon is four fold degenerate if it matches any pattern in FOUR_FOLD_DEGENERATE_PATTERNS
            if site3_polymorphic:
                codon = site1[0] + site2[0] + site3[0]
                for pattern in FOUR_FOLD_DEGENERATE_PATTERNS:
                    if re.match(pattern, codon):
                        #Update four fold degenerate SFS by adding values from local SFS
                        for maf, count in local_sfs.iteritems():
                            prev_occupations = four_fold_syn_sfs.get(maf, 0)
                            four_fold_syn_sfs[maf] = prev_occupations + count
            continue

        if not synonymous:
            if len(polymorph_site_usage) == len(translation_usage):
                #If all polymorphisms encode for different AA, we have multiple non-synonymous polymorphisms, where:
                #2 nucleotides = 1 polymorphism, 3 nucleotides = 2 polymorphisms, 4 nucleotides = 3 polymorphisms

                #Debug log statement
                log.info('local SFS: {0}'.format(local_sfs))

                #Update non synonymous SFS by adding values from local SFS
                for maf, count in local_sfs.iteritems():
                    prev_occupations = non_synonymous_sfs.get(maf, 0)
                    non_synonymous_sfs[maf] = prev_occupations + count
                continue
            else:
                #Some, but not all polymorphisms encode for different AA, making it unclear how this should be scored
                mixed_synonymous_polymorphisms += 1
                continue

    #Use Site Frequency Spectrum to calculate the number of synonymous and non synonymous polymorphisms
    #Note: this requires a complete SFS across synonymous & non_synonymous polymorphisms, be careful when updating code
    synonymous_polymorphisms = sum(synonymous_sfs.values())
    four_fold_syn_polymorphisms = sum(four_fold_syn_sfs.values())
    non_synonymous_polymorphisms = sum(non_synonymous_sfs.values())

    log.info('\n')

    log.info('synonymous_polymorphisms %i', synonymous_polymorphisms)
    log.info('synonymous_sfs:\t%s\n', _sfs2str(synonymous_sfs))

    log.info('four fold synonymous_polymorphisms %i', four_fold_syn_polymorphisms)
    log.info('four fold synonymous_sfs:\t%s\n', _sfs2str(four_fold_syn_sfs))

    log.info('non_synonymous_polymorphisms %i', non_synonymous_polymorphisms)
    log.info('non_synonymous_sfs:\t%s\n', _sfs2str(non_synonymous_sfs))

    log.info('mixed_synonymous_polymorphisms %i', mixed_synonymous_polymorphisms)
    log.info('multiple_site_polymorphisms %i', multiple_site_polymorphisms)

    #1. gene name
    gene_name = 'xyz'

    #2. number of strains (this will typically be the same for all genes)
    number_of_strains = len(alignment)

    #3. number of non-synonymous sites (see below for calculation) = LnP
    #TODOnumber_of_non_syn_sites = ?

    #4. SFS for non-synonymous polymorphsims
    non_synonymous_sfs

    #5. The sum of SFS for non-synonymous polymorphisms = Pn
    pn = sum(non_synonymous_sfs.values())

    #6. number of synonymous sites (see below for calculation) = LsP
    #TODO number_of_non_syn_sites = ?

    #7. SFS for synonymous prolymorphisms
    synonymous_sfs

    #8. The sum of the SFS for synonymous polymorphisms = Ps
    ps = sum(synonymous_sfs.values())

    #9. number of 4-fold synonymous sites = L4

    #10. SFS for 4-fold synonymous polymorphisms

    #11. Sum of the SFS for 4-folds = P4

    #12. number of non-synonymous sites for divergence from PAML (this might be different to 3, because this will come from two randomly chosen sequences) = LnD

    #13. number of non-synonymous substitutions from PAML = Dn

    #14. number of synonymous sites from PAML = LsD

    #15. number of synonymous substitutions from PAML = Ds

    #16. Direction of selection = Dn/(Dn+Ds) - Pn/(Pn+Ps)

    #I would calculate LnP and LsP as follows:
    #
    #LnP = L * LnD/(LnD+LsD)
    #LsP = L * LsD/(LnD+LsD)
    #
    #where L is the length of the sequence from which the polymorphism data is taken (i.e. 3* no of codons used)
    #
    #Finally, we might a line which gives the sum for columns 3 to 15 plus the mean of column 16

def main(args):
    """Main function called when run from command line or as part of pipeline."""
    usage = """
Usage: calculate_pnps.py
--alignment=FILE    alignment file in FASTA format
"""
    options = ['alignment']
    (alignment_file,) = parse_options(usage, options, args)

    alignment = AlignIO.read(alignment_file, 'fasta')
    _perform_calculations(alignment)

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


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

# Copyright (c) 2016 Ericsson AB
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from calvin.runtime.north.plugins.port.queue.common import QueueFull
from calvin.runtime.north.plugins.port.queue.fanout_base import FanoutBase

from calvin.utilities import calvinlogger

_log = calvinlogger.get_logger(__name__)

class FanoutOrderedFIFO(FanoutBase):

    """
    A FIFO which route tokens based on an ordering to peers

    """

    def __init__(self, port_properties, peer_port_properties):
        super(FanoutOrderedFIFO, self).__init__(port_properties, peer_port_properties)
        self.reader_turn = None
        self.turn_pos = 0
        self._set_turn()
        self._type = "dispatch:ordered"

    def _set_turn(self):
        # Get routing schedule based on info after ':' in type
        # Set the correct method
        self._update_turn = self._ordered
        if self.reader_turn is None:
            # Populate with alternating values up to N
            self.reader_turn = [p for n in range(self.N) for p in range(self.nbr_peers)][:self.N]
        # print "_set_turn(): "
        # print "    self._update_turn:", self._update_turn
        # print "    self.reader_turn:", self.reader_turn

    def _reset_turn(self):
        # Populate with alternating values up to N
        pos = self.reader_turn[self.turn_pos % self.N]
        self.reader_turn = [p for n in range(self.N) for p in range(self.nbr_peers)][pos:self.N+pos]
        self.turn_pos = 0

    def _ordered(self):
        reader = self.reader_turn[self.turn_pos % self.N]
        prev = self.reader_turn[(self.turn_pos - 1) % self.N]
        self.reader_turn[self.turn_pos % self.N] = (prev + 1) % self.nbr_peers
        self.turn_pos += 1
        # print "_ordered():"
        # print "    reader:", reader
        # print "    self.reader_turn:", self.reader_turn
        return reader

    def _state(self, remap=None):
        state = super(FanoutOrderedFIFO, self)._state(remap)
        if remap is None:
            state['reader_turn'] = self.reader_turn
            state['turn_pos'] = self.turn_pos
        else :
            state['reader_turn'] = None
            state['turn_pos'] = 0
        return state

    def _set_state(self, state):
        super(FanoutOrderedFIFO, self)._set_state(state)
        self.reader_turn = state['reader_turn']
        self.turn_pos = state['turn_pos']
        self._set_turn()

    def _set_port_order(self, order):
        if not set(order) == set(self.readers):
            # print order, self.readers
            raise Exception("Illegal port order list")
        self.readers = order

    def remove_reader(self, reader):
        removed = super(FanoutOrderedFIFO, self).remove_reader(reader)
        if removed:
            self._reset_turn()
        return removed

    def write(self, data, metadata):
        #_log.debug("WRITE1 %s" % metadata)
        if not self.slots_available(1, metadata):
            raise QueueFull()
        # Write token in peer's FIFO
        peer = self.readers[self._update_turn()]
        write_pos = self.write_pos[peer]
        #_log.debug("WRITE2 %s %s %d\n%s" % (metadata, peer, write_pos, str(map(str, self.fifo[peer]))))
        self.fifo[peer][write_pos % self.N] = data
        self.write_pos[peer] = write_pos + 1
        return True

    def slots_available(self, length, metadata):
        if length >= self.N:
            return False
        if length == 1:
            # shortcut for common special case
            peer = self.readers[self.reader_turn[self.turn_pos % self.N]]
            return self.write_pos[peer] - self.read_pos[peer] < self.N - 1
        # list of peer indexes that will be written to
        peers = [self.reader_turn[i % self.N] for i in range(self.turn_pos, self.turn_pos + length)]
        for p in peers:
            # How many slots for this peer
            c = peers.count(p)
            peer = self.readers[p]
            # If this peer does not have c slots then return False
            if (self.N - (self.write_pos[peer] - self.read_pos[peer]) - 1) < c:
                return False
        # ... otherwise all had enough slots
        return True
<commit_msg>queue: missed one<commit_after># -*- coding: utf-8 -*-

# Copyright (c) 2016 Ericsson AB
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from calvin.runtime.north.plugins.port.queue.common import QueueFull
from calvin.runtime.north.plugins.port.queue.fanout_base import FanoutBase

from calvin.utilities import calvinlogger

_log = calvinlogger.get_logger(__name__)

class FanoutOrderedFIFO(FanoutBase):

    """
    A FIFO which route tokens based on an ordering to peers

    """

    def __init__(self, port_properties, peer_port_properties):
        super(FanoutOrderedFIFO, self).__init__(port_properties, peer_port_properties)
        self.reader_turn = None
        self.turn_pos = 0
        self._set_turn()
        self._type = "dispatch:ordered"

    def _set_turn(self):
        # Get routing schedule based on info after ':' in type
        # Set the correct method
        self._update_turn = self._ordered
        if self.reader_turn is None:
            # Populate with alternating values up to N
            self.reader_turn = [p for n in range(self.N) for p in range(self.nbr_peers)][:self.N]
        # print "_set_turn(): "
        # print "    self._update_turn:", self._update_turn
        # print "    self.reader_turn:", self.reader_turn

    def _reset_turn(self):
        # Populate with alternating values up to N
        pos = self.reader_turn[self.turn_pos % self.N]
        self.reader_turn = [p for n in range(self.N) for p in range(self.nbr_peers)][pos:self.N+pos]
        self.turn_pos = 0

    def _ordered(self):
        reader = self.reader_turn[self.turn_pos % self.N]
        prev = self.reader_turn[(self.turn_pos - 1) % self.N]
        self.reader_turn[self.turn_pos % self.N] = (prev + 1) % self.nbr_peers
        self.turn_pos += 1
        # print "_ordered():"
        # print "    reader:", reader
        # print "    self.reader_turn:", self.reader_turn
        return reader

    def _state(self, remap=None):
        state = super(FanoutOrderedFIFO, self)._state(remap)
        if remap is None:
            state['reader_turn'] = self.reader_turn
            state['turn_pos'] = self.turn_pos
        else :
            state['reader_turn'] = None
            state['turn_pos'] = 0
        return state

    def _set_state(self, state):
        super(FanoutOrderedFIFO, self)._set_state(state)
        self.reader_turn = state['reader_turn']
        self.turn_pos = state['turn_pos']
        self._set_turn()

    def _set_port_order(self, order):
        if not set(order) == set(self.readers):
            # print order, self.readers
            raise Exception("Illegal port order list")
        self.readers = order

    def add_reader(self, reader, properties):
        super(FanoutOrderedFIFO, self).add_reader(reader, properties)
        self.readers.sort()

    def remove_reader(self, reader):
        removed = super(FanoutOrderedFIFO, self).remove_reader(reader)
        if removed:
            self._reset_turn()
        return removed

    def write(self, data, metadata):
        #_log.debug("WRITE1 %s" % metadata)
        if not self.slots_available(1, metadata):
            raise QueueFull()
        # Write token in peer's FIFO
        peer = self.readers[self._update_turn()]
        write_pos = self.write_pos[peer]
        #_log.debug("WRITE2 %s %s %d\n%s" % (metadata, peer, write_pos, str(map(str, self.fifo[peer]))))
        self.fifo[peer][write_pos % self.N] = data
        self.write_pos[peer] = write_pos + 1
        return True

    def slots_available(self, length, metadata):
        if length >= self.N:
            return False
        if length == 1:
            # shortcut for common special case
            peer = self.readers[self.reader_turn[self.turn_pos % self.N]]
            return self.write_pos[peer] - self.read_pos[peer] < self.N - 1
        # list of peer indexes that will be written to
        peers = [self.reader_turn[i % self.N] for i in range(self.turn_pos, self.turn_pos + length)]
        for p in peers:
            # How many slots for this peer
            c = peers.count(p)
            peer = self.readers[p]
            # If this peer does not have c slots then return False
            if (self.N - (self.write_pos[peer] - self.read_pos[peer]) - 1) < c:
                return False
        # ... otherwise all had enough slots
        return True
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- encoding: utf-8 -*-

import urlparse

def join(url, part):
    sep = "/"
    if url.endswith("/"):
        sep = ""
    return "%s%s%s" % (url, sep, part)

def make(url, path = None):
    u = ""

    if path is not None:
        u = urlparse.urlunparse((url.scheme, url.netloc, path, url.params,
                                 url.query, url.fragment))
    else:
        u = url.geturl()

    return u

def canonicalize(url, parent=None):
    if url.scheme:
        netloc_unauth = ('%s:%s' % (url.hostname, url.port) if url.port != 80
                         else url.hostname)
        return urlparse.urlunparse((url.scheme, netloc_unauth,
                                    url.path.replace('//', '/'),
                                    url.params, url.query, url.fragment))
    else:
        return urlparse.urljoin(parent.canonical_url,
                                url.path.replace('//', '/'))
<commit_msg>url.port can be None if it is not specified<commit_after>#!/usr/bin/env python
# -*- encoding: utf-8 -*-

import urlparse

def join(url, part):
    sep = "/"
    if url.endswith("/"):
        sep = ""
    return "%s%s%s" % (url, sep, part)

def make(url, path = None):
    u = ""

    if path is not None:
        u = urlparse.urlunparse((url.scheme, url.netloc, path, url.params,
                                 url.query, url.fragment))
    else:
        u = url.geturl()

    return u

def canonicalize(url, parent=None):
    if url.scheme:
        netloc_unauth = ('%s:%s' % (url.hostname, url.port)
                         if url.port not in (80, None)
                         else url.hostname)
        return urlparse.urlunparse((url.scheme, netloc_unauth,
                                    url.path.replace('//', '/'),
                                    url.params, url.query, url.fragment))
    else:
        return urlparse.urljoin(parent.canonical_url,
                                url.path.replace('//', '/'))
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- encoding: utf-8 -*-

import copy
import urlparse
import vobject
import StringIO
import uuid
from lxml import etree

from caldav.lib import error, vcal, url
from caldav.lib.namespace import ns
from caldav.elements import dav, cdav

class DAVObject(object):
    """
    Base class for all DAV objects.
    """
    id = None
    url = None
    client = None
    parent = None
    name = None

    def __init__(self, client, url=None, parent=None, name=None, id=None):
        """
        Default constructor.

        Parameters:
         * client: A DAVClient instance
         * url: The url for this object
         * parent: The parent object
         * name: A displayname
         * id: The resource id (UID for an Event)
        """

        self.client = client
        self.parent = parent
        self.name = name
        self.id = id
        if url is not None:
            self.url = urlparse.urlparse(url)

    @property
    def canonical_url(self):
        return url.canonicalize(self.url, self.parent)

    def children(self, type=None):
        """
        List children, using a propfind (resourcetype) on the parent object,
        at depth = 1.
        """
        c = []

        depth = 1
        properties = {}

        props = [dav.ResourceType(), ]

        prop = dav.Prop() + props
        root = dav.Propfind() + prop

        body = etree.tostring(root.xmlelement(), encoding="utf-8",
                              xml_declaration=True)

        response = self.client.propfind(self.url.path, body, depth)
        for r in response.tree.findall(dav.Response.tag):
            # We use canonicalized urls to index children
            href = urlparse.urlparse(r.find(dav.Href.tag).text)
            href = url.canonicalize(href, self)
            properties[href] = {}
            for p in props:
                t = r.find(".//" + p.tag)
                if len(list(t)) > 0:
                    if type is not None:
                        val = t.find(".//" + type)
                    else:
                        val = t.find(".//*")
                    if val is not None:
                        val = val.tag
                    else:
                        val = None
                else:
                    val = t.text
                properties[href][p.tag] = val

        for path in properties.keys():
            resource_type = properties[path][dav.ResourceType.tag]
            if resource_type == type or type is None:
                if path != self.canonical_url:
                    c.append((path, resource_type))

        return c

    def _get_properties(self, props=[], depth=0):
        properties = {}

        body = ""
        # build the propfind request
        if len(props) > 0:
            prop = dav.Prop() + props
            root = dav.Propfind() + prop

            body = etree.tostring(root.xmlelement(), encoding="utf-8",
                                  xml_declaration=True)

        response = self.client.propfind(self.url.path, body, depth)
        # All items should be in a <D:response> element
        for r in response.tree.findall(dav.Response.tag):
            href = r.find(dav.Href.tag).text
            properties[href] = {}
            for p in props:
                t = r.find(".//" + p.tag)
                if len(list(t)) > 0:
                    val = t.find(".//*")
                    if val is not None:
                        val = val.tag
                    else:
                        val = None
                else:
                    val = t.text
                properties[href][p.tag] = val

        return properties

    def get_properties(self, props=[], depth=0):
        """
        Get properties (PROPFIND) for this object. Works only for
        properties, that don't have complex types.

        Parameters:
         * props = [dav.ResourceType(), dav.DisplayName(), ...]

        Returns:
         * {proptag: value, ...}
        """
        rc = None
        properties =  self._get_properties(props, depth)
        path = self.url.path
        exchange_path = self.url.path + '/'

        if path in properties.keys():
            rc = properties[path]
        elif exchange_path in properties.keys():
            rc = properties[exchange_path]
        else:
            raise Exception("The CalDAV server you are using has a problem with path handling.")

        return rc

    def set_properties(self, props=[]):
        """
        Set properties (PROPPATCH) for this object.

        Parameters:
         * props = [dav.DisplayName('name'), ...]

        Returns:
         * self
        """
        prop = dav.Prop() + props
        set = dav.Set() + prop
        root = dav.PropertyUpdate() + set

        q = etree.tostring(root.xmlelement(), encoding="utf-8",
                           xml_declaration=True)
        r = self.client.proppatch(self.url.path, q)

        statuses = r.tree.findall(".//" + dav.Status.tag)
        for s in statuses:
            if not s.text.endswith("200 OK"):
                raise error.PropsetError(r.raw)

        return self

    def save(self):
        """
        Save the object. This is an abstract methed, that all classes
        derived .from DAVObject implement.

        Returns:
         * self
        """
        raise NotImplementedError()

    def delete(self):
        """
        Delete the object.
        """
        if self.url is not None:
            path = self.url.path
            r = self.client.delete(path)

            #TODO: find out why we get 404
            if r.status not in (200, 204, 404):
                raise error.DeleteError(r.raw)

class Principal(DAVObject):
    """
    This class represents a DAV Principal. It doesn't do much, except play
    the role of the parent to all calendar collections.
    """
    def __init__(self, client, url):
        """
        This object has a specific constructor, because its url is mandatory.
        """
        self.client = client
        self.url = urlparse.urlparse(url)

    def calendars(self):
        """
        List all calendar collections in this principal.

        Returns:
         * [Calendar(), ...]
        """
        cals = []

        data = self.children(cdav.Calendar.tag)
        for c_url, c_type in data:
            cals.append(Calendar(self.client, c_url, parent=self))

        return cals

class Calendar(DAVObject):
    """
    The `Calendar` object is used to represent a calendar collection.
    Refer to the RFC for details: http://www.ietf.org/rfc/rfc4791.txt
    """
    def _create(self, name, id=None):
        """
        Create a new calendar with display name `name` in `parent`.
        """
        path = None
        if id is None:
            id = str(uuid.uuid1())

        name = dav.DisplayName(name)
        cal = cdav.CalendarCollection()
        coll = dav.Collection() + cal
        type = dav.ResourceType() + coll

        prop = dav.Prop() + [type, name]
        set = dav.Set() + prop

        mkcol = dav.Mkcol() + set

        q = etree.tostring(mkcol.xmlelement(), encoding="utf-8",
                           xml_declaration=True)
        path = url.join(self.parent.url.path, id)

        r = self.client.mkcol(path, q)
        if r.status == 201:
            # XXX Should we use self.canonical_url ?
            path = url.make(self.parent.url, path)
        else:
            raise error.MkcolError(r.raw)

        return (id, path)

    def save(self):
        """
        The save method for a calendar is only used to create it, for now.
        We know we have to create it when we don't have a url.

        Returns:
         * self
        """
        if self.url is None:
            (id, path) = self._create(self.name, self.id)
            self.id = id
            self.url = urlparse.urlparse(path)
        return self

    def date_search(self, start, end=None):
        """
        Search events by date in the calendar. Recurring events are expanded
        if they have an occurence during the specified time frame.

        Parameters:
         * start = datetime.today().
         * end = same as above.

        Returns:
         * [Event(), ...]
        """
        matches = []

        # build the request
        expand = cdav.Expand(start, end)
        data = cdav.CalendarData() + expand
        prop = dav.Prop() + data

        range = cdav.TimeRange(start, end)
        vevent = cdav.CompFilter("VEVENT") + range
        vcal = cdav.CompFilter("VCALENDAR") + vevent
        filter = cdav.Filter() + vcal

        root = cdav.CalendarQuery() + [prop, filter]

        q = etree.tostring(root.xmlelement(), encoding="utf-8",
                           xml_declaration=True)
        response = self.client.report(self.url.path, q, 1)
        for r in response.tree.findall(".//" + dav.Response.tag):
            status = r.find(".//" + dav.Status.tag)
            if status.text.endswith("200 OK"):
                href = urlparse.urlparse(r.find(dav.Href.tag).text)
                href = url.canonicalize(href, self)
                data = r.find(".//" + cdav.CalendarData.tag).text
                e = Event(self.client, url=href, data=data, parent=self)
                matches.append(e)
            else:
                raise error.ReportError(response.raw)

        return matches

    def event(self, uid):
        """
        Get one event from the calendar.

        Parameters:
         * uid: the event uid

        Returns:
         * Event() or None
        """
        e = None

        data = cdav.CalendarData()
        prop = dav.Prop() + data

        match = cdav.TextMatch(uid)
        propf = cdav.PropFilter("UID") + match
        vevent = cdav.CompFilter("VEVENT") + propf
        vcal = cdav.CompFilter("VCALENDAR") + vevent
        filter = cdav.Filter() + vcal

        root = cdav.CalendarQuery() + [prop, filter]

        q = etree.tostring(root.xmlelement(), encoding="utf-8",
                           xml_declaration=True)
        response = self.client.report(self.url.path, q, 1)
        r = response.tree.find(".//" + dav.Response.tag)
        if r is not None:
            href = urlparse.urlparse(r.find(".//" + dav.Href.tag).text)
            href = url.canonicalize(href, self)
            data = r.find(".//" + cdav.CalendarData.tag).text
            e = Event(self.client, url=href, data=data, parent=self)
        else:
            raise error.NotFoundError(response.raw)

        return e

    def events(self):
        """
        List all events from the calendar.

        Returns:
         * [Event(), ...]
        """
        all = []

        data = self.children()
        for e_url, e_type in data:
            all.append(Event(self.client, e_url, parent=self))

        return all

    def __str__(self):
        return "Collection: %s" % self.canonical_url

class Event(DAVObject):
    """
    The `Event` object is used to represent an event.
    """
    _instance = None
    _data = None

    def __init__(self, client, url=None, data=None, parent=None, id=None):
        """
        Event has an additional parameter for its constructor:
         * data = "...", vCal data for the event
        """
        DAVObject.__init__(self, client=client, url=url, parent=parent, id=id)
        if data is not None:
            self.data = data

    def load(self):
        """
        Load the event from the caldav server.
        """
        r = self.client.request(self.url.path)
        self.data = vcal.fix(r.raw)
        return self

    def _create(self, data, id=None, path=None):
        if id is None:
            id = str(uuid.uuid1())
        if path is None:
            path = url.join(self.parent.url.path, id + ".ics")

        r = self.client.put(path, data,
                            {"Content-Type": 'text/calendar; charset="utf-8"'})
        if r.status == 204 or r.status == 201:
            # XXX Should we use self.canonical_url ?
            path = url.make(self.parent.url, path)
        else:
            raise error.PutError(r.raw)

        return (id, path)

    def save(self):
        """
        Save the event, can be used for creation and update.

        Returns:
         * self
        """
        if self._instance is not None:
            path = self.url.path if self.url else None
            (id, path) = self._create(self._instance.serialize(), self.id,
                     path)
            self.id = id
            self.url = urlparse.urlparse(path)
        return self

    def __str__(self):
        return "Event: %s" % self.canonical_url

    def set_data(self, data):
        self._data = vcal.fix(data)
        self._instance = vobject.readOne(StringIO.StringIO(self._data))
        return self

    def get_data(self):
        return self._data
    data = property(get_data, set_data,
                    doc="vCal representation of the event")

    def set_instance(self, inst):
        self._instance = inst
        self._data = inst.serialize()
        return self

    def get_instance(self):
        return self._instance
    instance = property(get_instance, set_instance,
                        doc="vobject instance of the event")
<commit_msg>Support passing in an already-parsed URL, to make copying objects easier<commit_after>#!/usr/bin/env python
# -*- encoding: utf-8 -*-

import copy
import urlparse
import vobject
import StringIO
import uuid
from lxml import etree

from caldav.lib import error, vcal, url
from caldav.lib.namespace import ns
from caldav.elements import dav, cdav

class DAVObject(object):
    """
    Base class for all DAV objects.
    """
    id = None
    url = None
    client = None
    parent = None
    name = None

    def __init__(self, client, url=None, parent=None, name=None, id=None):
        """
        Default constructor.

        Parameters:
         * client: A DAVClient instance
         * url: The url for this object
         * parent: The parent object
         * name: A displayname
         * id: The resource id (UID for an Event)
        """

        self.client = client
        self.parent = parent
        self.name = name
        self.id = id
        if isinstance(url, urlparse.ParseResult):
            self.url = url
        elif url is not None:
            self.url = urlparse.urlparse(url)

    @property
    def canonical_url(self):
        return url.canonicalize(self.url, self.parent)

    def children(self, type=None):
        """
        List children, using a propfind (resourcetype) on the parent object,
        at depth = 1.
        """
        c = []

        depth = 1
        properties = {}

        props = [dav.ResourceType(), ]

        prop = dav.Prop() + props
        root = dav.Propfind() + prop

        body = etree.tostring(root.xmlelement(), encoding="utf-8",
                              xml_declaration=True)

        response = self.client.propfind(self.url.path, body, depth)
        for r in response.tree.findall(dav.Response.tag):
            # We use canonicalized urls to index children
            href = urlparse.urlparse(r.find(dav.Href.tag).text)
            href = url.canonicalize(href, self)
            properties[href] = {}
            for p in props:
                t = r.find(".//" + p.tag)
                if len(list(t)) > 0:
                    if type is not None:
                        val = t.find(".//" + type)
                    else:
                        val = t.find(".//*")
                    if val is not None:
                        val = val.tag
                    else:
                        val = None
                else:
                    val = t.text
                properties[href][p.tag] = val

        for path in properties.keys():
            resource_type = properties[path][dav.ResourceType.tag]
            if resource_type == type or type is None:
                if path != self.canonical_url:
                    c.append((path, resource_type))

        return c

    def _get_properties(self, props=[], depth=0):
        properties = {}

        body = ""
        # build the propfind request
        if len(props) > 0:
            prop = dav.Prop() + props
            root = dav.Propfind() + prop

            body = etree.tostring(root.xmlelement(), encoding="utf-8",
                                  xml_declaration=True)

        response = self.client.propfind(self.url.path, body, depth)
        # All items should be in a <D:response> element
        for r in response.tree.findall(dav.Response.tag):
            href = r.find(dav.Href.tag).text
            properties[href] = {}
            for p in props:
                t = r.find(".//" + p.tag)
                if len(list(t)) > 0:
                    val = t.find(".//*")
                    if val is not None:
                        val = val.tag
                    else:
                        val = None
                else:
                    val = t.text
                properties[href][p.tag] = val

        return properties

    def get_properties(self, props=[], depth=0):
        """
        Get properties (PROPFIND) for this object. Works only for
        properties, that don't have complex types.

        Parameters:
         * props = [dav.ResourceType(), dav.DisplayName(), ...]

        Returns:
         * {proptag: value, ...}
        """
        rc = None
        properties =  self._get_properties(props, depth)
        path = self.url.path
        exchange_path = self.url.path + '/'

        if path in properties.keys():
            rc = properties[path]
        elif exchange_path in properties.keys():
            rc = properties[exchange_path]
        else:
            raise Exception("The CalDAV server you are using has a problem with path handling.")

        return rc

    def set_properties(self, props=[]):
        """
        Set properties (PROPPATCH) for this object.

        Parameters:
         * props = [dav.DisplayName('name'), ...]

        Returns:
         * self
        """
        prop = dav.Prop() + props
        set = dav.Set() + prop
        root = dav.PropertyUpdate() + set

        q = etree.tostring(root.xmlelement(), encoding="utf-8",
                           xml_declaration=True)
        r = self.client.proppatch(self.url.path, q)

        statuses = r.tree.findall(".//" + dav.Status.tag)
        for s in statuses:
            if not s.text.endswith("200 OK"):
                raise error.PropsetError(r.raw)

        return self

    def save(self):
        """
        Save the object. This is an abstract methed, that all classes
        derived .from DAVObject implement.

        Returns:
         * self
        """
        raise NotImplementedError()

    def delete(self):
        """
        Delete the object.
        """
        if self.url is not None:
            path = self.url.path
            r = self.client.delete(path)

            #TODO: find out why we get 404
            if r.status not in (200, 204, 404):
                raise error.DeleteError(r.raw)

class Principal(DAVObject):
    """
    This class represents a DAV Principal. It doesn't do much, except play
    the role of the parent to all calendar collections.
    """
    def __init__(self, client, url):
        """
        This object has a specific constructor, because its url is mandatory.
        """
        self.client = client
        self.url = urlparse.urlparse(url)

    def calendars(self):
        """
        List all calendar collections in this principal.

        Returns:
         * [Calendar(), ...]
        """
        cals = []

        data = self.children(cdav.Calendar.tag)
        for c_url, c_type in data:
            cals.append(Calendar(self.client, c_url, parent=self))

        return cals

class Calendar(DAVObject):
    """
    The `Calendar` object is used to represent a calendar collection.
    Refer to the RFC for details: http://www.ietf.org/rfc/rfc4791.txt
    """
    def _create(self, name, id=None):
        """
        Create a new calendar with display name `name` in `parent`.
        """
        path = None
        if id is None:
            id = str(uuid.uuid1())

        name = dav.DisplayName(name)
        cal = cdav.CalendarCollection()
        coll = dav.Collection() + cal
        type = dav.ResourceType() + coll

        prop = dav.Prop() + [type, name]
        set = dav.Set() + prop

        mkcol = dav.Mkcol() + set

        q = etree.tostring(mkcol.xmlelement(), encoding="utf-8",
                           xml_declaration=True)
        path = url.join(self.parent.url.path, id)

        r = self.client.mkcol(path, q)
        if r.status == 201:
            # XXX Should we use self.canonical_url ?
            path = url.make(self.parent.url, path)
        else:
            raise error.MkcolError(r.raw)

        return (id, path)

    def save(self):
        """
        The save method for a calendar is only used to create it, for now.
        We know we have to create it when we don't have a url.

        Returns:
         * self
        """
        if self.url is None:
            (id, path) = self._create(self.name, self.id)
            self.id = id
            self.url = urlparse.urlparse(path)
        return self

    def date_search(self, start, end=None):
        """
        Search events by date in the calendar. Recurring events are expanded
        if they have an occurence during the specified time frame.

        Parameters:
         * start = datetime.today().
         * end = same as above.

        Returns:
         * [Event(), ...]
        """
        matches = []

        # build the request
        expand = cdav.Expand(start, end)
        data = cdav.CalendarData() + expand
        prop = dav.Prop() + data

        range = cdav.TimeRange(start, end)
        vevent = cdav.CompFilter("VEVENT") + range
        vcal = cdav.CompFilter("VCALENDAR") + vevent
        filter = cdav.Filter() + vcal

        root = cdav.CalendarQuery() + [prop, filter]

        q = etree.tostring(root.xmlelement(), encoding="utf-8",
                           xml_declaration=True)
        response = self.client.report(self.url.path, q, 1)
        for r in response.tree.findall(".//" + dav.Response.tag):
            status = r.find(".//" + dav.Status.tag)
            if status.text.endswith("200 OK"):
                href = urlparse.urlparse(r.find(dav.Href.tag).text)
                href = url.canonicalize(href, self)
                data = r.find(".//" + cdav.CalendarData.tag).text
                e = Event(self.client, url=href, data=data, parent=self)
                matches.append(e)
            else:
                raise error.ReportError(response.raw)

        return matches

    def event(self, uid):
        """
        Get one event from the calendar.

        Parameters:
         * uid: the event uid

        Returns:
         * Event() or None
        """
        e = None

        data = cdav.CalendarData()
        prop = dav.Prop() + data

        match = cdav.TextMatch(uid)
        propf = cdav.PropFilter("UID") + match
        vevent = cdav.CompFilter("VEVENT") + propf
        vcal = cdav.CompFilter("VCALENDAR") + vevent
        filter = cdav.Filter() + vcal

        root = cdav.CalendarQuery() + [prop, filter]

        q = etree.tostring(root.xmlelement(), encoding="utf-8",
                           xml_declaration=True)
        response = self.client.report(self.url.path, q, 1)
        r = response.tree.find(".//" + dav.Response.tag)
        if r is not None:
            href = urlparse.urlparse(r.find(".//" + dav.Href.tag).text)
            href = url.canonicalize(href, self)
            data = r.find(".//" + cdav.CalendarData.tag).text
            e = Event(self.client, url=href, data=data, parent=self)
        else:
            raise error.NotFoundError(response.raw)

        return e

    def events(self):
        """
        List all events from the calendar.

        Returns:
         * [Event(), ...]
        """
        all = []

        data = self.children()
        for e_url, e_type in data:
            all.append(Event(self.client, e_url, parent=self))

        return all

    def __str__(self):
        return "Collection: %s" % self.canonical_url

class Event(DAVObject):
    """
    The `Event` object is used to represent an event.
    """
    _instance = None
    _data = None

    def __init__(self, client, url=None, data=None, parent=None, id=None):
        """
        Event has an additional parameter for its constructor:
         * data = "...", vCal data for the event
        """
        DAVObject.__init__(self, client=client, url=url, parent=parent, id=id)
        if data is not None:
            self.data = data

    def load(self):
        """
        Load the event from the caldav server.
        """
        r = self.client.request(self.url.path)
        self.data = vcal.fix(r.raw)
        return self

    def _create(self, data, id=None, path=None):
        if id is None:
            id = str(uuid.uuid1())
        if path is None:
            path = url.join(self.parent.url.path, id + ".ics")

        r = self.client.put(path, data,
                            {"Content-Type": 'text/calendar; charset="utf-8"'})
        if r.status == 204 or r.status == 201:
            # XXX Should we use self.canonical_url ?
            path = url.make(self.parent.url, path)
        else:
            raise error.PutError(r.raw)

        return (id, path)

    def save(self):
        """
        Save the event, can be used for creation and update.

        Returns:
         * self
        """
        if self._instance is not None:
            path = self.url.path if self.url else None
            (id, path) = self._create(self._instance.serialize(), self.id,
                     path)
            self.id = id
            self.url = urlparse.urlparse(path)
        return self

    def __str__(self):
        return "Event: %s" % self.canonical_url

    def set_data(self, data):
        self._data = vcal.fix(data)
        self._instance = vobject.readOne(StringIO.StringIO(self._data))
        return self

    def get_data(self):
        return self._data
    data = property(get_data, set_data,
                    doc="vCal representation of the event")

    def set_instance(self, inst):
        self._instance = inst
        self._data = inst.serialize()
        return self

    def get_instance(self):
        return self._instance
    instance = property(get_instance, set_instance,
                        doc="vobject instance of the event")
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding:utf-8 -*-

import os, json, re

def text2subpath(text):
  subpath = text

  # remove leading un-needed characters
  match = re.search(r'^[\d\s()-\.]+', subpath)
  if match:
    subpath = subpath[len(match.group()):]

  # remove trailing un-needed characters
  match = re.search(r'-\d$', subpath)
  if match:
    subpath = subpath[:-len(match.group())]

  subpath = subpath.lower()

  if subpath.endswith(u'pāḷi'):
    subpath = subpath[:-4]

  if subpath.endswith(u'nikāya'):
    subpath = subpath[:-6]

  if subpath.endswith(u'piṭaka'):
    subpath = subpath[:-6]

  return subpath

def traverseTreeAndSetSubpathProperty(node, space=0):
  if 'conflict' in node:
    subpath = text2subpath(node['text']) + node['conflict']
    del node['conflict']
  else:
    subpath = text2subpath(node['text'])
  #print(' '*space + '[text] => ' + node['text'])
  #print(' '*space + '[subpath]  => ' + subpath)
  # set 'subpath' property of node
  node['subpath'] = subpath

  if 'child' in node:
    subpaths = []
    conflictInfo = {}
    for child in node['child']:
      subpath = text2subpath(child['text'])
      if subpath in subpaths:
        # the same name of subpath already used

        if subpath in conflictInfo:
          conflictInfo[subpath] += 1
        else:
          conflictInfo[subpath] = 2

        child['conflict'] = unicode(conflictInfo[subpath])
        print('\nconflict: ' + subpath)
        #raw_input()

      subpaths.append(subpath)
      traverseTreeAndSetSubpathProperty(child, space+2)

    #raw_input()

if __name__ == '__main__':
  jsonPath = os.path.join(os.path.dirname(__file__), 'treeview.json')

  dstServerJsonPath = os.path.join(os.path.dirname(__file__), '../../tipitaka/gaelibs/json/treeviewAll.json')
  dstClientJsPath = os.path.join(os.path.dirname(__file__), '../../tipitaka/app/js/treeviewAllJson-service.js')

  with open(jsonPath, 'r') as f:
    treeviewJson = json.loads(f.read())

  # set subpath of Tipiṭaka, commentries, and sub-commentaries
  for child in treeviewJson['child']:
    traverseTreeAndSetSubpathProperty(child)

  # set subpath of Tipiṭaka (Mūla) as canon
  treeviewJson['child'][0]['subpath'] = 'canon'

  print(treeviewJson)

  if not os.path.exists(os.path.dirname(dstServerJsonPath)):
    os.makedirs(os.path.dirname(dstServerJsonPath))
  with open(dstServerJsonPath, 'w') as f:
    f.write(json.dumps(treeviewJson))

  with open(dstClientJsPath, 'w') as f:
    f.write("angular.module('pali.treeviewAllJson', []).\n")
    f.write("  factory('treeviewAllJson', [function() {\n")
    f.write("    var treeviewData = ")
    f.write(json.dumps(treeviewJson))
    f.write(";\n")
    f.write("    var serviceInstance = { all: treeviewData, tpk: treeviewData['child'][0] };\n")
    f.write("    return serviceInstance;\n")
    f.write("  }]);\n")
<commit_msg>covert text to subpath in commentries, sub-commentries, and anya<commit_after>#!/usr/bin/env python
# -*- coding:utf-8 -*-

import os, json, re

def text2subpath(text):
  subpath = text

  # remove leading un-needed characters
  match = re.search(r'^[\d\s()-\.]+', subpath)
  if match:
    subpath = subpath[len(match.group()):]

  # remove trailing un-needed characters
  match = re.search(r'-\d$', subpath)
  if match:
    subpath = subpath[:-len(match.group())]

  subpath = subpath.lower()

  if subpath.endswith(u'pāḷi'):
    subpath = subpath[:-4]

  if subpath.endswith(u'nikāya'):
    subpath = subpath[:-6]

  if subpath.endswith(u'piṭaka'):
    subpath = subpath[:-6]

  if subpath.endswith(u'piṭaka (aṭṭhakathā)'):
    subpath = subpath[:-19]

  if subpath.endswith(u'-aṭṭhakathā'):
    if subpath.endswith(u'kaṇḍa-aṭṭhakathā'):
      subpath = subpath[:-16]
    else:
      subpath = subpath[:-11]

  if subpath.endswith(u'nikāya (aṭṭhakathā)'):
    if subpath.endswith(u' nikāya (aṭṭhakathā)'):
      subpath = subpath[:-20]
    else:
      subpath = subpath[:-19]

  if subpath.endswith(u'piṭaka (ṭīkā)'):
    subpath = subpath[:-13]

  if subpath.endswith(u'nikāya (ṭīkā)'):
    subpath = subpath[:-13]

  if subpath.endswith(u'-mūlaṭīkā'):
    subpath = subpath[:-9]

  if subpath.endswith(u'-ṭīkā'):
    subpath = subpath[:-5]

  return subpath

def traverseTreeAndSetSubpathProperty(node, space=0):
  if 'conflict' in node:
    subpath = text2subpath(node['text']) + node['conflict']
    del node['conflict']
  else:
    subpath = text2subpath(node['text'])
  #print(' '*space + '[text] => ' + node['text'])
  #print(' '*space + '[subpath]  => ' + subpath)
  # set 'subpath' property of node
  node['subpath'] = subpath

  if 'child' in node:
    subpaths = []
    conflictInfo = {}
    for child in node['child']:
      subpath = text2subpath(child['text'])
      if subpath in subpaths:
        # the same name of subpath already used

        if subpath in conflictInfo:
          conflictInfo[subpath] += 1
        else:
          conflictInfo[subpath] = 2

        child['conflict'] = unicode(conflictInfo[subpath])
        print('\nconflict: ' + subpath)
        #raw_input()

      subpaths.append(subpath)
      traverseTreeAndSetSubpathProperty(child, space+2)

    #raw_input()

if __name__ == '__main__':
  jsonPath = os.path.join(os.path.dirname(__file__), 'treeview.json')

  dstServerJsonPath = os.path.join(os.path.dirname(__file__), '../../tipitaka/gaelibs/json/treeviewAll.json')
  dstClientJsPath = os.path.join(os.path.dirname(__file__), '../../tipitaka/app/js/treeviewAllJson-service.js')

  with open(jsonPath, 'r') as f:
    treeviewJson = json.loads(f.read())

  # set subpath of Tipiṭaka, commentries, and sub-commentaries
  for child in treeviewJson['child']:
    traverseTreeAndSetSubpathProperty(child)

  # set subpath of Tipiṭaka (Mūla) as canon
  treeviewJson['child'][0]['subpath'] = 'canon'

  print(treeviewJson)

  if not os.path.exists(os.path.dirname(dstServerJsonPath)):
    os.makedirs(os.path.dirname(dstServerJsonPath))
  with open(dstServerJsonPath, 'w') as f:
    f.write(json.dumps(treeviewJson))

  with open(dstClientJsPath, 'w') as f:
    f.write("angular.module('pali.treeviewAllJson', []).\n")
    f.write("  factory('treeviewAllJson', [function() {\n")
    f.write("    var treeviewData = ")
    f.write(json.dumps(treeviewJson))
    f.write(";\n")
    f.write("    var serviceInstance = { all: treeviewData, tpk: treeviewData['child'][0] };\n")
    f.write("    return serviceInstance;\n")
    f.write("  }]);\n")
<|endoftext|>