<commit_before>#!/usr/bin/env python3
from accessoryFunctions.accessoryFunctions import MetadataObject
from MLSTsippr.mlst import GeneSippr as MLSTSippr
from geneseekr.geneseekr import GeneSeekr
from geneseekr.parser import Parser
import multiprocessing
from glob import glob
import logging
import os

__author__ = 'adamkoziol'

class BLAST(object):

    def seekr(self):
        """
        Run the methods in the proper order
        """
        self.blast_db()
        self.run_blast()
        self.parse_results()
        self.create_reports()
        self.clean_object()
        logging.info('{at} analyses complete'.format(at=self.analysistype))

    def blast_db(self):
        """
        Make blast databases (if necessary)
        """
        logging.info('Creating {at} blast databases as required'.format(at=self.analysistype))
        for sample in self.metadata:
            self.geneseekr.makeblastdb(sample[self.analysistype].combinedtargets,
                                       self.program)

    def run_blast(self):
        """
        Perform BLAST analyses
        """
        logging.info('Performing {program} analyses on {at} targets'.format(program=self.program,
                                                                            at=self.analysistype))
        if 'mlst' in self.analysistype.lower():
            self.metadata = self.geneseekr.run_blast(self.metadata,
                                                     self.analysistype,
                                                     self.program,
                                                     self.outfmt,
                                                     evalue='1E-20',
                                                     num_threads=self.cpus,
                                                     perc_identity=99)
        elif 'sixteens' in self.analysistype:
            self.metadata = self.geneseekr.run_blast(self.metadata,
                                                     self.analysistype,
                                                     self.program,
                                                     self.outfmt,
                                                     evalue='1E-100',
                                                     num_threads=self.cpus,
                                                     num_alignments=1000,
                                                     perc_identity=98)
        elif self.analysistype == 'GDCS':
            self.metadata = self.geneseekr.run_blast(self.metadata,
                                                     self.analysistype,
                                                     self.program,
                                                     self.outfmt,
                                                     evalue=self.evalue,
                                                     num_threads=self.cpus,
                                                     task='blastn')
        else:
            self.metadata = self.geneseekr.run_blast(self.metadata,
                                                     self.analysistype,
                                                     self.program,
                                                     self.outfmt,
                                                     evalue=self.evalue,
                                                     num_threads=self.cpus)

    def parse_results(self):
        """
        Parse the output depending on whether unique results are desired
        """
        logging.info('Parsing {program} results for {at} targets'.format(program=self.program,
                                                                         at=self.analysistype))
        if 'sixteens' in self.analysistype:
            self.metadata = self.geneseekr.sixteens_parser(metadata=self.metadata,
                                                           analysistype=self.analysistype,
                                                           fieldnames=self.fieldnames,
                                                           cutoff=self.cutoff,
                                                           program=self.program)
        elif self.unique:
            # Run the unique blast parsing module
            self.metadata = self.geneseekr.unique_parse_blast(metadata=self.metadata,
                                                              analysistype=self.analysistype,
                                                              fieldnames=self.fieldnames,
                                                              cutoff=self.cutoff,
                                                              program=self.program)
            # Filter the unique hits
            self.metadata = self.geneseekr.filter_unique(metadata=self.metadata,
                                                         analysistype=self.analysistype)
        else:
            # Run the standard blast parsing module
            self.metadata = self.geneseekr.parse_blast(metadata=self.metadata,
                                                       analysistype=self.analysistype,
                                                       fieldnames=self.fieldnames,
                                                       cutoff=self.cutoff,
                                                       program=self.program)

    def create_reports(self):
        """
        Create reports
        """
        # Create dictionaries
        self.metadata = self.geneseekr.dict_initialise(self.metadata,
                                                       self.analysistype)
        # Create reports
        logging.info('Creating {at} reports'.format(at=self.analysistype))
        if 'resfinder' in self.analysistype:
            # ResFinder-specific report
            self.metadata = self.geneseekr.resfinder_reporter(metadata=self.metadata,
                                                              analysistype=self.analysistype,
                                                              reportpath=self.reportpath,
                                                              align=self.align,
                                                              targetfiles=self.targetfiles,
                                                              records=self.records,
                                                              program=self.program,
                                                              targetpath=self.targetpath)
        elif 'virulence' in self.analysistype:
            # VirulenceFinder-specific report
            self.geneseekr.virulencefinder_reporter(self.metadata,
                                                    self.analysistype,
                                                    self.reportpath)
        elif 'mlst' in self.analysistype.lower():
            # Adjust the analysis type to be consistently lowercase
            self.analysistype = self.analysistype.lower()
            # Create the necessary attributes for the MLST-typing method
            for sample in self.metadata:
                sample[self.analysistype].alleles = sorted(list(set(allele.split('_')[0]
                                                                    for allele in sample[self.analysistype]
                                                                    .targetnames)))
                # In order to work with the Enterobase cgMLST scheme that has underscores in the gene names (e.g.
                # AEJV01_03887, check for multiple underscores in the allele name, and act appropriately
                if len(sample[self.analysistype].alleles) > 53:
                    allele_set = set()
                    for allele in sample[self.analysistype].targetnames:
                        if len(allele.split('_')) == 3:
                            allele = '_'.join([allele.split('_')[0], allele.split('_')[1]])
                        else:
                            allele = allele.split('_')[0]
                        allele_set.add(allele)
                    sample[self.analysistype].alleles = sorted(list(allele_set))
                # Allele names attribute is apparently the same as the alleles attribute
                sample[self.analysistype].allelenames = sample[self.analysistype].alleles
                try:
                    sample[self.analysistype].profile = glob(os.path.join(sample[self.analysistype].targetpath,
                                                                          '*.txt'))[0]
                except IndexError:
                    sample[self.analysistype].profile = 'NA'
            # Create the typing object
            typing = MLST(args=self,
                          pipelinecommit='',
                          startingtime=self.start,
                          scriptpath='',
                          analysistype=self.analysistype,
                          cutoff=0.99,
                          pipeline=self.pipeline)
            # Perform typing, and create reports
            typing.reporter()
        elif 'sixteens' in self.analysistype:
            self.metadata = self.geneseekr.sixteens_reporter(metadata=self.metadata,
                                                             analysistype=self.analysistype,
                                                             reportpath=self.reportpath)
        elif self.analysistype == 'GDCS':
            self.metadata = self.geneseekr.gdcs_reporter(metadata=self.metadata,
                                                         analysistype=self.analysistype,
                                                         reportpath=self.reportpath)
        elif 'sero' in self.analysistype:
            self.metadata = self.geneseekr.sero_reporter(metadata=self.metadata,
                                                         analysistype=self.analysistype,
                                                         reportpath=self.reportpath)
        else:
            # GeneSeekr-specific report
            self.metadata = self.geneseekr.reporter(self.metadata,
                                                    self.analysistype,
                                                    self.reportpath,
                                                    self.align,
                                                    self.targetfiles,
                                                    self.records,
                                                    self.program)

    # noinspection PyNoneFunctionAssignment
    def clean_object(self):
        """
        Remove certain attributes from the object; they take up too much room on the .json report
        """
        self.metadata = self.geneseekr.clean_object(self.metadata,
                                                    self.analysistype)

    def __init__(self, args, analysistype='geneseekr', cutoff=70, program='blastn', genus_specific=False, unique=False,
                 evalue='1E-05', pipeline=True):
        try:
            args.program = args.program
        except AttributeError:
            args.program = program
        self.program = args.program
        try:
            self.cutoff = args.cutoff
        except AttributeError:
            self.cutoff = cutoff
        try:
            self.cpus = args.numthreads if args.numthreads else multiprocessing.cpu_count() - 1
        except AttributeError:
            self.cpus = args.cpus
        try:
            self.align = args.align
        except AttributeError:
            self.align = True
        if analysistype == 'geneseekr':
            try:
                self.analysistype = args.analysistype.lower()
            except AttributeError:
                self.analysistype = analysistype.lower()
                args.analysistype = analysistype.lower()
        elif analysistype == 'GDCS':
            self.analysistype = analysistype
        else:
            self.analysistype = analysistype.lower()
        try:
            self.resfinder = args.resfinder
        except AttributeError:
            self.resfinder = False
        try:
            self.virulencefinder = args.virulencefinder
        except AttributeError:
            self.virulencefinder = False
        try:
            self.typing = args.typing
        except AttributeError:
            self.typing = False
        # Automatically set self.unique to true for the appropriate analyses
        try:
            self.unique = True if unique or self.resfinder is True or self.virulencefinder is True or 'resfinder' in \
                                  self.analysistype or 'virulence' in self.analysistype or self.typing is True \
                                  or 'mlst' in self.analysistype.lower() or 'sixteens' in self.analysistype \
                                  else args.unique
        except AttributeError:
            self.unique = unique
        try:
            self.start = args.start
        except AttributeError:
            self.start = args.starttime
        try:
            self.evalue = args.evalue
        except AttributeError:
            self.evalue = evalue
        try:
            self.sequencepath = args.sequencepath
        except AttributeError:
            self.sequencepath = str()
        try:
            self.targetpath = os.path.join(args.reffilepath, self.analysistype)
        except AttributeError:
            self.targetpath = args.targetpath
        self.reportpath = args.reportpath
        self.genus_specific = genus_specific
        self.pipeline = pipeline
        try:
            self.metadata = args.runmetadata.samples
            parse = Parser(self)
            if not self.genus_specific:
                parse.target_find()
            parse.metadata_populate()
        except (AttributeError, KeyError):
            # Run the Parser class from the GeneSeekr methods script to create lists of the database targets, and
            # combined targets, fasta sequences, and metadata objects.
            parse = Parser(self)
            parse.main()
        # Extract the variables from the object
        self.reportpath = parse.reportpath
        self.targets = parse.targets
        self.strains = parse.strains
        self.combinedtargets = parse.combinedtargets
        self.metadata = parse.metadata
        # Fields used for custom outfmt 6 BLAST output:
        self.fieldnames = ['query_id', 'subject_id', 'positives', 'mismatches', 'gaps',
                           'evalue', 'bit_score', 'subject_length', 'alignment_length',
                           'query_start', 'query_end', 'query_sequence',
                           'subject_start', 'subject_end', 'subject_sequence']
        self.outfmt = "'6 qseqid sseqid positive mismatch gaps " \
                      "evalue bitscore slen length qstart qend qseq sstart send sseq'"
        self.targetfolders = set()
        self.targetfiles = list()
        self.records = dict()
        # Create the GeneSeekr object
        self.geneseekr = GeneSeekr()
        # Class variables required to use MLST
        self.path = self.sequencepath
        self.logfile = os.path.join(self.path, 'log')
        self.runmetadata = MetadataObject()
        self.runmetadata.samples = self.metadata

class MLST(MLSTSippr):
    def reporter(self):
        analysistype = 'rmlst' if 'rmlst' in self.analysistype.lower() else 'mlst'
        # Populate self.plusdict in order to reuse parsing code from an assembly-based method
        for sample in self.runmetadata.samples:
            if sample.general.bestassemblyfile != 'NA':
                for gene in sample[analysistype].allelenames:
                    for allele, percentidentity in sample[analysistype].blastresults.items():
                        if gene in allele:
                            # Split the allele number from the gene name using the appropriate delimiter
                            if '_' in allele:
                                splitter = '_'
                            elif '-' in allele:
                                splitter = '-'
                            else:
                                splitter = ''
                            # Create the plusdict dictionary as in the assembly-based (r)MLST method. Allows all the
                            # parsing and sequence typing code to be reused.
                            try:
                                self.plusdict[sample.name][gene][allele.split(splitter)[-1]][percentidentity] = 10
                            except IndexError:
                                pass
        self.profiler()
        self.sequencetyper()
        self.mlstreporter()
<commit_msg>Updates to assembly_typing functionality<commit_after>#!/usr/bin/env python3
from accessoryFunctions.accessoryFunctions import MetadataObject
from MLSTsippr.mlst import GeneSippr as MLSTSippr
from geneseekr.geneseekr import GeneSeekr
from geneseekr.parser import Parser
import multiprocessing
from glob import glob
import logging
import os

__author__ = 'adamkoziol'

class BLAST(object):

    def seekr(self):
        """
        Run the methods in the proper order
        """
        self.blast_db()
        self.run_blast()
        self.parse_results()
        self.create_reports()
        self.clean_object()
        logging.info('{at} analyses complete'.format(at=self.analysistype))

    def blast_db(self):
        """
        Make blast databases (if necessary)
        """
        logging.info('Creating {at} blast databases as required'.format(at=self.analysistype))
        for sample in self.metadata:
            self.geneseekr.makeblastdb(sample[self.analysistype].combinedtargets,
                                       self.program)

    def run_blast(self):
        """
        Perform BLAST analyses
        """
        logging.info('Performing {program} analyses on {at} targets'.format(program=self.program,
                                                                            at=self.analysistype))
        if 'mlst' in self.analysistype.lower():
            self.metadata = self.geneseekr.run_blast(self.metadata,
                                                     self.analysistype,
                                                     self.program,
                                                     self.outfmt,
                                                     evalue='1E-10',
                                                     num_threads=self.cpus,
                                                     perc_identity=self.cutoff)
        elif 'sixteens' in self.analysistype:
            self.metadata = self.geneseekr.run_blast(self.metadata,
                                                     self.analysistype,
                                                     self.program,
                                                     self.outfmt,
                                                     evalue='1E-100',
                                                     num_threads=self.cpus,
                                                     num_alignments=5000,
                                                     perc_identity=98)
        elif self.analysistype == 'GDCS':
            self.metadata = self.geneseekr.run_blast(self.metadata,
                                                     self.analysistype,
                                                     self.program,
                                                     self.outfmt,
                                                     evalue=self.evalue,
                                                     num_threads=self.cpus,
                                                     task='blastn')
        else:
            self.metadata = self.geneseekr.run_blast(self.metadata,
                                                     self.analysistype,
                                                     self.program,
                                                     self.outfmt,
                                                     evalue=self.evalue,
                                                     num_threads=self.cpus)

    def parse_results(self):
        """
        Parse the output depending on whether unique results are desired
        """
        logging.info('Parsing {program} results for {at} targets'.format(program=self.program,
                                                                         at=self.analysistype))
        if 'sixteens' in self.analysistype:
            self.metadata = self.geneseekr.sixteens_parser(metadata=self.metadata,
                                                           analysistype=self.analysistype,
                                                           fieldnames=self.fieldnames,
                                                           cutoff=self.cutoff,
                                                           program=self.program)
        elif self.unique:
            # Run the unique blast parsing module
            self.metadata = self.geneseekr.unique_parse_blast(metadata=self.metadata,
                                                              analysistype=self.analysistype,
                                                              fieldnames=self.fieldnames,
                                                              cutoff=self.cutoff,
                                                              program=self.program)
            # Filter the unique hits
            self.metadata = self.geneseekr.filter_unique(metadata=self.metadata,
                                                         analysistype=self.analysistype)
        else:
            # Run the standard blast parsing module
            self.metadata = self.geneseekr.parse_blast(metadata=self.metadata,
                                                       analysistype=self.analysistype,
                                                       fieldnames=self.fieldnames,
                                                       cutoff=self.cutoff,
                                                       program=self.program)

    def create_reports(self):
        """
        Create reports
        """
        # Create dictionaries
        self.metadata = self.geneseekr.dict_initialise(self.metadata,
                                                       self.analysistype)
        # Create reports
        logging.info('Creating {at} reports'.format(at=self.analysistype))
        if 'resfinder' in self.analysistype:
            # ResFinder-specific report
            self.metadata = self.geneseekr.resfinder_reporter(metadata=self.metadata,
                                                              analysistype=self.analysistype,
                                                              reportpath=self.reportpath,
                                                              align=self.align,
                                                              targetfiles=self.targetfiles,
                                                              records=self.records,
                                                              program=self.program,
                                                              targetpath=self.targetpath)
        elif 'virulence' in self.analysistype:
            # VirulenceFinder-specific report
            self.geneseekr.virulencefinder_reporter(self.metadata,
                                                    self.analysistype,
                                                    self.reportpath)
        elif 'mlst' in self.analysistype.lower():
            # Adjust the analysis type to be consistently lowercase
            self.analysistype = self.analysistype.lower()
            # Create the necessary attributes for the MLST-typing method
            for sample in self.metadata:
                sample[self.analysistype].alleles = sorted(list(set(allele.split('_')[0]
                                                                    for allele in sample[self.analysistype]
                                                                    .targetnames)))
                # In order to work with the Enterobase cgMLST scheme that has underscores in the gene names (e.g.
                # AEJV01_03887, check for multiple underscores in the allele name, and act appropriately
                if len(sample[self.analysistype].alleles) > 53:
                    allele_set = set()
                    for allele in sample[self.analysistype].targetnames:
                        if len(allele.split('_')) == 3:
                            allele = '_'.join([allele.split('_')[0], allele.split('_')[1]])
                        else:
                            allele = allele.split('_')[0]
                        allele_set.add(allele)
                    sample[self.analysistype].alleles = sorted(list(allele_set))
                # Allele names attribute is apparently the same as the alleles attribute
                sample[self.analysistype].allelenames = sample[self.analysistype].alleles
                try:
                    sample[self.analysistype].profile = glob(os.path.join(sample[self.analysistype].targetpath,
                                                                          '*.txt'))[0]
                except IndexError:
                    sample[self.analysistype].profile = 'NA'
            # Create the typing object
            typing = MLST(args=self,
                          pipelinecommit='',
                          startingtime=self.start,
                          scriptpath='',
                          analysistype=self.analysistype,
                          cutoff=0.99,
                          pipeline=self.pipeline)
            # Perform typing, and create reports
            typing.reporter()
        elif 'sixteens' in self.analysistype:
            self.metadata = self.geneseekr.sixteens_reporter(metadata=self.metadata,
                                                             analysistype=self.analysistype,
                                                             reportpath=self.reportpath)
        elif self.analysistype == 'GDCS':
            self.metadata = self.geneseekr.gdcs_reporter(metadata=self.metadata,
                                                         analysistype=self.analysistype,
                                                         reportpath=self.reportpath)
        elif 'sero' in self.analysistype:
            self.metadata = self.geneseekr.sero_reporter(metadata=self.metadata,
                                                         analysistype=self.analysistype,
                                                         reportpath=self.reportpath)
        else:
            # GeneSeekr-specific report
            self.metadata = self.geneseekr.reporter(self.metadata,
                                                    self.analysistype,
                                                    self.reportpath,
                                                    self.align,
                                                    self.targetfiles,
                                                    self.records,
                                                    self.program)

    # noinspection PyNoneFunctionAssignment
    def clean_object(self):
        """
        Remove certain attributes from the object; they take up too much room on the .json report
        """
        self.metadata = self.geneseekr.clean_object(self.metadata,
                                                    self.analysistype)

    def __init__(self, args, analysistype='geneseekr', cutoff=70, program='blastn', genus_specific=False, unique=False,
                 evalue='1E-05', pipeline=True):
        try:
            args.program = args.program
        except AttributeError:
            args.program = program
        self.program = args.program
        try:
            self.cutoff = args.cutoff
        except AttributeError:
            self.cutoff = cutoff
        try:
            self.cpus = args.numthreads if args.numthreads else multiprocessing.cpu_count() - 1
        except AttributeError:
            self.cpus = args.cpus
        try:
            self.align = args.align
        except AttributeError:
            self.align = True
        if analysistype == 'geneseekr':
            try:
                self.analysistype = args.analysistype.lower()
            except AttributeError:
                self.analysistype = analysistype.lower()
                args.analysistype = analysistype.lower()
        elif analysistype == 'GDCS':
            self.analysistype = analysistype
        else:
            self.analysistype = analysistype.lower()
        try:
            self.resfinder = args.resfinder
        except AttributeError:
            self.resfinder = False
        try:
            self.virulencefinder = args.virulencefinder
        except AttributeError:
            self.virulencefinder = False
        try:
            self.typing = args.typing
        except AttributeError:
            self.typing = False
        # Automatically set self.unique to true for the appropriate analyses
        try:
            self.unique = True if unique or self.resfinder is True or self.virulencefinder is True or 'resfinder' in \
                                  self.analysistype or 'virulence' in self.analysistype or self.typing is True \
                                  or 'mlst' in self.analysistype.lower() or 'sixteens' in self.analysistype \
                                  else args.unique
        except AttributeError:
            self.unique = unique
        try:
            self.start = args.start
        except AttributeError:
            self.start = args.starttime
        try:
            self.evalue = args.evalue
        except AttributeError:
            self.evalue = evalue
        try:
            self.sequencepath = args.sequencepath
        except AttributeError:
            self.sequencepath = str()
        try:
            self.targetpath = os.path.join(args.reffilepath, self.analysistype)
        except AttributeError:
            self.targetpath = args.targetpath
        self.reportpath = args.reportpath
        self.genus_specific = genus_specific
        self.pipeline = pipeline
        try:
            self.metadata = args.runmetadata.samples
            parse = Parser(args=self,
                           pipeline=self.pipeline)
            if not self.genus_specific:
                parse.target_find()
            parse.metadata_populate()
        except (AttributeError, KeyError):
            # Run the Parser class from the GeneSeekr methods script to create lists of the database targets, and
            # combined targets, fasta sequences, and metadata objects.
            parse = Parser(args=self,
                           pipeline=self.pipeline)
            parse.main()
        # Extract the variables from the object
        self.reportpath = parse.reportpath
        self.targets = parse.targets
        self.strains = parse.strains
        self.combinedtargets = parse.combinedtargets
        self.metadata = parse.metadata
        # Fields used for custom outfmt 6 BLAST output:
        self.fieldnames = ['query_id', 'subject_id', 'positives', 'mismatches', 'gaps',
                           'evalue', 'bit_score', 'subject_length', 'alignment_length',
                           'query_start', 'query_end', 'query_sequence',
                           'subject_start', 'subject_end', 'subject_sequence']
        self.outfmt = "'6 qseqid sseqid positive mismatch gaps " \
                      "evalue bitscore slen length qstart qend qseq sstart send sseq'"
        self.targetfolders = set()
        self.targetfiles = list()
        self.records = dict()
        # Create the GeneSeekr object
        self.geneseekr = GeneSeekr()
        # Class variables required to use MLST
        self.path = self.sequencepath
        self.logfile = os.path.join(self.path, 'log')
        self.runmetadata = MetadataObject()
        self.runmetadata.samples = self.metadata

class MLST(MLSTSippr):
    def reporter(self):
        analysistype = 'rmlst' if 'rmlst' in self.analysistype.lower() else 'mlst'
        # Populate self.plusdict in order to reuse parsing code from an assembly-based method
        for sample in self.runmetadata.samples:
            self.plusdict[sample.name] = dict()
            self.matchdict[sample.name] = dict()
            if sample.general.bestassemblyfile != 'NA':
                for gene in sample[analysistype].allelenames:
                    self.plusdict[sample.name][gene] = dict()
                    for allele, percentidentity in sample[analysistype].blastresults.items():
                        if gene in allele:
                            # Split the allele number from the gene name using the appropriate delimiter
                            if '_' in allele:
                                splitter = '_'
                            elif '-' in allele:
                                splitter = '-'
                            else:
                                splitter = ''
                            self.matchdict[sample.name].update({gene: allele.split(splitter)[-1]})
                            try:
                                self.plusdict[sample.name][gene][allele.split(splitter)[-1]][percentidentity] \
                                    = 10
                            except KeyError:
                                self.plusdict[sample.name][gene][allele.split(splitter)[-1]] = dict()
                                self.plusdict[sample.name][gene][allele.split(splitter)[-1]][percentidentity] \
                                    = 10
                    if gene not in self.matchdict[sample.name]:
                        self.matchdict[sample.name].update({gene: 'N'})
        self.profiler()
        self.sequencetyper()
        self.mlstreporter()
<|endoftext|>


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

import sys
import unittest

import stuff
stuff.append_src_to_pypath()

import pyrapi.rapi as rapi

class TestPyrapi(unittest.TestCase):
    def setUp(self):
        self.opts = rapi.opts()
        rapi.init(self.opts)

    def tearDown(self):
        rapi.shutdown()

    def test_init(self):
        r = rapi.init(rapi.opts())
        self.assertIsNone(r)

    def test_shutdown(self):
        rapi.init(rapi.opts())
        r = rapi.shutdown()
        self.assertIsNone(r)

    def test_aligner_name_vers(self):
        name = rapi.aligner_name()
        ver = rapi.aligner_version()
        self.assertGreater(len(name), 0)
        self.assertGreater(len(ver), 0)

class TestPyrapiRef(unittest.TestCase):
    def setUp(self):
        self.opts = rapi.opts()
        rapi.init(self.opts)
        self.ref = rapi.ref(stuff.MiniRef)

    def tearDown(self):
        self.ref.unload()
        rapi.shutdown()

    def test_bad_ref_path(self):
        self.assertRaises(RuntimeError, rapi.ref, 'bad/path')

    def test_ref_load_unload(self):
        self.assertEqual(stuff.MiniRef, self.ref.path)
        self.assertEqual(1, len(self.ref))

    def test_ref_len(self):
        self.assertEqual(1, len(self.ref))

    def test_ref_iteration(self):
        contigs = [ c for c in self.ref ]
        self.assertEquals(1, len(contigs))
        self.assertEquals(len(self.ref), len(contigs))

        c0 = contigs[0]
        self.assertEquals('chr1', c0.name)
        self.assertEquals(60000, c0.len)
        for v in (c0.assembly_identifier, c0.species, c0.uri, c0.md5):
            self.assertIsNone(v)

    def test_get_contig(self):
        c0 = self.ref.get_contig(0)
        self.assertEquals('chr1', c0.name)
        self.assertEquals(60000, c0.len)

    def test_get_contig_out_of_bounds(self):
        self.assertRaises(IndexError, self.ref.get_contig, 1)
        self.assertRaises(IndexError, self.ref.get_contig, -1)

    def test_get_contig_by_get_item(self):
        c0 = self.ref[0]
        self.assertEquals('chr1', c0.name)
        self.assertEquals(60000, c0.len)

    def test_get_contig_by_get_item_out_of_bounds(self):
        self.assertRaises(IndexError, self.ref.__getitem__, 1)
        self.assertRaises(IndexError, self.ref.__getitem__, -1)

class TestPyrapiReadBatch(unittest.TestCase):
    def setUp(self):
        self.opts = rapi.opts()
        rapi.init(self.opts)
        self.w = rapi.read_batch(2)

    def tearDown(self):
        rapi.shutdown()

    def test_hidden_members(self):
        self.assertFalse(hasattr(self.w, 'batch'))
        self.assertFalse(hasattr(self.w, 'fragment_num'))

    def test_create_bad_arg(self):
        self.assertRaises(ValueError, rapi.read_batch, -1)

    def test_create(self):
        self.assertEquals(2, self.w.n_reads_per_frag())
        self.assertEquals(0, len(self.w))
        self.assertEquals(0, self.w.n_fragments())

        w = rapi.read_batch(1)
        self.assertEquals(1, w.n_reads_per_frag())

    def test_reserve(self):
        self.w.reserve(10)
        self.assertGreaterEqual(10, self.w.capacity())
        self.w.reserve(100)
        self.assertGreaterEqual(100, self.w.capacity())
        self.assertEquals(0, self.w.n_fragments())
        self.assertEquals(0, len(self.w))
        self.assertRaises(ValueError, self.w.reserve, -1)

    def test_append_one(self):
        seq_pair = stuff.get_sequences()[0]
        # insert one read
        self.w.append(seq_pair[0], seq_pair[1], seq_pair[2], rapi.QENC_SANGER)
        self.assertEquals(1, len(self.w))
        read1 = self.w.get_read(0, 0)
        self.assertEquals(seq_pair[0], read1.id)
        self.assertEquals(seq_pair[1], read1.seq)
        self.assertEquals(seq_pair[2], read1.qual)
        self.assertEquals(len(seq_pair[1]), len(read1))
        self.assertEquals(0, self.w.n_fragments())

    def test_append_without_qual(self):
        seq_pair = stuff.get_sequences()[0]
        self.w.append(seq_pair[0], seq_pair[1], None, rapi.QENC_SANGER)
        self.assertEquals(1, len(self.w))
        read1 = self.w.get_read(0, 0)
        self.assertEquals(seq_pair[1], read1.seq)
        self.assertIsNone(read1.qual)

    def test_append_illumina(self):
        seq_pair = stuff.get_sequences()[0]
        # convert the base qualities from sanger to illumina encoding
        # (subtract sanger offset and add illumina offset)
        illumina_qualities = ''.join([ chr(ord(c) - 33 + 64) for c in seq_pair[2] ])
        self.w.append(seq_pair[0], seq_pair[1], illumina_qualities, rapi.QENC_ILLUMINA)
        read1 = self.w.get_read(0, 0)
        # Internally base qualities should be converted to sanger format
        self.assertEquals(seq_pair[2], read1.qual)

    def test_append_baseq_out_of_range(self):
        seq_pair = stuff.get_sequences()[0]

        new_q = chr(32)*len(seq_pair[1]) # sanger encoding goes down to 33
        self.assertRaises(ValueError, self.w.append, seq_pair[0], seq_pair[1], new_q, rapi.QENC_SANGER)

        new_q = chr(127)*len(seq_pair[1]) # and up to 126
        self.assertRaises(ValueError, self.w.append, seq_pair[0], seq_pair[1], new_q, rapi.QENC_SANGER)

        new_q = chr(63)*len(seq_pair[1]) # illumina encoding goes down to 64
        self.assertRaises(ValueError, self.w.append, seq_pair[0], seq_pair[1], new_q, rapi.QENC_ILLUMINA)

    def test_n_reads_per_frag(self):
        self.assertEquals(2, self.w.n_reads_per_frag())
        w = rapi.read_batch(1)
        self.assertEquals(1, w.n_reads_per_frag())

    def test_n_fragments(self):
        seq_pair = stuff.get_sequences()[0]
        self.assertEquals(0, self.w.n_fragments())
        # insert one read
        self.w.append(seq_pair[0], seq_pair[1], seq_pair[2], rapi.QENC_SANGER)
        self.assertEquals(0, self.w.n_fragments())
        self.w.append(seq_pair[0], seq_pair[3], seq_pair[4], rapi.QENC_SANGER)
        self.assertEquals(1, self.w.n_fragments())

    def test_get_read_out_of_bounds(self):
        self.assertEqual(0, len(self.w))
        self.assertRaises(IndexError, self.w.get_read, 0, 0)
        self.assertRaises(IndexError, self.w.get_read, -1, 0)
        self.assertRaises(IndexError, self.w.get_read, -1, -1)

        # now load a sequence and ensure we get exceptions if we access beyond the limits
        seq_pair = stuff.get_sequences()[0]
        self.w.append(seq_pair[0], seq_pair[1], seq_pair[2], rapi.QENC_SANGER)
        self.assertIsNotNone(self.w.get_read(0, 0))
        self.assertRaises(IndexError, self.w.get_read, 0, 1)
        self.assertRaises(IndexError, self.w.get_read, 1, 0)
        self.assertRaises(IndexError, self.w.get_read, 1, 1)

    def test_iteration_values(self):
        seqs = stuff.get_sequences()[0:2]
        for pair in seqs:
            self.w.append(pair[0], pair[1], pair[2], rapi.QENC_SANGER)
            self.w.append(pair[0], pair[3], pair[4], rapi.QENC_SANGER)

        for idx, fragment in enumerate(self.w):
            self.assertEquals(2, len(fragment))
            self.assertEquals(seqs[idx][0], fragment[0].id)
            self.assertEquals(seqs[idx][1], fragment[0].seq)
            self.assertEquals(seqs[idx][0], fragment[1].id)
            self.assertEquals(seqs[idx][3], fragment[1].seq)

    def test_iteration_completeness(self):
        # empty batch
        self.assertEquals(0, sum(1 for frag in self.w))
        # now put some sequences into it
        seqs = stuff.get_sequences()
        for pair in seqs:
            self.w.append(pair[0], pair[1], pair[2], rapi.QENC_SANGER)
            self.w.append(pair[0], pair[3], pair[4], rapi.QENC_SANGER)
        # test whether we iterate over all the fragments
        self.assertEquals(len(seqs), sum(1 for frag in self.w))

    def test_iteration_completeness_single_end(self):
        seqs = stuff.get_sequences()
        batch = rapi.read_batch(1)
        for pair in seqs:
            batch.append(pair[0], pair[1], pair[2], rapi.QENC_SANGER)
            batch.append(pair[0], pair[3], pair[4], rapi.QENC_SANGER)
        # test whether we iterate over all the fragments
        self.assertEquals(2 * len(seqs), sum(1 for frag in batch))
        it = iter(batch)
        fragment = next(it)
        self.assertEquals(1, len(fragment))

class TestPyrapiAlign(unittest.TestCase):
    def setUp(self):
        self.opts = rapi.opts()
        rapi.init(self.opts)
        self.ref = rapi.ref(stuff.MiniRef)

    def tearDown(self):
        self.ref.unload()
        rapi.shutdown()

    def test_align_pe(self):
        aligner = rapi.aligner(self.opts)
        batch = rapi.read_batch(2)
        reads = stuff.get_mini_ref_seqs()
        read_1 = reads[0]
        batch.append(read_1[0], read_1[1], read_1[2], rapi.QENC_SANGER)
        batch.append(read_1[0], read_1[3], read_1[4], rapi.QENC_SANGER)
        aligner.align_reads(self.ref, batch)

        rapi_read = batch.get_read(0, 0)
        self.assertTrue(rapi_read.n_alignments > 0)
        rapi_read = batch.get_read(0, 1)
        self.assertTrue(rapi_read.n_alignments > 0)

#    def test_align_se(self):
#        aligner = rapi.aligner(self.opts)
#        batch = rapi.read_batch(1)
#        reads = stuff.get_mini_ref_seqs()
#        batch.append(reads[0][0], reads[0][1], reads[0][2], rapi.QENC_SANGER)
#        batch.append(reads[1][0], reads[1][1], reads[1][2], rapi.QENC_SANGER)
#        aligner.align_reads(self.ref, batch)
#
#        rapi_read = batch.get_read(0, 0)
#        self.assertTrue(rapi_read.n_alignments > 0)
#        rapi_read = batch.get_read(1, 0)
#        self.assertTrue(rapi_read.n_alignments > 0)

def suite():
    s = unittest.TestLoader().loadTestsFromTestCase(TestPyrapi)
    s.addTests(unittest.TestLoader().loadTestsFromTestCase(TestPyrapiRef))
    s.addTests(unittest.TestLoader().loadTestsFromTestCase(TestPyrapiReadBatch))
    s.addTests(unittest.TestLoader().loadTestsFromTestCase(TestPyrapiAlign))
    return s

def main():
    result = unittest.TextTestRunner(verbosity=2).run(suite())
    return 0 if result.wasSuccessful() else 1

if __name__ == '__main__':
    sys.exit(main())
<commit_msg>Extend unit tests to cover part of the alignment functionality<commit_after>#!/usr/bin/env python

import re
import sys
import unittest

import stuff
stuff.append_src_to_pypath()

import pyrapi.rapi as rapi

class TestPyrapi(unittest.TestCase):
    def setUp(self):
        self.opts = rapi.opts()
        rapi.init(self.opts)

    def tearDown(self):
        rapi.shutdown()

    def test_init(self):
        r = rapi.init(rapi.opts())
        self.assertIsNone(r)

    def test_shutdown(self):
        rapi.init(rapi.opts())
        r = rapi.shutdown()
        self.assertIsNone(r)

    def test_aligner_name_vers(self):
        name = rapi.aligner_name()
        ver = rapi.aligner_version()
        self.assertGreater(len(name), 0)
        self.assertGreater(len(ver), 0)

class TestPyrapiRef(unittest.TestCase):
    def setUp(self):
        self.opts = rapi.opts()
        rapi.init(self.opts)
        self.ref = rapi.ref(stuff.MiniRef)

    def tearDown(self):
        self.ref.unload()
        rapi.shutdown()

    def test_bad_ref_path(self):
        self.assertRaises(RuntimeError, rapi.ref, 'bad/path')

    def test_ref_load_unload(self):
        self.assertEqual(stuff.MiniRef, self.ref.path)
        self.assertEqual(1, len(self.ref))

    def test_ref_len(self):
        self.assertEqual(1, len(self.ref))

    def test_ref_iteration(self):
        contigs = [ c for c in self.ref ]
        self.assertEquals(1, len(contigs))
        self.assertEquals(len(self.ref), len(contigs))

        c0 = contigs[0]
        self.assertEquals('chr1', c0.name)
        self.assertEquals(60000, c0.len)
        for v in (c0.assembly_identifier, c0.species, c0.uri, c0.md5):
            self.assertIsNone(v)

    def test_get_contig(self):
        c0 = self.ref.get_contig(0)
        self.assertEquals('chr1', c0.name)
        self.assertEquals(60000, c0.len)

    def test_get_contig_out_of_bounds(self):
        self.assertRaises(IndexError, self.ref.get_contig, 1)
        self.assertRaises(IndexError, self.ref.get_contig, -1)

    def test_get_contig_by_get_item(self):
        c0 = self.ref[0]
        self.assertEquals('chr1', c0.name)
        self.assertEquals(60000, c0.len)

    def test_get_contig_by_get_item_out_of_bounds(self):
        self.assertRaises(IndexError, self.ref.__getitem__, 1)
        self.assertRaises(IndexError, self.ref.__getitem__, -1)

class TestPyrapiReadBatch(unittest.TestCase):
    def setUp(self):
        self.opts = rapi.opts()
        rapi.init(self.opts)
        self.w = rapi.read_batch(2)

    def tearDown(self):
        rapi.shutdown()

    def test_hidden_members(self):
        self.assertFalse(hasattr(self.w, 'batch'))
        self.assertFalse(hasattr(self.w, 'fragment_num'))

    def test_create_bad_arg(self):
        self.assertRaises(ValueError, rapi.read_batch, -1)

    def test_create(self):
        self.assertEquals(2, self.w.n_reads_per_frag())
        self.assertEquals(0, len(self.w))
        self.assertEquals(0, self.w.n_fragments())

        w = rapi.read_batch(1)
        self.assertEquals(1, w.n_reads_per_frag())

    def test_reserve(self):
        self.w.reserve(10)
        self.assertGreaterEqual(10, self.w.capacity())
        self.w.reserve(100)
        self.assertGreaterEqual(100, self.w.capacity())
        self.assertEquals(0, self.w.n_fragments())
        self.assertEquals(0, len(self.w))
        self.assertRaises(ValueError, self.w.reserve, -1)

    def test_append_one(self):
        seq_pair = stuff.get_sequences()[0]
        # insert one read
        self.w.append(seq_pair[0], seq_pair[1], seq_pair[2], rapi.QENC_SANGER)
        self.assertEquals(1, len(self.w))
        read1 = self.w.get_read(0, 0)
        self.assertEquals(seq_pair[0], read1.id)
        self.assertEquals(seq_pair[1], read1.seq)
        self.assertEquals(seq_pair[2], read1.qual)
        self.assertEquals(len(seq_pair[1]), len(read1))
        self.assertEquals(0, self.w.n_fragments())

    def test_append_without_qual(self):
        seq_pair = stuff.get_sequences()[0]
        self.w.append(seq_pair[0], seq_pair[1], None, rapi.QENC_SANGER)
        self.assertEquals(1, len(self.w))
        read1 = self.w.get_read(0, 0)
        self.assertEquals(seq_pair[1], read1.seq)
        self.assertIsNone(read1.qual)

    def test_append_illumina(self):
        seq_pair = stuff.get_sequences()[0]
        # convert the base qualities from sanger to illumina encoding
        # (subtract sanger offset and add illumina offset)
        illumina_qualities = ''.join([ chr(ord(c) - 33 + 64) for c in seq_pair[2] ])
        self.w.append(seq_pair[0], seq_pair[1], illumina_qualities, rapi.QENC_ILLUMINA)
        read1 = self.w.get_read(0, 0)
        # Internally base qualities should be converted to sanger format
        self.assertEquals(seq_pair[2], read1.qual)

    def test_append_baseq_out_of_range(self):
        seq_pair = stuff.get_sequences()[0]

        new_q = chr(32)*len(seq_pair[1]) # sanger encoding goes down to 33
        self.assertRaises(ValueError, self.w.append, seq_pair[0], seq_pair[1], new_q, rapi.QENC_SANGER)

        new_q = chr(127)*len(seq_pair[1]) # and up to 126
        self.assertRaises(ValueError, self.w.append, seq_pair[0], seq_pair[1], new_q, rapi.QENC_SANGER)

        new_q = chr(63)*len(seq_pair[1]) # illumina encoding goes down to 64
        self.assertRaises(ValueError, self.w.append, seq_pair[0], seq_pair[1], new_q, rapi.QENC_ILLUMINA)

    def test_n_reads_per_frag(self):
        self.assertEquals(2, self.w.n_reads_per_frag())
        w = rapi.read_batch(1)
        self.assertEquals(1, w.n_reads_per_frag())

    def test_n_fragments(self):
        seq_pair = stuff.get_sequences()[0]
        self.assertEquals(0, self.w.n_fragments())
        # insert one read
        self.w.append(seq_pair[0], seq_pair[1], seq_pair[2], rapi.QENC_SANGER)
        self.assertEquals(0, self.w.n_fragments())
        self.w.append(seq_pair[0], seq_pair[3], seq_pair[4], rapi.QENC_SANGER)
        self.assertEquals(1, self.w.n_fragments())

    def test_get_read_out_of_bounds(self):
        self.assertEqual(0, len(self.w))
        self.assertRaises(IndexError, self.w.get_read, 0, 0)
        self.assertRaises(IndexError, self.w.get_read, -1, 0)
        self.assertRaises(IndexError, self.w.get_read, -1, -1)

        # now load a sequence and ensure we get exceptions if we access beyond the limits
        seq_pair = stuff.get_sequences()[0]
        self.w.append(seq_pair[0], seq_pair[1], seq_pair[2], rapi.QENC_SANGER)
        self.assertIsNotNone(self.w.get_read(0, 0))
        self.assertRaises(IndexError, self.w.get_read, 0, 1)
        self.assertRaises(IndexError, self.w.get_read, 1, 0)
        self.assertRaises(IndexError, self.w.get_read, 1, 1)

    def test_iteration_values(self):
        seqs = stuff.get_sequences()[0:2]
        for pair in seqs:
            self.w.append(pair[0], pair[1], pair[2], rapi.QENC_SANGER)
            self.w.append(pair[0], pair[3], pair[4], rapi.QENC_SANGER)

        for idx, fragment in enumerate(self.w):
            self.assertEquals(2, len(fragment))
            self.assertEquals(seqs[idx][0], fragment[0].id)
            self.assertEquals(seqs[idx][1], fragment[0].seq)
            self.assertEquals(seqs[idx][0], fragment[1].id)
            self.assertEquals(seqs[idx][3], fragment[1].seq)

    def test_iteration_completeness(self):
        # empty batch
        self.assertEquals(0, sum(1 for frag in self.w))
        # now put some sequences into it
        seqs = stuff.get_sequences()
        for pair in seqs:
            self.w.append(pair[0], pair[1], pair[2], rapi.QENC_SANGER)
            self.w.append(pair[0], pair[3], pair[4], rapi.QENC_SANGER)
        # test whether we iterate over all the fragments
        self.assertEquals(len(seqs), sum(1 for frag in self.w))

    def test_iteration_completeness_single_end(self):
        seqs = stuff.get_sequences()
        batch = rapi.read_batch(1)
        for pair in seqs:
            batch.append(pair[0], pair[1], pair[2], rapi.QENC_SANGER)
            batch.append(pair[0], pair[3], pair[4], rapi.QENC_SANGER)
        # test whether we iterate over all the fragments
        self.assertEquals(2 * len(seqs), sum(1 for frag in batch))
        it = iter(batch)
        fragment = next(it)
        self.assertEquals(1, len(fragment))

class TestPyrapiAlignment(unittest.TestCase):
    def setUp(self):
        self.opts = rapi.opts()
        rapi.init(self.opts)
        self.ref = rapi.ref(stuff.MiniRef)
        aligner = rapi.aligner(self.opts)
        self.batch = rapi.read_batch(2)
        reads = stuff.get_mini_ref_seqs()
        read_1 = reads[0]
        self.batch.append(read_1[0], read_1[1], read_1[2], rapi.QENC_SANGER)
        self.batch.append(read_1[0], read_1[3], read_1[4], rapi.QENC_SANGER)
        aligner.align_reads(self.ref, self.batch)

    def tearDown(self):
        self.ref.unload()
        rapi.shutdown()

    def test_alignment_struct(self):
        rapi_read = self.batch.get_read(0, 0)
        self.assertEqual("read_00", rapi_read.id)

        self.assertTrue(rapi_read.n_alignments > 0)

        aln = rapi_read.get_aln(0)
        self.assertEqual("chr1", aln.contig.name)
        self.assertEqual(32461, aln.pos)
        self.assertEqual("60M", aln.get_cigar_string())
        self.assertEqual( ( ('M', 60), ), aln.get_cigar_ops())
        self.assertEqual(60, aln.mapq)
        self.assertEqual(60, aln.score)
        # Flag 65 is 'p1'
        self.assertTrue(aln.paired)
        self.assertFalse(aln.prop_paired)
        self.assertTrue(aln.mapped)
        self.assertFalse(aln.reverse_strand)
        self.assertFalse(aln.secondary_aln)
        self.assertEqual(0, aln.n_mismatches)
        self.assertEqual(0, aln.n_gap_opens)
        self.assertEqual(0, aln.n_gap_extensions)

        expected_tags = dict(
                MD='60',
                XS=0)
        # the SAM has more tags (NM and AS), but RAPI makes that information
        # available through other members of the alignment structure.
        self.assertEqual(expected_tags, aln.get_tags())

    def test_get_aln_out_of_bounds(self):
        rapi_read = self.batch.get_read(0, 0)
        self.assertRaises(IndexError, rapi_read.get_aln, -1)
        self.assertRaises(IndexError, rapi_read.get_aln, rapi_read.n_alignments)

    def test_alignment_iter(self):
        return
        self.fail("Not implemented")

    def test_sam(self):
        # We ran this command line:
        #     bwa mem -p -T 0 -a mini_ref/mini_ref.fasta <(head -n 8 mini_ref/mini_ref_seqs.fastq )
        # First 8 lines == first two reads from fastq file.
        # BWA version: 0.7.8-r455
        # The BWA options only serve to tell it not to filter any reads from the output:
        #    * -p: first query file consists of interleaved paired-end sequences
        #    * -T 0: minimum score to output [30]
        #    * -a output all alignments for SE or unpaired PE
        #
        # We're using no options with RAPI, which is equivalent to that command line.
        # This is the output from BWA, without the SAM header.
        expected = [
            """read_00	65	chr1	32461	60	60M	=	32581	121	AAAACTGACCCACACAGAAAAACTAATTGTGAGAACCAATATTATACTAAATTCATTTGA	EEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEE	NM:i:0	MD:Z:60	AS:i:60	XS:i:0""",
            """read_00	129	chr1	32581	60	60M	=	32461	-121	CAAAAGTTAACCCATATGGAATGCAATGGAGGAAATCAATGACATATCAGATCTAGAAAC	EEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEE	NM:i:0	MD:Z:60	AS:i:60	XS:i:0"""
        ]

        rapi_sam = rapi.format_sam(self.batch, 0).split('\n')
        # The two outputs may not be identical because the order or the tags isn't defined.

        bwa_tag_start = [ re.search(r'\tNM:.*', line).start() for line in expected ]
        rapi_tag_start = [ re.search(r'\t[A-Z][A-Z]:[A-z]:.*', line).start() for line in rapi_sam ]

        bwa_no_tags = '\n'.join([ s[0:bwa_tag_start[i]] for i, s in enumerate(expected) ])
        rapi_no_tags = '\n'.join([ s[0:rapi_tag_start[i]] for i, s in enumerate(rapi_sam) ])
        self.assertEqual(bwa_no_tags, rapi_no_tags)

        # now verify the tags
        for read_num in xrange(len(bwa_tag_start)):
            # get the part of the string that contains the tags.  The '+1' is
            # because the regex we used to find the tags began with a tab, which
            # we want to skip.
            bwa_tags = expected[read_num][(bwa_tag_start[read_num]+1):]
            rapi_tags = rapi_sam[read_num][(rapi_tag_start[read_num]+1):]
            self.assertEquals(set(bwa_tags.split('\t')), set(rapi_tags.split('\t')))

#    def test_align_se(self):
#        aligner = rapi.aligner(self.opts)
#        batch = rapi.read_batch(1)
#        reads = stuff.get_mini_ref_seqs()
#        batch.append(reads[0][0], reads[0][1], reads[0][2], rapi.QENC_SANGER)
#        batch.append(reads[1][0], reads[1][1], reads[1][2], rapi.QENC_SANGER)
#        aligner.align_reads(self.ref, batch)
#
#        rapi_read = batch.get_read(0, 0)
#        self.assertTrue(rapi_read.n_alignments > 0)
#        rapi_read = batch.get_read(1, 0)
#        self.assertTrue(rapi_read.n_alignments > 0)

def suite():
    s = unittest.TestLoader().loadTestsFromTestCase(TestPyrapi)
    s.addTests(unittest.TestLoader().loadTestsFromTestCase(TestPyrapiRef))
    s.addTests(unittest.TestLoader().loadTestsFromTestCase(TestPyrapiReadBatch))
    s.addTests(unittest.TestLoader().loadTestsFromTestCase(TestPyrapiAlignment))
    return s

def main():
    result = unittest.TextTestRunner(verbosity=2).run(suite())
    return 0 if result.wasSuccessful() else 1

if __name__ == '__main__':
    sys.exit(main())
<|endoftext|>


<commit_before># Copyright (c) 2012-2013, Mark Peek <mark@peek.org>
# All rights reserved.
#
# See LICENSE file for full license.

import json
import re
import types

from . import validators

__version__ = "1.4.0"

# constants for DeletionPolicy
Delete = 'Delete'
Retain = 'Retain'
Snapshot = 'Snapshot'

# Pseudo Parameters
AWS_ACCOUNT_ID = 'AWS::AccountId'
AWS_NOTIFICATION_ARNS = 'AWS::NotificationARNs'
AWS_NO_VALUE = 'AWS::NoValue'
AWS_REGION = 'AWS::Region'
AWS_STACK_ID = 'AWS::StackId'
AWS_STACK_NAME = 'AWS::StackName'

valid_names = re.compile(r'^[a-zA-Z0-9]+$')

class BaseAWSObject(object):
    def __init__(self, title, template=None, **kwargs):
        self.title = title
        self.template = template
        # Cache the keys for validity checks
        self.propnames = self.props.keys()
        self.attributes = ['DependsOn', 'DeletionPolicy',
                           'Metadata', 'UpdatePolicy',
                           'Condition', 'CreationPolicy']

        # try to validate the title if its there
        if self.title is not None:
            self.validate_title()

        # Create the list of properties set on this object by the user
        self.properties = {}
        dictname = getattr(self, 'dictname', None)
        if dictname:
            self.resource = {
                dictname: self.properties,
            }
        else:
            self.resource = self.properties
        if hasattr(self, 'resource_type') and self.resource_type is not None:
            self.resource['Type'] = self.resource_type
        self.__initialized = True

        # Check for properties defined in the class
        for k, (_, required) in self.props.items():
            v = getattr(type(self), k, None)
            if v is not None and k not in kwargs:
                self.__setattr__(k, v)

        # Now that it is initialized, populate it with the kwargs
        for k, v in kwargs.items():
            self.__setattr__(k, v)

        # Bound it to template if we know it
        if self.template is not None:
            self.template.add_resource(self)

    def __getattr__(self, name):
        try:
            return self.properties.__getitem__(name)
        except KeyError:
            # Fall back to the name attribute in the object rather than
            # in the properties dict. This is for non-OpenStack backwards
            # compatibility since OpenStack objects use a "name" property.
            if name == 'name':
                return self.__getattribute__('title')
            raise AttributeError(name)

    def __setattr__(self, name, value):
        if name in self.__dict__.keys() \
                or '_BaseAWSObject__initialized' not in self.__dict__:
            return dict.__setattr__(self, name, value)
        elif name in self.attributes:
            self.resource[name] = value
            return None
        elif name in self.propnames:
            # Check the type of the object and compare against what we were
            # expecting.
            expected_type = self.props[name][0]

            # If the value is a AWSHelperFn we can't do much validation
            # we'll have to leave that to Amazon.  Maybe there's another way
            # to deal with this that we'll come up with eventually
            if isinstance(value, AWSHelperFn):
                return self.properties.__setitem__(name, value)

            # If it's a function, call it...
            elif isinstance(expected_type, types.FunctionType):
                value = expected_type(value)
                return self.properties.__setitem__(name, value)

            # If it's a list of types, check against those types...
            elif isinstance(expected_type, list):
                # If we're expecting a list, then make sure it is a list
                if not isinstance(value, list):
                    self._raise_type(name, value, expected_type)

                # Iterate over the list and make sure it matches our
                # type checks (as above accept AWSHelperFn because
                # we can't do the validation ourselves)
                for v in value:
                    if not isinstance(v, tuple(expected_type)) \
                       and not isinstance(v, AWSHelperFn):
                        self._raise_type(name, v, expected_type)
                # Validated so assign it
                return self.properties.__setitem__(name, value)

            # Single type so check the type of the object and compare against
            # what we were expecting. Special case AWS helper functions.
            elif isinstance(value, expected_type):
                return self.properties.__setitem__(name, value)
            else:
                self._raise_type(name, value, expected_type)

        type_name = getattr(self, 'resource_type', self.__class__.__name__)

        if type_name == 'AWS::CloudFormation::CustomResource' or \
                type_name.startswith('Custom::'):
            # Add custom resource arguments to the dict without any further
            # validation. The properties of a CustomResource is not known.
            return self.properties.__setitem__(name, value)

        raise AttributeError("%s object does not support attribute %s" %
                             (type_name, name))

    def _raise_type(self, name, value, expected_type):
        raise TypeError('%s: %s.%s is %s, expected %s' % (self.__class__,
                                                          self.title,
                                                          name,
                                                          type(value),
                                                          expected_type))

    def validate_title(self):
        if not valid_names.match(self.title):
            raise ValueError('Name "%s" not alphanumeric' % self.title)

    def validate(self):
        pass

    @classmethod
    def from_dict(cls, title, dict):
        obj = cls(title)
        obj.properties.update(dict)
        return obj

    def JSONrepr(self):
        for k, (_, required) in self.props.items():
            if required and k not in self.properties:
                rtype = getattr(self, 'resource_type', "<unknown type>")
                raise ValueError(
                    "Resource %s required in type %s" % (k, rtype))
        self.validate()
        # Mainly used to not have an empty "Properties".
        if self.properties:
            return self.resource
        elif hasattr(self, 'resource_type'):
            d = {}
            for k, v in self.resource.items():
                if k != 'Properties':
                    d[k] = v
            return d
        else:
            return {}

class AWSObject(BaseAWSObject):
    dictname = 'Properties'

class AWSDeclaration(BaseAWSObject):
    """
    Used for CloudFormation Resource Property objects
    http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/
    aws-product-property-reference.html
    """

    def __init__(self, title, **kwargs):
        super(AWSDeclaration, self).__init__(title, **kwargs)

class AWSProperty(BaseAWSObject):
    """
    Used for CloudFormation Resource Property objects
    http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/
    aws-product-property-reference.html
    """
    dictname = None

    def __init__(self, title=None, **kwargs):
        super(AWSProperty, self).__init__(title, **kwargs)

class AWSAttribute(BaseAWSObject):
    dictname = None

    """
    Used for CloudFormation Resource Attribute objects
    http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/
    aws-product-attribute-reference.html
    """

    def __init__(self, title=None, **kwargs):
        super(AWSAttribute, self).__init__(title, **kwargs)

def validate_pausetime(pausetime):
    if not pausetime.startswith('PT'):
        raise ValueError('PauseTime should look like PT#H#M#S')
    return pausetime

class UpdatePolicy(BaseAWSObject):
    def __init__(self, title, **kwargs):
        raise DeprecationWarning(
            "This UpdatePolicy class is deprecated, please switch to using "
            "the more general UpdatePolicy in troposphere.policies.\n"
        )

class AWSHelperFn(object):
    def getdata(self, data):
        if isinstance(data, BaseAWSObject):
            return data.title
        else:
            return data

class Base64(AWSHelperFn):
    def __init__(self, data):
        self.data = {'Fn::Base64': data}

    def JSONrepr(self):
        return self.data

class FindInMap(AWSHelperFn):
    def __init__(self, mapname, key, value):
        self.data = {'Fn::FindInMap': [self.getdata(mapname), key, value]}

    def JSONrepr(self):
        return self.data

class GetAtt(AWSHelperFn):
    def __init__(self, logicalName, attrName):
        self.data = {'Fn::GetAtt': [self.getdata(logicalName), attrName]}

    def JSONrepr(self):
        return self.data

class GetAZs(AWSHelperFn):
    def __init__(self, region=""):
        self.data = {'Fn::GetAZs': region}

    def JSONrepr(self):
        return self.data

class If(AWSHelperFn):
    def __init__(self, cond, true, false):
        self.data = {'Fn::If': [self.getdata(cond), true, false]}

    def JSONrepr(self):
        return self.data

class Equals(AWSHelperFn):
    def __init__(self, value_one, value_two):
        self.data = {'Fn::Equals': [value_one, value_two]}

    def JSONrepr(self):
        return self.data

class And(AWSHelperFn):
    def __init__(self, cond_one, cond_two, *conds):
        self.data = {'Fn::And': [cond_one, cond_two] + list(conds)}

    def JSONrepr(self):
        return self.data

class Or(AWSHelperFn):
    def __init__(self, cond_one, cond_two, *conds):
        self.data = {'Fn::Or': [cond_one, cond_two] + list(conds)}

    def JSONrepr(self):
        return self.data

class Not(AWSHelperFn):
    def __init__(self, cond):
        self.data = {'Fn::Not': [self.getdata(cond)]}

    def JSONrepr(self):
        return self.data

class Join(AWSHelperFn):
    def __init__(self, delimiter, values):
        self.data = {'Fn::Join': [delimiter, values]}

    def JSONrepr(self):
        return self.data

class Name(AWSHelperFn):
    def __init__(self, data):
        self.data = self.getdata(data)

    def JSONrepr(self):
        return self.data

class Select(AWSHelperFn):
    def __init__(self, indx, objects):
        self.data = {'Fn::Select': [indx, objects]}

    def JSONrepr(self):
        return self.data

class Ref(AWSHelperFn):
    def __init__(self, data):
        self.data = {'Ref': self.getdata(data)}

    def JSONrepr(self):
        return self.data

class Condition(AWSHelperFn):
    def __init__(self, data):
        self.data = {'Condition': self.getdata(data)}

    def JSONrepr(self):
        return self.data

class awsencode(json.JSONEncoder):
    def default(self, obj):
        if hasattr(obj, 'JSONrepr'):
            return obj.JSONrepr()
        return json.JSONEncoder.default(self, obj)

class Tags(AWSHelperFn):
    def __init__(self, **kwargs):
        self.tags = []
        for k, v in sorted(kwargs.iteritems()):
            self.tags.append({
                'Key': k,
                'Value': v,
            })

    # allow concatenation of the Tags object via '+' operator
    def __add__(self, newtags):
        newtags.tags = self.tags + newtags.tags
        return newtags

    def JSONrepr(self):
        return self.tags

class Template(object):
    props = {
        'AWSTemplateFormatVersion': (basestring, False),
        'Description': (basestring, False),
        'Parameters': (dict, False),
        'Mappings': (dict, False),
        'Resources': (dict, False),
        'Outputs': (dict, False),
    }

    def __init__(self):
        self.description = None
        self.metadata = {}
        self.conditions = {}
        self.mappings = {}
        self.outputs = {}
        self.parameters = {}
        self.resources = {}
        self.version = None

    def add_description(self, description):
        self.description = description

    def add_metadata(self, metadata):
        self.metadata = metadata

    def add_condition(self, name, condition):
        self.conditions[name] = condition

    def handle_duplicate_key(self, key):
        raise ValueError('duplicate key "%s" detected' % key)

    def _update(self, d, values):
        if isinstance(values, list):
            for v in values:
                if v.title in d:
                    self.handle_duplicate_key(v.title)
                d[v.title] = v
        else:
            if values.title in d:
                self.handle_duplicate_key(values.title)
            d[values.title] = values
        return values

    def add_output(self, output):
        return self._update(self.outputs, output)

    def add_mapping(self, name, mapping):
        self.mappings[name] = mapping

    def add_parameter(self, parameter):
        return self._update(self.parameters, parameter)

    def add_resource(self, resource):
        return self._update(self.resources, resource)

    def add_version(self, version=None):
        if version:
            self.version = version
        else:
            self.version = "2010-09-09"

    def to_json(self, indent=4, sort_keys=True, separators=(',', ': ')):
        t = {}
        if self.description:
            t['Description'] = self.description
        if self.metadata:
            t['Metadata'] = self.metadata
        if self.conditions:
            t['Conditions'] = self.conditions
        if self.mappings:
            t['Mappings'] = self.mappings
        if self.outputs:
            t['Outputs'] = self.outputs
        if self.parameters:
            t['Parameters'] = self.parameters
        if self.version:
            t['AWSTemplateFormatVersion'] = self.version
        t['Resources'] = self.resources

        return json.dumps(t, cls=awsencode, indent=indent,
                          sort_keys=sort_keys, separators=separators)

    def JSONrepr(self):
        return [self.parameters, self.mappings, self.resources]

class Output(AWSDeclaration):
    props = {
        'Description': (basestring, False),
        'Value': (basestring, True),
    }

class Parameter(AWSDeclaration):
    STRING_PROPERTIES = ['AllowedPattern', 'MaxLength', 'MinLength']
    NUMBER_PROPERTIES = ['MaxValue', 'MinValue']
    props = {
        'Type': (basestring, True),
        'Default': (basestring, False),
        'NoEcho': (bool, False),
        'AllowedValues': (list, False),
        'AllowedPattern': (basestring, False),
        'MaxLength': (validators.positive_integer, False),
        'MinLength': (validators.positive_integer, False),
        'MaxValue': (validators.integer, False),
        'MinValue': (validators.integer, False),
        'Description': (basestring, False),
        'ConstraintDescription': (basestring, False),
    }

    def validate(self):
        if self.properties['Type'] != 'String':
            for p in self.STRING_PROPERTIES:
                if p in self.properties:
                    raise ValueError("%s can only be used with parameters of "
                                     "the String type." % p)
        if self.properties['Type'] != 'Number':
            for p in self.NUMBER_PROPERTIES:
                if p in self.properties:
                    raise ValueError("%s can only be used with parameters of "
                                     "the Number type." % p)
<commit_msg>simplify if<commit_after># Copyright (c) 2012-2013, Mark Peek <mark@peek.org>
# All rights reserved.
#
# See LICENSE file for full license.

import json
import re
import types

from . import validators

__version__ = "1.4.0"

# constants for DeletionPolicy
Delete = 'Delete'
Retain = 'Retain'
Snapshot = 'Snapshot'

# Pseudo Parameters
AWS_ACCOUNT_ID = 'AWS::AccountId'
AWS_NOTIFICATION_ARNS = 'AWS::NotificationARNs'
AWS_NO_VALUE = 'AWS::NoValue'
AWS_REGION = 'AWS::Region'
AWS_STACK_ID = 'AWS::StackId'
AWS_STACK_NAME = 'AWS::StackName'

valid_names = re.compile(r'^[a-zA-Z0-9]+$')

class BaseAWSObject(object):
    def __init__(self, title, template=None, **kwargs):
        self.title = title
        self.template = template
        # Cache the keys for validity checks
        self.propnames = self.props.keys()
        self.attributes = ['DependsOn', 'DeletionPolicy',
                           'Metadata', 'UpdatePolicy',
                           'Condition', 'CreationPolicy']

        # try to validate the title if its there
        if self.title:
            self.validate_title()

        # Create the list of properties set on this object by the user
        self.properties = {}
        dictname = getattr(self, 'dictname', None)
        if dictname:
            self.resource = {
                dictname: self.properties,
            }
        else:
            self.resource = self.properties
        if hasattr(self, 'resource_type') and self.resource_type is not None:
            self.resource['Type'] = self.resource_type
        self.__initialized = True

        # Check for properties defined in the class
        for k, (_, required) in self.props.items():
            v = getattr(type(self), k, None)
            if v is not None and k not in kwargs:
                self.__setattr__(k, v)

        # Now that it is initialized, populate it with the kwargs
        for k, v in kwargs.items():
            self.__setattr__(k, v)

        # Bound it to template if we know it
        if self.template is not None:
            self.template.add_resource(self)

    def __getattr__(self, name):
        try:
            return self.properties.__getitem__(name)
        except KeyError:
            # Fall back to the name attribute in the object rather than
            # in the properties dict. This is for non-OpenStack backwards
            # compatibility since OpenStack objects use a "name" property.
            if name == 'name':
                return self.__getattribute__('title')
            raise AttributeError(name)

    def __setattr__(self, name, value):
        if name in self.__dict__.keys() \
                or '_BaseAWSObject__initialized' not in self.__dict__:
            return dict.__setattr__(self, name, value)
        elif name in self.attributes:
            self.resource[name] = value
            return None
        elif name in self.propnames:
            # Check the type of the object and compare against what we were
            # expecting.
            expected_type = self.props[name][0]

            # If the value is a AWSHelperFn we can't do much validation
            # we'll have to leave that to Amazon.  Maybe there's another way
            # to deal with this that we'll come up with eventually
            if isinstance(value, AWSHelperFn):
                return self.properties.__setitem__(name, value)

            # If it's a function, call it...
            elif isinstance(expected_type, types.FunctionType):
                value = expected_type(value)
                return self.properties.__setitem__(name, value)

            # If it's a list of types, check against those types...
            elif isinstance(expected_type, list):
                # If we're expecting a list, then make sure it is a list
                if not isinstance(value, list):
                    self._raise_type(name, value, expected_type)

                # Iterate over the list and make sure it matches our
                # type checks (as above accept AWSHelperFn because
                # we can't do the validation ourselves)
                for v in value:
                    if not isinstance(v, tuple(expected_type)) \
                       and not isinstance(v, AWSHelperFn):
                        self._raise_type(name, v, expected_type)
                # Validated so assign it
                return self.properties.__setitem__(name, value)

            # Single type so check the type of the object and compare against
            # what we were expecting. Special case AWS helper functions.
            elif isinstance(value, expected_type):
                return self.properties.__setitem__(name, value)
            else:
                self._raise_type(name, value, expected_type)

        type_name = getattr(self, 'resource_type', self.__class__.__name__)

        if type_name == 'AWS::CloudFormation::CustomResource' or \
                type_name.startswith('Custom::'):
            # Add custom resource arguments to the dict without any further
            # validation. The properties of a CustomResource is not known.
            return self.properties.__setitem__(name, value)

        raise AttributeError("%s object does not support attribute %s" %
                             (type_name, name))

    def _raise_type(self, name, value, expected_type):
        raise TypeError('%s: %s.%s is %s, expected %s' % (self.__class__,
                                                          self.title,
                                                          name,
                                                          type(value),
                                                          expected_type))

    def validate_title(self):
        if not valid_names.match(self.title):
            raise ValueError('Name "%s" not alphanumeric' % self.title)

    def validate(self):
        pass

    @classmethod
    def from_dict(cls, title, dict):
        obj = cls(title)
        obj.properties.update(dict)
        return obj

    def JSONrepr(self):
        for k, (_, required) in self.props.items():
            if required and k not in self.properties:
                rtype = getattr(self, 'resource_type', "<unknown type>")
                raise ValueError(
                    "Resource %s required in type %s" % (k, rtype))
        self.validate()
        # Mainly used to not have an empty "Properties".
        if self.properties:
            return self.resource
        elif hasattr(self, 'resource_type'):
            d = {}
            for k, v in self.resource.items():
                if k != 'Properties':
                    d[k] = v
            return d
        else:
            return {}

class AWSObject(BaseAWSObject):
    dictname = 'Properties'

class AWSDeclaration(BaseAWSObject):
    """
    Used for CloudFormation Resource Property objects
    http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/
    aws-product-property-reference.html
    """

    def __init__(self, title, **kwargs):
        super(AWSDeclaration, self).__init__(title, **kwargs)

class AWSProperty(BaseAWSObject):
    """
    Used for CloudFormation Resource Property objects
    http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/
    aws-product-property-reference.html
    """
    dictname = None

    def __init__(self, title=None, **kwargs):
        super(AWSProperty, self).__init__(title, **kwargs)

class AWSAttribute(BaseAWSObject):
    dictname = None

    """
    Used for CloudFormation Resource Attribute objects
    http://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/
    aws-product-attribute-reference.html
    """

    def __init__(self, title=None, **kwargs):
        super(AWSAttribute, self).__init__(title, **kwargs)

def validate_pausetime(pausetime):
    if not pausetime.startswith('PT'):
        raise ValueError('PauseTime should look like PT#H#M#S')
    return pausetime

class UpdatePolicy(BaseAWSObject):
    def __init__(self, title, **kwargs):
        raise DeprecationWarning(
            "This UpdatePolicy class is deprecated, please switch to using "
            "the more general UpdatePolicy in troposphere.policies.\n"
        )

class AWSHelperFn(object):
    def getdata(self, data):
        if isinstance(data, BaseAWSObject):
            return data.title
        else:
            return data

class Base64(AWSHelperFn):
    def __init__(self, data):
        self.data = {'Fn::Base64': data}

    def JSONrepr(self):
        return self.data

class FindInMap(AWSHelperFn):
    def __init__(self, mapname, key, value):
        self.data = {'Fn::FindInMap': [self.getdata(mapname), key, value]}

    def JSONrepr(self):
        return self.data

class GetAtt(AWSHelperFn):
    def __init__(self, logicalName, attrName):
        self.data = {'Fn::GetAtt': [self.getdata(logicalName), attrName]}

    def JSONrepr(self):
        return self.data

class GetAZs(AWSHelperFn):
    def __init__(self, region=""):
        self.data = {'Fn::GetAZs': region}

    def JSONrepr(self):
        return self.data

class If(AWSHelperFn):
    def __init__(self, cond, true, false):
        self.data = {'Fn::If': [self.getdata(cond), true, false]}

    def JSONrepr(self):
        return self.data

class Equals(AWSHelperFn):
    def __init__(self, value_one, value_two):
        self.data = {'Fn::Equals': [value_one, value_two]}

    def JSONrepr(self):
        return self.data

class And(AWSHelperFn):
    def __init__(self, cond_one, cond_two, *conds):
        self.data = {'Fn::And': [cond_one, cond_two] + list(conds)}

    def JSONrepr(self):
        return self.data

class Or(AWSHelperFn):
    def __init__(self, cond_one, cond_two, *conds):
        self.data = {'Fn::Or': [cond_one, cond_two] + list(conds)}

    def JSONrepr(self):
        return self.data

class Not(AWSHelperFn):
    def __init__(self, cond):
        self.data = {'Fn::Not': [self.getdata(cond)]}

    def JSONrepr(self):
        return self.data

class Join(AWSHelperFn):
    def __init__(self, delimiter, values):
        self.data = {'Fn::Join': [delimiter, values]}

    def JSONrepr(self):
        return self.data

class Name(AWSHelperFn):
    def __init__(self, data):
        self.data = self.getdata(data)

    def JSONrepr(self):
        return self.data

class Select(AWSHelperFn):
    def __init__(self, indx, objects):
        self.data = {'Fn::Select': [indx, objects]}

    def JSONrepr(self):
        return self.data

class Ref(AWSHelperFn):
    def __init__(self, data):
        self.data = {'Ref': self.getdata(data)}

    def JSONrepr(self):
        return self.data

class Condition(AWSHelperFn):
    def __init__(self, data):
        self.data = {'Condition': self.getdata(data)}

    def JSONrepr(self):
        return self.data

class awsencode(json.JSONEncoder):
    def default(self, obj):
        if hasattr(obj, 'JSONrepr'):
            return obj.JSONrepr()
        return json.JSONEncoder.default(self, obj)

class Tags(AWSHelperFn):
    def __init__(self, **kwargs):
        self.tags = []
        for k, v in sorted(kwargs.iteritems()):
            self.tags.append({
                'Key': k,
                'Value': v,
            })

    # allow concatenation of the Tags object via '+' operator
    def __add__(self, newtags):
        newtags.tags = self.tags + newtags.tags
        return newtags

    def JSONrepr(self):
        return self.tags

class Template(object):
    props = {
        'AWSTemplateFormatVersion': (basestring, False),
        'Description': (basestring, False),
        'Parameters': (dict, False),
        'Mappings': (dict, False),
        'Resources': (dict, False),
        'Outputs': (dict, False),
    }

    def __init__(self):
        self.description = None
        self.metadata = {}
        self.conditions = {}
        self.mappings = {}
        self.outputs = {}
        self.parameters = {}
        self.resources = {}
        self.version = None

    def add_description(self, description):
        self.description = description

    def add_metadata(self, metadata):
        self.metadata = metadata

    def add_condition(self, name, condition):
        self.conditions[name] = condition

    def handle_duplicate_key(self, key):
        raise ValueError('duplicate key "%s" detected' % key)

    def _update(self, d, values):
        if isinstance(values, list):
            for v in values:
                if v.title in d:
                    self.handle_duplicate_key(v.title)
                d[v.title] = v
        else:
            if values.title in d:
                self.handle_duplicate_key(values.title)
            d[values.title] = values
        return values

    def add_output(self, output):
        return self._update(self.outputs, output)

    def add_mapping(self, name, mapping):
        self.mappings[name] = mapping

    def add_parameter(self, parameter):
        return self._update(self.parameters, parameter)

    def add_resource(self, resource):
        return self._update(self.resources, resource)

    def add_version(self, version=None):
        if version:
            self.version = version
        else:
            self.version = "2010-09-09"

    def to_json(self, indent=4, sort_keys=True, separators=(',', ': ')):
        t = {}
        if self.description:
            t['Description'] = self.description
        if self.metadata:
            t['Metadata'] = self.metadata
        if self.conditions:
            t['Conditions'] = self.conditions
        if self.mappings:
            t['Mappings'] = self.mappings
        if self.outputs:
            t['Outputs'] = self.outputs
        if self.parameters:
            t['Parameters'] = self.parameters
        if self.version:
            t['AWSTemplateFormatVersion'] = self.version
        t['Resources'] = self.resources

        return json.dumps(t, cls=awsencode, indent=indent,
                          sort_keys=sort_keys, separators=separators)

    def JSONrepr(self):
        return [self.parameters, self.mappings, self.resources]

class Output(AWSDeclaration):
    props = {
        'Description': (basestring, False),
        'Value': (basestring, True),
    }

class Parameter(AWSDeclaration):
    STRING_PROPERTIES = ['AllowedPattern', 'MaxLength', 'MinLength']
    NUMBER_PROPERTIES = ['MaxValue', 'MinValue']
    props = {
        'Type': (basestring, True),
        'Default': (basestring, False),
        'NoEcho': (bool, False),
        'AllowedValues': (list, False),
        'AllowedPattern': (basestring, False),
        'MaxLength': (validators.positive_integer, False),
        'MinLength': (validators.positive_integer, False),
        'MaxValue': (validators.integer, False),
        'MinValue': (validators.integer, False),
        'Description': (basestring, False),
        'ConstraintDescription': (basestring, False),
    }

    def validate(self):
        if self.properties['Type'] != 'String':
            for p in self.STRING_PROPERTIES:
                if p in self.properties:
                    raise ValueError("%s can only be used with parameters of "
                                     "the String type." % p)
        if self.properties['Type'] != 'Number':
            for p in self.NUMBER_PROPERTIES:
                if p in self.properties:
                    raise ValueError("%s can only be used with parameters of "
                                     "the Number type." % p)
<|endoftext|>


<commit_before>import threading
import time
import multiprocessing
import Queue
import sys
import os
#sys.path.append(os.path.abspath("../core"))
#from stream import *

# Add lock for scheduled_agents

class ComputeEngine(object):
    """
    Manages the queue of agents scheduled for execution.
    When an agent has new data on which to operate, the
    agent is placed on a queue called
    q_agents.

    Parameters
    ----------
    input_queue: multiprocessing.Queue
       Elements for input streams of this thread are put
       in this queue.
    name_to_stream: dict
       key: stream name
       value: stream

    Attributes
    ----------
    q_agents: Queue.Queue() or multiprocessing.Queue()
       The queue of agents scheduled for execution.
    scheduled_agents: Set
       An agent is in the set if and only if it is
       in the queue. This set is used to ensure that
       each agent appears at most once in the queue.

    """
    def __init__(self, name=None):
        self.input_queue = multiprocessing.Queue()
        self.name_to_stream = None
        self.name = name
        self.q_agents = Queue.Queue()
        self.scheduled_agents = set()
        self.compute_thread = None
        self.started = False
        self.stopped = False
        self.lock = threading.Lock()
        self.ready = threading.Event()

    def put(self, a):
        """
        Puts the agent a into q_agents if the
        agent is not already in the queue.

        Parameters
        ----------
        a : Agent

        """
        with self.lock:
            if a not in self.scheduled_agents:
                self.scheduled_agents.add(a)
                self.q_agents.put(a)

    def get(self):
        """
        Waits until q_agents is non-empty,
        and then gets and returns the agent at the head
        of the queue.

        Returns
        -------
        a: Agent
           The agent at the head of the queue of scheduled
           agents.

        """
        with self.lock:
            a = self.q_agents.get()
            self.scheduled_agents.discard(a)
            return a

    def start(self):
        def execute_computation():
            self.ready.set()
            nnn = 0
            while not self.stopped:
                input_queue_was_empty = False
                try:
                    v = self.input_queue.get(timeout=0.5)
                    #v = self.input_queue.get()
                except:
                    input_queue_was_empty = True
                    print 'input queue empty.'
                    # Sleep for SLEEP_TIME seconds
                    # Stop after LIMIT number of empty gets
                    SLEEP_TIME = 0.5
                    LIMIT = 5
                    nnn += 1
                    if nnn > LIMIT:
                        self.stopped = True
                    time.sleep(SLEEP_TIME)
                if not input_queue_was_empty:
                    if v == 'closed':
                        self.stopped = True
                        break
                    out_stream_name, new_data_for_stream = v
                    out_stream = self.name_to_stream[out_stream_name]
                    out_stream.append(new_data_for_stream)
                # Process the new data for this stream until all agents
                # finish their work.
                self.step()
        self.compute_thread = threading.Thread(
            target=execute_computation, args=())
        self.compute_thread.start()

    def stop(self):
        self.stopped = True
        self.compute_thread.join()
    def join(self):
        self.compute_thread.join()

    def step(self):
        while self.scheduled_agents:
            a = self.q_agents.get()
            # This next discard is necessary!! Check.
            self.scheduled_agents.discard(a)
            a.next()
        return

# Creates one scheduler from compute_engine.py
#scheduler = ComputeEngine()

<commit_msg>Changes from FINAL<commit_after>import threading
import time
import multiprocessing
import Queue
import sys
import os
#sys.path.append(os.path.abspath("../core"))
#from stream import *

# Add lock for scheduled_agents

class ComputeEngine(object):
    """
    Manages the queue of agents scheduled for execution.
    When an agent has new data on which to operate, the
    agent is placed on a queue called
    q_agents.

    Parameters
    ----------
    input_queue: multiprocessing.Queue
       Elements for input streams of this thread are put
       in this queue.
    name_to_stream: dict
       key: stream name
       value: stream

    Attributes
    ----------
    q_agents: Queue.Queue() or multiprocessing.Queue()
       The queue of agents scheduled for execution.
    scheduled_agents: Set
       An agent is in the set if and only if it is
       in the queue. This set is used to ensure that
       each agent appears at most once in the queue.

    """
    def __init__(self, name=None):
        self.input_queue = multiprocessing.Queue()
        self.name_to_stream = {}
        self.name = name
        self.q_agents = Queue.Queue()
        self.scheduled_agents = set()
        self.compute_thread = None
        self.started = False
        self.stopped = False
        self.lock = threading.Lock()
        self.ready = threading.Event()

    def put(self, a):
        """
        Puts the agent a into q_agents if the
        agent is not already in the queue.

        Parameters
        ----------
        a : Agent

        """
        with self.lock:
            if a not in self.scheduled_agents:
                self.scheduled_agents.add(a)
                self.q_agents.put(a)

    def get(self):
        """
        Waits until q_agents is non-empty,
        and then gets and returns the agent at the head
        of the queue.

        Returns
        -------
        a: Agent
           The agent at the head of the queue of scheduled
           agents.

        """
        with self.lock:
            a = self.q_agents.get()
            self.scheduled_agents.discard(a)
            return a

    def start(self):
        def execute_computation():
            self.ready.set()
            nnn = 0
            while not self.stopped:
                input_queue_was_empty = False
                try:
                    v = self.input_queue.get(timeout=0.5)
                    #v = self.input_queue.get()
                except:
                    input_queue_was_empty = True
                    print 'input queue empty.'
                    # Sleep for SLEEP_TIME seconds
                    # Stop after LIMIT number of empty gets
                    SLEEP_TIME = 0.5
                    LIMIT = 5
                    nnn += 1
                    if nnn > LIMIT:
                        self.stopped = True
                    time.sleep(SLEEP_TIME)
                if not input_queue_was_empty:
                    if v == 'closed':
                        self.stopped = True
                        break
                    out_stream_name, new_data_for_stream = v
                    out_stream = self.name_to_stream[out_stream_name]
                    out_stream.append(new_data_for_stream)
                # Process the new data for this stream until all agents
                # finish their work.
                self.step()
        self.compute_thread = threading.Thread(
            target=execute_computation, args=())
        self.compute_thread.start()

    def stop(self):
        self.stopped = True
        self.compute_thread.join()
    def join(self):
        self.compute_thread.join()

    def step(self):
        while self.scheduled_agents:
            a = self.q_agents.get()
            # This next discard is necessary!! Check.
            self.scheduled_agents.discard(a)
            a.next()
        return

# Creates one scheduler from compute_engine.py
#scheduler = ComputeEngine()

<|endoftext|>


<commit_before>import unittest

from difference_of_squares import difference, square_of_sum, sum_of_squares

# Tests adapted from `problem-specifications//canonical-data.json` @ v1.1.0

class DifferenceOfSquaresTest(unittest.TestCase):
    def test_square_of_sum_1(self):
        self.assertEqual(square_of_sum(1), 1)

    def test_square_of_sum_5(self):
        self.assertEqual(square_of_sum(5), 225)

    def test_square_of_sum_100(self):
        self.assertEqual(square_of_sum(100), 25502500)

    def test_sum_of_squares_1(self):
        self.assertEqual(sum_of_squares(1), 1)

    def test_sum_of_squares_5(self):
        self.assertEqual(sum_of_squares(5), 55)

    def test_sum_of_squares_100(self):
        self.assertEqual(sum_of_squares(100), 338350)

    def test_difference_of_squares_1(self):
        self.assertEqual(difference(1), 0)

    def test_difference_of_squares_5(self):
        self.assertEqual(difference(5), 170)

    def test_difference_of_squares_100(self):
        self.assertEqual(difference(100), 25164150)

if __name__ == '__main__':
    unittest.main()
<commit_msg>difference-of-squares: update tests to v1.2.0 (#1277)<commit_after>import unittest

from difference_of_squares import difference, square_of_sum, sum_of_squares

# Tests adapted from `problem-specifications//canonical-data.json` @ v1.2.0

class DifferenceOfSquaresTest(unittest.TestCase):
    def test_square_of_sum_1(self):
        self.assertEqual(square_of_sum(1), 1)

    def test_square_of_sum_5(self):
        self.assertEqual(square_of_sum(5), 225)

    def test_square_of_sum_100(self):
        self.assertEqual(square_of_sum(100), 25502500)

    def test_sum_of_squares_1(self):
        self.assertEqual(sum_of_squares(1), 1)

    def test_sum_of_squares_5(self):
        self.assertEqual(sum_of_squares(5), 55)

    def test_sum_of_squares_100(self):
        self.assertEqual(sum_of_squares(100), 338350)

    def test_difference_of_squares_1(self):
        self.assertEqual(difference(1), 0)

    def test_difference_of_squares_5(self):
        self.assertEqual(difference(5), 170)

    def test_difference_of_squares_100(self):
        self.assertEqual(difference(100), 25164150)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>