<commit_before>import numpy as np

from selection.algorithms.softmax import nonnegative_softmax
import regreg.api as rr
from selection.bayesian.selection_probability_rr import cube_barrier_scaled, cube_gradient_scaled, cube_hessian_scaled,
nonnegative_softmax_scaled

def cube_subproblem_fs(argument,
                           c,
                           randomization_CGF_conjugate,
                           lagrange= 1., nstep=100,
                           initial=None,
                           lipschitz=0,
                           tol=1.e-10):
    '''
    Solve the subproblem
    $$
    \text{minimize}_{z} \Lambda_{-E}^*(u + c.z_{-E}) + b_{-E}(z)
    $$
    where $u$ is `argument`, $\Lambda_{-E}^*$ is the
    conjvex conjugate of the $-E$ coordinates of the
    randomization (assumes that randomization has independent
    coordinates) and
    $b_{-E}$ is a barrier approximation to
    the cube $\prod_{j \in -E} [-\lambda_j,\lambda_j]$ with
    $\lambda$ being `lagrange`.
    Returns the maximizer and the value of the convex conjugate.
    '''
    k = argument.shape[0]
    if initial is None:
        current = np.zeros(k, np.float)
    else:
        current = initial  # no copy

    current_value = np.inf

    conj_value = lambda x: randomization_CGF_conjugate.smooth_objective(x, 'func')
    conj_grad = lambda x: randomization_CGF_conjugate.smooth_objective(x, 'grad')

    step = np.ones(k, np.float)
    objective = lambda u: cube_barrier_scaled(u, lagrange) + conj_value(argument + c*u)

    for itercount in range(nstep):
        newton_step = ((cube_gradient_scaled(current, lagrange) +
                        (c*conj_grad(argument + c*current))) /
                       (cube_hessian_scaled(current, lagrange) + lipschitz))

        # make sure proposal is feasible

        count = 0
        while True:
            count += 1
            proposal = current - step * newton_step
            failing = (proposal > lagrange) + (proposal < - lagrange)
            if not failing.sum():
                break
            step *= 0.5 ** failing

            if count >= 40:
                raise ValueError('not finding a feasible point')

        # make sure proposal is a descent

        count = 0
        while True:
            proposal = current - step * newton_step
            proposed_value = objective(proposal)
            if proposed_value <= current_value:
                break
            step *= 0.5

        # stop if relative decrease is small

        if np.fabs(current_value - proposed_value) < tol * np.fabs(current_value):
            current = proposal
            current_value = proposed_value
            break

        current = proposal
        current_value = proposed_value

        if itercount % 4 == 0:
            step *= 2

    value = objective(current)
    return current, value

class cube_objective_fs(rr.smooth_atom):
    def __init__(self,
                 randomization_CGF_conjugate,
                 lagrange,
                 nstep=10,
                 tol=1.e-10,
                 initial=None,
                 coef=1.,
                 offset=None,
                 quadratic=None):

        (self.randomization_CGF_conjugate,
         self.lagrange,
         self.nstep,
         self.tol) = (randomization_CGF_conjugate,
                      lagrange,
                      nstep,
                      tol)

        rr.smooth_atom.__init__(self,
                                randomization_CGF_conjugate.shape,
                                initial=initial,
                                coef=coef,
                                offset=offset,
                                quadratic=quadratic)

    def smooth_objective(self, arg, mode='both', check_feasibility=False):

        arg = self.apply_offset(arg)

        arg_shape = arg.shape[0]

        c_bool = np.zeros(arg_shape, bool)

        c_bool[(arg_shape-1):] = 1

        z = arg[~c_bool]

        c = arg[c_bool]

        optimizer, value = cube_subproblem_fs(z,
                                              c,
                                              self.randomization_CGF_conjugate,
                                              self.lagrange,
                                              nstep=self.nstep,
                                              tol=self.tol)

        gradient_z = z + (c * optimizer)

        gradient_max_c = -np.true_divide((2* c* optimizer) + z, (c**2 + cube_hessian_scaled(optimizer, lagrange = 1.)))

        gradient_c = (c* z.T + cube_gradient_scaled(optimizer, lagrange = 1.).T + ((c**2)*optimizer.T)).\
            dot(gradient_max_c) + (c*np.power(optimizer, 2.)) + optimizer.T.dot(z)

        gradient = np.vstack([gradient_z, gradient_c])

        if mode == 'func':
            return self.scale(value)
        elif mode == 'grad':
            return self.scale(gradient)
        elif mode == 'both':
            return self.scale(value), self.scale(gradient)
        else:
            raise ValueError("mode incorrectly specified")

<commit_msg>setting lagrange to 1 in cube_solver<commit_after>import numpy as np

from selection.algorithms.softmax import nonnegative_softmax
import regreg.api as rr
from selection.bayesian.selection_probability_rr import cube_barrier_scaled, cube_gradient_scaled, cube_hessian_scaled,
nonnegative_softmax_scaled

def cube_subproblem_fs(argument,
                           c,
                           randomization_CGF_conjugate,
                           lagrange= 1., nstep=100,
                           initial=None,
                           lipschitz=0,
                           tol=1.e-10):
    '''
    Solve the subproblem
    $$
    \text{minimize}_{z} \Lambda_{-E}^*(u + c.z_{-E}) + b_{-E}(z)
    $$
    where $u$ is `argument`, $\Lambda_{-E}^*$ is the
    conjvex conjugate of the $-E$ coordinates of the
    randomization (assumes that randomization has independent
    coordinates) and
    $b_{-E}$ is a barrier approximation to
    the cube $\prod_{j \in -E} [-\lambda_j,\lambda_j]$ with
    $\lambda$ being `lagrange`.
    Returns the maximizer and the value of the convex conjugate.
    '''
    k = argument.shape[0]
    if initial is None:
        current = np.zeros(k, np.float)
    else:
        current = initial  # no copy

    current_value = np.inf

    conj_value = lambda x: randomization_CGF_conjugate.smooth_objective(x, 'func')
    conj_grad = lambda x: randomization_CGF_conjugate.smooth_objective(x, 'grad')

    step = np.ones(k, np.float)
    objective = lambda u: cube_barrier_scaled(u, lagrange) + conj_value(argument + c*u)

    for itercount in range(nstep):
        newton_step = ((cube_gradient_scaled(current, lagrange) +
                        (c*conj_grad(argument + c*current))) /
                       (cube_hessian_scaled(current, lagrange) + lipschitz))

        # make sure proposal is feasible

        count = 0
        while True:
            count += 1
            proposal = current - step * newton_step
            failing = (proposal > lagrange) + (proposal < - lagrange)
            if not failing.sum():
                break
            step *= 0.5 ** failing

            if count >= 40:
                raise ValueError('not finding a feasible point')

        # make sure proposal is a descent

        count = 0
        while True:
            proposal = current - step * newton_step
            proposed_value = objective(proposal)
            if proposed_value <= current_value:
                break
            step *= 0.5

        # stop if relative decrease is small

        if np.fabs(current_value - proposed_value) < tol * np.fabs(current_value):
            current = proposal
            current_value = proposed_value
            break

        current = proposal
        current_value = proposed_value

        if itercount % 4 == 0:
            step *= 2

    value = objective(current)
    return current, value

class cube_objective_fs(rr.smooth_atom):
    def __init__(self,
                 randomization_CGF_conjugate,
                 lagrange=1.,
                 nstep=10,
                 tol=1.e-10,
                 initial=None,
                 coef=1.,
                 offset=None,
                 quadratic=None):

        (self.randomization_CGF_conjugate,
         self.lagrange,
         self.nstep,
         self.tol) = (randomization_CGF_conjugate,
                      lagrange,
                      nstep,
                      tol)

        rr.smooth_atom.__init__(self,
                                randomization_CGF_conjugate.shape,
                                initial=initial,
                                coef=coef,
                                offset=offset,
                                quadratic=quadratic)

    def smooth_objective(self, arg, mode='both', check_feasibility=False):

        arg = self.apply_offset(arg)

        arg_shape = arg.shape[0]

        c_bool = np.zeros(arg_shape, bool)

        c_bool[(arg_shape-1):] = 1

        z = arg[~c_bool]

        c = arg[c_bool]

        optimizer, value = cube_subproblem_fs(z,
                                              c,
                                              self.randomization_CGF_conjugate,
                                              self.lagrange,
                                              nstep=self.nstep,
                                              tol=self.tol)

        gradient_z = z + (c * optimizer)

        gradient_max_c = -np.true_divide((2* c* optimizer) + z, (c**2 + cube_hessian_scaled(optimizer, lagrange = 1.)))

        gradient_c = (c* z.T + cube_gradient_scaled(optimizer, lagrange = 1.).T + ((c**2)*optimizer.T)).\
            dot(gradient_max_c) + (c*np.power(optimizer, 2.)) + optimizer.T.dot(z)

        gradient = np.vstack([gradient_z, gradient_c])

        if mode == 'func':
            return self.scale(value)
        elif mode == 'grad':
            return self.scale(gradient)
        elif mode == 'both':
            return self.scale(value), self.scale(gradient)
        else:
            raise ValueError("mode incorrectly specified")

class selection_probability_objective_fs(rr.smooth_atom):
    def __init__(self,
                 X,
                 feasible_point,
                 active,
                 active_signs,
                 lagrange,
                 mean_parameter,  # in R^n
                 noise_variance,
                 randomizer,
                 epsilon,
                 coef=1.,
                 offset=None,
                 quadratic=None,
                 nstep=10):

        """
        Objective function for $\beta_E$ (i.e. active) with $E$ the `active_set` optimization
        variables, and data $z \in \mathbb{R}^n$ (i.e. response).
        NEEDS UPDATING
        Above, $\beta_E^*$ is the `parameter`, $b_{\geq}$ is the softmax of the non-negative constraint,
        $$
        B_E = X^TX_E
        $$
        and
        $$
        \gamma_E = \begin{pmatrix} \lambda s_E\\ 0\end{pmatrix}
        $$
        with $\lambda$ being `lagrange`.
        Parameters
        ----------
        X : np.float
             Design matrix of shape (n,p)
        active : np.bool
             Boolean indicator of active set of shape (p,).
        active_signs : np.float
             Signs of active coefficients, of shape (active.sum(),).
        lagrange : np.float
             Array of lagrange penalties for LASSO of shape (p,)
        parameter : np.float
             Parameter $\beta_E^*$ for which we want to
             approximate the selection probability.
             Has shape (active_set.sum(),)
        randomization : np.float
             Variance of IID Gaussian noise
             that was added before selection.
        """

        n, p = X.shape
        E = active.sum()
        self._X = X
        self.active = active
        self.noise_variance = noise_variance
        self.randomization = randomizer

        self.inactive_conjugate = self.active_conjugate = randomizer.CGF_conjugate
        if self.active_conjugate is None:
            raise ValueError(
                'randomization must know its CGF_conjugate -- currently only isotropic_gaussian and laplace are implemented and are assumed to be randomization with IID coordinates')

        self.inactive_lagrange = lagrange[~active]

        initial = np.zeros(n + E, )
        initial[n:] = feasible_point

        rr.smooth_atom.__init__(self,
                                (n + E,),
                                offset=offset,
                                quadratic=quadratic,
                                initial=initial,
                                coef=coef)

        self.coefs[:] = initial

        self.active = active
        nonnegative = nonnegative_softmax_scaled(E)  # should there be a
        # scale to our softmax?
        opt_vars = np.zeros(n + E, bool)
        opt_vars[n:] = 1

        self._opt_selector = rr.selector(opt_vars, (n + E,))
        self.nonnegative_barrier = nonnegative.linear(self._opt_selector)
        self._response_selector = rr.selector(~opt_vars, (n + E,))

        X_E = self.X_E = X[:, active]
        B = X.T.dot(X_E)

        B_E = B[active]
        B_mE = B[~active]

        self.A_active = np.hstack([-X[:, active].T, (B_E + epsilon * np.identity(E)) * active_signs[None, :]])
        self.A_inactive = np.hstack([-X[:, ~active].T, (B_mE * active_signs[None, :])])

        self.offset_active = active_signs * lagrange[active]

        # defines \gamma and likelihood loss
        self.set_parameter(mean_parameter, noise_variance)

        self.inactive_subgrad = np.zeros(p - E)

        self.active_conj_loss = rr.affine_smooth(self.active_conjugate,
                                                 rr.affine_transform(self.A_active, self.offset_active))

        cube_obj = cube_objective(self.inactive_conjugate,
                                  lagrange[~active],
                                  nstep=nstep)

        self.cube_loss = rr.affine_smooth(cube_obj, self.A_inactive)

        self.total_loss = rr.smooth_sum([self.active_conj_loss,
                                         self.cube_loss,
                                         self.likelihood_loss,
                                         self.nonnegative_barrier])

    def set_parameter(self, mean_parameter, noise_variance):
        """
        Set $\beta_E^*$.
        """
        mean_parameter = np.squeeze(mean_parameter)
        likelihood_loss = rr.signal_approximator(mean_parameter, coef=1. / noise_variance)
        self.likelihood_loss = rr.affine_smooth(likelihood_loss, self._response_selector)

    def smooth_objective(self, param, mode='both', check_feasibility=False):
        """
        Evaluate the smooth objective, computing its value, gradient or both.
        Parameters
        ----------
        mean_param : ndarray
            The current parameter values.
        mode : str
            One of ['func', 'grad', 'both'].
        check_feasibility : bool
            If True, return `np.inf` when
            point is not feasible, i.e. when `mean_param` is not
            in the domain.
        Returns
        -------
        If `mode` is 'func' returns just the objective value
        at `mean_param`, else if `mode` is 'grad' returns the gradient
        else returns both.
        """

        param = self.apply_offset(param)

        if mode == 'func':
            f = self.total_loss.smooth_objective(param, 'func')
            return self.scale(f)
        elif mode == 'grad':
            g = self.total_loss.smooth_objective(param, 'grad')
            return self.scale(g)
        elif mode == 'both':
            f, g = self.total_loss.smooth_objective(param, 'both')
            return self.scale(f), self.scale(g)
        else:
            raise ValueError("mode incorrectly specified")

    def minimize(self, initial=None, min_its=10, max_its=50, tol=1.e-10):

        nonneg_con = self._opt_selector.output_shape[0]
        constraint = rr.separable(self.shape,
                                  [rr.nonnegative((nonneg_con,), offset=1.e-12 * np.ones(nonneg_con))],
                                  [self._opt_selector.index_obj])

        problem = rr.separable_problem.fromatom(constraint, self)
        problem.coefs[:] = 0.5
        soln = problem.solve(max_its=max_its, min_its=min_its, tol=tol)
        value = problem.objective(soln)
        return soln, value

    def minimize2(self, step=1, nstep=30, tol=1.e-8):

        n, p = self._X.shape

        current = self.coefs
        current_value = np.inf

        objective = lambda u: self.smooth_objective(u, 'func')
        grad = lambda u: self.smooth_objective(u, 'grad')

        for itercount in range(nstep):
            newton_step = grad(current) * self.noise_variance

            # make sure proposal is feasible

            count = 0
            while True:
                count += 1
                proposal = current - step * newton_step
                if np.all(proposal[n:] > 0):
                    break
                step *= 0.5
                if count >= 40:
                    raise ValueError('not finding a feasible point')

            # make sure proposal is a descent

            count = 0
            while True:
                proposal = current - step * newton_step
                proposed_value = objective(proposal)
                # print(current_value, proposed_value, 'minimize')
                if proposed_value <= current_value:
                    break
                step *= 0.5

            # stop if relative decrease is small

            if np.fabs(current_value - proposed_value) < tol * np.fabs(current_value):
                current = proposal
                current_value = proposed_value
                break

            current = proposal
            current_value = proposed_value

            if itercount % 4 == 0:
                step *= 2

        #print('iter', itercount)
        value = objective(current)
        return current, value

<|endoftext|>


<commit_before>import os
import sys
import re
from .standard_run_utilities import *
from .run_pipeline import CommandTemplate

from threading import Thread

#For matching the file names
import fnmatch

from optparse import OptionParser

## Subprocess import clause required for running commands on the shell##
import subprocess
import logging
logger = logging.getLogger(__name__)

class StandardRun(CommandTemplate):

    def is_script(self,fpath):
        return os.path.isfile(fpath)
    def is_exe(self,fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)
    #Stolen code from
    #http://stackoverflow.com/questions/377017/test-if-executable-exists-in-python
    def which(self,program,program_name):
        fpath, fname = os.path.split(program)
        if fpath:
            if self.is_exe(program):
                return program
            elif (self.is_script(program)):
                return program
        else:
            for path in os.environ["PATH"].split(os.pathsep):
                path = path.strip('"')
                exe_file = os.path.join(path, program)
                if self.is_exe(exe_file):
                    return exe_file
        logger.error(program_name +" path = " + fpath+" not locatable path or in the directory specified in your self.config file ")
        return None

    def check_executables_and_scripts_exist(self):
        executables=['plink','shapeit','impute','Rscript','python','ancestral_allele','indel_filter','multicore_ihh','qctool']
        if(self.which(self.config['plink']['plink_executable'],'plink')is None):
            return False
        if(self.which(self.config['shapeit']['shapeit_executable'],'shapeit') is None ):
            return False
        if(self.which(self.config['ancestral_allele']['ancestral_allele_script'],'ancestral_allele') is None):
            return False
        if(self.which(self.config['impute2']['impute_executable'],'impute2') is None):
            return False
        if(self.which(self.config['Rscript']['indel_filter'],'indel_filter') is None):
            return False
        if(self.which(self.config['Rscript']['rscript_executable'],'Rscript') is None):
            return False
        if(self.which(self.config['multicore_ihh']['multicore_ihh'],'multicore_ihh') is None):
            return False
        if(self.which(self.config['qctool']['qctool_executable'],'qctool') is None):
            return False
        return True

    def __init__(self,options,config,full_run=True):
        #Perform local executable check.#
        self.options = options
        self.config = config
        if(full_run):
            if( not self.check_executables_and_scripts_exist()):
                sys.exit(MISSING_EXECUTABLE_ERROR)
            self.threads=self.config['system']['cores_avaliable']

    def run_pipeline(self):
        if(self.options.phased_vcf):
            (haps,sample) = self.run_aa_annotate_haps(self.options.vcf_input,vcf=True)
            haps = self.indel_filter(haps)
            new_sample_file = self.fix_sample_file(sample)
            ihh = self.run_multi_coreihh(haps)
            tajimaSD = self.vcf_to_tajimas_d(vcf_input)
            haps2_haps = self.prepare_haps_for_variscan(new_sample_file)
            fayandwus = self.variscan_fayandwus(haps2_haps)
        else:
            (ped,map) = self.run_vcf_to_plink()
            (ped,map) = self.run_plink_filter(ped,map)
            (haps,sample) = self.run_shape_it(ped,map)
            if(self.options.imputation):
                (haps)= self.run_impute2(haps)
            haps = self.indel_filter(haps)
            new_sample_file = self.fix_sample_file(sample)
            vcf = self.haps_to_vcf(new_sample_file)
            vcf = self.fix_vcf_qctool(vcf)
            haps = self.run_aa_annotate_haps(haps)
            haps2_haps = self.prepare_haps_for_variscan(new_sample_file)
            fayandwus = self.variscan_fayandwus(self.config,haps2_haps)
            tajimaSD = self.vcf_to_tajimas_d(vcf)
        ihh = self.run_multi_coreihh(haps)

        ihs_file = ihh.split('.ihh')[0]+'.ihs'
        if not os.path.exists('results'):
            os.mkdir('results')
        os.rename(tajimaSD,'results/' + tajimaSD)
        os.rename(vcf,'results/' + vcf)
        os.rename(ihh,'results/' + ihh)
        os.rename(ihs_file,'results/'+ihs_file)
        os.rename(haps,'results/' + haps)
        os.rename(fayandwus,'results/' + fayandwus)
        if not os.path.exists('log'):
            os.mkdir('log')
        logger.info(self.options.log_file)
        os.rename(self.options.log_file,'log/' + self.options.log_file)
        if not self.options.no_clean_up:
            clean_folder('.')
        logger.info(tajimaSD)
        logger.info(vcf)
        logger.info(haps)
        logger.info(ihh)
        logger.info(fayandwus)
        logger.info("Pipeline completed successfully")
        logger.info("Goodbye :)")
    def run_vcf_to_plink(self):
        (cmd,prefix) = CommandTemplate.run_vcf_to_plink(self,self.options,self.config)
        run_subprocess(cmd,'vcftools')
        return(prefix + '.ped', prefix + '.map')

    # Calls a subprocess to run plink

    def run_plink_filter(self,ped,map):
        (cmd,prefix) = CommandTemplate.run_plink_filter(self,self.options,self.config,ped,map)
        run_subprocess(cmd,'plink')
        return(prefix+'.ped',prefix+'.map')

    # Calls a subprocess to run shape it

    def run_shape_it(self,ped,map):
        (cmd,prefix) = CommandTemplate.run_shape_it(self,self.options,self.config,ped,map)
        cmd.extend(['--thread',self.threads])
        run_subprocess(cmd,'shapeit')
        return(prefix + '.haps',prefix + '.sample')

    def haps_to_vcf(self,haps,new_sample_file):
        (cmd,output_name) = CommandTemplate.haps_to_vcf(self,self.options,self.config,haps,new_sample_file)
        run_subprocess(cmd,'hapstovcf')
        return(output_name)

    def join_impute2_files(self,output_prefix,no_commands):
        output_haps=open(output_prefix+'.haps','w')
        output_warnings=open(output_prefix+'.warnings','w')
        output_info=open(output_prefix+'.info','w')
        for i in range(no_commands):
            with open(output_prefix+'_'+str(i)+'.haps_haps','r') as h:
                with open(output_prefix + '_'+str(i)+'.warnings','r') as w:
                    with open(output_prefix + '_'+str(i) + '.info','r')as f:
                        output_haps.write(h.read())
                        output_warnings.write(w.read())
                        output_info.write(f.read())
        output_haps.close()
        output_warnings.close()
        output_info.close()
    def run_impute2(self,haps):
        (cmd_template,output_prefix) = CommandTemplate.run_impute2(self,self.options,self.config,haps)

        # change from megabases to bp which is what is
        # expected by the impute2 command line self.options
        distance=int(self.options.impute_split_size) * 1000000
        # Break files into 5 megabase regions.
        try:
            proc = subprocess.Popen("""tail -1 {0}| awk '{{print $3}}'""".format(haps),stdout=subprocess.PIPE,shell=True)
        except:
            logger.error("Tail command failed on haps file")
            sys.exit(SUBPROCESS_FAILED_EXIT)
        # get the start of the haps file
        try:
            head = subprocess.Popen("""head -1 {0}| awk '{{print $3}}'""".format(haps),stdout=subprocess.PIPE,shell=True)
        except:
            logger.error("Head command failed on haps file")
            sys.exit(SUBPROCESS_FAILED_EXIT)
        #Get the max position from your haps file#
        start_position = int(head.stdout.read())
        no_of_impute_jobs = (int(proc.stdout.read())-int(start_position))//distance + 1
        #create the command template
        #Get the max position from your haps file#

        # get the start of the first window
        #
        first_window = start_position // distance
        cmds = []
        for i in range(0,no_of_impute_jobs):
            individual_command=list(cmd_template)
            individual_command.extend(['-int',str((i+first_window)*distance),str((i+first_window+1)*distance)])
            individual_prefix=output_prefix + '_'+ str(i)
            individual_command.extend(['-o',individual_prefix+'.haps','-w',individual_prefix + '.warnings','-i',individual_prefix +'.info'])
            cmds.append(list(individual_command))
        queue_jobs(cmds,'impute2',self.config['system']['cores_avaliable'])
        join_impute2_files(self.options,self.config,output_prefix,no_of_impute_jobs)
        return(output_prefix+'.haps')

    def indel_filter(self,haps):
        (cmd,output_name) = CommandTemplate.indel_filter(self,self.options,self.config,haps)
        run_subprocess(cmd,'indel_filter')
        return(output_name)

    def run_aa_annotate_haps(self,haps,vcf=False):
        if(vcf):
            (cmd,output_name,sample_name) = CommandTemplate.run_aa_annotate_haps(self,self.options,self.config,haps,vcf)
            run_subprocess(cmd,'ancestral_annotation')
            return(output_name,sample_name)
        else:
            (cmd,output_name) = CommandTemplate.run_aa_annotate_haps(self,self.options,self.config,haps)
            run_subprocess(cmd,'ancestral_annotation')
            return(output_name)

    def run_multi_coreihh(self,haps):
        (cmd,output_name) = CommandTemplate.run_multi_coreihh(self,self.options,self.config,haps)
        cores=self.threads
        ihs_output=output_name.split('.ihh')[0]+'.ihs'
        cmd.extend(['--cores',cores])
        cmd.extend(['--working_dir','.'])
        cmd.extend(['--offset','1'])
        cmd.extend(['--ihs'])
        run_subprocess(cmd,'multcore_ihh')
        os.rename(self.options.population+'_chr_'+self.options.chromosome+"_wd_"+'.'+"_.ihh",output_name)
        os.rename(self.options.population+'_chr_'+self.options.chromosome+'_wd_'+'.'+"_.ihs",ihs_output)
        return output_name
    def fix_sample_file(self,sample_file):
        (cmd,output_name) = CommandTemplate.fix_sample_file(self,self.options,self.config,sample_file)
        new_sample_file=open(output_name,'w')
        run_subprocess(cmd,'fix sample file',stdout=new_sample_file)
        new_sample_file.close()
        return(output_name)
    def vcf_to_tajimas_d(self,vcf):
        (cmd,output_name) = CommandTemplate.vcf_to_tajimas_d(self,self.options,self.config,vcf)
        run_subprocess(cmd,'tajimas_d')
        taj_file =self.options.population + self.options.chromosome + '.taj_d'
        os.rename(output_name,taj_file)
        return(taj_file)
    def fix_vcf_qctool(self,vcf):
        (cmd,output_name) = CommandTemplate.fix_vcf_qctool(self,self.options,self.config,vcf)
        fixed_vcf = open(output_name,'w')
        run_subprocess(cmd,'fix vcf qctool',stdout=fixed_vcf)
        fixed_vcf.close()
        return(output_name)
    def prepare_haps_for_variscan(self,haps,sample):
        (cmd,output_name) = CommandTemplate.prepare_haps_for_variscan(self,self.options,self.config,haps,sample)
        run_subprocess(cmd,'haps to hapmap')
        return(output_name)
    def variscan_fayandwus(self,hap2):
        (cmd,output_name,varscan_conf) = CommandTemplate.variscan_fayandwus(self,self.options,self.config,hap2)
        output_variscan = open(output_name,'w')

        # get the start and the end of the hapmap2 file
        #
        varscan_config = open(varscan_conf ,'a')
        try:
            proc = subprocess.Popen("""tail -1 {0}| awk '{{print $4}}'""".format(hap2),stdout=subprocess.PIPE,shell=True)
        except:
            logger.error("Tail command failed on haps file")
            sys.exit(SUBPROCESS_FAILED_EXIT)
        # get the start of the haps file
        try:
            head = subprocess.Popen("""head -2 {0} | tail -1 | awk '{{print $4}}'""".format(hap2),stdout=subprocess.PIPE,shell=True)
        except:
            logger.error("Head command failed on haps file")
            sys.exit(SUBPROCESS_FAILED_EXIT)
        start_pos = head.stdout.read()
        start_position = int(start_pos)
        end_position = int(proc.stdout.read())

        # Write the start and end of the Fey and Wu's into the self.config file #
        varscan_config.write('StartPos = ' + str(start_position) + "\n")
        varscan_config.write('EndPos = ' + str(end_position)+ '\n')
        varscan_config.close()
        run_subprocess(cmd,'variscan',stdout=output_variscan)
        output_variscan.close()
        return(output_name)
<commit_msg>Started PEP-8 the code<commit_after>import os
import sys
from .standard_run_utilities import *
from .run_pipeline import CommandTemplate

## Subprocess import clause required for running commands on the shell##
import subprocess
import logging
logger = logging.getLogger(__name__)

class StandardRun(CommandTemplate):
    def is_script(self, fpath):
        return os.path.isfile(fpath)

    def is_exe(self, fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)
    #http://stackoverflow.com/questions/377017/test-if-executable-exists-in-python

    def which(self, program, program_name):
        fpath, fname = os.path.split(program)
        if fpath:
            if self.is_exe(program):
                return program
            elif (self.is_script(program)):
                return program
        else:
            for path in os.environ["PATH"].split(os.pathsep):
                path = path.strip('"')
                exe_file = os.path.join(path, program)
                if self.is_exe(exe_file):
                    return exe_file
        logger.error(program_name + " path = " + fpath +
                     " not locatable path or in the directory"
                     " specified in your self.config file")
        return None

    def check_executables_and_scripts_exist(self):
        if(self.which(
                self.config['plink']['plink_executable'],
                'plink')is None):
            return False
        if(self.which(
                self.config['shapeit']['shapeit_executable'],
                'shapeit') is None):
            return False
        if(self.which(
                self.config['ancestral_allele']['ancestral_allele_script'],
                'ancestral_allele'
                ) is None):
            return False
        if(self.which(
                self.config['impute2']['impute_executable'],
                'impute2'
                ) is None):
            return False
        if(self.which(
                self.config['Rscript']['indel_filter'],
                'indel_filter'
                ) is None):
            return False
        if(self.which(
                self.config['Rscript']['rscript_executable'],
                'Rscript'
                ) is None):
            return False
        if(self.which(
                self.config['multicore_ihh']['multicore_ihh'],
                'multicore_ihh'
                ) is None):
            return False
        if(self.which(
                self.config['qctool']['qctool_executable'],
                'qctool'
                ) is None):
            return False
        return True

    def __init__(self, options, config, full_run=True):
        #Perform local executable check.#
        self.options = options
        self.config = config
        if(full_run):
            if(not self.check_executables_and_scripts_exist()):
                sys.exit(MISSING_EXECUTABLE_ERROR)
            self.threads = self.config['system']['cores_avaliable']

    def run_pipeline(self):
        if(self.options.phased_vcf):
            (haps, sample) = self.run_aa_annotate_haps(
                self.options.vcf_input, vcf=True)
            haps = self.indel_filter(haps)
            new_sample_file = self.fix_sample_file(sample)
            ihh = self.run_multi_coreihh(haps)
            tajimaSD = self.vcf_to_tajimas_d(vcf_input)
            haps2_haps = self.prepare_haps_for_variscan(new_sample_file)
            fayandwus = self.variscan_fayandwus(haps2_haps)
        else:
            (ped, map) = self.run_vcf_to_plink()
            (ped, map) = self.run_plink_filter(ped, map)
            (haps, sample) = self.run_shape_it(ped, map)
            if(self.options.imputation):
                (haps) = self.run_impute2(haps)
            haps = self.indel_filter(haps)
            new_sample_file = self.fix_sample_file(sample)
            vcf = self.haps_to_vcf(new_sample_file)
            vcf = self.fix_vcf_qctool(vcf)
            haps = self.run_aa_annotate_haps(haps)
            haps2_haps = self.prepare_haps_for_variscan(new_sample_file)
            fayandwus = self.variscan_fayandwus(self.config, haps2_haps)
            tajimaSD = self.vcf_to_tajimas_d(vcf)
        ihh = self.run_multi_coreihh(haps)
        ihs_file = ihh.split('.ihh')[0]+'.ihs'
        if not os.path.exists('results'):
            os.mkdir('results')
        os.rename(tajimaSD, 'results/' + tajimaSD)
        os.rename(vcf, 'results/' + vcf)
        os.rename(ihh, 'results/' + ihh)
        os.rename(ihs_file, 'results/'+ihs_file)
        os.rename(haps, 'results/' + haps)
        os.rename(fayandwus, 'results/' + fayandwus)
        if not os.path.exists('log'):
            os.mkdir('log')
        logger.info(self.options.log_file)
        os.rename(self.options.log_file, 'log/' + self.options.log_file)
        if not self.options.no_clean_up:
            clean_folder('.')
        logger.info(tajimaSD)
        logger.info(vcf)
        logger.info(haps)
        logger.info(ihh)
        logger.info(fayandwus)
        logger.info("Pipeline completed successfully")
        logger.info("Goodbye :)")

    def run_vcf_to_plink(self):
        (cmd, prefix) = CommandTemplate.run_vcf_to_plink(self, self.options, self.config)
        run_subprocess(cmd, 'vcftools')
        return(prefix + '.ped', prefix + '.map')

    def run_plink_filter(self, ped, map):
        (cmd,prefix) = CommandTemplate.run_plink_filter(self, self.options, self.config, ped, map)
        run_subprocess(cmd,'plink')
        return(prefix+'.ped',prefix+'.map')

    # Calls a subprocess to run shape it

    def run_shape_it(self,ped,map):
        (cmd,prefix) = CommandTemplate.run_shape_it(self,self.options,self.config,ped,map)
        cmd.extend(['--thread',self.threads])
        run_subprocess(cmd,'shapeit')
        return(prefix + '.haps',prefix + '.sample')

    def haps_to_vcf(self,haps,new_sample_file):
        (cmd,output_name) = CommandTemplate.haps_to_vcf(self,self.options,self.config,haps,new_sample_file)
        run_subprocess(cmd,'hapstovcf')
        return(output_name)

    def join_impute2_files(self,output_prefix,no_commands):
        output_haps=open(output_prefix+'.haps','w')
        output_warnings=open(output_prefix+'.warnings','w')
        output_info=open(output_prefix+'.info','w')
        for i in range(no_commands):
            with open(output_prefix+'_'+str(i)+'.haps_haps','r') as h:
                with open(output_prefix + '_'+str(i)+'.warnings','r') as w:
                    with open(output_prefix + '_'+str(i) + '.info','r')as f:
                        output_haps.write(h.read())
                        output_warnings.write(w.read())
                        output_info.write(f.read())
        output_haps.close()
        output_warnings.close()
        output_info.close()
    def run_impute2(self,haps):
        (cmd_template,output_prefix) = CommandTemplate.run_impute2(self,self.options,self.config,haps)

        # change from megabases to bp which is what is
        # expected by the impute2 command line self.options
        distance=int(self.options.impute_split_size) * 1000000
        # Break files into 5 megabase regions.
        try:
            proc = subprocess.Popen("""tail -1 {0}| awk '{{print $3}}'""".format(haps),stdout=subprocess.PIPE,shell=True)
        except:
            logger.error("Tail command failed on haps file")
            sys.exit(SUBPROCESS_FAILED_EXIT)
        # get the start of the haps file
        try:
            head = subprocess.Popen("""head -1 {0}| awk '{{print $3}}'""".format(haps),stdout=subprocess.PIPE,shell=True)
        except:
            logger.error("Head command failed on haps file")
            sys.exit(SUBPROCESS_FAILED_EXIT)
        #Get the max position from your haps file#
        start_position = int(head.stdout.read())
        no_of_impute_jobs = (int(proc.stdout.read())-int(start_position))//distance + 1
        #create the command template
        #Get the max position from your haps file#

        # get the start of the first window
        #
        first_window = start_position // distance
        cmds = []
        for i in range(0,no_of_impute_jobs):
            individual_command=list(cmd_template)
            individual_command.extend(['-int',str((i+first_window)*distance),str((i+first_window+1)*distance)])
            individual_prefix=output_prefix + '_'+ str(i)
            individual_command.extend(['-o',individual_prefix+'.haps','-w',individual_prefix + '.warnings','-i',individual_prefix +'.info'])
            cmds.append(list(individual_command))
        queue_jobs(cmds,'impute2',self.config['system']['cores_avaliable'])
        join_impute2_files(self.options,self.config,output_prefix,no_of_impute_jobs)
        return(output_prefix+'.haps')

    def indel_filter(self,haps):
        (cmd,output_name) = CommandTemplate.indel_filter(self,self.options,self.config,haps)
        run_subprocess(cmd,'indel_filter')
        return(output_name)

    def run_aa_annotate_haps(self,haps,vcf=False):
        if(vcf):
            (cmd,output_name,sample_name) = CommandTemplate.run_aa_annotate_haps(self,self.options,self.config,haps,vcf)
            run_subprocess(cmd,'ancestral_annotation')
            return(output_name,sample_name)
        else:
            (cmd,output_name) = CommandTemplate.run_aa_annotate_haps(self,self.options,self.config,haps)
            run_subprocess(cmd,'ancestral_annotation')
            return(output_name)

    def run_multi_coreihh(self,haps):
        (cmd,output_name) = CommandTemplate.run_multi_coreihh(self,self.options,self.config,haps)
        cores=self.threads
        ihs_output=output_name.split('.ihh')[0]+'.ihs'
        cmd.extend(['--cores',cores])
        cmd.extend(['--working_dir','.'])
        cmd.extend(['--offset','1'])
        cmd.extend(['--ihs'])
        run_subprocess(cmd,'multcore_ihh')
        os.rename(self.options.population+'_chr_'+self.options.chromosome+"_wd_"+'.'+"_.ihh",output_name)
        os.rename(self.options.population+'_chr_'+self.options.chromosome+'_wd_'+'.'+"_.ihs",ihs_output)
        return output_name
    def fix_sample_file(self,sample_file):
        (cmd,output_name) = CommandTemplate.fix_sample_file(self,self.options,self.config,sample_file)
        new_sample_file=open(output_name,'w')
        run_subprocess(cmd,'fix sample file',stdout=new_sample_file)
        new_sample_file.close()
        return(output_name)
    def vcf_to_tajimas_d(self,vcf):
        (cmd,output_name) = CommandTemplate.vcf_to_tajimas_d(self,self.options,self.config,vcf)
        run_subprocess(cmd,'tajimas_d')
        taj_file =self.options.population + self.options.chromosome + '.taj_d'
        os.rename(output_name,taj_file)
        return(taj_file)
    def fix_vcf_qctool(self,vcf):
        (cmd,output_name) = CommandTemplate.fix_vcf_qctool(self,self.options,self.config,vcf)
        fixed_vcf = open(output_name,'w')
        run_subprocess(cmd,'fix vcf qctool',stdout=fixed_vcf)
        fixed_vcf.close()
        return(output_name)
    def prepare_haps_for_variscan(self,haps,sample):
        (cmd,output_name) = CommandTemplate.prepare_haps_for_variscan(self,self.options,self.config,haps,sample)
        run_subprocess(cmd,'haps to hapmap')
        return(output_name)
    def variscan_fayandwus(self,hap2):
        (cmd,output_name,varscan_conf) = CommandTemplate.variscan_fayandwus(self,self.options,self.config,hap2)
        output_variscan = open(output_name,'w')

        # get the start and the end of the hapmap2 file
        #
        varscan_config = open(varscan_conf ,'a')
        try:
            proc = subprocess.Popen("""tail -1 {0}| awk '{{print $4}}'""".format(hap2),stdout=subprocess.PIPE,shell=True)
        except:
            logger.error("Tail command failed on haps file")
            sys.exit(SUBPROCESS_FAILED_EXIT)
        # get the start of the haps file
        try:
            head = subprocess.Popen("""head -2 {0} | tail -1 | awk '{{print $4}}'""".format(hap2),stdout=subprocess.PIPE,shell=True)
        except:
            logger.error("Head command failed on haps file")
            sys.exit(SUBPROCESS_FAILED_EXIT)
        start_pos = head.stdout.read()
        start_position = int(start_pos)
        end_position = int(proc.stdout.read())

        # Write the start and end of the Fey and Wu's into the self.config file #
        varscan_config.write('StartPos = ' + str(start_position) + "\n")
        varscan_config.write('EndPos = ' + str(end_position)+ '\n')
        varscan_config.close()
        run_subprocess(cmd,'variscan',stdout=output_variscan)
        output_variscan.close()
        return(output_name)
<|endoftext|>


<commit_before>import os
import shutil
import sys
import time
from selenium import webdriver
from seleniumbase.config import settings
from seleniumbase.core.style_sheet import style
from seleniumbase.fixtures import page_actions

LATEST_REPORT_DIR = settings.LATEST_REPORT_DIR
ARCHIVE_DIR = settings.REPORT_ARCHIVE_DIR
HTML_REPORT = settings.HTML_REPORT
RESULTS_TABLE = settings.RESULTS_TABLE

def get_timestamp():
    return str(int(time.time() * 1000))

def process_successes(test, test_count):
    return(
        '"%s","%s","%s","%s","%s","%s","%s","%s","%s"' % (
            test_count,
            "Passed!",
            "*",
            "*",
            "*",
            test.browser,
            get_timestamp()[:-3],
            test.id(),
            "*"))

def process_failures(test, err, test_count, browser_type):
    bad_page_image = "failure_test_%s.jpg" % test_count
    bad_page_data = "failure_test_%s.txt" % test_count
    page_actions.save_screenshot(
        test.driver, bad_page_image, folder=LATEST_REPORT_DIR)
    page_actions.save_test_failure_data(
        test.driver, bad_page_data, browser_type, folder=LATEST_REPORT_DIR)
    exc_info = '(Unknown Failure)'
    exception = sys.exc_info()[1]
    if exception:
        if hasattr(exception, 'msg'):
            exc_info = exception.msg
        elif hasattr(exception, 'message'):
            exc_info = exception.message
        else:
            pass
    return(
        '"%s","%s","%s","%s","%s","%s","%s","%s","%s"' % (
            test_count,
            "FAILED!",
            bad_page_data,
            bad_page_image,
            test.driver.current_url,
            test.browser,
            get_timestamp()[:-3],
            test.id(),
            exc_info))

def clear_out_old_report_logs(archive_past_runs=True, get_log_folder=False):
        abs_path = os.path.abspath('.')
        file_path = abs_path + "/%s" % LATEST_REPORT_DIR
        if not os.path.exists(file_path):
            os.makedirs(file_path)

        if archive_past_runs:
            archive_timestamp = int(time.time())
            if not os.path.exists("%s/../%s/" % (file_path, ARCHIVE_DIR)):
                os.makedirs("%s/../%s/" % (file_path, ARCHIVE_DIR))
            archive_dir = "%s/../%s/report_%s" % (
                file_path, ARCHIVE_DIR, archive_timestamp)
            shutil.move(file_path, archive_dir)
            os.makedirs(file_path)
            if get_log_folder:
                return archive_dir
        else:
            # Just delete bad pages to make room for the latest run.
            filelist = [f for f in os.listdir(
                "./%s" % LATEST_REPORT_DIR) if f.startswith("failure_") or (
                f == HTML_REPORT) or (f.startswith("automation_failure")) or (
                f == RESULTS_TABLE)]
            for f in filelist:
                os.remove("%s/%s" % (file_path, f))

def add_bad_page_log_file(page_results_list):
    abs_path = os.path.abspath('.')
    file_path = abs_path + "/%s" % LATEST_REPORT_DIR
    log_file = "%s/%s" % (file_path, RESULTS_TABLE)
    f = open(log_file, 'w')
    h_p1 = '''"Num","Result","Failure Info","Screenshot",'''
    h_p2 = '''"URL","Browser","Epoch Time",'''
    h_p3 = '''"Test Case Address","Additional Info"\n'''
    page_header = h_p1 + h_p2 + h_p3
    f.write(page_header)
    for line in page_results_list:
        f.write("%s\n" % line)
    f.close()

def archive_new_report_logs():
    log_string = clear_out_old_report_logs(get_log_folder=True)
    log_folder = log_string.split('/')[-1]
    abs_path = os.path.abspath('.')
    file_path = abs_path + "/%s" % ARCHIVE_DIR
    report_log_path = "%s/%s" % (file_path, log_folder)
    return report_log_path

def add_results_page(html):
    abs_path = os.path.abspath('.')
    file_path = abs_path + "/%s" % LATEST_REPORT_DIR
    results_file_name = HTML_REPORT
    results_file = "%s/%s" % (file_path, results_file_name)
    f = open(results_file, 'w')
    f.write(html)
    f.close()
    return results_file

def build_report(report_log_path, page_results_list,
                 successes, failures, browser_type,
                 hide_report):

    web_log_path = "file://%s" % report_log_path
    successes_count = len(successes)
    failures_count = len(failures)
    total_test_count = successes_count + failures_count

    tf_color = "#11BB11"
    if failures_count > 0:
        tf_color = "#EE3A3A"

    new_data = '''

TEST REPORT SUMMARY

TESTS PASSING: %s

     TESTS FAILING: %s

TOTAL TESTS RUN: %s

''' % (successes_count,
                               tf_color,
                               failures_count,
                               total_test_count)

    new_view_1 = '''

        %s
''' % new_data

    log_link_shown = '../%s%s/' % (
        ARCHIVE_DIR, web_log_path.split(ARCHIVE_DIR)[1])
    csv_link = '%s/%s' % (web_log_path, RESULTS_TABLE)
    csv_link_shown = '%s' % RESULTS_TABLE
    new_view_2 = '''

LOG FILES LINK:  %s

RESULTS TABLE:  %s

''' % (
        web_log_path, log_link_shown, csv_link, csv_link_shown)

    new_view_3 = '

'
    any_screenshots = False
    for line in page_results_list:
        line = line.split(',')
        if line[1] == '"FAILED!"':
            if not any_screenshots:
                any_screenshots = True
                new_view_3 += '''

                    FAILURE DATA
                    SCREENSHOT
                    LOCATION OF FAILURE

'''
            line = '%s' % (
                "file://" + report_log_path + '/' + line[2], line[2]) + '''

                ''' + '%s' % (
                "file://" + report_log_path + '/' + line[3], line[3]) + '''

                ''' + '%s' % (line[4], line[4])
            line = line.replace('"', '')
            new_view_3 += '
%s
\n' % line
    new_view_3 += '

'

    new_view_4 = ''
    if failures:
        new_view_4 = '

'
        new_view_4 += '''

LIST OF FAILING TESTS

'''
        for failure in failures:
            new_view_4 += '
%s
\n' % failure
        new_view_4 += '

'

    new_view_5 = ''
    if successes:
        new_view_5 = '

'
        new_view_5 += '''

LIST OF PASSING TESTS

'''
        for success in successes:
            new_view_5 += '
%s
\n' % success
        new_view_5 += '

'

    new_view = '%s%s%s%s%s' % (
        new_view_1, new_view_2, new_view_3, new_view_4, new_view_5)
    results_content = '%s' % new_view
    new_source = '%s%s' % (
        style, results_content)
    results_file = add_results_page(new_source)
    archived_results_file = report_log_path + '/' + HTML_REPORT
    shutil.copyfile(results_file, archived_results_file)
    print "\n* The latest html report page is located at:\n" + results_file
    print "\n* Files saved for this report are located at:\n" + report_log_path
    print ""
    if not hide_report:
        if browser_type == 'chrome':
            browser = webdriver.Chrome()
        else:
            browser = webdriver.Firefox()
        browser.get("file://%s" % archived_results_file)
        print "\n*** Close the html report window to continue. ***"
        while len(browser.window_handles):
            time.sleep(0.1)
        browser.quit()
<commit_msg>Remove extra non-breaking spaces<commit_after>import os
import shutil
import sys
import time
from selenium import webdriver
from seleniumbase.config import settings
from seleniumbase.core.style_sheet import style
from seleniumbase.fixtures import page_actions

LATEST_REPORT_DIR = settings.LATEST_REPORT_DIR
ARCHIVE_DIR = settings.REPORT_ARCHIVE_DIR
HTML_REPORT = settings.HTML_REPORT
RESULTS_TABLE = settings.RESULTS_TABLE

def get_timestamp():
    return str(int(time.time() * 1000))

def process_successes(test, test_count):
    return(
        '"%s","%s","%s","%s","%s","%s","%s","%s","%s"' % (
            test_count,
            "Passed!",
            "*",
            "*",
            "*",
            test.browser,
            get_timestamp()[:-3],
            test.id(),
            "*"))

def process_failures(test, err, test_count, browser_type):
    bad_page_image = "failure_test_%s.jpg" % test_count
    bad_page_data = "failure_test_%s.txt" % test_count
    page_actions.save_screenshot(
        test.driver, bad_page_image, folder=LATEST_REPORT_DIR)
    page_actions.save_test_failure_data(
        test.driver, bad_page_data, browser_type, folder=LATEST_REPORT_DIR)
    exc_info = '(Unknown Failure)'
    exception = sys.exc_info()[1]
    if exception:
        if hasattr(exception, 'msg'):
            exc_info = exception.msg
        elif hasattr(exception, 'message'):
            exc_info = exception.message
        else:
            pass
    return(
        '"%s","%s","%s","%s","%s","%s","%s","%s","%s"' % (
            test_count,
            "FAILED!",
            bad_page_data,
            bad_page_image,
            test.driver.current_url,
            test.browser,
            get_timestamp()[:-3],
            test.id(),
            exc_info))

def clear_out_old_report_logs(archive_past_runs=True, get_log_folder=False):
        abs_path = os.path.abspath('.')
        file_path = abs_path + "/%s" % LATEST_REPORT_DIR
        if not os.path.exists(file_path):
            os.makedirs(file_path)

        if archive_past_runs:
            archive_timestamp = int(time.time())
            if not os.path.exists("%s/../%s/" % (file_path, ARCHIVE_DIR)):
                os.makedirs("%s/../%s/" % (file_path, ARCHIVE_DIR))
            archive_dir = "%s/../%s/report_%s" % (
                file_path, ARCHIVE_DIR, archive_timestamp)
            shutil.move(file_path, archive_dir)
            os.makedirs(file_path)
            if get_log_folder:
                return archive_dir
        else:
            # Just delete bad pages to make room for the latest run.
            filelist = [f for f in os.listdir(
                "./%s" % LATEST_REPORT_DIR) if f.startswith("failure_") or (
                f == HTML_REPORT) or (f.startswith("automation_failure")) or (
                f == RESULTS_TABLE)]
            for f in filelist:
                os.remove("%s/%s" % (file_path, f))

def add_bad_page_log_file(page_results_list):
    abs_path = os.path.abspath('.')
    file_path = abs_path + "/%s" % LATEST_REPORT_DIR
    log_file = "%s/%s" % (file_path, RESULTS_TABLE)
    f = open(log_file, 'w')
    h_p1 = '''"Num","Result","Failure Info","Screenshot",'''
    h_p2 = '''"URL","Browser","Epoch Time",'''
    h_p3 = '''"Test Case Address","Additional Info"\n'''
    page_header = h_p1 + h_p2 + h_p3
    f.write(page_header)
    for line in page_results_list:
        f.write("%s\n" % line)
    f.close()

def archive_new_report_logs():
    log_string = clear_out_old_report_logs(get_log_folder=True)
    log_folder = log_string.split('/')[-1]
    abs_path = os.path.abspath('.')
    file_path = abs_path + "/%s" % ARCHIVE_DIR
    report_log_path = "%s/%s" % (file_path, log_folder)
    return report_log_path

def add_results_page(html):
    abs_path = os.path.abspath('.')
    file_path = abs_path + "/%s" % LATEST_REPORT_DIR
    results_file_name = HTML_REPORT
    results_file = "%s/%s" % (file_path, results_file_name)
    f = open(results_file, 'w')
    f.write(html)
    f.close()
    return results_file

def build_report(report_log_path, page_results_list,
                 successes, failures, browser_type,
                 hide_report):

    web_log_path = "file://%s" % report_log_path
    successes_count = len(successes)
    failures_count = len(failures)
    total_test_count = successes_count + failures_count

    tf_color = "#11BB11"
    if failures_count > 0:
        tf_color = "#EE3A3A"

    new_data = '''

TEST REPORT SUMMARY

TESTS PASSING: %s

     TESTS FAILING: %s

TOTAL TESTS RUN: %s

''' % (successes_count,
                               tf_color,
                               failures_count,
                               total_test_count)

    new_view_1 = '''

        %s
''' % new_data

    log_link_shown = '../%s%s/' % (
        ARCHIVE_DIR, web_log_path.split(ARCHIVE_DIR)[1])
    csv_link = '%s/%s' % (web_log_path, RESULTS_TABLE)
    csv_link_shown = '%s' % RESULTS_TABLE
    new_view_2 = '''

LOG FILES LINK:  %s

RESULTS TABLE:  %s

''' % (
        web_log_path, log_link_shown, csv_link, csv_link_shown)

    new_view_3 = '

'
    any_screenshots = False
    for line in page_results_list:
        line = line.split(',')
        if line[1] == '"FAILED!"':
            if not any_screenshots:
                any_screenshots = True
                new_view_3 += '''

                    FAILURE DATA
                    SCREENSHOT
                    LOCATION OF FAILURE

'''
            line = '%s' % (
                "file://" + report_log_path + '/' + line[2], line[2]) + '''

                ''' + '%s' % (
                "file://" + report_log_path + '/' + line[3], line[3]) + '''

                ''' + '%s' % (line[4], line[4])
            line = line.replace('"', '')
            new_view_3 += '
%s
\n' % line
    new_view_3 += '

'

    new_view_4 = ''
    if failures:
        new_view_4 = '

'
        new_view_4 += '''

LIST OF FAILING TESTS

'''
        for failure in failures:
            new_view_4 += '
%s
\n' % failure
        new_view_4 += '

'

    new_view_5 = ''
    if successes:
        new_view_5 = '

'
        new_view_5 += '''

LIST OF PASSING TESTS

'''
        for success in successes:
            new_view_5 += '
%s
\n' % success
        new_view_5 += '

'

    new_view = '%s%s%s%s%s' % (
        new_view_1, new_view_2, new_view_3, new_view_4, new_view_5)
    results_content = '%s' % new_view
    new_source = '%s%s' % (
        style, results_content)
    results_file = add_results_page(new_source)
    archived_results_file = report_log_path + '/' + HTML_REPORT
    shutil.copyfile(results_file, archived_results_file)
    print "\n* The latest html report page is located at:\n" + results_file
    print "\n* Files saved for this report are located at:\n" + report_log_path
    print ""
    if not hide_report:
        if browser_type == 'chrome':
            browser = webdriver.Chrome()
        else:
            browser = webdriver.Firefox()
        browser.get("file://%s" % archived_results_file)
        print "\n*** Close the html report window to continue. ***"
        while len(browser.window_handles):
            time.sleep(0.1)
        browser.quit()
<|endoftext|>


<commit_before># Software License Agreement (BSD License)
#
# Copyright (c) 2009-2011, Eucalyptus Systems, Inc.
# All rights reserved.
#
# Redistribution and use of this software in source and binary forms, with or
# without modification, are permitted provided that the following conditions
# are met:
#
#   Redistributions of source code must retain the above
#   copyright notice, this list of conditions and the
#   following disclaimer.
#
#   Redistributions in binary form must reproduce the above
#   copyright notice, this list of conditions and the
#   following disclaimer in the documentation and/or other
#   materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
# Author: Neil Soman neil@eucalyptus.com
#         Mitch Garnaat mgarnaat@eucalyptus.com

import os
from boto.roboto.param import Param
import euca2ools.commands.eucacommand
import euca2ools.bundler
from euca2ools.exceptions import NotFoundError, CommandFailed

class Unbundle(euca2ools.commands.eucacommand.EucaCommand):

    Description = 'Unbundles a previously uploaded bundle.'
    Options = [Param(name='manifest_path',
                     short_name='m', long_name='manifest',
                     optional=False, ptype='file',
                     doc='Path to the manifest file.'),
               Param(name='private_key_path',
                     short_name='k', long_name='privatekey',
                     optional=True, ptype='file',
                     doc='Path to private key used to encrypt bundle.'),
               Param(name='destination_dir',
                     short_name='d', long_name='destination',
                     optional=True, ptype='dir', default='.',
                     doc="""Directory to store the image to.
                     Defaults to the current directory."""),
               Param(name='source_dir',
                     short_name='s', long_name='destination',
                     optional=True, ptype='dir',
                     doc="""Source directory for the bundled image parts.
                     Defaults to manifest directory.""")]

def main():
    if not self.source_dir:
        self.source_dir = self.get_file_path(self.manifest_path)
    if not self.private_key_path:
        self.private_key_path = self.get_environ('EC2_PRIVATE_KEY')
        if not os.path.isfile(self.private_key_path):
            msg = 'Private Key not found: %s' % self.private_key_path
            self.display_error_and_exit(msg)

    bundler = euca2ools.bundler.Bundler(self)
    (parts, encrypted_key, encrypted_iv) = \
        bundler.parse_manifest(self.manifest_path)
    image = bundler.assemble_parts(self.source_dir, self.directory,
                                   self.manifest_path, parts)
    print 'Decrypting image'
    decrypted_image = bundler.decrypt_image(image, encrypted_key,
                                            encrypted_iv, self.private_key_path)
    os.remove(image)
    print 'Uncompressing image'
    try:
        unencrypted_image = bundler.untarzip_image(self.directory,
                                                   decrypted_image)
    except NotFoundError:
        sys.exit(1)
    except CommandFailed:
        sys.exit(1)
    os.remove(decrypted_image)

def main_cli(self):
    self.main()

<commit_msg>Merging patch from Luca Invernizzi to fix LP 896481.<commit_after># Software License Agreement (BSD License)
#
# Copyright (c) 2009-2011, Eucalyptus Systems, Inc.
# All rights reserved.
#
# Redistribution and use of this software in source and binary forms, with or
# without modification, are permitted provided that the following conditions
# are met:
#
#   Redistributions of source code must retain the above
#   copyright notice, this list of conditions and the
#   following disclaimer.
#
#   Redistributions in binary form must reproduce the above
#   copyright notice, this list of conditions and the
#   following disclaimer in the documentation and/or other
#   materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
# Author: Neil Soman neil@eucalyptus.com
#         Mitch Garnaat mgarnaat@eucalyptus.com

import os
import sys
from boto.roboto.param import Param
import euca2ools.commands.eucacommand
import euca2ools.bundler
from euca2ools.exceptions import NotFoundError, CommandFailed

class Unbundle(euca2ools.commands.eucacommand.EucaCommand):

    Description = 'Unbundles a previously uploaded bundle.'
    Options = [Param(name='manifest_path',
                     short_name='m', long_name='manifest',
                     optional=False, ptype='file',
                     doc='Path to the manifest file.'),
               Param(name='private_key_path',
                     short_name='k', long_name='privatekey',
                     optional=True, ptype='file',
                     doc='Path to private key used to encrypt bundle.'),
               Param(name='destination_dir',
                     short_name='d', long_name='destination',
                     optional=True, ptype='dir', default='.',
                     doc="""Directory to store the image to.
                     Defaults to the current directory."""),
               Param(name='source_dir',
                     short_name='s', long_name='destination',
                     optional=True, ptype='dir',
                     doc="""Source directory for the bundled image parts.
                     Defaults to manifest directory.""")]

    def main(self):
        if not self.source_dir:
            self.source_dir = self.get_file_path(self.manifest_path)
        if not self.private_key_path:
            self.private_key_path = self.get_environ('EC2_PRIVATE_KEY')
            if not os.path.isfile(self.private_key_path):
                msg = 'Private Key not found: %s' % self.private_key_path
                self.display_error_and_exit(msg)

        bundler = euca2ools.bundler.Bundler(self)
        (parts, encrypted_key, encrypted_iv) = \
            bundler.parse_manifest(self.manifest_path)
        image = bundler.assemble_parts(self.source_dir, self.destination_dir,
                                       self.manifest_path, parts)
        print 'Decrypting image'
        decrypted_image = bundler.decrypt_image(image, encrypted_key,
                                                encrypted_iv, self.private_key_path)
        os.remove(image)
        print 'Uncompressing image'
        try:
            decrypted_image = bundler.untarzip_image(self.destination_dir,
                                                       decrypted_image)
        except NotFoundError:
            sys.exit(1)
        except CommandFailed:
            sys.exit(1)
        os.remove(decrypted_image)

    def main_cli(self):
        self.main()

<|endoftext|>


<commit_before>from dataclasses import dataclass, field, fields
from typing import Dict, List

from flask import current_app as app

from everyclass.server import logger
from everyclass.server.config import get_config
from everyclass.server.exceptions import RpcException
from everyclass.server.rpc.http import HttpRpc
from everyclass.server.utils.resource_identifier_encrypt import encrypt

def ensure_slots(cls, dct: Dict):
    """移除 dataclass 中不存在的key，预防 dataclass 的 __init__ 中 unexpected argument 的发生。"""
    _names = [x.name for x in fields(cls)]
    _del = []
    for key in dct:
        if key not in _names:
            _del.append(key)
    for key in _del:
        del dct[key]  # delete unexpected keys
        logger.warn("Unexpected field `{}` is removed when converting dict to dataclass `{}`".format(key, cls.__name__))
    return dct

@dataclass
class SearchResultStudentItem:
    student_id: str
    student_id_encoded: str
    name: str
    semesters: List[str]
    deputy: str
    klass: str

    @classmethod
    def make(cls, dct: Dict) -> "SearchResultStudentItem":
        dct['semesters'] = sorted(dct.pop("semester_list"))
        dct['student_id'] = dct.pop("student_code")  # rename
        dct['student_id_encoded'] = encrypt('student', dct['student_id'])
        dct['klass'] = dct.pop("class")
        del dct["type"]
        return cls(**ensure_slots(cls, dct))

@dataclass
class SearchResultTeacherItem:
    teacher_id: str
    teacher_id_encoded: str
    name: str
    semesters: List[str]
    unit: str
    title: str

    @classmethod
    def make(cls, dct: Dict) -> "SearchResultTeacherItem":
        dct['semesters'] = sorted(dct.pop("semester_list"))
        dct['teacher_id'] = dct.pop("teacher_code")  # rename
        dct['teacher_id_encoded'] = encrypt('teacher', dct['teacher_id'])
        del dct["type"]
        return cls(**ensure_slots(cls, dct))

@dataclass
class SearchResultClassroomItem:
    room_id: str
    room_id_encoded: str
    name: str
    semesters: List[str]
    campus: str
    building: str

    @classmethod
    def make(cls, dct: Dict) -> "SearchResultClassroomItem":
        dct['semesters'] = sorted(dct.pop("semester_list"))
        dct['room_id'] = dct.pop("room_code")  # rename
        dct['room_id_encoded'] = encrypt('room', dct['room_id'])
        del dct["type"]
        return cls(**ensure_slots(cls, dct))

@dataclass
class SearchResult:
    students: List[SearchResultStudentItem]
    teachers: List[SearchResultTeacherItem]
    classrooms: List[SearchResultClassroomItem]

    @classmethod
    def make(cls, dct: Dict) -> "SearchResult":
        del dct["status"]
        del dct["info"]
        dct["students"] = [SearchResultStudentItem.make(x) for x in dct['data'] if
                           'type' in x and x['type'] == 'student']
        dct["teachers"] = [SearchResultTeacherItem.make(x) for x in dct['data'] if
                           'type' in x and x['type'] == 'teacher']
        dct["classrooms"] = [SearchResultClassroomItem.make(x) for x in dct['data'] if
                             'type' in x and x['type'] == 'room']
        dct.pop("data")

        return cls(**ensure_slots(cls, dct))

    def append(self, to_append: Dict):
        """对于多页搜索结果，将第一页之后的结果追加到搜索结果对象"""
        new_result = self.__class__.make(to_append)
        self.students.extend(new_result.students)
        self.teachers.extend(new_result.teachers)
        self.classrooms.extend(new_result.classrooms)

@dataclass
class TeacherItem:
    teacher_id: str
    teacher_id_encoded: str
    name: str
    title: str

    @classmethod
    def make(cls, dct: Dict) -> "TeacherItem":
        dct['teacher_id'] = dct.pop("teacher_code")
        dct['teacher_id_encoded'] = encrypt('teacher', dct['teacher_id'])
        return cls(**ensure_slots(cls, dct))

@dataclass
class CardItem:
    name: str
    card_id: str
    card_id_encoded: str
    room: str
    room_id: str
    room_id_encoded: str
    weeks: List[int]
    week_string: str
    lesson: str
    teachers: List[TeacherItem]
    course_id: str

    @classmethod
    def make(cls, dct: Dict) -> "CardItem":
        dct["teachers"] = [TeacherItem.make(x) for x in dct.pop("teacher_list")]
        dct['room_id'] = dct.pop('room_code')
        dct['card_id'] = dct.pop('card_code')
        dct['weeks'] = dct.pop('week_list')
        dct['room_id_encoded'] = encrypt('room', dct['room_id'])
        dct['card_id_encoded'] = encrypt('klass', dct['card_id'])
        dct['course_id'] = dct.pop('course_code')
        return cls(**ensure_slots(cls, dct))

@dataclass
class ClassroomTimetableResult:
    room_id: str
    room_id_encoded: str
    name: str
    building: str
    campus: str
    semester: str
    semesters: List[str]
    cards: List[CardItem]

    @classmethod
    def make(cls, dct: Dict) -> "ClassroomTimetableResult":
        del dct["status"]
        dct['semesters'] = sorted(dct.pop('semester_list'))
        dct['room_id'] = dct.pop("room_code")
        dct['cards'] = [CardItem.make(x) for x in dct.pop('card_list')]

        dct['room_id_encoded'] = encrypt('room', dct['room_id'])
        return cls(**ensure_slots(cls, dct))

@dataclass
class CardResultTeacherItem:
    name: str
    teacher_id: str
    title: str
    unit: str

    @classmethod
    def make(cls, dct: Dict) -> "CardResultTeacherItem":
        dct['teacher_id'] = dct.pop('teacher_code')
        return cls(**ensure_slots(cls, dct))

@dataclass
class CardResultStudentItem:
    name: str
    student_id: str
    student_id_encoded: str
    klass: str
    deputy: str

    @classmethod
    def make(cls, dct: Dict) -> "CardResultStudentItem":
        dct["klass"] = dct.pop("class")
        dct["student_id"] = dct.pop("student_code")
        dct["student_id_encoded"] = encrypt("student", dct.get("student_id"))
        return cls(**ensure_slots(cls, dct))

@dataclass
class StudentResult:
    name: str
    student_id: str
    student_id_encoded: str
    campus: str
    deputy: str
    klass: str
    semesters: List[str] = field(default_factory=list)  # optional field

    @classmethod
    def make(cls, dct: Dict) -> "StudentResult":
        del dct["status"]
        dct["student_id"] = dct.pop("student_code")
        dct["student_id_encoded"] = encrypt("student", dct["student_id"])
        dct["klass"] = dct.pop("class")
        dct["semesters"] = dct.pop('semester_list')
        return cls(**ensure_slots(cls, dct))

@dataclass
class StudentTimetableResult:
    name: str  # 姓名
    student_id: str  # 学号
    student_id_encoded: str  # 编码后的学号
    campus: str  # 所在校区
    deputy: str  # 院系
    klass: str  # 班级
    cards: List[CardItem]  # card 列表
    semester: str  # 当前学期
    semesters: List[str] = field(default_factory=list)  # 学期列表

    @classmethod
    def make(cls, dct: Dict) -> "StudentTimetableResult":
        del dct["status"]
        dct["cards"] = [CardItem.make(x) for x in dct.pop("card_list")]
        dct["semesters"] = dct.pop("semester_list")
        dct["student_id"] = dct.pop("student_code")
        dct["student_id_encoded"] = encrypt("student", dct["student_id"])
        dct["klass"] = dct.pop("class")
        return cls(**ensure_slots(cls, dct))

@dataclass
class TeacherTimetableResult:
    name: str  # 姓名
    teacher_id: str  # 教工号
    teacher_id_encoded: str  # 编码后的教工号
    degree: str  # 教师学历（可能为空）
    title: str  # 职称
    unit: str  # 所在单位
    cards: List[CardItem]  # card 列表
    semester: str  # 当前学期
    semesters: List[str] = field(default_factory=list)  # 所有学期

    @classmethod
    def make(cls, dct: Dict) -> "TeacherTimetableResult":
        del dct["status"]
        dct["cards"] = [CardItem.make(x) for x in dct.pop("card_list")]
        dct['semesters'] = sorted(dct.pop('semester_list'))
        dct['teacher_id'] = dct.pop('teacher_code')
        dct["teacher_id_encoded"] = encrypt("teacher", dct["teacher_id"])
        return cls(**ensure_slots(cls, dct))

@dataclass
class CardResult:
    name: str  # 课程名
    card_id: str  # card id
    card_id_encoded: str  # 编码后的 card id
    semester: str  # 学期
    union_name: str  # 合班（教学班）名称
    hour: int  # 课时
    lesson: str  # 上课时间，如10506
    type: str  # 课程类型
    picked: int  # 选课人数
    room: str  # 教室名
    room_id: str  # 教室 ID
    room_id_encoded: str  # 编码后的教室 ID
    students: List[CardResultStudentItem]  # 学生列表
    teachers: List[CardResultTeacherItem]  # 老师列表
    weeks: List[int]  # 周次列表
    week_string: str  # 周次字符串表示
    course_id: str  # 课程 ID

    @classmethod
    def make(cls, dct: Dict) -> "CardResult":
        del dct["status"]
        dct["teachers"] = [CardResultTeacherItem.make(x) for x in dct.pop("teacher_list")]
        dct["students"] = [CardResultStudentItem.make(x) for x in dct.pop("student_list")]
        dct['card_id'] = dct.pop('card_code')
        dct["card_id_encoded"] = encrypt("klass", dct["card_id"])
        dct['room_id'] = dct.pop('room_code')
        dct['room_id_encoded'] = encrypt("room", dct["room_id"])
        dct['weeks'] = dct.pop("week_list")
        dct['course_id'] = dct.pop('course_code')
        dct['union_name'] = dct.pop('tea_class')
        return cls(**ensure_slots(cls, dct))

def teacher_list_to_str(teachers: List[CardResultTeacherItem]) -> str:
    """CardResultTeacherItem 列表转换为老师列表字符串"""
    string = ''
    for teacher in teachers:
        string = string + teacher.name + teacher.title + '、'
    return string[:len(string) - 1]

class APIServer:
    @classmethod
    def search(cls, keyword: str) -> SearchResult:
        """在 API Server 上搜索

        :param keyword: 需要搜索的关键词
        :return: 搜索结果列表
        """
        keyword = keyword.replace("/", "")

        resp = HttpRpc.call(method="GET",
                            url='{}/search/query?key={}&page_size={}'.format(app.config['API_SERVER_BASE_URL'],
                                                                             keyword,
                                                                             100),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        page_num = resp['info']['page_num']
        search_result = SearchResult.make(resp)

        # 多页结果
        if page_num > 1:
            for page_index in range(2, page_num + 1):
                resp = HttpRpc.call(method="GET",
                                    url='{}/search/query?key={}&page_size={}&page_index={}'.format(
                                            app.config['API_SERVER_BASE_URL'],
                                            keyword,
                                            100, page_index),
                                    retry=True,
                                    headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
                if resp["status"] != "success":
                    raise RpcException('API Server returns non-success status')
                search_result.append(resp)

        return search_result

    @classmethod
    def get_student(cls, student_id: str):
        """
        根据学号获得学生课表

        :param student_id: 学号
        :return:
        """
        resp = HttpRpc.call(method="GET",
                            url='{}/student/{}'.format(app.config['API_SERVER_BASE_URL'],
                                                       student_id),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        search_result = StudentResult.make(resp)
        return search_result

    @classmethod
    def get_student_timetable(cls, student_id: str, semester: str):
        """
        根据学期和学号获得学生课表

        :param student_id: 学号
        :param semester: 学期，如 2018-2019-1
        :return:
        """
        resp = HttpRpc.call(method="GET",
                            url='{}/student/{}/timetable/{}'.format(app.config['API_SERVER_BASE_URL'],
                                                                    student_id,
                                                                    semester),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        search_result = StudentTimetableResult.make(resp)
        return search_result

    @classmethod
    def get_teacher_timetable(cls, teacher_id: str, semester: str):
        """
        根据学期和教工号获得老师课表

        :param teacher_id: 教工号
        :param semester: 学期，如 2018-2019-1
        :return:
        """
        resp = HttpRpc.call(method="GET",
                            url='{}/teacher/{}/timetable/{}'.format(app.config['API_SERVER_BASE_URL'],
                                                                    teacher_id,
                                                                    semester),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        search_result = TeacherTimetableResult.make(resp)
        return search_result

    @classmethod
    def get_classroom_timetable(cls, semester: str, room_id: str):
        """
        根据学期和教室ID获得教室课表
        :param semester: 学期，如 2018-2019-1
        :param room_id: 教室ID
        :return:
        """
        resp = HttpRpc.call(method="GET",
                            url='{}/room/{}/timetable/{}'.format(app.config['API_SERVER_BASE_URL'],
                                                                 room_id,
                                                                 semester),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        search_result = ClassroomTimetableResult.make(resp)
        return search_result

    @classmethod
    def get_card(cls, semester: str, card_id: str) -> CardResult:
        """
        根据学期和card ID获得card
        :param semester: 学期，如 2018-2019-1
        :param card_id: card ID
        :return:
        """
        resp = HttpRpc.call(method="GET",
                            url='{}/card/{}/timetable/{}'.format(app.config['API_SERVER_BASE_URL'], card_id,
                                                                 semester),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        search_result = CardResult.make(resp)
        return search_result
<commit_msg>增加pattern字段<commit_after>from dataclasses import dataclass, field, fields
from typing import Dict, List

from flask import current_app as app

from everyclass.server import logger
from everyclass.server.config import get_config
from everyclass.server.exceptions import RpcException
from everyclass.server.rpc.http import HttpRpc
from everyclass.server.utils.resource_identifier_encrypt import encrypt

def ensure_slots(cls, dct: Dict):
    """移除 dataclass 中不存在的key，预防 dataclass 的 __init__ 中 unexpected argument 的发生。"""
    _names = [x.name for x in fields(cls)]
    _del = []
    for key in dct:
        if key not in _names:
            _del.append(key)
    for key in _del:
        del dct[key]  # delete unexpected keys
        logger.warn("Unexpected field `{}` is removed when converting dict to dataclass `{}`".format(key, cls.__name__))
    return dct

@dataclass
class SearchResultStudentItem:
    student_id: str
    student_id_encoded: str
    name: str
    semesters: List[str]
    deputy: str
    klass: str
    pattern: str

    @classmethod
    def make(cls, dct: Dict) -> "SearchResultStudentItem":
        dct['semesters'] = sorted(dct.pop("semester_list"))
        dct['student_id'] = dct.pop("student_code")  # rename
        dct['student_id_encoded'] = encrypt('student', dct['student_id'])
        dct['klass'] = dct.pop("class")
        del dct["type"]
        return cls(**ensure_slots(cls, dct))

@dataclass
class SearchResultTeacherItem:
    teacher_id: str
    teacher_id_encoded: str
    name: str
    semesters: List[str]
    unit: str
    title: str
    pattern: str

    @classmethod
    def make(cls, dct: Dict) -> "SearchResultTeacherItem":
        dct['semesters'] = sorted(dct.pop("semester_list"))
        dct['teacher_id'] = dct.pop("teacher_code")  # rename
        dct['teacher_id_encoded'] = encrypt('teacher', dct['teacher_id'])
        del dct["type"]
        return cls(**ensure_slots(cls, dct))

@dataclass
class SearchResultClassroomItem:
    room_id: str
    room_id_encoded: str
    name: str
    semesters: List[str]
    campus: str
    building: str
    pattern: str  # 搜索结果来源

    @classmethod
    def make(cls, dct: Dict) -> "SearchResultClassroomItem":
        dct['semesters'] = sorted(dct.pop("semester_list"))
        dct['room_id'] = dct.pop("room_code")  # rename
        dct['room_id_encoded'] = encrypt('room', dct['room_id'])
        del dct["type"]
        return cls(**ensure_slots(cls, dct))

@dataclass
class SearchResult:
    students: List[SearchResultStudentItem]
    teachers: List[SearchResultTeacherItem]
    classrooms: List[SearchResultClassroomItem]

    @classmethod
    def make(cls, dct: Dict) -> "SearchResult":
        del dct["status"]
        del dct["info"]
        dct["students"] = [SearchResultStudentItem.make(x) for x in dct['data'] if
                           'type' in x and x['type'] == 'student']
        dct["teachers"] = [SearchResultTeacherItem.make(x) for x in dct['data'] if
                           'type' in x and x['type'] == 'teacher']
        dct["classrooms"] = [SearchResultClassroomItem.make(x) for x in dct['data'] if
                             'type' in x and x['type'] == 'room']
        dct.pop("data")

        return cls(**ensure_slots(cls, dct))

    def append(self, to_append: Dict):
        """对于多页搜索结果，将第一页之后的结果追加到搜索结果对象"""
        new_result = self.__class__.make(to_append)
        self.students.extend(new_result.students)
        self.teachers.extend(new_result.teachers)
        self.classrooms.extend(new_result.classrooms)

@dataclass
class TeacherItem:
    teacher_id: str
    teacher_id_encoded: str
    name: str
    title: str

    @classmethod
    def make(cls, dct: Dict) -> "TeacherItem":
        dct['teacher_id'] = dct.pop("teacher_code")
        dct['teacher_id_encoded'] = encrypt('teacher', dct['teacher_id'])
        return cls(**ensure_slots(cls, dct))

@dataclass
class CardItem:
    name: str
    card_id: str
    card_id_encoded: str
    room: str
    room_id: str
    room_id_encoded: str
    weeks: List[int]
    week_string: str
    lesson: str
    teachers: List[TeacherItem]
    course_id: str

    @classmethod
    def make(cls, dct: Dict) -> "CardItem":
        dct["teachers"] = [TeacherItem.make(x) for x in dct.pop("teacher_list")]
        dct['room_id'] = dct.pop('room_code')
        dct['card_id'] = dct.pop('card_code')
        dct['weeks'] = dct.pop('week_list')
        dct['room_id_encoded'] = encrypt('room', dct['room_id'])
        dct['card_id_encoded'] = encrypt('klass', dct['card_id'])
        dct['course_id'] = dct.pop('course_code')
        return cls(**ensure_slots(cls, dct))

@dataclass
class ClassroomTimetableResult:
    room_id: str
    room_id_encoded: str
    name: str
    building: str
    campus: str
    semester: str
    semesters: List[str]
    cards: List[CardItem]

    @classmethod
    def make(cls, dct: Dict) -> "ClassroomTimetableResult":
        del dct["status"]
        dct['semesters'] = sorted(dct.pop('semester_list'))
        dct['room_id'] = dct.pop("room_code")
        dct['cards'] = [CardItem.make(x) for x in dct.pop('card_list')]

        dct['room_id_encoded'] = encrypt('room', dct['room_id'])
        return cls(**ensure_slots(cls, dct))

@dataclass
class CardResultTeacherItem:
    name: str
    teacher_id: str
    title: str
    unit: str

    @classmethod
    def make(cls, dct: Dict) -> "CardResultTeacherItem":
        dct['teacher_id'] = dct.pop('teacher_code')
        return cls(**ensure_slots(cls, dct))

@dataclass
class CardResultStudentItem:
    name: str
    student_id: str
    student_id_encoded: str
    klass: str
    deputy: str

    @classmethod
    def make(cls, dct: Dict) -> "CardResultStudentItem":
        dct["klass"] = dct.pop("class")
        dct["student_id"] = dct.pop("student_code")
        dct["student_id_encoded"] = encrypt("student", dct.get("student_id"))
        return cls(**ensure_slots(cls, dct))

@dataclass
class StudentResult:
    name: str
    student_id: str
    student_id_encoded: str
    campus: str
    deputy: str
    klass: str
    semesters: List[str] = field(default_factory=list)  # optional field

    @classmethod
    def make(cls, dct: Dict) -> "StudentResult":
        del dct["status"]
        dct["student_id"] = dct.pop("student_code")
        dct["student_id_encoded"] = encrypt("student", dct["student_id"])
        dct["klass"] = dct.pop("class")
        dct["semesters"] = dct.pop('semester_list')
        return cls(**ensure_slots(cls, dct))

@dataclass
class StudentTimetableResult:
    name: str  # 姓名
    student_id: str  # 学号
    student_id_encoded: str  # 编码后的学号
    campus: str  # 所在校区
    deputy: str  # 院系
    klass: str  # 班级
    cards: List[CardItem]  # card 列表
    semester: str  # 当前学期
    semesters: List[str] = field(default_factory=list)  # 学期列表

    @classmethod
    def make(cls, dct: Dict) -> "StudentTimetableResult":
        del dct["status"]
        dct["cards"] = [CardItem.make(x) for x in dct.pop("card_list")]
        dct["semesters"] = dct.pop("semester_list")
        dct["student_id"] = dct.pop("student_code")
        dct["student_id_encoded"] = encrypt("student", dct["student_id"])
        dct["klass"] = dct.pop("class")
        return cls(**ensure_slots(cls, dct))

@dataclass
class TeacherTimetableResult:
    name: str  # 姓名
    teacher_id: str  # 教工号
    teacher_id_encoded: str  # 编码后的教工号
    degree: str  # 教师学历（可能为空）
    title: str  # 职称
    unit: str  # 所在单位
    cards: List[CardItem]  # card 列表
    semester: str  # 当前学期
    semesters: List[str] = field(default_factory=list)  # 所有学期

    @classmethod
    def make(cls, dct: Dict) -> "TeacherTimetableResult":
        del dct["status"]
        dct["cards"] = [CardItem.make(x) for x in dct.pop("card_list")]
        dct['semesters'] = sorted(dct.pop('semester_list'))
        dct['teacher_id'] = dct.pop('teacher_code')
        dct["teacher_id_encoded"] = encrypt("teacher", dct["teacher_id"])
        return cls(**ensure_slots(cls, dct))

@dataclass
class CardResult:
    name: str  # 课程名
    card_id: str  # card id
    card_id_encoded: str  # 编码后的 card id
    semester: str  # 学期
    union_name: str  # 合班（教学班）名称
    hour: int  # 课时
    lesson: str  # 上课时间，如10506
    type: str  # 课程类型
    picked: int  # 选课人数
    room: str  # 教室名
    room_id: str  # 教室 ID
    room_id_encoded: str  # 编码后的教室 ID
    students: List[CardResultStudentItem]  # 学生列表
    teachers: List[CardResultTeacherItem]  # 老师列表
    weeks: List[int]  # 周次列表
    week_string: str  # 周次字符串表示
    course_id: str  # 课程 ID

    @classmethod
    def make(cls, dct: Dict) -> "CardResult":
        del dct["status"]
        dct["teachers"] = [CardResultTeacherItem.make(x) for x in dct.pop("teacher_list")]
        dct["students"] = [CardResultStudentItem.make(x) for x in dct.pop("student_list")]
        dct['card_id'] = dct.pop('card_code')
        dct["card_id_encoded"] = encrypt("klass", dct["card_id"])
        dct['room_id'] = dct.pop('room_code')
        dct['room_id_encoded'] = encrypt("room", dct["room_id"])
        dct['weeks'] = dct.pop("week_list")
        dct['course_id'] = dct.pop('course_code')
        dct['union_name'] = dct.pop('tea_class')
        return cls(**ensure_slots(cls, dct))

def teacher_list_to_str(teachers: List[CardResultTeacherItem]) -> str:
    """CardResultTeacherItem 列表转换为老师列表字符串"""
    string = ''
    for teacher in teachers:
        string = string + teacher.name + teacher.title + '、'
    return string[:len(string) - 1]

class APIServer:
    @classmethod
    def search(cls, keyword: str) -> SearchResult:
        """在 API Server 上搜索

        :param keyword: 需要搜索的关键词
        :return: 搜索结果列表
        """
        keyword = keyword.replace("/", "")

        resp = HttpRpc.call(method="GET",
                            url='{}/search/query?key={}&page_size={}'.format(app.config['API_SERVER_BASE_URL'],
                                                                             keyword,
                                                                             100),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        page_num = resp['info']['page_num']
        search_result = SearchResult.make(resp)

        # 多页结果
        if page_num > 1:
            for page_index in range(2, page_num + 1):
                resp = HttpRpc.call(method="GET",
                                    url='{}/search/query?key={}&page_size={}&page_index={}'.format(
                                            app.config['API_SERVER_BASE_URL'],
                                            keyword,
                                            100, page_index),
                                    retry=True,
                                    headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
                if resp["status"] != "success":
                    raise RpcException('API Server returns non-success status')
                search_result.append(resp)

        return search_result

    @classmethod
    def get_student(cls, student_id: str):
        """
        根据学号获得学生课表

        :param student_id: 学号
        :return:
        """
        resp = HttpRpc.call(method="GET",
                            url='{}/student/{}'.format(app.config['API_SERVER_BASE_URL'],
                                                       student_id),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        search_result = StudentResult.make(resp)
        return search_result

    @classmethod
    def get_student_timetable(cls, student_id: str, semester: str):
        """
        根据学期和学号获得学生课表

        :param student_id: 学号
        :param semester: 学期，如 2018-2019-1
        :return:
        """
        resp = HttpRpc.call(method="GET",
                            url='{}/student/{}/timetable/{}'.format(app.config['API_SERVER_BASE_URL'],
                                                                    student_id,
                                                                    semester),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        search_result = StudentTimetableResult.make(resp)
        return search_result

    @classmethod
    def get_teacher_timetable(cls, teacher_id: str, semester: str):
        """
        根据学期和教工号获得老师课表

        :param teacher_id: 教工号
        :param semester: 学期，如 2018-2019-1
        :return:
        """
        resp = HttpRpc.call(method="GET",
                            url='{}/teacher/{}/timetable/{}'.format(app.config['API_SERVER_BASE_URL'],
                                                                    teacher_id,
                                                                    semester),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        search_result = TeacherTimetableResult.make(resp)
        return search_result

    @classmethod
    def get_classroom_timetable(cls, semester: str, room_id: str):
        """
        根据学期和教室ID获得教室课表
        :param semester: 学期，如 2018-2019-1
        :param room_id: 教室ID
        :return:
        """
        resp = HttpRpc.call(method="GET",
                            url='{}/room/{}/timetable/{}'.format(app.config['API_SERVER_BASE_URL'],
                                                                 room_id,
                                                                 semester),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        search_result = ClassroomTimetableResult.make(resp)
        return search_result

    @classmethod
    def get_card(cls, semester: str, card_id: str) -> CardResult:
        """
        根据学期和card ID获得card
        :param semester: 学期，如 2018-2019-1
        :param card_id: card ID
        :return:
        """
        resp = HttpRpc.call(method="GET",
                            url='{}/card/{}/timetable/{}'.format(app.config['API_SERVER_BASE_URL'], card_id,
                                                                 semester),
                            retry=True,
                            headers={'X-Auth-Token': get_config().API_SERVER_TOKEN})
        if resp["status"] != "success":
            raise RpcException('API Server returns non-success status')
        search_result = CardResult.make(resp)
        return search_result
<|endoftext|>