<commit_before>"""
Genotype a set of SV calls in a given set of BAMs.
"""
import argparse
from collections import defaultdict
from joblib import Parallel, delayed
import logging
import numpy as np
import pybedtools
import pysam
from scipy import stats

np.random.seed(1)

# create logger
logger = logging.getLogger("calculate_depth")
logger.setLevel(logging.INFO)

# create console handler and set level to debug
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)

# create formatter
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

# add formatter to ch
ch.setFormatter(formatter)

# add ch to logger
logger.addHandler(ch)

# Input file looks like this with SV coordinates in local assembly at columns 6-8.
# chr1    350793  350794  insertion       40      chr1-350756-362784|utg7180000000000|merged      37      77
CHROMOSOME=0
START=1
END=2
EVENT_TYPE=3
EVENT_LENGTH=4
CONTIG_NAME=8
CONTIG_START=9
CONTIG_END=10

def get_min_depth_for_region(bam_fields, breakpoints):
    """
    Return read depth for concordant reads in the given dictionary
    ``bam_fields`` with a pysam.AlignmentFile in the field "bam" and the given
    pybedtools.BedTool of regions, ``regions``.

    The depth calculated across the given regions depends on the type of region
    which is "control" by default. Other acceptable values for the region type
    are "insertion" or "deletion". If the region type is not "control" then the
    coordinates given by ``breakpoints`` are used to evaluate concordant and
    discordant read pairs.
    """
    # Convert breakpoints to a list for easier indexing.
    region = "%s:%s-%s" % list(breakpoints)[0]
    region_size = breakpoints[0][2] - breakpoints[0][1]
    logger.debug("Min depth for region %s (size: %s)", region, region_size)
    depths = [int(line.rstrip().split("\t")[-1])
              for line in pysam.depth("-Q", "20", "-q", "20", "-r", region, bam_fields["filename"])]

    # If there are not depths for every base of the given region, fill those
    # missing values with enough additional zeroes to equal the region size.
    if len(depths) < region_size:
        depths = depths + [0] * (region_size - len(depths))

    logger.debug("Found %i depths for region %s: %s", len(depths), region, depths)
    logger.debug("Median depths: %s", np.median(depths))
    logger.debug("Standard deviation depths: %s", np.std(depths))
    logger.debug("Standard error depths: %s", stats.sem(depths))
    high_quality_depth = max(int(np.round(np.median(depths) - np.std(depths))), 0)

    return high_quality_depth

def get_depth_for_sv_call(sv_call, bams_by_name, chromosome_sizes):
    SLOP_FOR_BREAKPOINTS = 50

    if sv_call[EVENT_TYPE] == "deletion":
        breakpoint_intervals = [
            (sv_call[CONTIG_NAME], int(sv_call[CONTIG_START]), int(sv_call[CONTIG_END]) + SLOP_FOR_BREAKPOINTS),
        ]
        reference_intervals = [
            (sv_call[CHROMOSOME], int(sv_call[START]), int(sv_call[END]))
        ]
        reference_call_type = "insertion"
    else:
        breakpoint_intervals = [
            (sv_call[CONTIG_NAME], int(sv_call[CONTIG_START]), int(sv_call[CONTIG_END]))
        ]
        reference_intervals = [
            (sv_call[CHROMOSOME], int(sv_call[START]), int(sv_call[START]) + SLOP_FOR_BREAKPOINTS),
        ]
        reference_call_type = "deletion"

    logger.debug("Breakpoint intervals: %s", breakpoint_intervals)

    for bam_name, bam in bams_by_name.iteritems():
        breakpoint_concordant_depth = get_min_depth_for_region(bam, breakpoint_intervals)
        breakpoint_discordant_depth = get_min_depth_for_region(bam, reference_intervals)
        logger.debug("Found concordant depth for %s: %s", sv_call[EVENT_TYPE], breakpoint_concordant_depth)
        logger.debug("Found discordant depth for %s: %s", sv_call[EVENT_TYPE], breakpoint_discordant_depth)

        print("\t".join(map(str, (
            bam["sample"],
            sv_call[CHROMOSOME],
            sv_call[START],
            sv_call[END],
            sv_call[EVENT_TYPE],
            sv_call[CONTIG_NAME],
            sv_call[CONTIG_START],
            sv_call[CONTIG_END],
            breakpoint_concordant_depth,
            breakpoint_discordant_depth
        ))))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("sv_calls", help="BED file of SV calls with reference coordinates in columns 1-3, SV type in 4, length in 5, and contig coordinates in 6-8")
    parser.add_argument("bams", nargs="+", help="one or more BAMs per sample to genotype; read groups must be present in BAM")
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--nproc", type=int, default=1, help="number of processes to run in parallel")
    args = parser.parse_args()

    if args.debug:
        logger.setLevel(logging.DEBUG)
        ch.setLevel(logging.DEBUG)
        verbosity = 5
    else:
        verbosity = 0

    chromosome_sizes = None
    bams_by_name = defaultdict(dict)
    for bam in args.bams:
        bam_file = pysam.AlignmentFile(bam, "rb")

        # Prepare chromosome lengths based on one BAM's header.
        if chromosome_sizes is None:
            chromosome_sizes = dict(zip(bam_file.references, [(0, length) for length in bam_file.lengths]))
            logger.debug("Got %i chromosome sizes", len(chromosome_sizes))

        bams_by_name[bam]["file"] = None
        bams_by_name[bam]["filename"] = bam
        bams_by_name[bam]["sample"] = bam_file.header["RG"][0]["SM"]
        bam_file.close()

    columns = ("sample", "chr", "start", "end", "sv_call", "contig", "contig_start", "contig_end", "concordant", "discordant", "total_concordant", "total_discordant")
    print("\t".join(columns))

    sv_calls = pybedtools.BedTool(args.sv_calls)
    # for sv_call in sv_calls:
    #     get_depth_for_sv_call(sv_call, bams_by_name, chromosome_sizes)

    Parallel(n_jobs=args.nproc, verbose=verbosity)(delayed(get_depth_for_sv_call)(sv_call, bams_by_name, chromosome_sizes) for sv_call in sv_calls)
<commit_msg>removed header columns for data columns that don't exist anymore.<commit_after>"""
Genotype a set of SV calls in a given set of BAMs.
"""
import argparse
from collections import defaultdict
from joblib import Parallel, delayed
import logging
import numpy as np
import pybedtools
import pysam
from scipy import stats

np.random.seed(1)

# create logger
logger = logging.getLogger("calculate_depth")
logger.setLevel(logging.INFO)

# create console handler and set level to debug
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)

# create formatter
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

# add formatter to ch
ch.setFormatter(formatter)

# add ch to logger
logger.addHandler(ch)

# Input file looks like this with SV coordinates in local assembly at columns 6-8.
# chr1    350793  350794  insertion       40      chr1-350756-362784|utg7180000000000|merged      37      77
CHROMOSOME=0
START=1
END=2
EVENT_TYPE=3
EVENT_LENGTH=4
CONTIG_NAME=8
CONTIG_START=9
CONTIG_END=10

def get_min_depth_for_region(bam_fields, breakpoints):
    """
    Return read depth for concordant reads in the given dictionary
    ``bam_fields`` with a pysam.AlignmentFile in the field "bam" and the given
    pybedtools.BedTool of regions, ``regions``.

    The depth calculated across the given regions depends on the type of region
    which is "control" by default. Other acceptable values for the region type
    are "insertion" or "deletion". If the region type is not "control" then the
    coordinates given by ``breakpoints`` are used to evaluate concordant and
    discordant read pairs.
    """
    # Convert breakpoints to a list for easier indexing.
    region = "%s:%s-%s" % list(breakpoints)[0]
    region_size = breakpoints[0][2] - breakpoints[0][1]
    logger.debug("Min depth for region %s (size: %s)", region, region_size)
    depths = [int(line.rstrip().split("\t")[-1])
              for line in pysam.depth("-Q", "20", "-q", "20", "-r", region, bam_fields["filename"])]

    # If there are not depths for every base of the given region, fill those
    # missing values with enough additional zeroes to equal the region size.
    if len(depths) < region_size:
        depths = depths + [0] * (region_size - len(depths))

    logger.debug("Found %i depths for region %s: %s", len(depths), region, depths)
    logger.debug("Median depths: %s", np.median(depths))
    logger.debug("Standard deviation depths: %s", np.std(depths))
    logger.debug("Standard error depths: %s", stats.sem(depths))
    high_quality_depth = max(int(np.round(np.median(depths) - np.std(depths))), 0)

    return high_quality_depth

def get_depth_for_sv_call(sv_call, bams_by_name, chromosome_sizes):
    SLOP_FOR_BREAKPOINTS = 50

    if sv_call[EVENT_TYPE] == "deletion":
        breakpoint_intervals = [
            (sv_call[CONTIG_NAME], int(sv_call[CONTIG_START]), int(sv_call[CONTIG_END]) + SLOP_FOR_BREAKPOINTS),
        ]
        reference_intervals = [
            (sv_call[CHROMOSOME], int(sv_call[START]), int(sv_call[END]))
        ]
        reference_call_type = "insertion"
    else:
        breakpoint_intervals = [
            (sv_call[CONTIG_NAME], int(sv_call[CONTIG_START]), int(sv_call[CONTIG_END]))
        ]
        reference_intervals = [
            (sv_call[CHROMOSOME], int(sv_call[START]), int(sv_call[START]) + SLOP_FOR_BREAKPOINTS),
        ]
        reference_call_type = "deletion"

    logger.debug("Breakpoint intervals: %s", breakpoint_intervals)

    for bam_name, bam in bams_by_name.iteritems():
        breakpoint_concordant_depth = get_min_depth_for_region(bam, breakpoint_intervals)
        breakpoint_discordant_depth = get_min_depth_for_region(bam, reference_intervals)
        logger.debug("Found concordant depth for %s: %s", sv_call[EVENT_TYPE], breakpoint_concordant_depth)
        logger.debug("Found discordant depth for %s: %s", sv_call[EVENT_TYPE], breakpoint_discordant_depth)

        print("\t".join(map(str, (
            bam["sample"],
            sv_call[CHROMOSOME],
            sv_call[START],
            sv_call[END],
            sv_call[EVENT_TYPE],
            sv_call[CONTIG_NAME],
            sv_call[CONTIG_START],
            sv_call[CONTIG_END],
            breakpoint_concordant_depth,
            breakpoint_discordant_depth
        ))))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("sv_calls", help="BED file of SV calls with reference coordinates in columns 1-3, SV type in 4, length in 5, and contig coordinates in 6-8")
    parser.add_argument("bams", nargs="+", help="one or more BAMs per sample to genotype; read groups must be present in BAM")
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--nproc", type=int, default=1, help="number of processes to run in parallel")
    args = parser.parse_args()

    if args.debug:
        logger.setLevel(logging.DEBUG)
        ch.setLevel(logging.DEBUG)
        verbosity = 5
    else:
        verbosity = 0

    chromosome_sizes = None
    bams_by_name = defaultdict(dict)
    for bam in args.bams:
        bam_file = pysam.AlignmentFile(bam, "rb")

        # Prepare chromosome lengths based on one BAM's header.
        if chromosome_sizes is None:
            chromosome_sizes = dict(zip(bam_file.references, [(0, length) for length in bam_file.lengths]))
            logger.debug("Got %i chromosome sizes", len(chromosome_sizes))

        bams_by_name[bam]["file"] = None
        bams_by_name[bam]["filename"] = bam
        bams_by_name[bam]["sample"] = bam_file.header["RG"][0]["SM"]
        bam_file.close()

    columns = ("sample", "chr", "start", "end", "sv_call", "contig", "contig_start", "contig_end", "concordant", "discordant")
    print("\t".join(columns))

    sv_calls = pybedtools.BedTool(args.sv_calls)
    # for sv_call in sv_calls:
    #     get_depth_for_sv_call(sv_call, bams_by_name, chromosome_sizes)

    Parallel(n_jobs=args.nproc, verbose=verbosity)(delayed(get_depth_for_sv_call)(sv_call, bams_by_name, chromosome_sizes) for sv_call in sv_calls)
<|endoftext|>


<commit_before>
import pytest
import pandas as pd

pyspark = pytest.importorskip('pyspark')
from pyspark.sql import SchemaRDD, Row, SQLContext
from pyspark.sql import ArrayType, StructField, StructType, IntegerType
from pyspark.sql import StringType

from datashape import dshape, int64
from into import into, discover
from into.utils import tmpfile
from into.backends.sparksql import schema_to_dshape, dshape_to_schema

data = [['Alice', 100.0, 1],
        ['Bob', 200.0, 2],
        ['Alice', 50.0, 3]]

df = pd.DataFrame(data, columns=['name', 'amount', 'id'])

@pytest.yield_fixture(scope='module')
def people(sc):
    with tmpfile('.txt') as fn:
        df.to_csv(fn, header=False, index=False)
        raw = sc.textFile(fn)
        parts = raw.map(lambda line: line.split(','))
        yield parts.map(lambda person: Row(name=person[0],
                                           amount=float(person[1]),
                                           id=int(person[2])))

@pytest.fixture(scope='module')
def ctx(sqlctx, people):
    schema = sqlctx.inferSchema(people)
    schema.registerTempTable('t')
    return sqlctx

def test_pyspark_to_sparksql(ctx, people):
    sdf = into(ctx, data)
    assert isinstance(sdf, SchemaRDD)
    assert into(list, people) == into(list, sdf)

def test_into_sparksql_from_other(ctx):
    sdf = into(ctx, df)
    assert isinstance(sdf, SchemaRDD)
    assert into(list, sdf) == into(list, df)

def test_discover_context(ctx):
    result = discover(ctx)
    assert result is not None

def test_schema_to_dshape():
    assert schema_to_dshape(IntegerType()) == int64

    assert schema_to_dshape(
        ArrayType(IntegerType(), False)) == dshape("var * int64")

    assert schema_to_dshape(
        ArrayType(IntegerType(), True)) == dshape("var * ?int64")

    assert schema_to_dshape(StructType([
        StructField('name', StringType(), False),
        StructField('amount', IntegerType(), True)])) \
        == dshape("{name: string, amount: ?int64}")

def test_dshape_to_schema():
    assert dshape_to_schema('int32') == IntegerType()

    assert dshape_to_schema('5 * int32') == ArrayType(IntegerType(), False)

    assert dshape_to_schema('5 * ?int32') == ArrayType(IntegerType(), True)

    assert dshape_to_schema('{name: string, amount: int32}') == \
        StructType([StructField('name', StringType(), False),
                    StructField('amount', IntegerType(), False)])

    assert dshape_to_schema('10 * {name: string, amount: ?int32}') == \
        ArrayType(StructType(
            [StructField('name', StringType(), False),
             StructField('amount', IntegerType(), True)]),
        False)
<commit_msg>cleanup imports<commit_after>
import pytest
import pandas as pd

pyspark = pytest.importorskip('pyspark')

from pyspark.sql import SchemaRDD, Row
from pyspark.sql import ArrayType, StructField, StructType, IntegerType
from pyspark.sql import StringType

import datashape
from datashape import dshape
from into import into, discover
from into.utils import tmpfile
from into.backends.sparksql import schema_to_dshape, dshape_to_schema

data = [['Alice', 100.0, 1],
        ['Bob', 200.0, 2],
        ['Alice', 50.0, 3]]

df = pd.DataFrame(data, columns=['name', 'amount', 'id'])

@pytest.yield_fixture(scope='module')
def people(sc):
    with tmpfile('.txt') as fn:
        df.to_csv(fn, header=False, index=False)
        raw = sc.textFile(fn)
        parts = raw.map(lambda line: line.split(','))
        yield parts.map(lambda person: Row(name=person[0],
                                           amount=float(person[1]),
                                           id=int(person[2])))

@pytest.fixture(scope='module')
def ctx(sqlctx, people):
    schema = sqlctx.inferSchema(people)
    schema.registerTempTable('t')
    return sqlctx

def test_pyspark_to_sparksql(ctx, people):
    sdf = into(ctx, data)
    assert isinstance(sdf, SchemaRDD)
    assert into(list, people) == into(list, sdf)

def test_into_sparksql_from_other(ctx):
    sdf = into(ctx, df)
    assert isinstance(sdf, SchemaRDD)
    assert into(list, sdf) == into(list, df)

def test_discover_context(ctx):
    result = discover(ctx)
    assert result is not None

def test_schema_to_dshape():
    assert schema_to_dshape(IntegerType()) == int64

    assert schema_to_dshape(
        ArrayType(IntegerType(), False)) == dshape("var * int64")

    assert schema_to_dshape(
        ArrayType(IntegerType(), True)) == dshape("var * ?int64")

    assert schema_to_dshape(StructType([
        StructField('name', StringType(), False),
        StructField('amount', IntegerType(), True)])) \
        == dshape("{name: string, amount: ?int64}")

def test_dshape_to_schema():
    assert dshape_to_schema('int32') == IntegerType()

    assert dshape_to_schema('5 * int32') == ArrayType(IntegerType(), False)

    assert dshape_to_schema('5 * ?int32') == ArrayType(IntegerType(), True)

    assert dshape_to_schema('{name: string, amount: int32}') == \
        StructType([StructField('name', StringType(), False),
                    StructField('amount', IntegerType(), False)])

    assert dshape_to_schema('10 * {name: string, amount: ?int32}') == \
        ArrayType(StructType(
            [StructField('name', StringType(), False),
             StructField('amount', IntegerType(), True)]),
        False)
<|endoftext|>


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

'''
@author Luke Campbell <lcampbell@asascience.com>
@file ion/processes/data/replay_process.py
@description Implementation for the Replay Agent
'''
from gevent.greenlet import Greenlet
from gevent.coros import RLock
from interface.objects import BlogBase, StreamGranuleContainer
from pyon.datastore.datastore import DataStore
from pyon.ion.endpoint import StreamPublisherRegistrar
from pyon.public import log
from interface.services.dm.ireplay_process import BaseReplayProcess
class ReplayProcess(BaseReplayProcess):
    process_type="standalone"
    def __init__(self, *args, **kwargs):
        super(ReplayProcess, self).__init__(*args,**kwargs)
        #@todo Init stuff
        # mutex for shared resources between threads
        self.lock = RLock()

    def on_start(self):
        '''
        Creates a publisher for each stream_id passed in as publish_streams
        Creates an attribute with the name matching the stream name which corresponds to the publisher
        ex: say we have publish_streams:{'output': my_output_stream_id }
          then the instance has an attribute output which corresponds to the publisher for the stream
          in my_output_stream_id
        '''
        self.stream_publisher_registrar = StreamPublisherRegistrar(process=self,node=self.container.node)

        # Get the stream(s)
        streams = self.CFG.get('process',{}).get('publish_streams',{})

        # Get the query
        self.query = self.CFG.get('process',{}).get('query',{})

        # Get the delivery_format
        self.delivery_format = self.CFG.get('process',{}).get('delivery_format',{}) or {}

        self.datastore_name = self.CFG.get('process',{}).get('datastore_name','dm_datastore')
        self.view_name = self.CFG.get('process',{}).get('view_name')
        self.key_id = self.CFG.get('process',{}).get('key_id')
        # Get a stream_id for this process
        self.stream_id = self.CFG.get('process',{}).get('publish_streams',{}).get('output')

        if not (self.stream_id and hasattr(self,'output')):
            raise RuntimeError('The replay agent requires an output stream publisher named output. Invalid configuration!')

    def _publish_query(self, results):
        '''
        Callback to publish the specified results
        '''
        #-----------------------
        # Iteration
        #-----------------------
        #  - Go through the results, if the user had include_docs=True in the options field
        #    then the full document is in result.doc; however if the query did not include_docs,
        #    then only the doc_id is provided in the result.value.
        #
        #  - What this allows us to do is limit the amount of traffic in information for large queries.
        #    If we only are making a query in a sequence of queries (such as map and reduce) then we don't
        #    care about the full document, yet, we only care about the doc id and will retrieve the document later.
        #  - Example:
        #      Imagine the blogging example, we want the latest blog by author George and all the comments for that blog
        #      The series of queries would go, post_by_updated -> posts_by_author -> posts_join_comments and then
        #      in the last query we'll set include_docs to true and parse the docs.
        #-----------------------

        #@todo: Add thread sync here because self.output is shared and deadlocks COULD occur
        log.warn('results: %s', results)

        for result in results:
            log.warn('Result: %s' % result)
            if 'doc' in result:
                log.debug('Result contains document.')
                replay_obj_msg = result['doc']
                if isinstance(replay_obj_msg, BlogBase):
                    replay_obj_msg.is_replay = True
                else:
                    # Override the resource_stream_id so ingestion doesn't reingest, also this is a NEW stream (replay)
                    replay_obj_msg.stream_resource_id = self.stream_id
            else:
                replay_obj_msg = result['value'] # Document ID, not a document

            # Handle delivery options
            # Case: Chopping granules
            if isinstance(replay_obj_msg, StreamGranuleContainer) and self.delivery_format.get('chop', False):
                for identifiable in replay_obj_msg.identifiables:
                    self.lock.acquire()
                    self.output.publish(identifiable)
                    self.lock.release()
            # Default: Publish
            else:
                self.lock.acquire()
                self.output.publish(replay_obj_msg)
                self.lock.release()

        #@todo: log when there are not results
        if results is None:
            log.warn('No results found in replay query!')
        else:
            log.debug('Published replay!')

    def execute_replay(self):
        log.debug('(Replay Agent %s)', self.name)

        # Handle the query
        datastore_name = self.datastore_name
        key_id = self.key_id

        # Got the post ID, pull the post and the comments
        view_name = self.view_name
        opts = {
            'start_key':[key_id, 0],
            'end_key':[key_id,2],
            'include_docs': True
        }
        g = Greenlet(self._query,datastore_name=datastore_name, view_name=view_name, opts=opts,
            callback=lambda results: self._publish_query(results))
        g.start()

    def _query(self,datastore_name='dm_datastore', view_name='posts/posts_by_id', opts={}, callback=None):
        '''
        Performs the query action
        '''
        log.debug('Couch Query:\n\t%s\n\t%s\n\t%s', datastore_name, view_name, opts)
        #@todo: Fix this datastore management profile with correct data profile in near future
        db = self.container.datastore_manager.get_datastore(datastore_name, DataStore.DS_PROFILE.EXAMPLES, self.CFG)

        ret = db.query_view(view_name=view_name,opts=opts)

        callback(ret)

<commit_msg>Luke: Replay - Refactored get -> get_safe<commit_after>#!/usr/bin/env python

'''
@author Luke Campbell <lcampbell@asascience.com>
@file ion/processes/data/replay_process.py
@description Implementation for the Replay Agent
'''
from gevent.greenlet import Greenlet
from gevent.coros import RLock
from interface.objects import BlogBase, StreamGranuleContainer
from pyon.datastore.datastore import DataStore
from pyon.ion.endpoint import StreamPublisherRegistrar
from pyon.public import log
from interface.services.dm.ireplay_process import BaseReplayProcess
from pyon.util.containers import DotDict

class ReplayProcess(BaseReplayProcess):
    process_type="standalone"
    def __init__(self, *args, **kwargs):
        super(ReplayProcess, self).__init__(*args,**kwargs)
        #@todo Init stuff
        # mutex for shared resources between threads
        self.lock = RLock()

    def on_start(self):
        '''
        Creates a publisher for each stream_id passed in as publish_streams
        Creates an attribute with the name matching the stream name which corresponds to the publisher
        ex: say we have publish_streams:{'output': my_output_stream_id }
          then the instance has an attribute output which corresponds to the publisher for the stream
          in my_output_stream_id
        '''
        self.stream_publisher_registrar = StreamPublisherRegistrar(process=self,node=self.container.node)

        # Get the query
        self.query = self.CFG.get_safe('process.query',{})

        # Get the delivery_format
        self.delivery_format = self.CFG.get_safe('process.delivery_format',{})
        self.datastore_name = self.CFG.get_safe('process.datastore_name','dm_datastore')

        self.view_name = self.CFG.get_safe('process.view_name','datasets/dataset_by_id')
        self.key_id = self.CFG.get_safe('process.key_id')
        # Get a stream_id for this process
        self.stream_id = self.CFG.get_safe('process.publish_streams.output',{})

        if not (self.stream_id and hasattr(self,'output')):
            raise RuntimeError('The replay agent requires an output stream publisher named output. Invalid configuration!')

    def _publish_query(self, results):
        '''
        Callback to publish the specified results
        '''
        #-----------------------
        # Iteration
        #-----------------------
        #  - Go through the results, if the user had include_docs=True in the options field
        #    then the full document is in result.doc; however if the query did not include_docs,
        #    then only the doc_id is provided in the result.value.
        #
        #  - What this allows us to do is limit the amount of traffic in information for large queries.
        #    If we only are making a query in a sequence of queries (such as map and reduce) then we don't
        #    care about the full document, yet, we only care about the doc id and will retrieve the document later.
        #  - Example:
        #      Imagine the blogging example, we want the latest blog by author George and all the comments for that blog
        #      The series of queries would go, post_by_updated -> posts_by_author -> posts_join_comments and then
        #      in the last query we'll set include_docs to true and parse the docs.
        #-----------------------

        #@todo: Add thread sync here because self.output is shared and deadlocks COULD occur
        log.warn('results: %s', results)

        for result in results:
            log.warn('Result: %s' % result)
            if 'doc' in result:
                log.debug('Result contains document.')
                replay_obj_msg = result['doc']
                if isinstance(replay_obj_msg, BlogBase):
                    replay_obj_msg.is_replay = True
                else:
                    # Override the resource_stream_id so ingestion doesn't reingest, also this is a NEW stream (replay)
                    replay_obj_msg.stream_resource_id = self.stream_id
            else:
                replay_obj_msg = result['value'] # Document ID, not a document

            # Handle delivery options
            # Case: Chopping granules
            if isinstance(replay_obj_msg, StreamGranuleContainer) and self.delivery_format.get('chop', False):
                for identifiable in replay_obj_msg.identifiables:
                    self.lock.acquire()
                    self.output.publish(identifiable)
                    self.lock.release()
            # Default: Publish
            else:
                self.lock.acquire()
                self.output.publish(replay_obj_msg)
                self.lock.release()

        #@todo: log when there are not results
        if results is None:
            log.warn('No results found in replay query!')
        else:
            log.debug('Published replay!')

    def execute_replay(self):
        log.debug('(Replay Agent %s)', self.name)

        # Handle the query
        datastore_name = self.datastore_name
        key_id = self.key_id

        # Got the post ID, pull the post and the comments
        view_name = self.view_name
        opts = {
            'start_key':[key_id, 0],
            'end_key':[key_id,2],
            'include_docs': True
        }
        g = Greenlet(self._query,datastore_name=datastore_name, view_name=view_name, opts=opts,
            callback=lambda results: self._publish_query(results))
        g.start()

    def _query(self,datastore_name='dm_datastore', view_name='posts/posts_by_id', opts={}, callback=None):
        '''
        Performs the query action
        '''
        log.debug('Couch Query:\n\t%s\n\t%s\n\t%s', datastore_name, view_name, opts)
        #@todo: Fix this datastore management profile with correct data profile in near future
        db = self.container.datastore_manager.get_datastore(datastore_name, DataStore.DS_PROFILE.EXAMPLES, self.CFG)

        ret = db.query_view(view_name=view_name,opts=opts)

        callback(ret)

<|endoftext|>


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

"""Test Markdown parser and reader."""

#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------

from ...core import format_manager, convert
from ._utils import (_test_reader, _test_writer, _diff, _show_outputs,
                     _exec_test_file, _read_test_file)

#------------------------------------------------------------------------------
# Test Markdown parser
#------------------------------------------------------------------------------

def _test_markdown_reader(basename):
    """Check that (test cells) and (test contents ==> cells) are the same."""
    converted, expected = _test_reader(basename, 'markdown')
    assert converted == expected

def _test_markdown_writer(basename):
    """Check that (test contents) and (test cells ==> contents) are the same.
    """
    converted, expected = _test_writer(basename, 'markdown')
    assert _diff(converted, expected) == ''

def _test_markdown_markdown(basename):
    """Check that the double conversion is the identity."""

    contents = _read_test_file(basename, 'markdown')
    cells = convert(contents, from_='markdown')
    converted = convert(cells, to='markdown')

    assert _diff(contents, converted) == ''

def test_markdown_reader():
    _test_markdown_reader('ex1')
    _test_markdown_reader('ex2')

def test_markdown_writer():
    _test_markdown_writer('ex1')
    _test_markdown_writer('ex2')

def test_markdown_markdown():
    _test_markdown_markdown('ex1')
    _test_markdown_markdown('ex2')

def test_decorator():

    source = "\n".join(("@decorator",
                        "def f():",
                        "    # pass",
                        "    ",
                        "    return",
                        ))
    cells = [{'cell_type': 'code',
              'input': source,
              'output': 'blah\nblah'}]
    markdown = convert(cells, to='markdown')
    cells_bis = convert(markdown, from_='markdown')
    markdown_bis = convert(cells_bis, to='markdown')

    print("***")
    print(markdown_bis)
    print("***")
<commit_msg>WIP: adding failing test.<commit_after># -*- coding: utf-8 -*-

"""Test Markdown parser and reader."""

#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------

from ...core import format_manager, convert
from ._utils import (_test_reader, _test_writer, _diff, _show_outputs,
                     _exec_test_file, _read_test_file)

#------------------------------------------------------------------------------
# Test Markdown parser
#------------------------------------------------------------------------------

def _test_markdown_reader(basename):
    """Check that (test cells) and (test contents ==> cells) are the same."""
    converted, expected = _test_reader(basename, 'markdown')
    assert converted == expected

def _test_markdown_writer(basename):
    """Check that (test contents) and (test cells ==> contents) are the same.
    """
    converted, expected = _test_writer(basename, 'markdown')
    assert _diff(converted, expected) == ''

def _test_markdown_markdown(basename):
    """Check that the double conversion is the identity."""

    contents = _read_test_file(basename, 'markdown')
    cells = convert(contents, from_='markdown')
    converted = convert(cells, to='markdown')

    assert _diff(contents, converted) == ''

def test_markdown_reader():
    _test_markdown_reader('ex1')
    _test_markdown_reader('ex2')

def test_markdown_writer():
    _test_markdown_writer('ex1')
    _test_markdown_writer('ex2')

def test_markdown_markdown():
    _test_markdown_markdown('ex1')
    _test_markdown_markdown('ex2')

def test_decorator():

    markdown = '\n'.join(('```python',
                          '>>> @decorator',
                          '... def f():',
                          '...     """Docstring."""',
                          '...',
                          '...     # Comment.',
                          '...     pass',
                          '...',
                          '...     # Comment.',
                          '...     pass',
                          '...     pass',
                          'blah',
                          'blah',
                          '```'))

    cells = convert(markdown, from_='markdown')
    assert cells[0]['output'] == 'blah\nblah'
<|endoftext|>


<commit_before>#!/usr/bin/python

# Configuration file generator for epidemic model

from lxml import etree as ET
import sys
import random as randNumLocations
import random as randNumPersons
import random as randTTCH
import random as randDTI
import random as randSuscep
import random as randIsVacc
import random as randIS

###### Settings go here ######

# Diffusion parameters
# model - "FullyConnected" or "WattsStrogatz"
MODEL                           = "WattsStrogatz"
DIFFUSION_SEED                  = 101

# Watts-Strogatz model parameters
K                               = 8
BETA                            = 0.1

# Disease parameters
TRANSMISSIBILITY                = 0.12
LATENT_DWELL_TIME               = 200
LATENT_INFECTIVITY              = 0
INCUBATING_DWELL_TIME           = 100
INCUBATING_INFECTIVITY          = 0.3
INFECTIOUS_DWELL_TIME           = 400
INFECTIOUS_INFECTIVITY          = 1.0
ASYMPT_DWELL_TIME               = 200
ASYMPT_INFECTIVITY              = 0.5
PROB_UL_U                       = 0.2
PROB_UL_V                       = 0.9
PROB_UR_V                       = 0.5
PROB_UI_V                       = 0.1
PROB_UI_U                       = 0.3
LOCATION_STATE_REFRESH_INTERVAL = 50
DISEASE_SEED                    = 90

# Data capture
IS_DATA_CAPTURE_NEEDED          = "no"
DATA_CAPTURE_FILENAME           = "epidemic_data.csv"

# Regions
NUMBER_OF_REGIONS               = 50
MIN_NUM_OF_LOCS_PER_REGION      = 100
MAX_NUM_OF_LOCS_PER_REGION      = 120

# Location
MIN_NUM_OF_PERSONS_PER_LOC      = 10
MAX_NUM_OF_PERSONS_PER_LOC      = 30
MIN_TRAVEL_TIME_TO_CENTRAL_HUB  = 50
MAX_TRAVEL_TIME_TO_CENTRAL_HUB  = 400
MIN_DIFFUSION_TRIG_INTERVAL     = 500
MAX_DIFFUSION_TRIG_INTERVAL     = 1500

# Person
SUSCEPTIBILITY_PRECISION        = 100.0
INFECTION_STATES                = ["uninfected", "latent", "incubating", "infectious", "asympt", "recovered"]
VACCINATION_STATUS              = ["yes", "no"]

###### Don't edit below here ######

def main():
    outFileName = sys.argv[1]

    randNumLocations.seed()
    randNumPersons.seed()
    randTTCH.seed()
    randDTI.seed()
    randSuscep.seed()
    randIsVacc.seed()
    randIS.seed()

    root = ET.Element("epidemic_configuration")

    doc = ET.SubElement(root, "diffusion")
    field = ET.SubElement(doc, "model")
    field.text = str(MODEL)
    field = ET.SubElement(doc, "seed")
    field.text = str(DIFFUSION_SEED)

    if MODEL == "WattsStrogatz":
        doc = ET.SubElement(doc, "watts_strogatz")
        field = ET.SubElement(doc, "k")
        field.text = str(K)
        field = ET.SubElement(doc, "beta")
        field.text = str(BETA)

    doc = ET.SubElement(root, "disease")
    field = ET.SubElement(doc, "transmissibility")
    field.text = str(TRANSMISSIBILITY)
    field = ET.SubElement(doc, "latent_dwell_time")
    field.text = str(LATENT_DWELL_TIME)
    field = ET.SubElement(doc, "latent_infectivity")
    field.text = str(LATENT_INFECTIVITY)
    field = ET.SubElement(doc, "incubating_dwell_time")
    field.text = str(INFECTIOUS_DWELL_TIME)
    field = ET.SubElement(doc, "incubating_infectivity")
    field.text = str(INCUBATING_INFECTIVITY)
    field = ET.SubElement(doc, "infectious_dwell_time")
    field.text = str(INFECTIOUS_DWELL_TIME)
    field = ET.SubElement(doc, "infectious_infectivity")
    field.text = str(INFECTIOUS_INFECTIVITY)
    field = ET.SubElement(doc, "asympt_dwell_time")
    field.text = str(ASYMPT_DWELL_TIME)
    field = ET.SubElement(doc, "asympt_infectivity")
    field.text = str(ASYMPT_INFECTIVITY)
    field = ET.SubElement(doc, "prob_ul_u")
    field.text = str(PROB_UL_U)
    field = ET.SubElement(doc, "prob_ul_v")
    field.text = str(PROB_UL_V)
    field = ET.SubElement(doc, "prob_ur_v")
    field.text = str(PROB_UR_V)
    field = ET.SubElement(doc, "prob_ui_v")
    field.text = str(PROB_UI_V)
    field = ET.SubElement(doc, "prob_ui_u")
    field.text = str(PROB_UI_U)
    field = ET.SubElement(doc, "location_state_refresh_interval")
    field.text = str(LOCATION_STATE_REFRESH_INTERVAL)
    field = ET.SubElement(doc, "seed")
    field.text = str(DISEASE_SEED)

    doc = ET.SubElement(root, "data_capture")
    field = ET.SubElement(doc, "is_needed")
    field.text = str(IS_DATA_CAPTURE_NEEDED)
    if IS_DATA_CAPTURE_NEEDED == "yes":
        field = ET.SubElement(doc, "capture_file")
        field.text = str(DATA_CAPTURE_FILENAME)

    doc = ET.SubElement(root, "number_of_regions")
    doc.text = str(NUMBER_OF_REGIONS)

    pid = 1
    for regionNum in range(NUMBER_OF_REGIONS):
        doc = ET.SubElement(root, "region")
        field = ET.SubElement(doc, "region_name")
        field.text = "Region"+str(regionNum)
        field = ET.SubElement(doc, "number_of_locations")
        numberOfLocations = randNumLocations.randrange(MIN_NUM_OF_LOCS_PER_REGION, MAX_NUM_OF_LOCS_PER_REGION+1)
        field.text = str(numberOfLocations)
        for locationNum in range(numberOfLocations):
            doc1 = ET.SubElement(doc, "location")
            field = ET.SubElement(doc1, "location_name")
            field.text = "Location"+str(locationNum)
            field = ET.SubElement(doc1, "number_of_persons")
            numberOfPersons = randNumPersons.randrange(MIN_NUM_OF_PERSONS_PER_LOC, MAX_NUM_OF_PERSONS_PER_LOC+1)
            field.text = str(numberOfPersons)
            field = ET.SubElement(doc1, "travel_time_to_central_hub")
            field.text = str(randTTCH.randrange(MIN_TRAVEL_TIME_TO_CENTRAL_HUB, MAX_TRAVEL_TIME_TO_CENTRAL_HUB+1))
            field = ET.SubElement(doc1, "diffusion_trigger_interval")
            field.text = str(randDTI.randrange(MIN_DIFFUSION_TRIG_INTERVAL, MAX_DIFFUSION_TRIG_INTERVAL+1))
            for numPerson in range(numberOfPersons):
                doc2 = ET.SubElement(doc1, "person")
                field = ET.SubElement(doc2, "pid")
                field.text = str(pid)
                pid += 1
                field = ET.SubElement(doc2, "susceptibility")
                field.text = str(randSuscep.randrange(0, SUSCEPTIBILITY_PRECISION)/SUSCEPTIBILITY_PRECISION)
                field = ET.SubElement(doc2, "is_vaccinated")
                field.text = str(VACCINATION_STATUS[randIsVacc.randrange(0, len(VACCINATION_STATUS))])
                field = ET.SubElement(doc2, "infection_state")
                field.text = str(INFECTION_STATES[randIS.randrange(0, len(INFECTION_STATES))])

    tree = ET.ElementTree(root)
    tree.write(outFileName, pretty_print=True)

if __name__ == "__main__":
    main()

<commit_msg>Switched to standard library modules<commit_after>#!/usr/bin/python

# Configuration file generator for epidemic model

import xml.etree.ElementTree as ET
import xml.dom.minidom as minidom
import sys
import random as randNumLocations
import random as randNumPersons
import random as randTTCH
import random as randDTI
import random as randSuscep
import random as randIsVacc
import random as randIS

###### Settings go here ######

# Diffusion parameters
# model - "FullyConnected" or "WattsStrogatz"
MODEL                           = "WattsStrogatz"
DIFFUSION_SEED                  = 101

# Watts-Strogatz model parameters
K                               = 8
BETA                            = 0.1

# Disease parameters
TRANSMISSIBILITY                = 0.12
LATENT_DWELL_TIME               = 200
LATENT_INFECTIVITY              = 0
INCUBATING_DWELL_TIME           = 100
INCUBATING_INFECTIVITY          = 0.3
INFECTIOUS_DWELL_TIME           = 400
INFECTIOUS_INFECTIVITY          = 1.0
ASYMPT_DWELL_TIME               = 200
ASYMPT_INFECTIVITY              = 0.5
PROB_UL_U                       = 0.2
PROB_UL_V                       = 0.9
PROB_UR_V                       = 0.5
PROB_UI_V                       = 0.1
PROB_UI_U                       = 0.3
LOCATION_STATE_REFRESH_INTERVAL = 50
DISEASE_SEED                    = 90

# Data capture
IS_DATA_CAPTURE_NEEDED          = "no"
DATA_CAPTURE_FILENAME           = "epidemic_data.csv"

# Regions
NUMBER_OF_REGIONS               = 50
MIN_NUM_OF_LOCS_PER_REGION      = 100
MAX_NUM_OF_LOCS_PER_REGION      = 120

# Location
MIN_NUM_OF_PERSONS_PER_LOC      = 10
MAX_NUM_OF_PERSONS_PER_LOC      = 30
MIN_TRAVEL_TIME_TO_CENTRAL_HUB  = 50
MAX_TRAVEL_TIME_TO_CENTRAL_HUB  = 400
MIN_DIFFUSION_TRIG_INTERVAL     = 500
MAX_DIFFUSION_TRIG_INTERVAL     = 1500

# Person
SUSCEPTIBILITY_PRECISION        = 100.0
INFECTION_STATES                = ["uninfected", "latent", "incubating", "infectious", "asympt", "recovered"]
VACCINATION_STATUS              = ["yes", "no"]

###### Don't edit below here ######

def main():
    outFileName = sys.argv[1]

    randNumLocations.seed()
    randNumPersons.seed()
    randTTCH.seed()
    randDTI.seed()
    randSuscep.seed()
    randIsVacc.seed()
    randIS.seed()

    root = ET.Element("epidemic_configuration")

    doc = ET.SubElement(root, "diffusion")
    field = ET.SubElement(doc, "model")
    field.text = str(MODEL)
    field = ET.SubElement(doc, "seed")
    field.text = str(DIFFUSION_SEED)

    if MODEL == "WattsStrogatz":
        doc = ET.SubElement(doc, "watts_strogatz")
        field = ET.SubElement(doc, "k")
        field.text = str(K)
        field = ET.SubElement(doc, "beta")
        field.text = str(BETA)

    doc = ET.SubElement(root, "disease")
    field = ET.SubElement(doc, "transmissibility")
    field.text = str(TRANSMISSIBILITY)
    field = ET.SubElement(doc, "latent_dwell_time")
    field.text = str(LATENT_DWELL_TIME)
    field = ET.SubElement(doc, "latent_infectivity")
    field.text = str(LATENT_INFECTIVITY)
    field = ET.SubElement(doc, "incubating_dwell_time")
    field.text = str(INFECTIOUS_DWELL_TIME)
    field = ET.SubElement(doc, "incubating_infectivity")
    field.text = str(INCUBATING_INFECTIVITY)
    field = ET.SubElement(doc, "infectious_dwell_time")
    field.text = str(INFECTIOUS_DWELL_TIME)
    field = ET.SubElement(doc, "infectious_infectivity")
    field.text = str(INFECTIOUS_INFECTIVITY)
    field = ET.SubElement(doc, "asympt_dwell_time")
    field.text = str(ASYMPT_DWELL_TIME)
    field = ET.SubElement(doc, "asympt_infectivity")
    field.text = str(ASYMPT_INFECTIVITY)
    field = ET.SubElement(doc, "prob_ul_u")
    field.text = str(PROB_UL_U)
    field = ET.SubElement(doc, "prob_ul_v")
    field.text = str(PROB_UL_V)
    field = ET.SubElement(doc, "prob_ur_v")
    field.text = str(PROB_UR_V)
    field = ET.SubElement(doc, "prob_ui_v")
    field.text = str(PROB_UI_V)
    field = ET.SubElement(doc, "prob_ui_u")
    field.text = str(PROB_UI_U)
    field = ET.SubElement(doc, "location_state_refresh_interval")
    field.text = str(LOCATION_STATE_REFRESH_INTERVAL)
    field = ET.SubElement(doc, "seed")
    field.text = str(DISEASE_SEED)

    doc = ET.SubElement(root, "data_capture")
    field = ET.SubElement(doc, "is_needed")
    field.text = str(IS_DATA_CAPTURE_NEEDED)
    if IS_DATA_CAPTURE_NEEDED == "yes":
        field = ET.SubElement(doc, "capture_file")
        field.text = str(DATA_CAPTURE_FILENAME)

    doc = ET.SubElement(root, "number_of_regions")
    doc.text = str(NUMBER_OF_REGIONS)

    pid = 1
    for regionNum in range(NUMBER_OF_REGIONS):
        doc = ET.SubElement(root, "region")
        field = ET.SubElement(doc, "region_name")
        field.text = "Region"+str(regionNum)
        field = ET.SubElement(doc, "number_of_locations")
        numberOfLocations = randNumLocations.randrange(MIN_NUM_OF_LOCS_PER_REGION, MAX_NUM_OF_LOCS_PER_REGION+1)
        field.text = str(numberOfLocations)
        for locationNum in range(numberOfLocations):
            doc1 = ET.SubElement(doc, "location")
            field = ET.SubElement(doc1, "location_name")
            field.text = "Location"+str(locationNum)
            field = ET.SubElement(doc1, "number_of_persons")
            numberOfPersons = randNumPersons.randrange(MIN_NUM_OF_PERSONS_PER_LOC, MAX_NUM_OF_PERSONS_PER_LOC+1)
            field.text = str(numberOfPersons)
            field = ET.SubElement(doc1, "travel_time_to_central_hub")
            field.text = str(randTTCH.randrange(MIN_TRAVEL_TIME_TO_CENTRAL_HUB, MAX_TRAVEL_TIME_TO_CENTRAL_HUB+1))
            field = ET.SubElement(doc1, "diffusion_trigger_interval")
            field.text = str(randDTI.randrange(MIN_DIFFUSION_TRIG_INTERVAL, MAX_DIFFUSION_TRIG_INTERVAL+1))
            for numPerson in range(numberOfPersons):
                doc2 = ET.SubElement(doc1, "person")
                field = ET.SubElement(doc2, "pid")
                field.text = str(pid)
                pid += 1
                field = ET.SubElement(doc2, "susceptibility")
                field.text = str(randSuscep.randrange(0, SUSCEPTIBILITY_PRECISION)/SUSCEPTIBILITY_PRECISION)
                field = ET.SubElement(doc2, "is_vaccinated")
                field.text = str(VACCINATION_STATUS[randIsVacc.randrange(0, len(VACCINATION_STATUS))])
                field = ET.SubElement(doc2, "infection_state")
                field.text = str(INFECTION_STATES[randIS.randrange(0, len(INFECTION_STATES))])

    with open(outFileName, 'w') as f:
        s = ET.tostring(root)
        minidom.parseString(s).writexml(f, addindent='  ' * 4, newl='\n')

if __name__ == "__main__":
    main()

<|endoftext|>


<commit_before>from plenum.common.messages.node_messages import InstanceChange, ViewChangeDone
from plenum.test.delayers import icDelay, vcd_delay
from plenum.test.helper import waitForViewChange
from plenum.test.node_catchup.helper import ensure_all_nodes_have_same_data
from plenum.test.test_node import ensureElectionsDone

nodeCount = 7

def check_vcd_msgs(node, expected_count, expected_view_no):
    if expected_view_no not in node.msgsForFutureViews:
        raise AssertionError('no messages for expected view')
    vcd_count = sum(1 for msg in node.msgsForFutureViews[expected_view_no] if isinstance(msg[0], ViewChangeDone))
    assert expected_count == vcd_count

def check_no_ic_msgs(node, expected_view_no):
    if expected_view_no not in node.msgsForFutureViews:
        raise AssertionError('no messages for expected view')
    ic_count = sum(1 for msg in node.msgsForFutureViews[expected_view_no] if isinstance(msg[0], InstanceChange))
    assert ic_count == 0

def test_delayed_instance_changes_after_vcd_for_next_view(looper, txnPoolNodeSet):
    '''
    A node is doing view change to view=1, while the other nodes already finished view change to view=2.
    The node receives a quorum of VCD messages for view=2 before a quorum of InstanceChange messages for view=2.
    Nevertheless, the node should not start a view change to view=2 without a quorum of InstanceChanges,
    that is it should not go to propagate primary mode since it's already in view chanage state.
    The node should eventually finish view change to view=2 once receives all VCD and IS msgs for view=2
    '''
    nodes = txnPoolNodeSet
    slow_node = nodes[-1]
    fast_nodes = [n for n in nodes if n != slow_node]
    slow_stasher = slow_node.nodeIbStasher

    # 1. DO FIRST VIEW CHANGE

    # delay VCD for the first ViewChange
    slow_stasher.delay(vcd_delay())

    # Trigger view change
    for n in nodes:
        n.view_changer.on_master_degradation()
    waitForViewChange(looper, nodes, expectedViewNo=1)

    # make sure view change is finished on all nodes except the slow one
    ensureElectionsDone(looper, fast_nodes, numInstances=3)

    # 2. DO SECOND VIEW CHANGE

    # delay Instance Changes and drop all VCD to view=1
    # so that the slow node receives VCD for view=2 before
    # a quorum of InstanceChanges for that view while still doing view change to view=1
    slow_stasher.drop_delayeds()
    slow_stasher.resetDelays()
    slow_stasher.delay(icDelay())

    # Trigger view change
    for n in nodes:
        n.view_changer.on_master_degradation()
    waitForViewChange(looper, fast_nodes, expectedViewNo=2)

    # make sure view change is finished on all nodes except the slow one
    ensureElectionsDone(looper, fast_nodes, numInstances=3)

    # slow node is still on view=1
    assert slow_node.viewNo == 1
    assert slow_node.view_change_in_progress

    # make sure that the slow node receives VCD msgs for view=2
    # and didn't receive IS msgs for view=2
    check_vcd_msgs(slow_node, expected_view_no=2, expected_count=len(fast_nodes), )
    check_no_ic_msgs(slow_node, expected_view_no=2)

    # 3. RESET DELAYS AND CHECK
    slow_stasher.reset_delays_and_process_delayeds()

    waitForViewChange(looper, nodes, expectedViewNo=2)
    ensureElectionsDone(looper, nodes, numInstances=3)
    ensure_all_nodes_have_same_data(looper, nodes=nodes)
<commit_msg>improve the test<commit_after>from plenum.common.messages.node_messages import InstanceChange, ViewChangeDone
from plenum.test.delayers import icDelay, vcd_delay
from plenum.test.helper import waitForViewChange
from plenum.test.node_catchup.helper import ensure_all_nodes_have_same_data
from plenum.test.test_node import ensureElectionsDone

nodeCount = 7

def check_vcd_msgs(node, expected_count, expected_view_no):
    if expected_view_no not in node.msgsForFutureViews:
        raise AssertionError('no messages for expected view')
    vcd_count = sum(1 for msg in node.msgsForFutureViews[expected_view_no] if isinstance(msg[0], ViewChangeDone))
    assert expected_count == vcd_count

def check_no_ic_msgs(node, expected_view_no):
    if expected_view_no not in node.msgsForFutureViews:
        raise AssertionError('no messages for expected view')
    ic_count = sum(1 for msg in node.msgsForFutureViews[expected_view_no] if isinstance(msg[0], InstanceChange))
    assert ic_count == 0

def test_delayed_instance_changes_after_vcd_for_next_view(looper, txnPoolNodeSet):
    '''
    A node is doing view change to view=1, while the other nodes already finished view change to view=2.
    The node receives a quorum of VCD messages for view=2 before a quorum of InstanceChange messages for view=2.
    Nevertheless, the node should not start a view change to view=2 without a quorum of InstanceChanges,
    that is it should not go to propagate primary mode since it's already in view chanage state.
    The node should eventually finish view change to view=2 once receives all VCD and IS msgs for view=2
    '''
    nodes = txnPoolNodeSet
    slow_node = nodes[-1]
    fast_nodes = [n for n in nodes if n != slow_node]
    slow_stasher = slow_node.nodeIbStasher

    # 1. DO FIRST VIEW CHANGE

    # delay VCD for the first ViewChange
    slow_stasher.delay(vcd_delay())

    # Trigger view change
    for n in nodes:
        n.view_changer.on_master_degradation()
    waitForViewChange(looper, nodes, expectedViewNo=1)

    # make sure view change is finished on all nodes except the slow one
    ensureElectionsDone(looper, fast_nodes, numInstances=3)

    # 2. DO SECOND VIEW CHANGE

    # delay Instance Changes and drop all VCD to view=1
    # so that the slow node receives VCD for view=2 before
    # a quorum of InstanceChanges for that view while still doing view change to view=1
    slow_stasher.drop_delayeds()
    slow_stasher.resetDelays()
    slow_stasher.delay(icDelay())

    # Trigger view change
    for n in nodes:
        n.view_changer.on_master_degradation()
    waitForViewChange(looper, fast_nodes, expectedViewNo=2)

    # make sure view change is finished on all nodes except the slow one
    ensureElectionsDone(looper, fast_nodes, numInstances=3)

    # slow node is still on view=1
    assert slow_node.viewNo == 1
    assert slow_node.view_change_in_progress

    # make sure that the slow node receives VCD msgs for view=2
    # and didn't receive IS msgs for view=2
    check_vcd_msgs(slow_node, expected_view_no=2, expected_count=len(fast_nodes), )
    check_no_ic_msgs(slow_node, expected_view_no=2)

    # 3. RESET DELAYS AND CHECK
    slow_stasher.reset_delays_and_process_delayeds()

    waitForViewChange(looper, nodes, expectedViewNo=2)
    ensureElectionsDone(looper, nodes)
    assert not slow_node.view_change_in_progress
    ensure_all_nodes_have_same_data(looper, nodes=nodes)
<|endoftext|>


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

__author__ = 'idanmo'

import unittest
import os
from os import path
from worker_installer import init_worker_installer
from worker_installer import DEFAULT_MIN_WORKERS, DEFAULT_MAX_WORKERS
from worker_installer import FabricRunner
from worker_installer.tasks import create_celery_configuration
from worker_installer.tasks import CELERY_INIT_PATH, CELERY_CONFIG_PATH
from cloudify.mocks import MockCloudifyContext
from cloudify.context import BootstrapContext

@init_worker_installer
def m(ctx, runner, worker_config, **kwargs):
    return worker_config

class CeleryWorkerConfigurationTest(unittest.TestCase):

    def setUp(self):
        os.environ['MANAGEMENT_USER'] = 'user'

    def test_deployment_config(self):
        ctx = MockCloudifyContext(deployment_id='deployment_id')
        conf = m(ctx)
        self.assertTrue('base_dir' in conf)
        self.assertTrue('init_file' in conf)
        self.assertTrue('config_file' in conf)

    def test_vm_config_validation(self):
        ctx = MockCloudifyContext(node_id='node',
                                  properties={'worker_config': {
                                      'distro': 'Ubuntu', }})
        self.assertRaises(ValueError, m, ctx)
        ctx = MockCloudifyContext(node_id='node',
                                  properties={
                                      'worker_config': {
                                          'distro': 'Ubuntu', },
                                      'ip': '192.168.0.1'
                                  })
        self.assertRaises(ValueError, m, ctx)
        ctx = MockCloudifyContext(node_id='node',
                                  properties={
                                      'worker_config': {
                                          'user': 'user',
                                          'distro': 'Ubuntu', },
                                      'ip': '192.168.0.1'
                                  })
        self.assertRaises(ValueError, m, ctx)
        ctx = MockCloudifyContext(node_id='node',
                                  properties={
                                      'worker_config': {
                                          'user': 'user',
                                          'key': 'key.pem',
                                          'distro': 'Ubuntu',
                                      },
                                      'ip': '192.168.0.1'
                                  })
        m(ctx)

    def test_worker_config(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'distro': 'Ubuntu',
                }
            }
        )
        conf = m(ctx)
        self.assertTrue('base_dir' in conf)
        self.assertTrue('init_file' in conf)
        self.assertTrue('config_file' in conf)
        self.assertTrue('includes_file' in conf)

    def test_disable_requiretty_config(self):
        self._test_disable_requiretty_config('true', True)
        self._test_disable_requiretty_config('false', False)
        self._test_disable_requiretty_config('true', True)
        self._test_disable_requiretty_config('true', True)
        self._test_disable_requiretty_config(True, True)
        self._test_disable_requiretty_config(False, False)
        self._test_disable_requiretty_config(value=None,
                                             should_raise_exception=True)
        self._test_disable_requiretty_config(value='1234',
                                             should_raise_exception=True)

    def _test_disable_requiretty_config(self,
                                        value=None,
                                        expected=None,
                                        should_raise_exception=False):
        ctx = MockCloudifyContext(
            deployment_id='test',
            properties={
                'worker_config': {
                    'disable_requiretty': value
                }
            }
        )
        if should_raise_exception:
            self.assertRaises(ValueError, m, ctx)
        else:
            conf = m(ctx)
            self.assertEqual(expected, conf['disable_requiretty'])

    def test_autoscale_configuration(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'distro': 'Ubuntu',
                }
            }
        )
        conf = m(ctx)
        self.assertEqual(conf['min_workers'], DEFAULT_MIN_WORKERS)
        self.assertEqual(conf['max_workers'], DEFAULT_MAX_WORKERS)
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'min_workers': 2,
                    'max_workers': 5,
                    'distro': 'Ubuntu',
                }
            }
        )
        conf = m(ctx)
        self.assertEqual(conf['min_workers'], 2)
        self.assertEqual(conf['max_workers'], 5)

    def test_illegal_autoscale_configuration(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'min_workers': 10,
                    'max_workers': 5,
                    'distro': 'Ubuntu',
                }
            }
        )
        self.assertRaises(ValueError, m, ctx)
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'min_workers': 'aaa',
                    'max_workers': 5,
                    'distro': 'Ubuntu',
                }
            }
        )
        self.assertRaises(ValueError, m, ctx)

    def test_autoscale_from_bootstrap_context(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'distro': 'Ubuntu',
                }
            },
            bootstrap_context=BootstrapContext({
                'cloudify_agent': {
                    'min_workers': 2,
                    'max_workers': 5,
                }
            })
        )
        conf = m(ctx)
        self.assertEqual(conf['min_workers'], 2)
        self.assertEqual(conf['max_workers'], 5)

    def test_key_from_bootstrap_context(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                }
            },
            bootstrap_context=BootstrapContext({
                'cloudify_agent': {
                    'agent_key_path': 'here'
                }
            })
        )
        conf = m(ctx)
        self.assertEqual(conf['key'], 'here')

    def test_user_from_bootstrap_context(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            bootstrap_context=BootstrapContext({
                'cloudify_agent': {
                    'agent_key_path': 'here',
                    'user': 'john doe'

                }
            })
        )
        conf = m(ctx)
        self.assertEqual(conf['user'], 'john doe')

    def test_ssh_port_from_bootstrap_context(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            bootstrap_context=BootstrapContext({
                'cloudify_agent': {
                    'agent_key_path': 'here',
                    'user': 'john doe',
                    'remote_execution_port': 2222

                }
            })
        )
        conf = m(ctx)
        self.assertEqual(conf['port'], 2222)

    def test_workflows_worker_config(self):
        ctx = MockCloudifyContext(
            deployment_id='test',
            properties={
                'worker_config': {
                    'workflows_worker': 'true',
                    'distro': 'Ubuntu',
                }
            },
            runtime_properties={
                'ip': '192.168.0.1'
            }
        )
        conf = m(ctx)
        self.assertEqual(conf['name'], 'test_workflows')

class MockFabricRunner(FabricRunner):

    def __init__(self):
        self.put_files = {}

    def put(self, file_path, content, sudo=False):
        self.put_files[file_path] = content

class ConfigurationCreationTest(unittest.TestCase):

    def setUp(self):
        os.environ['MANAGEMENT_USER'] = 'user'
        os.environ['MANAGER_REST_PORT'] = '8100'
        os.environ['MANAGEMENT_IP'] = '192.168.0.1'
        os.environ['AGENT_IP'] = '192.168.0.2'

    def read_file(self, file_name):
        file_path = path.join(path.dirname(__file__), file_name)
        with open(file_path, 'r') as f:
            return f.read()

    def get_resource(self, resource_name):
        if CELERY_INIT_PATH in resource_name:
            return self.read_file('celeryd-cloudify.init.jinja2')
        elif CELERY_CONFIG_PATH in resource_name:
            return self.read_file('celeryd-cloudify.conf.jinja2')
        return None

    def test_prepare_configuration(self):
        ctx = MockCloudifyContext(deployment_id='deployment_id')
        worker_config = m(ctx)
        runner = MockFabricRunner()
        create_celery_configuration(ctx,
                                    runner,
                                    worker_config,
                                    self.get_resource)
        self.assertEquals(3, len(runner.put_files))
        self.assertTrue(worker_config['init_file'] in runner.put_files)
        self.assertTrue(worker_config['config_file'] in runner.put_files)
        self.assertTrue(worker_config['includes_file'] in runner.put_files)
<commit_msg>CFY-781 fixed tests<commit_after>#########
# Copyright (c) 2014 GigaSpaces Technologies Ltd. All rights reserved
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
#  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  * See the License for the specific language governing permissions and
#  * limitations under the License.

__author__ = 'idanmo'

import unittest
import os
from os import path
from worker_installer import init_worker_installer
from worker_installer import DEFAULT_MIN_WORKERS, DEFAULT_MAX_WORKERS
from worker_installer import FabricRunner
from worker_installer.tasks import create_celery_configuration
from worker_installer.tasks import CELERY_INIT_PATH, CELERY_CONFIG_PATH
from cloudify.mocks import MockCloudifyContext
from cloudify.context import BootstrapContext

@init_worker_installer
def m(ctx, runner, worker_config, **kwargs):
    return worker_config

class CeleryWorkerConfigurationTest(unittest.TestCase):

    def setUp(self):
        os.environ['MANAGEMENT_USER'] = 'user'

    def test_deployment_config(self):
        ctx = MockCloudifyContext(deployment_id='deployment_id')
        conf = m(ctx)
        self.assertTrue('base_dir' in conf)
        self.assertTrue('init_file' in conf)
        self.assertTrue('config_file' in conf)

    def test_vm_config_validation(self):
        ctx = MockCloudifyContext(node_id='node',
                                  properties={'worker_config': {
                                      'distro': 'Ubuntu', }})
        self.assertRaises(ValueError, m, ctx)
        ctx = MockCloudifyContext(node_id='node',
                                  properties={
                                      'worker_config': {
                                          'distro': 'Ubuntu', },
                                      'ip': '192.168.0.1'
                                  })
        self.assertRaises(ValueError, m, ctx)
        ctx = MockCloudifyContext(node_id='node',
                                  properties={
                                      'worker_config': {
                                          'user': 'user',
                                          'distro': 'Ubuntu', },
                                      'ip': '192.168.0.1'
                                  })
        self.assertRaises(ValueError, m, ctx)
        ctx = MockCloudifyContext(node_id='node',
                                  properties={
                                      'worker_config': {
                                          'user': 'user',
                                          'key': 'key.pem',
                                          'distro': 'Ubuntu',
                                      },
                                      'ip': '192.168.0.1'
                                  })
        m(ctx)

    def test_worker_config(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'distro': 'Ubuntu',
                }
            }
        )
        conf = m(ctx)
        self.assertTrue('base_dir' in conf)
        self.assertTrue('init_file' in conf)
        self.assertTrue('config_file' in conf)
        self.assertTrue('includes_file' in conf)

    def test_disable_requiretty_config(self):
        self._test_disable_requiretty_config('true', True)
        self._test_disable_requiretty_config('false', False)
        self._test_disable_requiretty_config('true', True)
        self._test_disable_requiretty_config('true', True)
        self._test_disable_requiretty_config(True, True)
        self._test_disable_requiretty_config(False, False)
        self._test_disable_requiretty_config(value=None,
                                             should_raise_exception=True)
        self._test_disable_requiretty_config(value='1234',
                                             should_raise_exception=True)

    def _test_disable_requiretty_config(self,
                                        value=None,
                                        expected=None,
                                        should_raise_exception=False):
        ctx = MockCloudifyContext(
            deployment_id='test',
            properties={
                'worker_config': {
                    'disable_requiretty': value
                }
            }
        )
        if should_raise_exception:
            self.assertRaises(ValueError, m, ctx)
        else:
            conf = m(ctx)
            self.assertEqual(expected, conf['disable_requiretty'])

    def test_autoscale_configuration(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'distro': 'Ubuntu',
                }
            }
        )
        conf = m(ctx)
        self.assertEqual(conf['min_workers'], DEFAULT_MIN_WORKERS)
        self.assertEqual(conf['max_workers'], DEFAULT_MAX_WORKERS)
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'min_workers': 2,
                    'max_workers': 5,
                    'distro': 'Ubuntu',
                }
            }
        )
        conf = m(ctx)
        self.assertEqual(conf['min_workers'], 2)
        self.assertEqual(conf['max_workers'], 5)

    def test_illegal_autoscale_configuration(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'min_workers': 10,
                    'max_workers': 5,
                    'distro': 'Ubuntu',
                }
            }
        )
        self.assertRaises(ValueError, m, ctx)
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'min_workers': 'aaa',
                    'max_workers': 5,
                    'distro': 'Ubuntu',
                }
            }
        )
        self.assertRaises(ValueError, m, ctx)

    def test_autoscale_from_bootstrap_context(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'key': 'key.pem',
                    'distro': 'Ubuntu',
                }
            },
            bootstrap_context=BootstrapContext({
                'cloudify_agent': {
                    'min_workers': 2,
                    'max_workers': 5,
                }
            })
        )
        conf = m(ctx)
        self.assertEqual(conf['min_workers'], 2)
        self.assertEqual(conf['max_workers'], 5)

    def test_key_from_bootstrap_context(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'user': 'user',
                    'distro': 'Ubuntu',
                }
            },
            bootstrap_context=BootstrapContext({
                'cloudify_agent': {
                    'agent_key_path': 'here'
                }
            })
        )
        conf = m(ctx)
        self.assertEqual(conf['key'], 'here')

    def test_user_from_bootstrap_context(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'distro': 'Ubuntu',
                },
            },
            bootstrap_context=BootstrapContext({
                'cloudify_agent': {
                    'agent_key_path': 'here',
                    'user': 'john doe'

                }
            })
        )
        conf = m(ctx)
        self.assertEqual(conf['user'], 'john doe')

    def test_ssh_port_from_bootstrap_context(self):
        node_id = 'node_id'
        ctx = MockCloudifyContext(
            deployment_id='test',
            node_id=node_id,
            runtime_properties={
                'ip': '192.168.0.1'
            },
            properties={
                'worker_config': {
                    'distro': 'Ubuntu',
                },
            },
            bootstrap_context=BootstrapContext({
                'cloudify_agent': {
                    'agent_key_path': 'here',
                    'user': 'john doe',
                    'remote_execution_port': 2222

                }
            })
        )
        conf = m(ctx)
        self.assertEqual(conf['port'], 2222)

    def test_workflows_worker_config(self):
        ctx = MockCloudifyContext(
            deployment_id='test',
            properties={
                'worker_config': {
                    'workflows_worker': 'true',
                    'distro': 'Ubuntu',
                }
            },
            runtime_properties={
                'ip': '192.168.0.1'
            }
        )
        conf = m(ctx)
        self.assertEqual(conf['name'], 'test_workflows')

class MockFabricRunner(FabricRunner):

    def __init__(self):
        self.put_files = {}

    def put(self, file_path, content, sudo=False):
        self.put_files[file_path] = content

class ConfigurationCreationTest(unittest.TestCase):

    def setUp(self):
        os.environ['MANAGEMENT_USER'] = 'user'
        os.environ['MANAGER_REST_PORT'] = '8100'
        os.environ['MANAGEMENT_IP'] = '192.168.0.1'
        os.environ['AGENT_IP'] = '192.168.0.2'

    def read_file(self, file_name):
        file_path = path.join(path.dirname(__file__), file_name)
        with open(file_path, 'r') as f:
            return f.read()

    def get_resource(self, resource_name):
        if CELERY_INIT_PATH in resource_name:
            return self.read_file('celeryd-cloudify.init.jinja2')
        elif CELERY_CONFIG_PATH in resource_name:
            return self.read_file('celeryd-cloudify.conf.jinja2')
        return None

    def test_prepare_configuration(self):
        ctx = MockCloudifyContext(deployment_id='deployment_id')
        worker_config = m(ctx)
        runner = MockFabricRunner()
        create_celery_configuration(ctx,
                                    runner,
                                    worker_config,
                                    self.get_resource)
        self.assertEquals(3, len(runner.put_files))
        self.assertTrue(worker_config['init_file'] in runner.put_files)
        self.assertTrue(worker_config['config_file'] in runner.put_files)
        self.assertTrue(worker_config['includes_file'] in runner.put_files)
<|endoftext|>


<commit_before># Created: 13.03.2011
# Copyright (c) 2011-2020, Manfred Moitzi
# License: MIT License
from typing import TYPE_CHECKING, Iterable, List, Iterator
from itertools import chain

from ezdxf.lldxf.tags import DXFStructureError
from ezdxf.entities import entity_linker

if TYPE_CHECKING:
    from ezdxf.eztypes import (
        TagWriter, Drawing, DXFEntity, Tags, DXFTagStorage, DXFGraphic,
        BlockRecord,
    )

class StoredSection:
    def __init__(self, entities: List['Tags']):
        self.entities = entities

    def export_dxf(self, tagwriter: 'TagWriter'):
        # (0, SECTION) (2, NAME) is stored in entities
        for entity in self.entities:
            tagwriter.write_tags(entity)
        # ENDSEC not stored in entities !!!
        tagwriter.write_str('  0\nENDSEC\n')

class EntitySection:
    """ :class:`EntitiesSection` is just a proxy for :class:`Modelspace` and
    active :class:`Paperspace` linked together.
    """

    def __init__(self, doc: 'Drawing' = None,
                 entities: Iterable['DXFEntity'] = None):
        self.doc = doc
        if entities is not None:
            self._build(iter(entities))

    def __iter__(self) -> Iterable['DXFEntity']:
        """ Iterable for all entities of modelspace and active paperspace. """
        layouts = self.doc.layouts
        for entity in chain(layouts.modelspace(), layouts.active_layout()):
            yield entity

    def __len__(self) -> int:
        """ Returns count of all entities of modelspace and active paperspace.
        """
        layouts = self.doc.layouts
        return len(layouts.modelspace()) + len(layouts.active_layout())

    # none public interface

    def _build(self, entities: Iterator['DXFEntity']) -> None:
        section_head: 'DXFTagStorage' = next(entities)
        if section_head.dxftype() != 'SECTION' or section_head.base_class[
            1] != (2, 'ENTITIES'):
            raise DXFStructureError(
                "Critical structure error in ENTITIES section.")

        def add(entity: 'DXFGraphic'):
            handle = entity.dxf.owner
            # higher priority for owner handle
            if handle == msp_layout_key:
                paperspace = 0
            elif handle == psp_layout_key:
                paperspace = 1
            else:  # paperspace flag as fallback
                paperspace = entity.dxf.paperspace

            if paperspace:
                psp.add_entity(entity)
            else:
                msp.add_entity(entity)

        msp: 'BlockRecord' = self.doc.block_records.get('*Model_Space')
        psp: 'BlockRecord' = self.doc.block_records.get('*Paper_Space')
        msp_layout_key = msp.dxf.handle
        psp_layout_key = psp.dxf.handle
        linked_entities = entity_linker()
        # Don't store linked entities (VERTEX, ATTRIB, SEQEND) in entity space
        for entity in entities:
            if not linked_entities(entity):
                add(entity)

    def export_dxf(self, tagwriter: 'TagWriter') -> None:
        layouts = self.doc.layouts
        tagwriter.write_str("  0\nSECTION\n  2\nENTITIES\n")
        # Just write *Model_Space and the active *Paper_Space into the
        # ENTITIES section.
        layouts.modelspace().entity_space.export_dxf(tagwriter)
        layouts.active_layout().entity_space.export_dxf(tagwriter)
        tagwriter.write_tag2(0, "ENDSEC")
<commit_msg>reformat code by black<commit_after># Copyright (c) 2011-2021, Manfred Moitzi
# License: MIT License
from typing import TYPE_CHECKING, Iterable, List, Iterator
from itertools import chain

from ezdxf.lldxf.tags import DXFStructureError
from ezdxf.entities import entity_linker

if TYPE_CHECKING:
    from ezdxf.eztypes import (
        TagWriter,
        Drawing,
        DXFEntity,
        Tags,
        DXFTagStorage,
        DXFGraphic,
        BlockRecord,
    )

class StoredSection:
    def __init__(self, entities: List["Tags"]):
        self.entities = entities

    def export_dxf(self, tagwriter: "TagWriter"):
        # (0, SECTION) (2, NAME) is stored in entities
        for entity in self.entities:
            tagwriter.write_tags(entity)
        # ENDSEC not stored in entities !!!
        tagwriter.write_str("  0\nENDSEC\n")

class EntitySection:
    """:class:`EntitiesSection` is just a proxy for :class:`Modelspace` and
    active :class:`Paperspace` linked together.
    """

    def __init__(
        self, doc: "Drawing" = None, entities: Iterable["DXFEntity"] = None
    ):
        self.doc = doc
        if entities is not None:
            self._build(iter(entities))

    def __iter__(self) -> Iterable["DXFEntity"]:
        """Iterable for all entities of modelspace and active paperspace."""
        layouts = self.doc.layouts
        for entity in chain(layouts.modelspace(), layouts.active_layout()):
            yield entity

    def __len__(self) -> int:
        """Returns count of all entities of modelspace and active paperspace."""
        layouts = self.doc.layouts
        return len(layouts.modelspace()) + len(layouts.active_layout())

    # none public interface

    def _build(self, entities: Iterator["DXFEntity"]) -> None:
        section_head: "DXFTagStorage" = next(entities)
        if section_head.dxftype() != "SECTION" or section_head.base_class[
            1
        ] != (2, "ENTITIES"):
            raise DXFStructureError(
                "Critical structure error in ENTITIES section."
            )

        def add(entity: "DXFGraphic"):
            handle = entity.dxf.owner
            # higher priority for owner handle
            if handle == msp_layout_key:
                paperspace = 0
            elif handle == psp_layout_key:
                paperspace = 1
            else:  # paperspace flag as fallback
                paperspace = entity.dxf.paperspace

            if paperspace:
                psp.add_entity(entity)
            else:
                msp.add_entity(entity)

        msp: "BlockRecord" = self.doc.block_records.get("*Model_Space")
        psp: "BlockRecord" = self.doc.block_records.get("*Paper_Space")
        msp_layout_key = msp.dxf.handle
        psp_layout_key = psp.dxf.handle
        linked_entities = entity_linker()
        # Don't store linked entities (VERTEX, ATTRIB, SEQEND) in entity space
        for entity in entities:
            if not linked_entities(entity):
                add(entity)

    def export_dxf(self, tagwriter: "TagWriter") -> None:
        layouts = self.doc.layouts
        tagwriter.write_str("  0\nSECTION\n  2\nENTITIES\n")
        # Just write *Model_Space and the active *Paper_Space into the
        # ENTITIES section.
        layouts.modelspace().entity_space.export_dxf(tagwriter)
        layouts.active_layout().entity_space.export_dxf(tagwriter)
        tagwriter.write_tag2(0, "ENDSEC")
<|endoftext|>


<commit_before>import os
import logging
import tensorflow as tf

from FEExpresser import FEExpresser
from util.tf_util import weight_variable, bias_variable, conv2d, max_pool_2x2

class FEGhoshCNN(FEExpresser):
    ''' Implements a the architecture described in
    http://ieeexplore.ieee.org/xpl/login.jsp?tp=&arnumber=7344632&url=http%3A%2F%2Fieeexplore.ieee.org%2Fiel7%2F7332343%2F7344542%2F07344632.pdf%3Farnumber%3D7344632
    '''

    def __init__(self, image_size, codes, repo=None):
        """ Intializes

        :param image_size: The rectangular length the model should expect
        :type image_size: int
        :param codes: List of action units to train on and recognize
        :type codes: [int]
        :param repo: The data repository to pull training examples from. Not necessary to provide on pre-trained models
        :type repo: data.repositories.FACRepository
        """

        logging.info("Initializing FEMultiLabelCNN")
        self.repo = repo
        self.codes = codes
        self.image_size = image_size
        self.isTrained = False

        #Model placholders
        self.x = tf.placeholder(tf.float32, shape=(None, self.image_size, self.image_size, 1), name='x_var')
        self.y_ = tf.placeholder(tf.float32, shape=(None, len(self.codes)), name='y_var')
        self.keep_prob = tf.placeholder(tf.float32, name='keep_prob')

        #Graph variables
        logging.info("Creating network variables")
        self.trainer, self.output, self.ops = self._create_network()

        self.saver = tf.train.Saver()
        self.session = tf.Session()
        self.session.run(self.ops)

    def train(self):
        '''Trains the expression recognition model
            Must be overriden by any extendng class

            :raises: NotImplementedError
        '''
        raise NotImplementedError

    def predict(self, face):
        ''' Given a face image, predicts the corresponding expression

        :param face: Image which matches in size with the parameters of the model
        :type face: OpenCV Mat (numpy ndarray would also probably work)
        :returns: list of expression data
        :raises: NotImplementedError
        '''
        raise NotImplementedError

    def _create_network(self):

        print("Begin creating graph")

    def _multilabel_error(self, label, output):
        p_hat = tf.exp(output) / tf.sum(tf.exp(output))
        return -tf.reduce_mean(tf.sum(label) * tf.log(p_hat))

    def get_image_size(self):

        raise 96<commit_msg>Implemented Ghosh graph<commit_after>import os
import logging
import tensorflow as tf

from FEExpresser import FEExpresser
from util.tf_util import weight_variable, bias_variable, conv2d, max_pool_2x2

KERNEL_SIZE=5
LEARNING_RATE=0.01
WEIGHT_STD = 0.3

FULL_SIZE_1 = 500
FULL_SIZE_2 = 100
FULL_SIZE_3 = 10

class FEGhoshCNN(FEExpresser):
    ''' Implements a the architecture described in
    http://ieeexplore.ieee.org/xpl/login.jsp?tp=&arnumber=7344632&url=http%3A%2F%2Fieeexplore.ieee.org%2Fiel7%2F7332343%2F7344542%2F07344632.pdf%3Farnumber%3D7344632
    '''

    def __init__(self, image_size, codes, repo=None):
        """ Intializes

        :param image_size: The rectangular length the model should expect
        :type image_size: int
        :param codes: List of action units to train on and recognize
        :type codes: [int]
        :param repo: The data repository to pull training examples from. Not necessary to provide on pre-trained models
        :type repo: data.repositories.FACRepository
        """

        logging.info("Initializing FEMultiLabelCNN")
        self.repo = repo
        self.codes = codes
        self.image_size = image_size
        self.isTrained = False

        #Model placholders
        self.x = tf.placeholder(tf.float32, shape=(None, self.image_size, self.image_size, 1), name='x_var')
        self.y_ = tf.placeholder(tf.float32, shape=(None, len(self.codes)), name='y_var')
        self.keep_prob = tf.placeholder(tf.float32, name='keep_prob')

        #Graph variables
        logging.info("Creating network variables")
        self.trainer, self.output, self.ops = self._create_network()

        self.saver = tf.train.Saver()
        self.session = tf.Session()
        self.session.run(self.ops)

    def train(self):
        '''Trains the expression recognition model
            Must be overriden by any extendng class

            :raises: NotImplementedError
        '''
        raise NotImplementedError

    def predict(self, face):
        ''' Given a face image, predicts the corresponding expression

        :param face: Image which matches in size with the parameters of the model
        :type face: OpenCV Mat (numpy ndarray would also probably work)
        :returns: list of expression data
        '''
        raise NotImplementedError

    def _create_network(self):
        logging.info("Begin creating Ghosh graph")

        #First "mega" layer
        conv_W_1 = weight_variable([KERNEL_SIZE, KERNEL_SIZE, 1, 70], name='conv_W_1')
        conv_b_1 = bias_variable([70], name="conv_b_1")
        conv_1 = conv2d(self.x, conv_W_1, 'conv_layer_1')
        relu_1 = tf.nn.relu(conv_1 + conv_b_1, 'relu_1')
        drop_1 = tf.nn.dropout(relu_1, self.keep_prob, name='drop_1')
        max_pool_1 = max_pool_2x2(drop_1, name='pool_1')

        #Second "mega" layer
        conv_W_2 = weight_variable([KERNEL_SIZE, KERNEL_SIZE, 1, 10], name='conv_W_2')
        conv_b_2 = bias_variable([70], name="conv_b_2")
        conv_2 = conv2d(max_pool_1, conv_W_2, 'conv_layer_2')
        relu_2 = tf.nn.relu(conv_2 + conv_b_2, 'relu_2')
        drop_2 = tf.nn.dropout(relu_2, self.keep_prob, name='drop_2')
        max_pool_2 = max_pool_2x2(drop_2, name='pool_2')

        #Fully connected layers

        flattened = tf.reshape(max_pool_2, [-1], name='layer3_pool')

        #1st fully connected
        flat_W_1 = weight_variable(shape=[FULL_SIZE_1, size(flattened)], name='flat_W_1')
        flat_b_1 = bias_variable(shape=[FULL_SIZE_1], name='flat_b_1')
        flat_1 = tf.nn.relu(tf.matmul(flat_W_1, flattened) + flat_b_1, name='flat_1')

        #2nd
        flat_W_2 = weight_variable(shape=[FULL_SIZE_2, FULL_SIZE_1], name='flat_W_2')
        flat_b_2 = bias_variable(shape=[FULL_SIZE_2], name='flat_b_2')
        flat_2 = tf.nn.relu(tf.matmul(flat_W_2, flat_1) + flat_b_2, name='flat_2')

        #3rd
        flat_W_3 = weight_variable(shape=[FULL_SIZE_3, FULL_SIZE_2], name='flat_W_3')
        flat_b_3 = bias_variable(shape=[FULL_SIZE_3], name='flat_b_3')
        flat_3 = tf.nn.relu(tf.matmul(flat_W_3, flat_2) + flat_b_3, name='flat_3')

        cost = self._multilabel_error(self.y_, flat_3)

        tf.scalar_summary('cost/summary', cost)
        train_step = tf.train.AdamOptimizer(LEARNING_RATE).minimize(cost)

        #initialize the variables of the graph
        init_ops = tf.initialize_all_variables()

        return train_step, flat_3, init_ops

    def _multilabel_error(self, label, output):
        p_hat = tf.exp(output) / tf.sum(tf.exp(output))
        return -tf.reduce_mean(tf.sum(label) * tf.log(p_hat))

    def get_image_size(self):
        raise 96<|endoftext|>


<commit_before>from django.core.management.base import BaseCommand

from judge.models import *

class Command(BaseCommand):
    help = 'creates a user'

    def add_arguments(self, parser):
        parser.add_argument('name', help='username')
        parser.add_argument('email', help='email, not necessary to be resolvable')
        parser.add_argument('language', default='PY2', help='default language ID for user')

        parser.add_argument('--superuser', action='store_true', default=False,
                            help="if specified, creates user with superuser privileges")
        parser.add_argument('--staff', action='store_true', default=False,
                            help="if specified, creates user with staff privileges")

    def handle(self, *args, **options):
        usr = User(name=options['name'], email=options['email'], is_active=True)
        usr.is_superuser = options['superuser']
        usr.is_staff = options['staff']
        usr.save()

        profile = Profile(user=usr)
        profile.language = Language.objects.get(key=options['language'])
        profile.save()
<commit_msg>Fix adduser command<commit_after>from django.core.management.base import BaseCommand
from django.contrib.auth.models import User

from judge.models import *

class Command(BaseCommand):
    help = 'creates a user'

    def add_arguments(self, parser):
        parser.add_argument('name', help='username')
        parser.add_argument('email', help='email, not necessary to be resolvable')
        parser.add_argument('password', help='password for the user')
        parser.add_argument('language', nargs='?', default='PY2', help='default language ID for user')

        parser.add_argument('--superuser', action='store_true', default=False,
                            help="if specified, creates user with superuser privileges")
        parser.add_argument('--staff', action='store_true', default=False,
                            help="if specified, creates user with staff privileges")

    def handle(self, *args, **options):
        usr = User(username=options['name'], email=options['email'], is_active=True)
        usr.set_password(options['password'])
        usr.is_superuser = options['superuser']
        usr.is_staff = options['staff']
        usr.save()

        profile = Profile(user=usr)
        profile.language = Language.objects.get(key=options['language'])
        profile.save()
<|endoftext|>