<commit_before>from __future__ import print_function
import os, sys
# we assume (and assert) that this script is running from the virus directory, i.e. inside H7N9 or zika
sys.path.append(os.path.join(os.path.dirname(__file__), '../..'))
import base.process
from base.process import process
import argparse
from dengue_titers import titer_model, titer_export

##### Define clades #####
genotypes = {
'denv1': {'I': [('E', 461, 'V'), ('E', 484, 'L'), ('M', 107, 'T')],
           'II': [('E', 345, 'A'), ('E', 432, 'M'), ('E', 439, 'V')],
           'IV': [('E', 339, 'S'), ('M', 72, 'E'), ('E', 88, 'T')],
           'V': [('E', 297, 'T'), ('NS5', 135, 'M')]},
 'denv2': {'AMERICAN': [('E', 71, 'D'), ('E', 81, 'T'), ('E', 129, 'I')],
           'ASIANAMERICAN': [('E', 491, 'A'), ('M', 15, 'G'), ('M', 39, 'I')],
           'ASIANI': [('E', 484, 'I'), ('NS5', 688, 'I'), ('NS1', 222, 'N')],
           'COSMOPOLITAN': [('E', 71, 'A'), ('E', 149, 'N'), ('E', 462, 'V')],
           'SYLVATIC': [('E', 59, 'F'), ('E', 236, 'M'), ('E', 432, 'V')]},
 'denv3': {'I': [('E', 233, 'K'), ('M', 128, 'F'), ('E', 68, 'V')],
           'II': [('M', 57, 'A'), ('NS5', 750, 'K')],
           'III': [('E', 303, 'T'), ('E', 454, 'V'), ('E', 132, 'Y')],
           'IV': [('E', 22, 'E'), ('E', 50, 'V'), ('E', 62, 'G')]},
 'denv4': {'I': [('E', 494, 'H'), ('NS1', 5, 'A')],
           'II': [('E', 265, 'A'), ('E', 46, 'T'), ('NS1', 246, 'S')],
           'SYLVATIC': [('E', 132, 'V'), ('E', 154, 'S'), ('E', 162, 'T')]},
 'all': {}
 }

for i in ['denv1', 'denv2', 'denv3', 'denv4']:
    for k,v in genotypes[i].iteritems():
        genotypes['all'][i.upper()+'_'+k] = v

regions = ['africa', 'europe', 'north_america', 'china', 'south_asia',
            'japan_korea', 'south_pacific', 'oceania', 'south_america',
            'southeast_asia', 'west_asia']

def collect_args():
    """Returns a dengue-specific argument parser.
    """
    parser = base.process.collect_args()

    # parser.add_argument('-j', '--json', default=None, nargs='+', type=str, help="Accepts path to prepared JSON(s); overrides -s argument")
    parser.add_argument('-s', '--serotypes', default=["multiple"], nargs='+', type=str, choices=['denv1', 'denv2', 'denv3', 'denv4', 'all', 'multiple'],
    help="Look for prepared JSON(s) like ./prepared/dengue_SEROTYPE.json; 'multiple' will run all five builds. Default='multiple'")
    parser.add_argument('--no_mut_freqs', default=False, action='store_true', help="skip mutation frequencies")
    parser.add_argument('--no_tree_freqs', default=False, action='store_true', help="skip tree (clade) frequencies")
    parser.add_argument('--no_titers', default=False, action='store_true', help="skip titer models")
    parser.set_defaults(json = None)

    return parser

def make_config (prepared_json, args):
    return {
        "dir": "dengue",
        "in": prepared_json,
        "geo_inference": ['region'], # what traits to perform this on
        "auspice": { ## settings for auspice JSON export
            "extra_attr": ['serum'],
            "color_options": {
                # "country":{"key":"country", "legendTitle":"Country", "menuItem":"country", "type":"discrete"},
                "region":{"key":"region", "legendTitle":"Region", "menuItem":"region", "type":"discrete"},
                "gt":{"key":"genotype", "legendTitle":"Genotype", "menuItem":"genotype", "type":"discrete"},
            },
            "controls": {'authors':['authors']},
            "defaults": {'geoResolution': 'region', 'colorBy': 'region', 'distanceMeasure': 'div', 'mapTriplicate': True}
            },

        "timetree_options": {"Tc": False},
        "fit_titer_model": not args.no_titers,
        "titers": {
            "lam_avi":3.0,
            "lam_pot":0.5,
            "lam_drop":1.0,
            "training_fraction":0.9
        },
        "estimate_mutation_frequencies": not args.no_mut_freqs,
        "estimate_tree_frequencies": not args.no_tree_freqs,
        "clean": args.clean,
        "pivot_spacing": 1.0/4, # is this n timepoints per year?,
        "newick_tree_options":{
            "raxml": not args.no_raxml
        }
    }

if __name__=="__main__":
    parser = collect_args()
    args = parser.parse_args()

    if args.json: # Specified JSON path overrides serotype argument
        args.json = [args.json]
    else: # Look for serotype-specific JSONs in the ./prepared/ directory
        if 'multiple' in args.serotypes: # "multiple" = run all 5 builds
            args.serotypes = ['denv1', 'denv2', 'denv3', 'denv4', 'all']
        else:
            args.serotypes = args.serotypes
        args.json = ['./prepared/dengue_%s.json'%s for s in args.serotypes] # Look for ./prepared/dengue_SEROTYPE.json if no file paths given

    for j in args.json:            # validate input JSONs exist
        assert os.path.isfile(j)

    for prepared_json in args.json:
        print("Processing %s"%prepared_json)
        runner = process(make_config(prepared_json, args))
        runner.align()
        runner.build_tree()
        runner.timetree_setup_filter_run()
        runner.run_geo_inference()

        # estimate mutation frequencies here.
        if runner.config["estimate_mutation_frequencies"]:
            pivots = runner.get_pivots_via_spacing()
            runner.estimate_mutation_frequencies(pivots=pivots, min_freq=0.02, inertia=np.exp(-1.0/12), stiffness=0.8*12)

        # estimate tree frequencies here.
        if runner.config["estimate_tree_frequencies"]:
            pivots = runner.get_pivots_via_spacing()
            runner.estimate_tree_frequencies(pivots=pivots, stiffness=2)

            for region in ['southeast_asia', 'south_america']: #regions:
                try:
                    runner.estimate_tree_frequencies(region=region, stiffness=2)
                except:
                    continue
        # titers
        if runner.config["fit_titer_model"] and runner.config["titers"]:
            titer_model(runner,
                        lam_pot = runner.config['titers']['lam_pot'],
                        lam_avi = runner.config['titers']['lam_avi'],
                        lam_drop = runner.config['titers']['lam_drop'],
                        training_fraction = runner.config['titers']['training_fraction'])
            titer_export(runner)

        # runner.matchClades(genotypes[runner.info['lineage']])

        runner.auspice_export()
<commit_msg>Turn off mutation frequencies<commit_after>from __future__ import print_function
import os, sys
# we assume (and assert) that this script is running from the virus directory, i.e. inside H7N9 or zika
sys.path.append(os.path.join(os.path.dirname(__file__), '../..'))
import base.process
from base.process import process
import argparse
from dengue_titers import titer_model, titer_export

##### Define clades #####
genotypes = {
'denv1': {'I': [('E', 461, 'V'), ('E', 484, 'L'), ('M', 107, 'T')],
           'II': [('E', 345, 'A'), ('E', 432, 'M'), ('E', 439, 'V')],
           'IV': [('E', 339, 'S'), ('M', 72, 'E'), ('E', 88, 'T')],
           'V': [('E', 297, 'T'), ('NS5', 135, 'M')]},
 'denv2': {'AMERICAN': [('E', 71, 'D'), ('E', 81, 'T'), ('E', 129, 'I')],
           'ASIANAMERICAN': [('E', 491, 'A'), ('M', 15, 'G'), ('M', 39, 'I')],
           'ASIANI': [('E', 484, 'I'), ('NS5', 688, 'I'), ('NS1', 222, 'N')],
           'COSMOPOLITAN': [('E', 71, 'A'), ('E', 149, 'N'), ('E', 462, 'V')],
           'SYLVATIC': [('E', 59, 'F'), ('E', 236, 'M'), ('E', 432, 'V')]},
 'denv3': {'I': [('E', 233, 'K'), ('M', 128, 'F'), ('E', 68, 'V')],
           'II': [('M', 57, 'A'), ('NS5', 750, 'K')],
           'III': [('E', 303, 'T'), ('E', 454, 'V'), ('E', 132, 'Y')],
           'IV': [('E', 22, 'E'), ('E', 50, 'V'), ('E', 62, 'G')]},
 'denv4': {'I': [('E', 494, 'H'), ('NS1', 5, 'A')],
           'II': [('E', 265, 'A'), ('E', 46, 'T'), ('NS1', 246, 'S')],
           'SYLVATIC': [('E', 132, 'V'), ('E', 154, 'S'), ('E', 162, 'T')]},
 'all': {}
 }

for i in ['denv1', 'denv2', 'denv3', 'denv4']:
    for k,v in genotypes[i].iteritems():
        genotypes['all'][i.upper()+'_'+k] = v

regions = ['africa', 'europe', 'north_america', 'china', 'south_asia',
            'japan_korea', 'south_pacific', 'oceania', 'south_america',
            'southeast_asia', 'west_asia']

def collect_args():
    """Returns a dengue-specific argument parser.
    """
    parser = base.process.collect_args()

    # parser.add_argument('-j', '--json', default=None, nargs='+', type=str, help="Accepts path to prepared JSON(s); overrides -s argument")
    parser.add_argument('-s', '--serotypes', default=["multiple"], nargs='+', type=str, choices=['denv1', 'denv2', 'denv3', 'denv4', 'all', 'multiple'],
    help="Look for prepared JSON(s) like ./prepared/dengue_SEROTYPE.json; 'multiple' will run all five builds. Default='multiple'")
    parser.add_argument('--no_mut_freqs', default=True, action='store_true', help="skip mutation frequencies")
    parser.add_argument('--no_tree_freqs', default=False, action='store_true', help="skip tree (clade) frequencies")
    parser.add_argument('--no_titers', default=False, action='store_true', help="skip titer models")
    parser.set_defaults(json = None)
    return parser

def make_config (prepared_json, args):
    return {
        "dir": "dengue",
        "in": prepared_json,
        "geo_inference": ['region'], # what traits to perform this on
        "auspice": { ## settings for auspice JSON export
            "extra_attr": ['serum'],
            "color_options": {
                # "country":{"key":"country", "legendTitle":"Country", "menuItem":"country", "type":"discrete"},
                "region":{"key":"region", "legendTitle":"Region", "menuItem":"region", "type":"discrete"},
                "gt":{"key":"genotype", "legendTitle":"Genotype", "menuItem":"genotype", "type":"discrete"},
            },
            "controls": {'authors':['authors']},
            "defaults": {'geoResolution': 'region', 'colorBy': 'region', 'distanceMeasure': 'div', 'mapTriplicate': True}
            },

        "timetree_options": {"Tc": False},
        "fit_titer_model": not args.no_titers,
        "titers": {
            "lam_avi":3.0,
            "lam_pot":0.5,
            "lam_drop":1.0,
            "training_fraction":0.9
        },
        "estimate_mutation_frequencies": not args.no_mut_freqs,
        "estimate_tree_frequencies": not args.no_tree_freqs,
        "clean": args.clean,
        "pivot_spacing": 1.0/4, # is this n timepoints per year?,
        "newick_tree_options":{
            "raxml": not args.no_raxml
        }
    }

if __name__=="__main__":
    parser = collect_args()
    args = parser.parse_args()

    if args.json: # Specified JSON path overrides serotype argument
        args.json = [args.json]
    else: # Look for serotype-specific JSONs in the ./prepared/ directory
        if 'multiple' in args.serotypes: # "multiple" = run all 5 builds
            args.serotypes = ['denv1', 'denv2', 'denv3', 'denv4', 'all']
        else:
            args.serotypes = args.serotypes
        args.json = ['./prepared/dengue_%s.json'%s for s in args.serotypes] # Look for ./prepared/dengue_SEROTYPE.json if no file paths given

    for j in args.json:            # validate input JSONs exist
        assert os.path.isfile(j)

    for prepared_json in args.json:
        print("Processing %s"%prepared_json)
        runner = process(make_config(prepared_json, args))
        runner.align()
        runner.build_tree()
        runner.timetree_setup_filter_run()
        runner.run_geo_inference()

        # estimate mutation frequencies here.
        if runner.config["estimate_mutation_frequencies"]:
            pivots = runner.get_pivots_via_spacing()
            runner.estimate_mutation_frequencies(pivots=pivots, min_freq=0.02, inertia=np.exp(-1.0/12), stiffness=0.8*12)

        # estimate tree frequencies here.
        if runner.config["estimate_tree_frequencies"]:
            pivots = runner.get_pivots_via_spacing()
            runner.estimate_tree_frequencies(pivots=pivots, stiffness=2)

            for region in ['southeast_asia', 'south_america']: #regions:
                try:
                    runner.estimate_tree_frequencies(region=region, stiffness=2)
                except:
                    continue
        # titers
        if runner.config["fit_titer_model"] and runner.config["titers"]:
            titer_model(runner,
                        lam_pot = runner.config['titers']['lam_pot'],
                        lam_avi = runner.config['titers']['lam_avi'],
                        lam_drop = runner.config['titers']['lam_drop'],
                        training_fraction = runner.config['titers']['training_fraction'])
            titer_export(runner)

        # runner.matchClades(genotypes[runner.info['lineage']])

        runner.auspice_export()
<|endoftext|>


<commit_before>from __future__ import print_function
import os, sys
# we assume (and assert) that this script is running from the virus directory, i.e. inside H7N9 or zika
sys.path.append(os.path.join(os.path.dirname(__file__), '../..'))
import base.process
from base.process import process
import argparse
import numpy as np
from dengue_titers import titer_model, titer_export ## Set up and parameterize the titer model separately for tidiness

##### Define references and metadata #####

regions = ['africa', 'europe', 'north_america', 'china', 'south_asia',
            'japan_korea', 'south_pacific', 'oceania', 'south_america',
            'southeast_asia', 'west_asia']

##### Parse args, set up config #####
def collect_args():
    """Returns a dengue-specific argument parser."""
    parser = base.process.collect_args()

    # parser.add_argument('-j', '--json', default=None, nargs='+', type=str, help="Accepts path to prepared JSON(s); overrides -s argument")
    parser.add_argument('-s', '--serotypes', default=["multiple"], nargs='+', type=str, choices=['denv1', 'denv2', 'denv3', 'denv4', 'all', 'multiple'],
    help="Look for prepared JSON(s) like ./prepared/dengue_SEROTYPE.json; 'multiple' will run all five builds. Default='multiple'")
    parser.add_argument('--no_mut_freqs', default=True, action='store_true', help="skip mutation frequencies")
    parser.add_argument('--no_tree_freqs', default=False, action='store_true', help="skip tree (clade) frequencies")
    parser.add_argument('--no_titers', default=False, action='store_true', help="skip titer models")
    parser.set_defaults(json = None)
    return parser

def make_config (prepared_json, args):
    """
    Configure your analysis here.
    Parsed as a function to enable running multiple builds with one cmd.
    """
    return {
        "dir": "dengue",
        "in": prepared_json,
        "geo_inference": ['region'], # what traits to perform this on; don't run country (too many demes, too few sequences per deme to be reliable)
        "auspice": { ## settings for auspice JSON export
            "extra_attr": ['serum'],
            "color_options": { # which traits to color the tree by in auspice; titer colorbys are added in dengue_titers
                "country":{"key":"country", "legendTitle":"Country", "menuItem":"country", "type":"discrete"},
                "region":{"key":"region", "legendTitle":"Region", "menuItem":"region", "type":"discrete"},
                "gt":{"key":"genotype", "legendTitle":"Genotype", "menuItem":"genotype", "type":"discrete"},
            },
            "controls": {'authors':['authors']},
            "defaults": {'geoResolution': 'region', 'colorBy': 'region', 'distanceMeasure': 'div', 'mapTriplicate': True}
            },

        "timetree_options": {"Tc": False},
        "fit_titer_model": not args.no_titers,
        "titers": { # regularization parameter values and cross-validation fraction
            "lam_avi":0.0,
            "lam_pot":0.5,
            "lam_drop":1.0,
            "training_fraction":0.9,
        },
        "estimate_mutation_frequencies": not args.no_mut_freqs,
        "estimate_tree_frequencies": not args.no_tree_freqs,
        "clean": args.clean,
        "pivot_spacing": 1.0/4, # pivots = time points; 1/N timepoints per year
        "newick_tree_options":{
            "raxml": not args.no_raxml # for dev work only
        }
    }

##### Parse input files/params and run #####
if __name__=="__main__":
    parser = collect_args()
    args = parser.parse_args()

    ### Find the right input files ###
    if args.json: # If provided, a specified JSON path overrides serotype argument
        args.json = [args.json]
    else: # Look for serotype-specific JSONs in the ./prepared/ directory
        if 'multiple' in args.serotypes: # "multiple" = run all 5 builds
            args.serotypes = ['denv1', 'denv2', 'denv3', 'denv4', 'all']
        else:
            args.serotypes = args.serotypes
        args.json = ['./prepared/dengue_%s.json'%s for s in args.serotypes] # Look for ./prepared/dengue_SEROTYPE.json if no file paths given

    for j in args.json:            # validate input JSONs exist
        assert os.path.isfile(j)

    ### Run analyses ###
    for prepared_json in args.json:
        try:
            print("Processing %s"%prepared_json)
            runner = process(make_config(prepared_json, args)) # parse
            runner.align() # run alignment with mafft
            runner.build_tree() # build tree with fasttree -> raxml
            runner.timetree_setup_filter_run() # infer ML ancestral states (geo traits, node dates, mutations)
            runner.run_geo_inference() # run mugration model to infer transmissions

            # estimate mutation frequencies here.
            if runner.config["estimate_mutation_frequencies"]:
                pivots = runner.get_pivots_via_spacing()
                runner.estimate_mutation_frequencies(pivots=pivots, min_freq=0.02, inertia=np.exp(-1.0/12), stiffness=2)

            # estimate tree frequencies here.
            if runner.config["estimate_tree_frequencies"]: # methods @ [ref]
                pivots = runner.get_pivots_via_spacing()
                runner.estimate_tree_frequencies(pivots=pivots, stiffness=2) # stiffness ~= amount of smoothing

                for region in ['southeast_asia', 'south_america']: #regions:
                    try:
                    runner.estimate_tree_frequencies(region=region, stiffness=2)
                    except:
                        continue
            # titers
            if runner.config["fit_titer_model"] and runner.config["titers"]: # methods @ Neher et al., PNAS 2016
                titer_model(runner,
                            lam_pot = runner.config['titers']['lam_pot'],
                            lam_avi = runner.config['titers']['lam_avi'],
                            lam_drop = runner.config['titers']['lam_drop'],
                            training_fraction = runner.config['titers']['training_fraction'],
                            plot=False,
                            criterium = lambda node: True) # calculate dTiter for all branches
                titer_export(runner)

            ### Export for visualization in auspice
            runner.auspice_export()

        except:
            continue

##### Extra code bank #####
'''
genotypes = {
'denv1': {'I': [('E', 461, 'V'), ('E', 484, 'L'), ('M', 107, 'T')],
           'II': [('E', 345, 'A'), ('E', 432, 'M'), ('E', 439, 'V')],
           'IV': [('E', 339, 'S'), ('M', 72, 'E'), ('E', 88, 'T')],
           'V': [('E', 297, 'T'), ('NS5', 135, 'M')]},
 'denv2': {'AMERICAN': [('E', 71, 'D'), ('E', 81, 'T'), ('E', 129, 'I')],
           'ASIANAMERICAN': [('E', 491, 'A'), ('M', 15, 'G'), ('M', 39, 'I')],
           'ASIANI': [('E', 484, 'I'), ('NS5', 688, 'I'), ('NS1', 222, 'N')],
           'COSMOPOLITAN': [('E', 71, 'A'), ('E', 149, 'N'), ('E', 462, 'V')],
           'SYLVATIC': [('E', 59, 'F'), ('E', 236, 'M'), ('E', 432, 'V')]},
 'denv3': {'I': [('E', 233, 'K'), ('M', 128, 'F'), ('E', 68, 'V')],
           'II': [('M', 57, 'A'), ('NS5', 750, 'K')],
           'III': [('E', 303, 'T'), ('E', 454, 'V'), ('E', 132, 'Y')],
           'IV': [('E', 22, 'E'), ('E', 50, 'V'), ('E', 62, 'G')]},
 'denv4': {'I': [('E', 494, 'H'), ('NS1', 5, 'A')],
           'II': [('E', 265, 'A'), ('E', 46, 'T'), ('NS1', 246, 'S')],
           'SYLVATIC': [('E', 132, 'V'), ('E', 154, 'S'), ('E', 162, 'T')]},
 'all': {}
 }

for i in ['denv1', 'denv2', 'denv3', 'denv4']:
    for k,v in genotypes[i].iteritems():
        genotypes['all'][i.upper()+'_'+k] = v

# Label named clades based on mutations/genotypes at defining sites
runner.matchClades(genotypes[runner.info['lineage']])
# this is tricky with dengue because the canonical genotypes
# don't really represent the present-day viral diversity.
# I'll get around to redefining these soon-ish hopefully.

### Comparison: force dTiter values to be non-zero only on interserotype brances
    def is_interserotype(node):
        descendents = node.get_terminals()
        serotypes = [k.name.split('/')[0] for k in descendents if 'DENV' in k.name]
        serotypes = [s for s in serotypes if s != 'DENV']
        return len(set(serotypes)) > 1

    interserotype_branches = []
    for node in runner.tree.tree.find_clades():
        if is_interserotype(node):
            interserotype_branches.append(node)
            for child in node.clades:
                interserotype_branches.append(child)
    for node in runner.tree.tree.find_clades():
        if node in interserotype_branches:
            node.interserotype = True
        else:
            node.interserotype = False

    titer_model(runner,
                lam_pot = runner.config['titers']['lam_pot'],
                lam_avi = runner.config['titers']['lam_avi'],
                lam_drop = runner.config['titers']['lam_drop'],
                training_fraction = runner.config['titers']['training_fraction'],
                plot=False,
                criterium = lambda node: node.interserotype == True,
                csv_fname='~/Users/Sidney/Dropbox/dengue/data/titer-model/interserotype-branch-effects/model_predictions.csv')
'''
<commit_msg>Calc and export sanofi vaccine strain dTiter vals<commit_after>from __future__ import print_function
import os, sys
# we assume (and assert) that this script is running from the virus directory, i.e. inside H7N9 or zika
sys.path.append(os.path.join(os.path.dirname(__file__), '../..'))
import base.process
from base.process import process
import argparse
import numpy as np
from dengue_titers import titer_model, titer_export ## Set up and parameterize the titer model separately for tidiness

##### Define references and metadata #####
sanofi_vaccine_strains = {
    'denv1': 'DENV1/THAILAND/PUO359/1980',
    'denv2': 'DENV2/THAILAND/PUO218/1980',
    'denv3': 'DENV3/THAILAND/PAH88188/1988',
    'denv4': 'DENV4/INDONESIA/S1228/1978',
    'all': None}

regions = ['africa', 'europe', 'north_america', 'china', 'south_asia',
            'japan_korea', 'south_pacific', 'oceania', 'south_america',
            'southeast_asia', 'west_asia']

##### Parse args, set up config #####
def collect_args():
    """Returns a dengue-specific argument parser."""
    parser = base.process.collect_args()

    # parser.add_argument('-j', '--json', default=None, nargs='+', type=str, help="Accepts path to prepared JSON(s); overrides -s argument")
    parser.add_argument('-s', '--serotypes', default=["multiple"], nargs='+', type=str, choices=['denv1', 'denv2', 'denv3', 'denv4', 'all', 'multiple'],
    help="Look for prepared JSON(s) like ./prepared/dengue_SEROTYPE.json; 'multiple' will run all five builds. Default='multiple'")
    parser.add_argument('--no_mut_freqs', default=True, action='store_true', help="skip mutation frequencies")
    parser.add_argument('--no_tree_freqs', default=False, action='store_true', help="skip tree (clade) frequencies")
    parser.add_argument('--no_titers', default=False, action='store_true', help="skip titer models")
    parser.set_defaults(json = None)
    return parser

def make_config (prepared_json, args):
    """
    Configure your analysis here.
    Parsed as a function to enable running multiple builds with one cmd.
    """
    return {
        "dir": "dengue",
        "in": prepared_json,
        "geo_inference": ['region'], # what traits to perform this on; don't run country (too many demes, too few sequences per deme to be reliable)
        "auspice": { ## settings for auspice JSON export
            "extra_attr": ['serum', 'clade', 'dTiter_sanofi'], # keys from tree.tree.clade['attr'] to include in export
            "color_options": { # which traits to color the tree by in auspice; titer colorbys are added in dengue_titers
                "country":{"key":"country", "legendTitle":"Country", "menuItem":"country", "type":"discrete"},
                "region":{"key":"region", "legendTitle":"Region", "menuItem":"region", "type":"discrete"},
                "gt":{"key":"genotype", "legendTitle":"Genotype", "menuItem":"genotype", "type":"discrete"},
            },
            "controls": {'authors':['authors']},
            "defaults": {'geoResolution': 'region', 'colorBy': 'region', 'distanceMeasure': 'div', 'mapTriplicate': True}
            },

        "timetree_options": {"Tc": False},
        "fit_titer_model": not args.no_titers,
        "titers": { # regularization parameter values and cross-validation fraction
            "lam_avi":0.0,
            "lam_pot":0.5,
            "lam_drop":1.0,
            "training_fraction":0.9,
        },
        "estimate_mutation_frequencies": not args.no_mut_freqs,
        "estimate_tree_frequencies": not args.no_tree_freqs,
        "clean": args.clean,
        "pivot_spacing": 1.0/4, # pivots = time points; 1/N timepoints per year
        "newick_tree_options":{
            "raxml": not args.no_raxml # for dev work only
        }
    }

##### Parse input files/params and run #####
if __name__=="__main__":
    parser = collect_args()
    args = parser.parse_args()

    ### Find the right input files ###
    if args.json: # If provided, a specified JSON path overrides serotype argument
        args.json = [args.json]
    else: # Look for serotype-specific JSONs in the ./prepared/ directory
        if 'multiple' in args.serotypes: # "multiple" = run all 5 builds
            args.serotypes = ['denv1', 'denv2', 'denv3', 'denv4', 'all']
        else:
            args.serotypes = args.serotypes
        args.json = ['./prepared/dengue_%s.json'%s for s in args.serotypes] # Look for ./prepared/dengue_SEROTYPE.json if no file paths given

    for j in args.json:            # validate input JSONs exist
        assert os.path.isfile(j)

    ### Run analyses ###
    for prepared_json in args.json:
        try:
            print("Processing %s"%prepared_json)
            runner = process(make_config(prepared_json, args)) # parse
            runner.align() # run alignment with mafft
            runner.build_tree() # build tree with fasttree -> raxml
            runner.timetree_setup_filter_run() # infer ML ancestral states (geo traits, node dates, mutations)
            runner.run_geo_inference() # run mugration model to infer transmissions

            # estimate mutation frequencies here.
            if runner.config["estimate_mutation_frequencies"]:
                pivots = runner.get_pivots_via_spacing()
                runner.estimate_mutation_frequencies(pivots=pivots, min_freq=0.02, inertia=np.exp(-1.0/12), stiffness=2)

            # estimate tree frequencies here.
            if runner.config["estimate_tree_frequencies"]: # methods @ [ref]
                pivots = runner.get_pivots_via_spacing()
                runner.estimate_tree_frequencies(pivots=pivots, stiffness=2) # stiffness ~= amount of smoothing

                for region in ['southeast_asia', 'south_america']: #regions:
                    try:
                        runner.estimate_tree_frequencies(region=region, stiffness=2)
                    except:
                        continue
            # titers
            if runner.config["fit_titer_model"] and runner.config["titers"]: # methods @ Neher et al., PNAS 2016
                titer_model(runner,
                            lam_pot = runner.config['titers']['lam_pot'],
                            lam_avi = runner.config['titers']['lam_avi'],
                            lam_drop = runner.config['titers']['lam_drop'],
                            training_fraction = runner.config['titers']['training_fraction'],
                            sanofi_strain = sanofi_vaccine_strains[runner.info['lineage']], # vaccine strain for each serotype-specific build
                            plot=False,
                            criterium = lambda node: True) # calculate dTiter for all branches
                titer_export(runner)

            ### Export for visualization in auspice
            runner.auspice_export()

        except:
            continue

##### Extra code bank #####
'''
genotypes = {
'denv1': {'I': [('E', 461, 'V'), ('E', 484, 'L'), ('M', 107, 'T')],
           'II': [('E', 345, 'A'), ('E', 432, 'M'), ('E', 439, 'V')],
           'IV': [('E', 339, 'S'), ('M', 72, 'E'), ('E', 88, 'T')],
           'V': [('E', 297, 'T'), ('NS5', 135, 'M')]},
 'denv2': {'AMERICAN': [('E', 71, 'D'), ('E', 81, 'T'), ('E', 129, 'I')],
           'ASIANAMERICAN': [('E', 491, 'A'), ('M', 15, 'G'), ('M', 39, 'I')],
           'ASIANI': [('E', 484, 'I'), ('NS5', 688, 'I'), ('NS1', 222, 'N')],
           'COSMOPOLITAN': [('E', 71, 'A'), ('E', 149, 'N'), ('E', 462, 'V')],
           'SYLVATIC': [('E', 59, 'F'), ('E', 236, 'M'), ('E', 432, 'V')]},
 'denv3': {'I': [('E', 233, 'K'), ('M', 128, 'F'), ('E', 68, 'V')],
           'II': [('M', 57, 'A'), ('NS5', 750, 'K')],
           'III': [('E', 303, 'T'), ('E', 454, 'V'), ('E', 132, 'Y')],
           'IV': [('E', 22, 'E'), ('E', 50, 'V'), ('E', 62, 'G')]},
 'denv4': {'I': [('E', 494, 'H'), ('NS1', 5, 'A')],
           'II': [('E', 265, 'A'), ('E', 46, 'T'), ('NS1', 246, 'S')],
           'SYLVATIC': [('E', 132, 'V'), ('E', 154, 'S'), ('E', 162, 'T')]},
 'all': {}
 }

for i in ['denv1', 'denv2', 'denv3', 'denv4']:
    for k,v in genotypes[i].iteritems():
        genotypes['all'][i.upper()+'_'+k] = v

# Label named clades based on mutations/genotypes at defining sites
runner.matchClades(genotypes[runner.info['lineage']])
# this is tricky with dengue because the canonical genotypes
# don't really represent the present-day viral diversity.
# I'll get around to redefining these soon-ish hopefully.

### Comparison: force dTiter values to be non-zero only on interserotype brances
    def is_interserotype(node):
        descendents = node.get_terminals()
        serotypes = [k.name.split('/')[0] for k in descendents if 'DENV' in k.name]
        serotypes = [s for s in serotypes if s != 'DENV']
        return len(set(serotypes)) > 1

    interserotype_branches = []
    for node in runner.tree.tree.find_clades():
        if is_interserotype(node):
            interserotype_branches.append(node)
            for child in node.clades:
                interserotype_branches.append(child)
    for node in runner.tree.tree.find_clades():
        if node in interserotype_branches:
            node.interserotype = True
        else:
            node.interserotype = False

    titer_model(runner,
                lam_pot = runner.config['titers']['lam_pot'],
                lam_avi = runner.config['titers']['lam_avi'],
                lam_drop = runner.config['titers']['lam_drop'],
                training_fraction = runner.config['titers']['training_fraction'],
                plot=False,
                criterium = lambda node: node.interserotype == True,
                csv_fname='~/Users/Sidney/Dropbox/dengue/data/titer-model/interserotype-branch-effects/model_predictions.csv')
'''
<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
#################################################################################
#                                                                               #
#    Model module for OpenERP                                                   #
#    Copyright (C) 2010 Sébastien BEAU <sebastien.beau@akretion.comr>           #
#                                                                               #
#    This program is free software: you can redistribute it and/or modify       #
#    it under the terms of the GNU Affero General Public License as             #
#    published by the Free Software Foundation, either version 3 of the         #
#    License, or (at your option) any later version.                            #
#                                                                               #
#    This program is distributed in the hope that it will be useful,            #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of             #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              #
#    GNU Affero General Public License for more details.                        #
#                                                                               #
#    You should have received a copy of the GNU Affero General Public License   #
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.      #
#                                                                               #
#################################################################################

from osv import osv, fields
import netsvc

class ir_cron(osv.osv, netsvc.Agent):
    _inherit = "ir.cron"

    _columns = {
        'email_to' : fields.char('TO', size=256, help="If an error occure with this scheduler an email will be send"),
        'message' : fields.text('Message', help="If an error occure with this scheduler this message will be send via email"),
        'send_email' : fields.boolean('Active'),
        'email_account' : fields.many2one('poweremail.core_accounts', 'FROM')
    }

    def _handle_callback_exception(self, cr, uid, model, func, args, job_id, job_exception):
        super(ir_cron, self)._handle_callback_exception(cr, uid, model, func, args, job_id, job_exception)
        job = self.read(cr, uid, job_id, ['send_email?', 'message', 'email_TO', 'email_account', 'name'])
        #TODO USE POWEREMAIL TEMPLATE
        if job['send_email?']:
            addresses = {'To' : job['email_TO']}
            mail_obj = self.pool.get('poweremail.mailbox')
            id = mail_obj.create(cr, uid, {
                                            'pem_to' : job['email_TO'],
                                            'pem_subject' : "OPENERP : error when excecuting scheduler " + job["name"],
                                            'pem_body_text' : job['message'],
                                            'pem_account_id' : job['email_account'][0],
                                            'mail_type' : 'text/plain',
                                            'folder' : 'outbox',
                                            'state' :'na',
                                        })
            mail_obj.send_this_mail(cr, uid, [id])
        return res
ir_cron()

<commit_msg>[FIX] fix bug and port on v6<commit_after># -*- encoding: utf-8 -*-
#################################################################################
#                                                                               #
#    Model module for OpenERP                                                   #
#    Copyright (C) 2010 Sébastien BEAU <sebastien.beau@akretion.comr>           #
#                                                                               #
#    This program is free software: you can redistribute it and/or modify       #
#    it under the terms of the GNU Affero General Public License as             #
#    published by the Free Software Foundation, either version 3 of the         #
#    License, or (at your option) any later version.                            #
#                                                                               #
#    This program is distributed in the hope that it will be useful,            #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of             #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              #
#    GNU Affero General Public License for more details.                        #
#                                                                               #
#    You should have received a copy of the GNU Affero General Public License   #
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.      #
#                                                                               #
#################################################################################

from osv import osv, fields
import netsvc

class ir_cron(osv.osv, netsvc.Agent):
    _inherit = "ir.cron"

    _columns = {
        'email_to' : fields.char('TO', size=256, help="If an error occure with this scheduler an email will be send"),
        'message' : fields.text('Message', help="If an error occure with this scheduler this message will be send via email"),
        'send_email' : fields.boolean('Active'),
        'email_account' : fields.many2one('poweremail.core_accounts', 'FROM')
    }

    def _handle_callback_exception(self, cr, uid, model, func, args, job_id, job_exception):
        res = super(ir_cron, self)._handle_callback_exception(cr, uid, model, func, args, job_id, job_exception)
        job = self.read(cr, uid, job_id, ['send_email', 'message', 'email_TO', 'email_account', 'name'])
        #TODO USE POWEREMAIL TEMPLATE
        if job['send_email']:
            addresses = {'To' : job['email_TO']}
            mail_obj = self.pool.get('poweremail.mailbox')
            id = mail_obj.create(cr, uid, {
                                            'pem_to' : job['email_TO'],
                                            'pem_subject' : "OPENERP : error when excecuting scheduler " + job["name"],
                                            'pem_body_text' : job['message'],
                                            'pem_account_id' : job['email_account'][0],
                                            'mail_type' : 'text/plain',
                                            'folder' : 'outbox',
                                            'state' :'na',
                                        })
            mail_obj.send_this_mail(cr, uid, [id])
        return res
ir_cron()

<|endoftext|>


<commit_before>#
# Copyright (c) 2017, Arista Networks, Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#   Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
#
#   Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in the
#   documentation and/or other materials provided with the distribution.
#
#   Neither the name of Arista Networks nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL ARISTA NETWORKS
# BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
# IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
import sys
import os
import unittest

sys.path.append(os.path.join(os.path.dirname(__file__), '../lib'))

from testlib import get_fixture, function
from testlib import EapiConfigUnitTest

import pyeapi.api.vrfs

class TestApiVrfs(EapiConfigUnitTest):

    def __init__(self, *args, **kwargs):
        super(TestApiVrfs, self).__init__(*args, **kwargs)
        self.instance = pyeapi.api.vrfs.instance(None)
        self.config = open(get_fixture('running_config.vrf')).read()

    def test_get(self):
        result = self.instance.get('blah')
        vrf = dict(rd='10:10', vrf_name='blah', description='blah desc',
                   ipv4_routing=True, ipv6_routing=False)
        self.assertEqual(vrf, result)
        result2 = self.instance.get('test')
        vrf2 = dict(rd='200:500', vrf_name='test', description='!',
                    ipv4_routing=False, ipv6_routing=True)
        self.assertEqual(vrf2, result2)

    def test_get_not_configured(self):
        self.assertIsNone(self.instance.get('notthere'))

    def test_getall(self):
        result = self.instance.getall()
        self.assertIsInstance(result, dict)
        self.assertEqual(len(result), 3)

    def test_vrf_functions(self):
        for name in ['create', 'delete', 'default']:
            vrf_name = 'testvrf'
            if name == 'create':
                cmds = ['vrf definition %s' % vrf_name]
            elif name == 'delete':
                cmds = 'no vrf definition %s' % vrf_name
            elif name == 'default':
                cmds = 'default vrf definition %s' % vrf_name
            func = function(name, vrf_name)
            self.eapi_positive_config_test(func, cmds)

    def test_vrf_create_with_rd(self):
        vrf_name = 'testvrfrd'
        rd = '10:10'
        cmds = ['vrf definition %s' % vrf_name, 'rd %s' % rd]
        func = function('create', vrf_name, rd=rd)
        self.eapi_positive_config_test(func, cmds)

    def test_set_rd(self):
        vrf_name = 'testrdvrf'
        rd = '10:10'
        cmds = ['vrf definition %s' % vrf_name, 'rd %s' % rd]
        func = function('set_rd', vrf_name, rd)
        self.eapi_positive_config_test(func, cmds)

    def test_set_description(self):
        for state in ['config', 'negate', 'default']:
            vrf_name = 'testdescvrf'
            if state == 'config':
                description = 'testing'
                cmds = ['vrf definition %s' % vrf_name,
                        'description %s' % description]
                func = function('set_description', vrf_name, description)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'negate':
                cmds = ['vrf definition %s' % vrf_name, 'no description']
                func = function('set_description', vrf_name, disable=True)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'default':
                cmds = ['vrf definition %s' % vrf_name, 'default description']
                func = function('set_description', vrf_name, default=True)
                self.eapi_positive_config_test(func, cmds)

    def test_set_ipv4_routing(self):
        for state in ['config', 'negate', 'default']:
            vrf_name = 'testipv4vrf'
            if state == 'config':
                cmds = ['ip routing vrf %s' % vrf_name]
                func = function('set_ipv4_routing', vrf_name)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'negate':
                cmds = ['no ip routing vrf %s' % vrf_name]
                func = function('set_ipv4_routing', vrf_name, disable=True)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'default':
                cmds = ['default ip routing vrf %s' % vrf_name]
                func = function('set_ipv4_routing', vrf_name, default=True)
                self.eapi_positive_config_test(func, cmds)

    def test_set_ipv6_routing(self):
        for state in ['config', 'negate', 'default']:
            vrf_name = 'testipv6vrf'
            if state == 'config':
                cmds = ['ipv6 unicast-routing vrf %s' % vrf_name]
                func = function('set_ipv6_routing', vrf_name)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'negate':
                cmds = ['no ipv6 unicast-routing vrf %s' % vrf_name]
                func = function('set_ipv6_routing', vrf_name, disable=True)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'default':
                cmds = ['default ipv6 unicast-routing vrf %s' % vrf_name]
                func = function('set_ipv6_routing', vrf_name, default=True)
                self.eapi_positive_config_test(func, cmds)

    def test_set_interface(self):
        for state in ['config', 'negate', 'default']:
            vrf_name = 'testintvrf'
            interface = 'Ethernet1'
            if state == 'config':
                cmds = ['interface %s' % interface,
                        'vrf forwarding %s' % vrf_name]
                func = function('set_interface', vrf_name, interface)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'negate':
                cmds = ['interface %s' % interface, 'no vrf forwarding']
                func = function('set_interface', vrf_name, interface,
                                disable=True)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'default':
                cmds = ['interface %s' % interface, 'default vrf forwarding']
                func = function('set_interface', vrf_name, interface,
                                default=True)
                self.eapi_positive_config_test(func, cmds)

if __name__ == '__main__':
    unittest.main()
<commit_msg>fix bad test<commit_after>#
# Copyright (c) 2017, Arista Networks, Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#   Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
#
#   Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in the
#   documentation and/or other materials provided with the distribution.
#
#   Neither the name of Arista Networks nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL ARISTA NETWORKS
# BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
# IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
import sys
import os
import unittest

sys.path.append(os.path.join(os.path.dirname(__file__), '../lib'))

from testlib import get_fixture, function
from testlib import EapiConfigUnitTest

import pyeapi.api.vrfs

class TestApiVrfs(EapiConfigUnitTest):

    def __init__(self, *args, **kwargs):
        super(TestApiVrfs, self).__init__(*args, **kwargs)
        self.instance = pyeapi.api.vrfs.instance(None)
        self.config = open(get_fixture('running_config.vrf')).read()

    def test_get(self):
        result = self.instance.get('blah')
        vrf = dict(rd='10:10', vrf_name='blah', description='blah desc',
                   ipv4_routing=True, ipv6_routing=False)
        self.assertEqual(vrf, result)
        result2 = self.instance.get('test')
        vrf2 = dict(rd='200:500', vrf_name='test', description='',
                    ipv4_routing=False, ipv6_routing=True)
        self.assertEqual(vrf2, result2)

    def test_get_not_configured(self):
        self.assertIsNone(self.instance.get('notthere'))

    def test_getall(self):
        result = self.instance.getall()
        self.assertIsInstance(result, dict)
        self.assertEqual(len(result), 3)

    def test_vrf_functions(self):
        for name in ['create', 'delete', 'default']:
            vrf_name = 'testvrf'
            if name == 'create':
                cmds = ['vrf definition %s' % vrf_name]
            elif name == 'delete':
                cmds = 'no vrf definition %s' % vrf_name
            elif name == 'default':
                cmds = 'default vrf definition %s' % vrf_name
            func = function(name, vrf_name)
            self.eapi_positive_config_test(func, cmds)

    def test_vrf_create_with_rd(self):
        vrf_name = 'testvrfrd'
        rd = '10:10'
        cmds = ['vrf definition %s' % vrf_name, 'rd %s' % rd]
        func = function('create', vrf_name, rd=rd)
        self.eapi_positive_config_test(func, cmds)

    def test_set_rd(self):
        vrf_name = 'testrdvrf'
        rd = '10:10'
        cmds = ['vrf definition %s' % vrf_name, 'rd %s' % rd]
        func = function('set_rd', vrf_name, rd)
        self.eapi_positive_config_test(func, cmds)

    def test_set_description(self):
        for state in ['config', 'negate', 'default']:
            vrf_name = 'testdescvrf'
            if state == 'config':
                description = 'testing'
                cmds = ['vrf definition %s' % vrf_name,
                        'description %s' % description]
                func = function('set_description', vrf_name, description)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'negate':
                cmds = ['vrf definition %s' % vrf_name, 'no description']
                func = function('set_description', vrf_name, disable=True)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'default':
                cmds = ['vrf definition %s' % vrf_name, 'default description']
                func = function('set_description', vrf_name, default=True)
                self.eapi_positive_config_test(func, cmds)

    def test_set_ipv4_routing(self):
        for state in ['config', 'negate', 'default']:
            vrf_name = 'testipv4vrf'
            if state == 'config':
                cmds = ['ip routing vrf %s' % vrf_name]
                func = function('set_ipv4_routing', vrf_name)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'negate':
                cmds = ['no ip routing vrf %s' % vrf_name]
                func = function('set_ipv4_routing', vrf_name, disable=True)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'default':
                cmds = ['default ip routing vrf %s' % vrf_name]
                func = function('set_ipv4_routing', vrf_name, default=True)
                self.eapi_positive_config_test(func, cmds)

    def test_set_ipv6_routing(self):
        for state in ['config', 'negate', 'default']:
            vrf_name = 'testipv6vrf'
            if state == 'config':
                cmds = ['ipv6 unicast-routing vrf %s' % vrf_name]
                func = function('set_ipv6_routing', vrf_name)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'negate':
                cmds = ['no ipv6 unicast-routing vrf %s' % vrf_name]
                func = function('set_ipv6_routing', vrf_name, disable=True)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'default':
                cmds = ['default ipv6 unicast-routing vrf %s' % vrf_name]
                func = function('set_ipv6_routing', vrf_name, default=True)
                self.eapi_positive_config_test(func, cmds)

    def test_set_interface(self):
        for state in ['config', 'negate', 'default']:
            vrf_name = 'testintvrf'
            interface = 'Ethernet1'
            if state == 'config':
                cmds = ['interface %s' % interface,
                        'vrf forwarding %s' % vrf_name]
                func = function('set_interface', vrf_name, interface)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'negate':
                cmds = ['interface %s' % interface, 'no vrf forwarding']
                func = function('set_interface', vrf_name, interface,
                                disable=True)
                self.eapi_positive_config_test(func, cmds)
            elif state == 'default':
                cmds = ['interface %s' % interface, 'default vrf forwarding']
                func = function('set_interface', vrf_name, interface,
                                default=True)
                self.eapi_positive_config_test(func, cmds)

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


<commit_before>from rethinkdb import r
import json

class rwrapper(object):

  id            = None

  _limit        = 0
  _order_by     = None
  _meta         = None
  _changed      = False

  def __init__(self, **kwargs):
    self._meta = {}
    self._fmeta = {}
    for key in dir(self):
      try:
        if key.startswith('_'):
          continue
        value = getattr(self, key)
        if not value.__rfield__ and not value.__rffield__:
          continue
        self._meta[key] = value
        self._meta[key].name = key
        default = None
        if not value.default == '_rwrapper_default':
          default = value.default
        setattr(self, key, default)
      except:
        continue
    for key in kwargs:
      setattr(self, key, kwargs[key])

  def __json__(self):
    d = {}
    for key in dir(self):
      try:
        if not key == None and not key.startswith('_') and not hasattr(getattr(self, key), '__call__'):
          d[key] = getattr(self, key)
      except:
        continue
    return d

  def __setattr__(self, name, value):
    try:
      if not name.startswith('_') and not value == getattr(self, name):
        self._changed = True
    except:
      pass
    return object.__setattr__(self, name, value)

  def __getattribute__(self, name):
    if name == '__dict__':
      return self.__json__()
    return object.__getattribute__(self, name)

  def evaluate_insert(self, result):
    if result['errors']>1:
      raise IOError(json.dumps(result))
    elif result['inserted'] == 1:
      self.id = result['generated_keys'][0]
    return self.id

  def evaluate_update(self, result):
    if result['updated'] == 0:
      raise ValueError(json.dumps(result))
    if result['errors'] > 0:
      raise IOError(json.dumps(result))
    return result

  def rq(self, filter=False):
    if not filter:
      filter = self._filter()
    rq = r.table(self._db_table)
    if len(filter)>0:
      rq = rq.filter(filter)
    if not self._order_by == None:
      rq = rq.order_by(*tuple([order if not order[:1] == '-' else r.desc(order[1:]) for order in list(self._order_by) ]))
    if not self._limit == 0:
      rq = rq.limit(int(self._limit))
    return rq

  def _filter(self):
    filter = {}
    for key, value in self.__dict__.iteritems():
      if key in self._meta and self._meta[key].default == value:
        continue
      if not value == None:
        filter[key] = value
    return filter

  def get(self, o=False, exception=False):
    try:
      result = self.rq().limit(1).run()[0]
      result = result if o == False else o(**result)
      if o:
        result.changed(False)
      return result
    except:
      if exception == False:
        return None
      raise ValueError('Row not found in table.')

  def save(self):
    # Try and be lazy about saving. Only save if our values have actually
    # changed
    if not self._changed:
      return False

    # Validate any defined fields and set any defaults
    doc = self.__dict__
    if len(self._meta) > 0:
      for key in self._meta.keys():
        setattr(self, key, self._meta[key].validate(doc[key]))

    # id being none means we should insert
    if self.id == None:
      if 'id' in doc:
        del doc['id']
      self.changed(False)
      return self.evaluate_insert(r.table(self._db_table).insert(doc).run())

    # id found; update
    self.changed(False)
    return self.evaluate_update(r.table(self._db_table).filter({'id': self.id}).update(self._filter()).run())

  def changed(self, value):
    if value == True or value == False:
      self._changed = False
    return self

  def order_by(self, *args):
    self._order_by = args
    return self

  def limit(self, amount):
    self._limit = amount
    return self

  def count(self, filter=False):
    return self.rq(filter).count().run()

  def all(self, o=False):
    return [ row if o == False else o(**row) for row in self.rq().run() ]

  def delete(self, filter=False):
    return self.rq(filter).delete().run()
<commit_msg>Save including None types<commit_after>from rethinkdb import r
import json

class rwrapper(object):

  id            = None

  _limit        = 0
  _order_by     = None
  _meta         = None
  _changed      = False

  def __init__(self, **kwargs):
    self._meta = {}
    self._fmeta = {}
    for key in dir(self):
      try:
        if key.startswith('_'):
          continue
        value = getattr(self, key)
        if not value.__rfield__ and not value.__rffield__:
          continue
        self._meta[key] = value
        self._meta[key].name = key
        default = None
        if not value.default == '_rwrapper_default':
          default = value.default
        setattr(self, key, default)
      except:
        continue
    for key in kwargs:
      setattr(self, key, kwargs[key])

  def __json__(self):
    d = {}
    for key in dir(self):
      try:
        if not key == None and not key.startswith('_') and not hasattr(getattr(self, key), '__call__'):
          d[key] = getattr(self, key)
      except:
        continue
    return d

  def __setattr__(self, name, value):
    try:
      if not name.startswith('_') and not value == getattr(self, name):
        self._changed = True
    except:
      pass
    return object.__setattr__(self, name, value)

  def __getattribute__(self, name):
    if name == '__dict__':
      return self.__json__()
    return object.__getattribute__(self, name)

  def evaluate_insert(self, result):
    if result['errors']>1:
      raise IOError(json.dumps(result))
    elif result['inserted'] == 1:
      self.id = result['generated_keys'][0]
    return self.id

  def evaluate_update(self, result):
    if result['updated'] == 0:
      raise ValueError(json.dumps(result))
    if result['errors'] > 0:
      raise IOError(json.dumps(result))
    return result

  def rq(self, filter=False):
    if not filter:
      filter = self._filter()
    rq = r.table(self._db_table)
    if len(filter)>0:
      rq = rq.filter(filter)
    if not self._order_by == None:
      rq = rq.order_by(*tuple([order if not order[:1] == '-' else r.desc(order[1:]) for order in list(self._order_by) ]))
    if not self._limit == 0:
      rq = rq.limit(int(self._limit))
    return rq

  def _filter(self):
    filter = {}
    for key, value in self.__dict__.iteritems():
      if key in self._meta and self._meta[key].default == value:
        continue
      if not value == None:
        filter[key] = value
    return filter

  def get(self, o=False, exception=False):
    try:
      result = self.rq().limit(1).run()[0]
      result = result if o == False else o(**result)
      if o:
        result.changed(False)
      return result
    except:
      if exception == False:
        return None
      raise ValueError('Row not found in table.')

  def save(self):
    # Try and be lazy about saving. Only save if our values have actually
    # changed
    if not self._changed:
      return False

    # Validate any defined fields and set any defaults
    doc = self.__dict__
    if len(self._meta) > 0:
      for key in self._meta.keys():
        setattr(self, key, self._meta[key].validate(doc[key]))

    # id being none means we should insert
    if self.id == None:
      if 'id' in doc:
        del doc['id']
      self.changed(False)
      return self.evaluate_insert(r.table(self._db_table).insert(doc).run())

    # id found; update
    self.changed(False)
    return self.evaluate_update(r.table(self._db_table).filter({'id': self.id}).update(self.__dict__).run())

  def changed(self, value):
    if value == True or value == False:
      self._changed = False
    return self

  def order_by(self, *args):
    self._order_by = args
    return self

  def limit(self, amount):
    self._limit = amount
    return self

  def count(self, filter=False):
    return self.rq(filter).count().run()

  def all(self, o=False):
    return [ row if o == False else o(**row) for row in self.rq().run() ]

  def delete(self, filter=False):
    return self.rq(filter).delete().run()
<|endoftext|>


<commit_before>from threading import Thread
import sys
import serial

class GPS(Thread):
    def is_valid(self):
        return self.signal_validity

    def __init__(self):
        Thread.__init__(self)
        self.port = '/dev/ttyS0'
        self.speed = 9600
        self.signal_validity = False
        self.latitude = 0.
        self.longitude = 0.
        self.speed = 0.

    @staticmethod
    def convert_degrees(degrees, minutes):
        return degrees + minutes/60.

    def calculate_decimal(self, value, direction):
        degrees = int(value/100)
        minutes = float(value) - (100*degrees)
        abs_degree = self.convert_degrees(degrees, minutes)
        if direction == 'W' or direction == 'S':
            abs_degree *= -1.
        return abs_degree

    def parse_gprmc(self, msg):
        params = msg.split(',')

        if params[2] == 'V':  # checks validity of signal
            self.signal_validity = False
            return
        else:
            self.signal_validity = True

        # tries to parse all values, but if failed, set signal_validity to false
        try:
            lat_raw = float(params[3])
            lng_raw = float(params[5])
            speed_raw = float(params[7])

            # calculates latitude, longitude and speed
            self.latitude = round(self.calculate_decimal(lat_raw, params[4]), 5)
            self.longitude = round(self.calculate_decimal(lng_raw, params[6]), 5)
            self.speed = round(speed_raw*1.852, 2)  # converts from knots to km/h
            self.signal_validity = True
        except ValueError:
            self.signal_validity = False

        return self.latitude, self.longitude, self.speed

    def getLatitude(self):
        return self.latitude

    def getLongitude(self):
        return self.longitude

    def getSpeed(self):
        return self.speed

    def run(self):
        try:
            my_serial = serial.Serial(self.port, self.speed)

            while True:
                data = my_serial.readline()
                if data.startswith("$GPRMC"):
                    self.parse_gprmc(data)

        except Exception as e:
            self.run_event.clear()
            # sys.stderr.write('Erro reading serial port %s: %s\n' % (type(e).__name__, e))
<commit_msg>Remove old variable run_event in gps.py<commit_after>from threading import Thread
import sys
import serial

class GPS(Thread):
    def is_valid(self):
        return self.signal_validity

    def __init__(self):
        Thread.__init__(self)
        self.port = '/dev/ttyS0'
        self.speed = 9600
        self.signal_validity = False
        self.latitude = 0.
        self.longitude = 0.
        self.speed = 0.

    @staticmethod
    def convert_degrees(degrees, minutes):
        return degrees + minutes/60.

    def calculate_decimal(self, value, direction):
        degrees = int(value/100)
        minutes = float(value) - (100*degrees)
        abs_degree = self.convert_degrees(degrees, minutes)
        if direction == 'W' or direction == 'S':
            abs_degree *= -1.
        return abs_degree

    def parse_gprmc(self, msg):
        params = msg.split(',')

        if params[2] == 'V':  # checks validity of signal
            self.signal_validity = False
            return
        else:
            self.signal_validity = True

        # tries to parse all values, but if failed, set signal_validity to false
        try:
            lat_raw = float(params[3])
            lng_raw = float(params[5])
            speed_raw = float(params[7])

            # calculates latitude, longitude and speed
            self.latitude = round(self.calculate_decimal(lat_raw, params[4]), 5)
            self.longitude = round(self.calculate_decimal(lng_raw, params[6]), 5)
            self.speed = round(speed_raw*1.852, 2)  # converts from knots to km/h
            self.signal_validity = True
        except ValueError:
            self.signal_validity = False

        return self.latitude, self.longitude, self.speed

    def getLatitude(self):
        return self.latitude

    def getLongitude(self):
        return self.longitude

    def getSpeed(self):
        return self.speed

    def run(self):
        try:
            my_serial = serial.Serial(self.port, self.speed)

            while True:
                data = my_serial.readline()
                if data.startswith("$GPRMC"):
                    self.parse_gprmc(data)

        except Exception as e:
            return
<|endoftext|>


<commit_before> # Custom widgets / Complex controls required
from django.forms import Textarea, Widget, MultiWidget
from django.forms import widgets
from registry.utils import get_static_url
from django.utils.safestring import mark_safe
from django.core.urlresolvers import reverse_lazy

import logging
logger = logging.getLogger("registry_log")

import pycountry

class BadCustomFieldWidget(Textarea):
    """
    Widget to use instead if a custom widget is defined and fails on creation
    """
    pass

class CustomWidgetC18583(Widget):
    def render(self, name, value, attrs=None):
        return "
%s
" % value

class DatatypeWidgetAlphanumericxxx(Textarea):
    def render(self, name, value, attrs=None):
        html = super(DatatypeWidgetAlphanumericxxx, self).render(name, value, attrs)
        return "
%s
" % html

class OtherPleaseSpecifyWidget(MultiWidget):
    def __init__(self, main_choices, other_please_specify_value, unset_value, attrs=None):
        self.main_choices = main_choices
        self.other_please_specify_value = other_please_specify_value
        self.unset_value = unset_value

        _widgets = (
            widgets.Select(attrs=attrs, choices=self.main_choices),
            widgets.TextInput(attrs=attrs)
        )

        super(OtherPleaseSpecifyWidget, self).__init__(_widgets, attrs)

    def format_output(self, rendered_widgets):
        output = u'
'.join(rendered_widgets)
        return output

    def decompress(self, value):
        """
        :param value: value from db or None
        :return: values to be supplied to the select widget and text widget:
        If no value, we show unset for dropdown and nothing in text
        If a value is supplied outside of the dropdown range we provide it to the text box
        and set the select widget to the indicator for "Other please specify"
        otherwise we provide the selected value to the select box and an empty string to
        the textbox
        """
        if not value:
            return [self.unset_value, ""]

        if value not in [choice[0] for choice in self.main_choices]:
            return [self.other_please_specify_value, value]
        else:
            return [value, ""]

    def value_from_datadict(self, data, files, name):
        logger.debug("value from datadict: data = %s name = %s" % (data, name))
        if name in data:
            return data[name]
        else:
            option_selected = data.get(name + "_0", self.unset_value)
            text_entered = data.get(name + "_1","")

            if option_selected == self.other_please_specify_value:
                return text_entered
            else:
                return option_selected

    def render(self, name, value, attrs=None):
        select_id = "id_" + name + "_0"
        specified_value_textbox_id = "id_" + name + "_1"
        script = """

        """ % (select_id, self.other_please_specify_value, specified_value_textbox_id, specified_value_textbox_id, select_id)

        return super(OtherPleaseSpecifyWidget, self).render(name, value, attrs) + script

class HgvsValidation(widgets.TextInput):

    def __init__(self, attrs={}):
        self.attrs = attrs
        super(HgvsValidation, self).__init__(attrs=attrs)

    def render(self, name, value, attrs):
        return mark_safe(u'''<input type="text" name="%s" id="id_%s" value="%s">

                         ''' % (name, name, value or '', name, name, name))

class CalculatedFieldWidget(widgets.TextInput):
    def __init__(self, script, attrs={}):
        attrs['readonly'] = 'readonly'
        self.script = script
        super(CalculatedFieldWidget, self).__init__(attrs=attrs)

    def render(self, name, value, attrs):
        #attrs['readonly'] = 'readonly'
        return super(CalculatedFieldWidget, self).render(name, value, attrs) + self.script

class ExtensibleListWidget(MultiWidget):
    def __init__(self, prototype_widget, attrs={}):
        self.widget_count = 1
        self.prototype_widget = prototype_widget
        super(ExtensibleListWidget, self).__init__([prototype_widget], attrs)

    def _buttons_html(self):
        return """<button type="button" onclick="alert('todo')">Click Me!</button>"""

    def decompress(self, data):
        """

        :param data: dictionary contains items key with a list of row data for each widget
        We create as many widgets on the fly so that render method can iterate
        data  must not be a list else render won't call decompress ...
        :return: a list of data for the widgets to render
        """
        from copy import copy
        if not data:
            self.widgets = [ copy(self.prototype_widget) ]
            return [None]
        else:
            items = data["items"]
            num_widgets = len(items)
            self.widgets = [ copy(self.prototype_widget) for i in range(num_widgets) ]
            return data

    def render(self, name, value):
        html = super(ExtensibleListWidget, self).render(name, value)
        return html + self._buttons_html()

class LookupWidget(widgets.TextInput):
    SOURCE_URL = ""

    def render(self, name, value, attrs):
        return """
            <input type="text" name="%s" id="id_%s" value="%s">

        """ % (name, name, value or '', name, self.SOURCE_URL)

class GeneLookupWidget(LookupWidget):
    SOURCE_URL = reverse_lazy('gene_source')

class LaboratoryLookupWidget(LookupWidget):
    SOURCE_URL = reverse_lazy('laboratory_source')

class DateWidget(widgets.TextInput):

    def render(self, name, value, attrs):
        def just_date(value):
            if value:
                if hasattr(value, 'date'):
                    d = value.date()
                    return "%s-%s-%s" % (d.day, d.month, d.year)
                else:
                    return value
            else:
                return value
        return """
            <input type="text" name="%s" id="id_%s" value="%s" class="datepicker" readonly>
        """ % (name, name, just_date(value) or '')

class CountryWidget(widgets.Select):

    def render(self, name, value, attrs):
        if not value:
            value = self.attrs['default']

        countries = pycountry.countries

        output = ["<select onChange='select_country(this);' id='%s' name='%s'>" % (name, name),]
        for country in countries:
            if value == country.alpha2:
                output.append("<option value='%s' selected>%s</option>" % (country.alpha2, country.name))
            else:
                output.append("<option value='%s'>%s</option>" % (country.alpha2, country.name))
        output.append("</select>")
        return mark_safe('\n'.join(output))

class StateWidget(widgets.Select):

    def render(self, name, value, attrs):
        if not value:
            value = self.attrs['default']

        try:
            state = pycountry.subdivisions.get(code=value)
        except KeyError:
            state = pycountry.subdivisions.get(code=self.attrs['default'])

        country_states = pycountry.subdivisions.get(country_code=state.country.alpha2)

        output = ["<select id='%s' name='%s'>" % (name, name),]
        for state in country_states:
            if value == state.code:
                output.append("<option value='%s' selected>%s</option>" % (state.code, state.name))
            else:
                output.append("<option value='%s'>%s</option>" % (state.code, state.name))
        output.append("</select>")
        return mark_safe('\n'.join(output))
<commit_msg>linting<commit_after> # Custom widgets / Complex controls required
from django.forms import Textarea, Widget, MultiWidget
from django.forms import widgets
from registry.utils import get_static_url
from django.utils.safestring import mark_safe
from django.core.urlresolvers import reverse_lazy

import logging
logger = logging.getLogger("registry_log")

import pycountry

class BadCustomFieldWidget(Textarea):
    """
    Widget to use instead if a custom widget is defined and fails on creation
    """
    pass

class CustomWidgetC18583(Widget):
    def render(self, name, value, attrs=None):
        return "
%s
" % value

class DatatypeWidgetAlphanumericxxx(Textarea):
    def render(self, name, value, attrs=None):
        html = super(DatatypeWidgetAlphanumericxxx, self).render(name, value, attrs)
        return "
%s
" % html

class OtherPleaseSpecifyWidget(MultiWidget):
    def __init__(self, main_choices, other_please_specify_value, unset_value, attrs=None):
        self.main_choices = main_choices
        self.other_please_specify_value = other_please_specify_value
        self.unset_value = unset_value

        _widgets = (
            widgets.Select(attrs=attrs, choices=self.main_choices),
            widgets.TextInput(attrs=attrs)
        )

        super(OtherPleaseSpecifyWidget, self).__init__(_widgets, attrs)

    def format_output(self, rendered_widgets):
        output = u'
'.join(rendered_widgets)
        return output

    def decompress(self, value):
        """
        :param value: value from db or None
        :return: values to be supplied to the select widget and text widget:
        If no value, we show unset for dropdown and nothing in text
        If a value is supplied outside of the dropdown range we provide it to the text box
        and set the select widget to the indicator for "Other please specify"
        otherwise we provide the selected value to the select box and an empty string to
        the textbox
        """
        if not value:
            return [self.unset_value, ""]

        if value not in [choice[0] for choice in self.main_choices]:
            return [self.other_please_specify_value, value]
        else:
            return [value, ""]

    def value_from_datadict(self, data, files, name):
        logger.debug("value from datadict: data = %s name = %s" % (data, name))
        if name in data:
            return data[name]
        else:
            option_selected = data.get(name + "_0", self.unset_value)
            text_entered = data.get(name + "_1", "")

            if option_selected == self.other_please_specify_value:
                return text_entered
            else:
                return option_selected

    def render(self, name, value, attrs=None):
        select_id = "id_" + name + "_0"
        specified_value_textbox_id = "id_" + name + "_1"
        script = """

        """ % (select_id, self.other_please_specify_value, specified_value_textbox_id, specified_value_textbox_id, select_id)

        return super(OtherPleaseSpecifyWidget, self).render(name, value, attrs) + script

class HgvsValidation(widgets.TextInput):

    def __init__(self, attrs={}):
        self.attrs = attrs
        super(HgvsValidation, self).__init__(attrs=attrs)

    def render(self, name, value, attrs):
        return mark_safe(u'''<input type="text" name="%s" id="id_%s" value="%s">

                         ''' % (name, name, value or '', name, name, name))

class CalculatedFieldWidget(widgets.TextInput):
    def __init__(self, script, attrs={}):
        attrs['readonly'] = 'readonly'
        self.script = script
        super(CalculatedFieldWidget, self).__init__(attrs=attrs)

    def render(self, name, value, attrs):
        #attrs['readonly'] = 'readonly'
        return super(CalculatedFieldWidget, self).render(name, value, attrs) + self.script

class ExtensibleListWidget(MultiWidget):
    def __init__(self, prototype_widget, attrs={}):
        self.widget_count = 1
        self.prototype_widget = prototype_widget
        super(ExtensibleListWidget, self).__init__([prototype_widget], attrs)

    def _buttons_html(self):
        return """<button type="button" onclick="alert('todo')">Click Me!</button>"""

    def decompress(self, data):
        """

        :param data: dictionary contains items key with a list of row data for each widget
        We create as many widgets on the fly so that render method can iterate
        data  must not be a list else render won't call decompress ...
        :return: a list of data for the widgets to render
        """
        from copy import copy
        if not data:
            self.widgets = [copy(self.prototype_widget)]
            return [None]
        else:
            items = data["items"]
            num_widgets = len(items)
            self.widgets = [copy(self.prototype_widget) for i in range(num_widgets)]
            return data

    def render(self, name, value):
        html = super(ExtensibleListWidget, self).render(name, value)
        return html + self._buttons_html()

class LookupWidget(widgets.TextInput):
    SOURCE_URL = ""

    def render(self, name, value, attrs):
        return """
            <input type="text" name="%s" id="id_%s" value="%s">

        """ % (name, name, value or '', name, self.SOURCE_URL)

class GeneLookupWidget(LookupWidget):
    SOURCE_URL = reverse_lazy('gene_source')

class LaboratoryLookupWidget(LookupWidget):
    SOURCE_URL = reverse_lazy('laboratory_source')

class DateWidget(widgets.TextInput):

    def render(self, name, value, attrs):
        def just_date(value):
            if value:
                if hasattr(value, 'date'):
                    d = value.date()
                    return "%s-%s-%s" % (d.day, d.month, d.year)
                else:
                    return value
            else:
                return value
        return """
            <input type="text" name="%s" id="id_%s" value="%s" class="datepicker" readonly>
        """ % (name, name, just_date(value) or '')

class CountryWidget(widgets.Select):

    def render(self, name, value, attrs):
        if not value:
            value = self.attrs['default']

        countries = pycountry.countries

        output = ["<select onChange='select_country(this);' id='%s' name='%s'>" % (name, name)]
        for country in countries:
            if value == country.alpha2:
                output.append("<option value='%s' selected>%s</option>" % (country.alpha2, country.name))
            else:
                output.append("<option value='%s'>%s</option>" % (country.alpha2, country.name))
        output.append("</select>")
        return mark_safe('\n'.join(output))

class StateWidget(widgets.Select):

    def render(self, name, value, attrs):
        if not value:
            value = self.attrs['default']

        try:
            state = pycountry.subdivisions.get(code=value)
        except KeyError:
            state = pycountry.subdivisions.get(code=self.attrs['default'])

        country_states = pycountry.subdivisions.get(country_code=state.country.alpha2)

        output = ["<select id='%s' name='%s'>" % (name, name)]
        for state in country_states:
            if value == state.code:
                output.append("<option value='%s' selected>%s</option>" % (state.code, state.name))
            else:
                output.append("<option value='%s'>%s</option>" % (state.code, state.name))
        output.append("</select>")
        return mark_safe('\n'.join(output))
<|endoftext|>


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

import sys
import os
import re
import shutil
import subprocess
import argparse
import json

# Images comparison
from PIL import Image, ImageChops
import math
import operator

class cd:
  """Context manager for changing the current working directory"""
  def __init__(self, newPath):
    self.newPath = newPath

  def __enter__(self):
    self.savedPath = os.getcwd()
    os.chdir(self.newPath)

  def __exit__(self, etype, value, traceback):
    os.chdir(self.savedPath)

class bcolors:
  OKGREEN = '\033[92m'
  OKBLUE = '\033[94m'
  FAIL = '\033[91m'
  WARNING = '\033[93m'
  ENDC = '\033[0m'

  def disable(self):
    self.OKGREEN = ''
    self.OKBLUE = ''
    self.WARNING = ''
    self.FAIL = ''
    self.ENDC = ''

def success(message):
  print bcolors.OKGREEN + message + bcolors.ENDC

def info(message):
  print "[INFO] " + message

def error(message):
  print bcolors.FAIL + message + bcolors.ENDC

def debug(message):
  if verbose:
    print message

def warn(message):
  print bcolors.WARNING + message + bcolors.ENDC

def croak_and_slam_the_door(message):
  error(message)
  error("Exiting with status 1 (no hard feelings).")
  sys.exit(1)

def sanitize_test_cases():
  with cd('test_cases'):
    tests = os.listdir('.')

    errors = {}
    dirs = []

    for test in tests:
      if test == '.tmp':
        continue

      if sanitize_test_case(test):
        dirs.append(os.path.abspath(test))
      else:
        croak_and_slam_the_door("Invalid test case.")

    return dirs

def sanitize_test_case(test):
  if not has_desc(test):
    warn("[WARN] " + test + " does not contain 'desc' file.")

  if not has_expected_image(test):
    warn("[WARN] " + test + " does not contain 'expected.png' image.")

  missings = are_files_missing(os.path.abspath(test))
  if len(missings):
    error("Testcase '" + test + "' is missing the following files : " + ', '.join(missings))
    return False

  return True

def has_desc(path):
  return os.path.isfile(path + '/desc')

def has_expected_image(path):
  return os.path.isfile(path + '/expected.png')

def are_files_missing(path):
  mandatories = ['html', 'javascript']
  return filter(lambda m:not os.path.isfile(path + '/' + m), mandatories)

def create_temp_dir():
  target = ".tmp"

  if os.path.exists(target):
    info("Directory .tmp already exists, removing it")
    shutil.rmtree(target)

  os.mkdir(target)
  info("Creating .tmp directory")

  return target

def get_first_line(path):
  if os.path.isfile(path):
    with open(path, "r") as source:
      return source.readline().rstrip()
  else:
    return ""

def get_content(path):
  content = ""
  with open(path, "r") as source:
    content = source.read()

  return content

def generate_test_files(dirs, project_path, template_path):
  info('Creating test files')
  with cd('.tmp'):
    for test in dirs:
      name = generate_test_file(test, project_path, template_path + "/template.html")
      generate_phantom_file(test, project_path, template_path + "/capture.js")
      debug("  created '" + name + "' test case")

  debug('...done')

def generate_test_file(test, project_path, template_path):
  name = os.path.basename(test)
  html = get_content(test + '/html')
  javascript = get_content(test + '/javascript')

  with open(name + ".html", 'w') as target:
    with open(template_path, 'r') as src:
      line = src.read()
      line = re.sub(r'%html%', html, line)
      line = re.sub(r'%javascript%', javascript, line)
      line = re.sub(r'%project_path%', project_path, line)

      target.write(line)
  return name

def generate_phantom_file(test, project_path, template_path):
  name = os.path.basename(test)

  phantomjs = 'function(){}'
  if os.path.isfile(test + '/phantomjs'):
    phantomjs = get_content(test + '/phantomjs')

  with open(name + ".js", 'w') as target:
    with open(template_path, 'r') as src:
      line = src.read()
      line = re.sub(r'%phantomjs%', phantomjs, line)

      target.write(line)
  return name

def capture_tests(project_path):
  with cd('.tmp'):
    files = os.listdir('.')

    info('Capturing images')
    html_files = [f for f in files if os.path.splitext(f)[1] == '.html']

    for file in html_files:
      img = file.split('.')[0] + '.png'
      js = file.split('.')[0] + '.js'

      subprocess.call([
        project_path + 'node_modules/phantomjs/bin/phantomjs',
        js,
        file,
        img
      ])

      debug('  captured ' + img)

    debug('...done')

def compare(dirs):
  images = {}

  with cd('.tmp'):
    files = os.listdir('.')

    for file in files:
      sp = file.split('.')
      if sp[1] == 'png':
        images[sp[0]] = {
          'computed': os.path.abspath(file)
        }
  errors = 0
  for name in images:
    o = images[name]

    with cd('test_cases/' + name):
      o['expected'] = os.path.abspath('expected.png')
      o['desc'] = get_first_line(os.path.abspath('desc'))

    if not os.path.isfile(o['expected']):
      o['comment'] = "No expected image"
      o['success'] = False
      errors += 1
      croak_and_slam_the_door('No expected image, aborting (please use "run.py -u -o ' + name + '" to init the test case')
    elif not os.path.isfile(o['computed']):
      o['comment'] = "No computed image"
      o['success'] = False
      errors += 1
      croak_and_slam_the_door('No computed image, aborting.')
    else:
      sc = histo_diff(o['expected'], o['computed'])
      o['score'] = sc
      if sc > 100:
        o['success'] = False
        o['comment'] = "The two images appear to be different (result > 100)"
        errors += 1
      else:
        o['success'] = True

      o['comment'] = None
      pretty_print_result(o, name)

  return images, errors

def read_image (source):
  return Image.open(source)

def histo_diff (expected, computed):
  h0 = read_image(expected).histogram()
  h1 = read_image(computed).histogram()

  return math.sqrt(reduce(operator.add, map(lambda a,b: (a-b)**2, h0, h1))/len(h0))

def pretty_print_result(result, name):
  if result['comment'] is not None:
    error('[FAIL] ' + name + ' : ' + result['desc'])
    error('       ' + result['comment'])
  else:
    if result['success']:
      success('[PASS] ' + name + ' : ' + result['desc'])
    else:
      error('[FAIL] ' + name + ' : ' + result['desc'])
      error('       diff too high : ' + str(result['score']))

def bootstrap():
    parser = argparse.ArgumentParser(prog='line-chart visual regression tool')

    parser.add_argument('-v', '--verbose', action="store_true")
    parser.add_argument('-o', '--only')
    parser.add_argument('-u', '--update', action="store_true")

    return parser

def copy_computed_to_tests(dirs):
  with cd('.tmp'):
    for test in dirs:
      name = os.path.basename(test)
      debug("Copying capture for " + name)
      shutil.copy(name + ".png", test + "/expected.png")

    info(str(len(dirs)) + " screenshots copied as references")

def prepare_for_git(results):
  os.mkdir('.tmp/ready_for_git')
  with cd('.tmp/ready_for_git'):

    data = {
      "build": {
        "number": os.getenv('TRAVIS_BUILD_NUMBER', None),
        "branch": os.getenv('TRAVIS_BRANCH', None),
        "is_a_pull_request": os.getenv('TRAVIS_PULL_REQUEST', None)
      },
      "results": []
    }

    for key in results:
      result = results[key]

      data['results'].append({
        "name": key,
        "comment": result["comment"],
        "description": result["desc"],
        "score": result["score"],
        "success": result["success"],
        "expected_image": key + '_expected.png',
        "computed_image": key + '_computed.png'
      })

      shutil.copy(result['expected'], key + '_expected.png')
      shutil.copy(result['computed'], key + '_computed.png')

    with open("test_results.json", "w") as f:
      f.write(json.dumps(data))

args = bootstrap().parse_args()
verbose = args.verbose

scriptsDir = os.path.abspath(os.path.dirname(__file__))
visual = os.path.normpath(scriptsDir + '/../')

with cd(visual):
  project_path = "../../../"

  dirs = sanitize_test_cases()

  if args.only:
    dirs = filter(lambda d: os.path.basename(d) == args.only, dirs)

  create_temp_dir()
  generate_test_files(dirs, project_path, os.path.abspath("scripts"))

  capture_tests(project_path)

  if args.update:
    copy_computed_to_tests(dirs)
  else:
    results, errors = compare(dirs)
    prepare_for_git(results)

    if errors > 0:
      croak_and_slam_the_door(str(errors) + " failure(s)")
    else:
      sys.exit(0)
<commit_msg>Losened visual regression tests<commit_after>#!/usr/bin/env python

import sys
import os
import re
import shutil
import subprocess
import argparse
import json

# Images comparison
from PIL import Image, ImageChops
import math
import operator

class cd:
  """Context manager for changing the current working directory"""
  def __init__(self, newPath):
    self.newPath = newPath

  def __enter__(self):
    self.savedPath = os.getcwd()
    os.chdir(self.newPath)

  def __exit__(self, etype, value, traceback):
    os.chdir(self.savedPath)

class bcolors:
  OKGREEN = '\033[92m'
  OKBLUE = '\033[94m'
  FAIL = '\033[91m'
  WARNING = '\033[93m'
  ENDC = '\033[0m'

  def disable(self):
    self.OKGREEN = ''
    self.OKBLUE = ''
    self.WARNING = ''
    self.FAIL = ''
    self.ENDC = ''

def success(message):
  print bcolors.OKGREEN + message + bcolors.ENDC

def info(message):
  print "[INFO] " + message

def error(message):
  print bcolors.FAIL + message + bcolors.ENDC

def debug(message):
  if verbose:
    print message

def warn(message):
  print bcolors.WARNING + message + bcolors.ENDC

def croak_and_slam_the_door(message):
  error(message)
  error("Exiting with status 1 (no hard feelings).")
  sys.exit(1)

def sanitize_test_cases():
  with cd('test_cases'):
    tests = os.listdir('.')

    errors = {}
    dirs = []

    for test in tests:
      if test == '.tmp':
        continue

      if sanitize_test_case(test):
        dirs.append(os.path.abspath(test))
      else:
        croak_and_slam_the_door("Invalid test case.")

    return dirs

def sanitize_test_case(test):
  if not has_desc(test):
    warn("[WARN] " + test + " does not contain 'desc' file.")

  if not has_expected_image(test):
    warn("[WARN] " + test + " does not contain 'expected.png' image.")

  missings = are_files_missing(os.path.abspath(test))
  if len(missings):
    error("Testcase '" + test + "' is missing the following files : " + ', '.join(missings))
    return False

  return True

def has_desc(path):
  return os.path.isfile(path + '/desc')

def has_expected_image(path):
  return os.path.isfile(path + '/expected.png')

def are_files_missing(path):
  mandatories = ['html', 'javascript']
  return filter(lambda m:not os.path.isfile(path + '/' + m), mandatories)

def create_temp_dir():
  target = ".tmp"

  if os.path.exists(target):
    info("Directory .tmp already exists, removing it")
    shutil.rmtree(target)

  os.mkdir(target)
  info("Creating .tmp directory")

  return target

def get_first_line(path):
  if os.path.isfile(path):
    with open(path, "r") as source:
      return source.readline().rstrip()
  else:
    return ""

def get_content(path):
  content = ""
  with open(path, "r") as source:
    content = source.read()

  return content

def generate_test_files(dirs, project_path, template_path):
  info('Creating test files')
  with cd('.tmp'):
    for test in dirs:
      name = generate_test_file(test, project_path, template_path + "/template.html")
      generate_phantom_file(test, project_path, template_path + "/capture.js")
      debug("  created '" + name + "' test case")

  debug('...done')

def generate_test_file(test, project_path, template_path):
  name = os.path.basename(test)
  html = get_content(test + '/html')
  javascript = get_content(test + '/javascript')

  with open(name + ".html", 'w') as target:
    with open(template_path, 'r') as src:
      line = src.read()
      line = re.sub(r'%html%', html, line)
      line = re.sub(r'%javascript%', javascript, line)
      line = re.sub(r'%project_path%', project_path, line)

      target.write(line)
  return name

def generate_phantom_file(test, project_path, template_path):
  name = os.path.basename(test)

  phantomjs = 'function(){}'
  if os.path.isfile(test + '/phantomjs'):
    phantomjs = get_content(test + '/phantomjs')

  with open(name + ".js", 'w') as target:
    with open(template_path, 'r') as src:
      line = src.read()
      line = re.sub(r'%phantomjs%', phantomjs, line)

      target.write(line)
  return name

def capture_tests(project_path):
  with cd('.tmp'):
    files = os.listdir('.')

    info('Capturing images')
    html_files = [f for f in files if os.path.splitext(f)[1] == '.html']

    for file in html_files:
      img = file.split('.')[0] + '.png'
      js = file.split('.')[0] + '.js'

      subprocess.call([
        project_path + 'node_modules/phantomjs/bin/phantomjs',
        js,
        file,
        img
      ])

      debug('  captured ' + img)

    debug('...done')

def compare(dirs):
  images = {}

  with cd('.tmp'):
    files = os.listdir('.')

    for file in files:
      sp = file.split('.')
      if sp[1] == 'png':
        images[sp[0]] = {
          'computed': os.path.abspath(file)
        }
  errors = 0
  warnings = 0
  for name in images:
    o = images[name]

    with cd('test_cases/' + name):
      o['expected'] = os.path.abspath('expected.png')
      o['desc'] = get_first_line(os.path.abspath('desc'))

    if not os.path.isfile(o['expected']):
      o['comment'] = "No expected image"
      o['status'] = 'error'
      errors += 1
      croak_and_slam_the_door('No expected image, aborting (please use "run.py -u -o ' + name + '" to init the test case')
    elif not os.path.isfile(o['computed']):
      o['comment'] = "No computed image"
      o['status'] = 'error'
      errors += 1
      croak_and_slam_the_door('No computed image, aborting.')
    else:
      sc = histo_diff(o['expected'], o['computed'])
      o['score'] = sc
      if sc > 100:
        o['status'] = 'warning'
        o['comment'] = "The two images appear to be different (diff > 100)"
        warnings += 1
      else:
        o['status'] = 'success'

      o['comment'] = None
      pretty_print_result(o, name)

  return images, errors, warnings

def read_image (source):
  return Image.open(source)

def histo_diff (expected, computed):
  h0 = read_image(expected).histogram()
  h1 = read_image(computed).histogram()

  return math.sqrt(reduce(operator.add, map(lambda a,b: (a-b)**2, h0, h1))/len(h0))

def pretty_print_result(result, name):
  if result['comment'] is not None:
    error('[FAIL] ' + name + ' : ' + result['desc'])
    error('       ' + result['comment'])
  else:
    if result['status'] == 'success' :
      success('[PASS] ' + name + ' : ' + result['desc'])
    elif result['status'] == 'warning' :
      warn('[WARN] ' + name + ' : ' + result['desc'])
      warn('       diff maybe too high : ' + str(result['score']))

def bootstrap():
    parser = argparse.ArgumentParser(prog='line-chart visual regression tool')

    parser.add_argument('-v', '--verbose', action="store_true")
    parser.add_argument('-o', '--only')
    parser.add_argument('-u', '--update', action="store_true")

    return parser

def copy_computed_to_tests(dirs):
  with cd('.tmp'):
    for test in dirs:
      name = os.path.basename(test)
      debug("Copying capture for " + name)
      shutil.copy(name + ".png", test + "/expected.png")

    info(str(len(dirs)) + " screenshots copied as references")

def prepare_for_git(results):
  os.mkdir('.tmp/ready_for_git')
  with cd('.tmp/ready_for_git'):

    data = {
      "build": {
        "number": os.getenv('TRAVIS_BUILD_NUMBER', None),
        "branch": os.getenv('TRAVIS_BRANCH', None),
        "is_a_pull_request": os.getenv('TRAVIS_PULL_REQUEST', None)
      },
      "results": []
    }

    for key in results:
      result = results[key]

      data['results'].append({
        "name": key,
        "comment": result["comment"],
        "description": result["desc"],
        "score": result["score"],
        "status": result["status"],
        "expected_image": key + '_expected.png',
        "computed_image": key + '_computed.png'
      })

      shutil.copy(result['expected'], key + '_expected.png')
      shutil.copy(result['computed'], key + '_computed.png')

    with open("test_results.json", "w") as f:
      f.write(json.dumps(data))

args = bootstrap().parse_args()
verbose = args.verbose

scriptsDir = os.path.abspath(os.path.dirname(__file__))
visual = os.path.normpath(scriptsDir + '/../')

with cd(visual):
  project_path = "../../../"

  dirs = sanitize_test_cases()

  if args.only:
    dirs = filter(lambda d: os.path.basename(d) == args.only, dirs)

  create_temp_dir()
  generate_test_files(dirs, project_path, os.path.abspath("scripts"))

  capture_tests(project_path)

  if args.update:
    copy_computed_to_tests(dirs)
  else:
    results, errors, warnings = compare(dirs)
    prepare_for_git(results)

    if warnings > 0:
      warn(str(warnings) + " warning(s)")

    if errors > 0:
      croak_and_slam_the_door(str(errors) + " failure(s)")
    else:
      sys.exit(0)
<|endoftext|>