<commit_before>"""
Copyright [2009-2017] EMBL-European Bioinformatics Institute
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.
"""

import operator as op
import itertools as it
from collections import Counter

from caching.base import CachingMixin, CachingManager
from django.core.exceptions import ObjectDoesNotExist
from django.core.urlresolvers import reverse
from django.db import models
from django.db.models import Prefetch, Min, Max, Q
from django.utils.functional import cached_property

from database import Database
from genomic_coordinates import GenomicCoordinates
from modification import Modification
from rna_precomputed import RnaPrecomputed
from reference import Reference
from xref import Xref
from .rfam import RfamHit, RfamAnalyzedSequences
from .accession import Accession

from formatters import Gff3Formatter, GffFormatter, _xref_to_bed_format
from portal.utils import descriptions as desc
from portal.rfam_matches import check_issues

class Rna(CachingMixin, models.Model):
    id = models.IntegerField(db_column='id')
    upi = models.CharField(max_length=13, db_index=True, primary_key=True)
    timestamp = models.DateField()
    userstamp = models.CharField(max_length=30)
    crc64 = models.CharField(max_length=16)
    length = models.IntegerField(db_column='len')
    seq_short = models.CharField(max_length=4000)
    seq_long = models.TextField()
    md5 = models.CharField(max_length=32, unique=True, db_index=True)

    objects = CachingManager()

    class Meta:
        db_table = 'rna'

    def get_absolute_url(self):
        """Get a URL for an RNA object. Used for generating sitemaps."""
        return reverse('unique-rna-sequence', kwargs={'upi': self.upi})

    def get_publications(self, taxid=None):
        """
        Get all publications associated with a Unique RNA Sequence.
        Use raw SQL query for better performance.
        """
        query = """
        SELECT b.id, b.location, b.title, b.pmid as pubmed, b.doi, b.authors
        FROM
            (SELECT DISTINCT t3.id
            FROM xref t1, rnc_reference_map t2, RNC_REFERENCES t3
            WHERE t1.ac = t2.accession AND
                  t1.upi = %s AND
                  {taxid_clause}
                  t2.reference_id = t3.id) a
        JOIN
            rnc_references b
        ON a.id = b.id
        ORDER BY b.title
        """

        query = query.format(taxid_clause='t1.taxid = %s AND' % taxid) if taxid else query.format(taxid_clause='')

        return Reference.objects.raw(query, [self.upi])

    def is_active(self):
        """A sequence is considered active if it has at least one active cross_reference."""
        values_list = self.xrefs.values_list('deleted', flat=True)
        print "self.xrefs.values_list('deleted', flat=True) = %s" % values_list
        return 'N' in self.xrefs.values_list('deleted', flat=True).distinct()  # deleted xrefs are marked with N

    def has_genomic_coordinates(self, taxid=None):
        """Return True if at least one cross-reference has genomic coordinates."""
        xrefs = self.xrefs
        if taxid:
            xrefs = xrefs.filter(taxid=taxid)
        chromosomes = xrefs.all().values_list('accession__coordinates__chromosome', flat=True)
        for chromosome in chromosomes:
            if chromosome:
                return True
        return False

    def get_sequence(self):
        """
        Sequences of up to 4000 nucleotides are stored in seq_short, while the
        longer ones are in stored in seq_long.
        This was due to Oracle column size restrictions.
        """
        if self.seq_short:
            sequence = self.seq_short
        else:
            sequence = self.seq_long
        return sequence.replace('T', 'U').upper()

    def count_symbols(self):
        """
        Returns the number of occurrences of all symbols in RNA,
        including non-canonical nucleotides and random garbage.
        :return: dict {'A': 1, 'T': 2, 'C': 3, 'G': 4, 'N': 5, 'I': 6, '*': 7}
        """
        return dict(Counter(self.get_sequence()))

    def get_xrefs(self, taxid=None):
        """Get all xrefs, show non-ENA annotations first."""
        # Exclude source ENA entries that are associated with other expert db entries.
        # For example, only fetch Vega xrefs and don't retrieve the ENA entries they are based on.
        expert_db_projects = Database.objects.exclude(project_id__isnull=True)\
                                             .values_list('project_id', flat=True)

        xrefs = self.xrefs.filter(deleted='N', upi=self.upi)\
                          .filter(~Q(accession__project__in=expert_db_projects, db__id=1) | Q(accession__project__isnull=True))\
                          .order_by('-db__id')\
                          .select_related()\
                          .prefetch_related(
                              Prefetch(
                                  'modifications',
                                  queryset=Modification.objects.select_related('modification_id')
                              )
                          )\
                          .prefetch_related(
                              Prefetch(
                                  'accession__coordinates',
                                  queryset=GenomicCoordinates.objects.filter(chromosome__isnull=False)
                              )
                          )

        if taxid:
            xrefs = xrefs.filter(taxid=taxid)

        # Sometimes xrefs are deleted from databases (e.g. when by mistake they were
        # annotated as RNA being in fact protein-coding sequences). If our xrefs list
        # doesn't contain proper RNA sequences, we should at least return these
        # wrong annotations to hard-links to deleted sequences accessible from web.
        if xrefs.exists():
            return xrefs
        else:
            return self.xrefs.filter(deleted='Y').select_related()

    def count_xrefs(self, taxid=None):
        """Count the number of cross-references associated with the sequence."""
        xrefs = self.xrefs.filter(db__project_id__isnull=True, deleted='N')
        if taxid:
            xrefs = xrefs.filter(taxid=taxid)
        return xrefs.count()

    @cached_property
    def count_distinct_organisms(self):
        """Count the number of distinct taxids referenced by the sequence."""
        queryset = self.xrefs.values('accession__species')
        results = queryset.filter(deleted='N').distinct().count()
        if not results:
            results = queryset.distinct().count()
        return results

    def get_distinct_database_names(self, taxid=None):
        """Get a non-redundant list of databases referencing the sequence."""
        databases = self.xrefs.filter(deleted='N')
        if taxid:
            databases = databases.filter(taxid=taxid)
        databases = list(databases.values_list('db__display_name', flat=True).distinct())
        databases = sorted(databases, key=lambda s: s.lower())  # case-insensitive
        return databases

    @cached_property
    def first_seen(self):
        """Return the earliest release the sequence is referenced in."""
        data = self.xrefs.aggregate(first_seen=Min('created__release_date'))
        return data['first_seen']

    @cached_property
    def last_seen(self):
        """Like `first_seen` but with reversed order."""
        data = self.xrefs.aggregate(last_seen=Max('last__release_date'))
        return data['last_seen']

    def get_sequence_fasta(self):
        """Split long sequences by a fixed number of characters per line."""
        max_column = 80
        seq = self.get_sequence()
        split_seq = ''
        i = 0
        while i < len(seq):
            split_seq += seq[i:i+max_column] + "\n"
            i += max_column
        description = self.get_description()
        fasta = ">%s %s\n%s" % (self.upi, description, split_seq)
        return fasta

    def get_gff(self):
        """
        Format genomic coordinates from all xrefs into a single file in GFF2 format.
        To reduce redundancy, keep only xrefs from the source entries,
        not the entries added from the DR lines.
        """
        xrefs = self.xrefs.filter(db__project_id__isnull=True).all()
        gff = ''
        for xref in xrefs:
            gff += GffFormatter(xref)()
        return gff

    def get_gff3(self):
        """Format genomic coordinates from all xrefs into a single file in GFF3 format."""
        xrefs = self.xrefs.filter(deleted='N').all()
        gff = '##gff-version 3\n'
        for xref in xrefs:
            gff += Gff3Formatter(xref)()
        return gff

    def get_ucsc_bed(self):
        """
        Format genomic coordinates from all xrefs into a single file in UCSC BED format.
        Example:
        chr1    29554    31097    RNA000063C361    0    +   29554    31097    255,0,0    3    486,104,122    0,1009,1421
        """
        xrefs = self.xrefs.filter(db__project_id__isnull=True).all()
        bed = ''
        for xref in xrefs:
            bed += _xref_to_bed_format(xref)
        return bed

    def get_rna_type(self, taxid=None, recompute=False):
        """Determine the rna type for the given sequence. This will use the
        precomuted data if possible. If not asked to recompute it will do so.
        Providing an taxid will compute the rna_type for the given taxon only.
        This means it will determine the rna_type for only that organism.

        Parameters
        ----------
        taxid : int, None
            The taxon id, if any to use for finding the rna_type.
        recompute : bool, False
            Flag to indicate if this should use the already stored data, or
            recompute it.

        Returns
        -------
        rna_type : str
            The rna type computed for this sequence and possibly taxon id.
        """

        if not recompute:
            queryset = RnaPrecomputed.objects.filter(taxid=taxid)
            try:
                rna_type = queryset.get(upi=self.upi).rna_type
                if rna_type is None:
                    xrefs = self.find_valid_xrefs(taxid=taxid)
                    return desc.get_rna_type(self, xrefs, taxid=taxid)
                return rna_type
            except ObjectDoesNotExist:
                pass

        xrefs = self.find_valid_xrefs(taxid=taxid)
        return desc.get_rna_type(self, xrefs, taxid=taxid)

    def get_description(self, taxid=None, recompute=False):
        """
        Compute the description of this sequence. The description is intented
        to be a good description of the sequence and is based upon the xrefs
        this sequence has. If given a taxid this will produce a description
        that is specific to that species, otherwise it will return a
        description that is general for all species this has been observed in.

        Normally, this will simply lookup a stored description in
        rna_precomputed, however, if recompute=True is given then this will
        recompute the description.

        Parameters
        ----------
        taxid : int, None
            The taxon id to create a description for.

        recompute : bool, False
            If this should compute the description or simply look it up.

        Returns
        -------
        description : str
            The description of this sequence.
        """
        if not recompute:
            if taxid:
                queryset = RnaPrecomputed.objects.filter(taxid=taxid)
            else:
                queryset = RnaPrecomputed.objects.filter(taxid__isnull=True)

            try:
                obj = queryset.get(upi=self.upi)
                return obj.description
            except ObjectDoesNotExist:
                pass

        xrefs = self.find_valid_xrefs(taxid=taxid)
        return desc.get_description(self, xrefs, taxid=taxid)

    def find_valid_xrefs(self, taxid=None):
        """
        Determine the valid xrefs for this sequence and taxid. This will
        attempt to get all active (not deleted) xrefs for the given sequence
        and taxon id. If there are no active xrefs then this will switch to use
        the deleted xrefs.

        taxid : int, None
            The taxon id to use. None indicates no species constraint,
            otherwise the taxid is the taxid of the xref to limit to.

        Returns
        -------
        xrefs : queryset
            The collection of xrefs that are valid for the sequence and taxid.
        """

        base = Xref.objects.filter(upi=self.upi)
        xrefs = base.filter(deleted='N')
        if taxid is not None:
            xrefs = xrefs.filter(taxid=taxid)

        if not xrefs.exists():
            xrefs = base
            if taxid is not None:
                xrefs = xrefs.filter(taxid=taxid)

        return xrefs.select_related('accession', 'db')

    def has_rfam_hits(self):
        """
        Check if this has any Rfam hits.

        :returns bool: True if there are any Rfam matches.
        """
        return bool(self.get_rfam_hits())

    def get_rfam_hits(self, allow_suppressed=True):
        """
        This gets all matches of this Rna to any Rfam hit. Note that this does
        not exclude any families which are supressed. If two families from the
        same clan hit the same sequence then this will selec the hit with the
        lowest e-value and highest score.

        :returns list: A list of all Rfam hits to this sequence.
        """

        query = RfamHit.objects.select_related('rfam_model').\
                                filter(upi=self.upi)
        if not allow_suppressed:
            query = query.filter(rfam_model__is_suppressed=False)
        return query.order_by('rfam_model_id', 'sequence_start')

    def grouped_rfam_hits(self, allow_suppressed=True):
        hits = it.groupby(
            self.get_rfam_hits(allow_suppressed=allow_suppressed),
            op.attrgetter('rfam_model_id')
        )
        results = []
        for _, hits in hits:
            hits = list(hits)
            ranges = []
            for hit in hits:
                ranges.append((
                    hit.sequence_start,
                    hit.sequence_stop,
                    hit.model_completeness,
                ))

            results.append({
                'raw': hits,
                'ranges': ranges,
                'rfam_model': hits[0].rfam_model,
                'rfam_model_id': hits[0].rfam_model_id,
            })
        return sorted(results, key=lambda h: h['ranges'][0][0])

    def get_rfam_status(self, taxid=None):
        return check_issues(self, taxid=taxid)

    def was_rfam_analyzed(self):
        """
        Check if this Rna was analyzed with Rfam. This does not check if there
        are matches, because some sequences will have no match even when
        analyzed, so this looks at the RfamAnalyzedSequences table for this
        information.

        :returns bool: True if this was ever analyzed with Rfam.
        """

        has = bool(RfamAnalyzedSequences.objects.get(upi=self.upi))
        return has

    def get_domains(self, taxid=None, ignore_unclassified=False):
        """
        Get all domains this sequence has been found in. If taxid is given then
        only the domain for all accessions from the given taxid will be
        returned. If ignore_unclassified is given then this will exclude any
        domains where the organism comes from an enviromental sample or is
        uncultured.
        """

        domains = set()
        accessions = Accession.objects.filter(xrefs__upi=self.upi)
        if taxid:
            accessions = accessions.filter(xrefs__taxid=taxid)
        for accession in accessions:
            classification = accession.classification
            if ignore_unclassified:
                if 'uncultured' in classification or \
                        'environmental' in classification:
                    continue
            domains.add(classification.split(';')[0])
        return domains

    def get_rfam_hit_families(self, **kwargs):
        hits = self.get_rfam_hits(**kwargs)
        return sorted(set(hit.rfam_model for hit in hits))
<commit_msg>Use Xref.default_objects<commit_after>"""
Copyright [2009-2017] EMBL-European Bioinformatics Institute
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.
"""

import operator as op
import itertools as it
from collections import Counter

from caching.base import CachingMixin, CachingManager
from django.core.exceptions import ObjectDoesNotExist
from django.core.urlresolvers import reverse
from django.db import models
from django.db.models import Prefetch, Min, Max, Q
from django.utils.functional import cached_property

from database import Database
from genomic_coordinates import GenomicCoordinates
from modification import Modification
from rna_precomputed import RnaPrecomputed
from reference import Reference
from xref import Xref
from .rfam import RfamHit, RfamAnalyzedSequences
from .accession import Accession

from formatters import Gff3Formatter, GffFormatter, _xref_to_bed_format
from portal.utils import descriptions as desc
from portal.rfam_matches import check_issues

class Rna(CachingMixin, models.Model):
    id = models.IntegerField(db_column='id')
    upi = models.CharField(max_length=13, db_index=True, primary_key=True)
    timestamp = models.DateField()
    userstamp = models.CharField(max_length=30)
    crc64 = models.CharField(max_length=16)
    length = models.IntegerField(db_column='len')
    seq_short = models.CharField(max_length=4000)
    seq_long = models.TextField()
    md5 = models.CharField(max_length=32, unique=True, db_index=True)

    objects = CachingManager()

    class Meta:
        db_table = 'rna'

    def get_absolute_url(self):
        """Get a URL for an RNA object. Used for generating sitemaps."""
        return reverse('unique-rna-sequence', kwargs={'upi': self.upi})

    def get_publications(self, taxid=None):
        """
        Get all publications associated with a Unique RNA Sequence.
        Use raw SQL query for better performance.
        """
        query = """
        SELECT b.id, b.location, b.title, b.pmid as pubmed, b.doi, b.authors
        FROM
            (SELECT DISTINCT t3.id
            FROM xref t1, rnc_reference_map t2, RNC_REFERENCES t3
            WHERE t1.ac = t2.accession AND
                  t1.upi = %s AND
                  {taxid_clause}
                  t2.reference_id = t3.id) a
        JOIN
            rnc_references b
        ON a.id = b.id
        ORDER BY b.title
        """

        query = query.format(taxid_clause='t1.taxid = %s AND' % taxid) if taxid else query.format(taxid_clause='')

        return Reference.objects.raw(query, [self.upi])

    def is_active(self):
        """A sequence is considered active if it has at least one active cross_reference."""
        values_list = self.xrefs.values_list('deleted', flat=True)
        print "self.xrefs.values_list('deleted', flat=True) = %s" % values_list
        return 'N' in self.xrefs.values_list('deleted', flat=True).distinct()  # deleted xrefs are marked with N

    def has_genomic_coordinates(self, taxid=None):
        """Return True if at least one cross-reference has genomic coordinates."""
        xrefs = self.xrefs
        if taxid:
            xrefs = xrefs.filter(taxid=taxid)
        chromosomes = xrefs.all().values_list('accession__coordinates__chromosome', flat=True)
        for chromosome in chromosomes:
            if chromosome:
                return True
        return False

    def get_sequence(self):
        """
        Sequences of up to 4000 nucleotides are stored in seq_short, while the
        longer ones are in stored in seq_long.
        This was due to Oracle column size restrictions.
        """
        if self.seq_short:
            sequence = self.seq_short
        else:
            sequence = self.seq_long
        return sequence.replace('T', 'U').upper()

    def count_symbols(self):
        """
        Returns the number of occurrences of all symbols in RNA,
        including non-canonical nucleotides and random garbage.
        :return: dict {'A': 1, 'T': 2, 'C': 3, 'G': 4, 'N': 5, 'I': 6, '*': 7}
        """
        return dict(Counter(self.get_sequence()))

    def get_xrefs(self, taxid=None):
        """Get all xrefs, show non-ENA annotations first."""
        # Exclude source ENA entries that are associated with other expert db entries.
        # For example, only fetch Vega xrefs and don't retrieve the ENA entries they are based on.
        expert_db_projects = Database.objects.exclude(project_id__isnull=True)\
                                             .values_list('project_id', flat=True)

        xrefs = self.xrefs.filter(deleted='N', upi=self.upi)\
                          .filter(~Q(accession__project__in=expert_db_projects, db__id=1) | Q(accession__project__isnull=True))\
                          .order_by('-db__id')\
                          .select_related()\
                          .prefetch_related(
                              Prefetch(
                                  'modifications',
                                  queryset=Modification.objects.select_related('modification_id')
                              )
                          )\
                          .prefetch_related(
                              Prefetch(
                                  'accession__coordinates',
                                  queryset=GenomicCoordinates.objects.filter(chromosome__isnull=False)
                              )
                          )

        if taxid:
            xrefs = xrefs.filter(taxid=taxid)

        # Sometimes xrefs are deleted from databases (e.g. when by mistake they were
        # annotated as RNA being in fact protein-coding sequences). If our xrefs list
        # doesn't contain proper RNA sequences, we should at least return these
        # wrong annotations to hard-links to deleted sequences accessible from web.
        if xrefs.exists():
            return xrefs
        else:
            return self.xrefs.filter(deleted='Y').select_related()

    def count_xrefs(self, taxid=None):
        """Count the number of cross-references associated with the sequence."""
        xrefs = self.xrefs.filter(db__project_id__isnull=True, deleted='N')
        if taxid:
            xrefs = xrefs.filter(taxid=taxid)
        return xrefs.count()

    @cached_property
    def count_distinct_organisms(self):
        """Count the number of distinct taxids referenced by the sequence."""
        queryset = self.xrefs.values('accession__species')
        results = queryset.filter(deleted='N').distinct().count()
        if not results:
            results = queryset.distinct().count()
        return results

    def get_distinct_database_names(self, taxid=None):
        """Get a non-redundant list of databases referencing the sequence."""
        databases = self.xrefs.filter(deleted='N')
        if taxid:
            databases = databases.filter(taxid=taxid)
        databases = list(databases.values_list('db__display_name', flat=True).distinct())
        databases = sorted(databases, key=lambda s: s.lower())  # case-insensitive
        return databases

    @cached_property
    def first_seen(self):
        """Return the earliest release the sequence is referenced in."""
        data = self.xrefs.aggregate(first_seen=Min('created__release_date'))
        return data['first_seen']

    @cached_property
    def last_seen(self):
        """Like `first_seen` but with reversed order."""
        data = self.xrefs.aggregate(last_seen=Max('last__release_date'))
        return data['last_seen']

    def get_sequence_fasta(self):
        """Split long sequences by a fixed number of characters per line."""
        max_column = 80
        seq = self.get_sequence()
        split_seq = ''
        i = 0
        while i < len(seq):
            split_seq += seq[i:i+max_column] + "\n"
            i += max_column
        description = self.get_description()
        fasta = ">%s %s\n%s" % (self.upi, description, split_seq)
        return fasta

    def get_gff(self):
        """
        Format genomic coordinates from all xrefs into a single file in GFF2 format.
        To reduce redundancy, keep only xrefs from the source entries,
        not the entries added from the DR lines.
        """
        xrefs = self.xrefs.filter(db__project_id__isnull=True).all()
        gff = ''
        for xref in xrefs:
            gff += GffFormatter(xref)()
        return gff

    def get_gff3(self):
        """Format genomic coordinates from all xrefs into a single file in GFF3 format."""
        xrefs = self.xrefs.filter(deleted='N').all()
        gff = '##gff-version 3\n'
        for xref in xrefs:
            gff += Gff3Formatter(xref)()
        return gff

    def get_ucsc_bed(self):
        """
        Format genomic coordinates from all xrefs into a single file in UCSC BED format.
        Example:
        chr1    29554    31097    RNA000063C361    0    +   29554    31097    255,0,0    3    486,104,122    0,1009,1421
        """
        xrefs = self.xrefs.filter(db__project_id__isnull=True).all()
        bed = ''
        for xref in xrefs:
            bed += _xref_to_bed_format(xref)
        return bed

    def get_rna_type(self, taxid=None, recompute=False):
        """Determine the rna type for the given sequence. This will use the
        precomuted data if possible. If not asked to recompute it will do so.
        Providing an taxid will compute the rna_type for the given taxon only.
        This means it will determine the rna_type for only that organism.

        Parameters
        ----------
        taxid : int, None
            The taxon id, if any to use for finding the rna_type.
        recompute : bool, False
            Flag to indicate if this should use the already stored data, or
            recompute it.

        Returns
        -------
        rna_type : str
            The rna type computed for this sequence and possibly taxon id.
        """

        if not recompute:
            queryset = RnaPrecomputed.objects.filter(taxid=taxid)
            try:
                rna_type = queryset.get(upi=self.upi).rna_type
                if rna_type is None:
                    xrefs = self.find_valid_xrefs(taxid=taxid)
                    return desc.get_rna_type(self, xrefs, taxid=taxid)
                return rna_type
            except ObjectDoesNotExist:
                pass

        xrefs = self.find_valid_xrefs(taxid=taxid)
        return desc.get_rna_type(self, xrefs, taxid=taxid)

    def get_description(self, taxid=None, recompute=False):
        """
        Compute the description of this sequence. The description is intented
        to be a good description of the sequence and is based upon the xrefs
        this sequence has. If given a taxid this will produce a description
        that is specific to that species, otherwise it will return a
        description that is general for all species this has been observed in.

        Normally, this will simply lookup a stored description in
        rna_precomputed, however, if recompute=True is given then this will
        recompute the description.

        Parameters
        ----------
        taxid : int, None
            The taxon id to create a description for.

        recompute : bool, False
            If this should compute the description or simply look it up.

        Returns
        -------
        description : str
            The description of this sequence.
        """
        if not recompute:
            if taxid:
                queryset = RnaPrecomputed.objects.filter(taxid=taxid)
            else:
                queryset = RnaPrecomputed.objects.filter(taxid__isnull=True)

            try:
                obj = queryset.get(upi=self.upi)
                return obj.description
            except ObjectDoesNotExist:
                pass

        xrefs = self.find_valid_xrefs(taxid=taxid)
        return desc.get_description(self, xrefs, taxid=taxid)

    def find_valid_xrefs(self, taxid=None):
        """
        Determine the valid xrefs for this sequence and taxid. This will
        attempt to get all active (not deleted) xrefs for the given sequence
        and taxon id. If there are no active xrefs then this will switch to use
        the deleted xrefs.

        taxid : int, None
            The taxon id to use. None indicates no species constraint,
            otherwise the taxid is the taxid of the xref to limit to.

        Returns
        -------
        xrefs : queryset
            The collection of xrefs that are valid for the sequence and taxid.
        """

        base = Xref.default_objects.filter(upi=self.upi)
        xrefs = base.filter(deleted='N')
        if taxid is not None:
            xrefs = xrefs.filter(taxid=taxid)

        if not xrefs.exists():
            xrefs = base
            if taxid is not None:
                xrefs = xrefs.filter(taxid=taxid)

        return xrefs.select_related('accession', 'db')

    def has_rfam_hits(self):
        """
        Check if this has any Rfam hits.

        :returns bool: True if there are any Rfam matches.
        """
        return bool(self.get_rfam_hits())

    def get_rfam_hits(self, allow_suppressed=True):
        """
        This gets all matches of this Rna to any Rfam hit. Note that this does
        not exclude any families which are supressed. If two families from the
        same clan hit the same sequence then this will selec the hit with the
        lowest e-value and highest score.

        :returns list: A list of all Rfam hits to this sequence.
        """

        query = RfamHit.objects.select_related('rfam_model').\
                                filter(upi=self.upi)
        if not allow_suppressed:
            query = query.filter(rfam_model__is_suppressed=False)
        return query.order_by('rfam_model_id', 'sequence_start')

    def grouped_rfam_hits(self, allow_suppressed=True):
        hits = it.groupby(
            self.get_rfam_hits(allow_suppressed=allow_suppressed),
            op.attrgetter('rfam_model_id')
        )
        results = []
        for _, hits in hits:
            hits = list(hits)
            ranges = []
            for hit in hits:
                ranges.append((
                    hit.sequence_start,
                    hit.sequence_stop,
                    hit.model_completeness,
                ))

            results.append({
                'raw': hits,
                'ranges': ranges,
                'rfam_model': hits[0].rfam_model,
                'rfam_model_id': hits[0].rfam_model_id,
            })
        return sorted(results, key=lambda h: h['ranges'][0][0])

    def get_rfam_status(self, taxid=None):
        return check_issues(self, taxid=taxid)

    def was_rfam_analyzed(self):
        """
        Check if this Rna was analyzed with Rfam. This does not check if there
        are matches, because some sequences will have no match even when
        analyzed, so this looks at the RfamAnalyzedSequences table for this
        information.

        :returns bool: True if this was ever analyzed with Rfam.
        """

        has = bool(RfamAnalyzedSequences.objects.get(upi=self.upi))
        return has

    def get_domains(self, taxid=None, ignore_unclassified=False):
        """
        Get all domains this sequence has been found in. If taxid is given then
        only the domain for all accessions from the given taxid will be
        returned. If ignore_unclassified is given then this will exclude any
        domains where the organism comes from an enviromental sample or is
        uncultured.
        """

        domains = set()
        accessions = Accession.objects.filter(xrefs__upi=self.upi)
        if taxid:
            accessions = accessions.filter(xrefs__taxid=taxid)
        for accession in accessions:
            classification = accession.classification
            if ignore_unclassified:
                if 'uncultured' in classification or \
                        'environmental' in classification:
                    continue
            domains.add(classification.split(';')[0])
        return domains

    def get_rfam_hit_families(self, **kwargs):
        hits = self.get_rfam_hits(**kwargs)
        return sorted(set(hit.rfam_model for hit in hits))
<|endoftext|>


<commit_before>#
# Copyright (C) 2013 RoboIME
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
import pygame
from numpy import abs

from .. import Skill

class Joystick(Skill):
    """
    This tactic enables joystick control for a robot.
    It uses the library 'pygame'.
    """

    speed_ratio = 2.0
    turbo_ratio = 2.0  # used to multiply speed_ratio
    angle_ratio = 180.0
    power_ratio = 1.0
    deadzone = 0.2

    def __init__(self, robot, relative=True, **kwargs):
        super(Joystick, self).__init__(robot, deterministic=True, **kwargs)
        pygame.init()
        pygame.joystick.init()
        self.check_joysticks()
        self.relative = relative

    def __del__(self):
        pygame.quit()

    def check_joysticks(self, index=0):
        if pygame.joystick.get_count() == 0:
            self.joystick_found = False
            # raise RuntimeError('No joysticks found.')
            print 'WARNING: No joysticks found.'
            print 'Setting keyboard as main controller.'
            pygame.joystick.quit()

            #key_control = pygame.display.set_mode((150, 150))
            pygame.display.set_mode((150, 150))
            pygame.display.set_caption('Keyboard Controller')

        else:
            if index < pygame.joystick.get_count():
                # Available joysticks: ADD NEW JOYSTICK TEMPLATES HERE!
                # name_of_joystick = (numaxes, numhats, numbuttons)
                available_templates = {
                    (5, 1, 10): 'xbox',
                    (6, 0, 15): 'xbox',
                    (3, 0, 11): 'attack3',
                    (5, 1, 12): 'maxprint',
                }

                # Yes, we do have a joystick.
                self.joystick_found = True
                # Assigns a joystick to the corresponding index and initializes it.
                self.joystick = pygame.joystick.Joystick(index)
                self.joystick.init()
                # Not yet in use:
                # self.fishingrod = FishingRod(self.robot, b_angle=0)

                # Acquires the actual template of the joystick
                joystick_topology = (self.joystick.get_numaxes(), self.joystick.get_numhats(), self.joystick.get_numbuttons())
                template = available_templates.get(joystick_topology)
                print 'template:', template

                # Available joysticks: ADD NEW JOYSTICK KEYMAPS HERE!
                if template == 'xbox':
                    # XBOX Template:
                    # --- Buttons:
                    # A: 0
                    # B: 1
                    # X: 2
                    # Y: 3
                    # LB: 4
                    # RB: 5
                    # BACK: 6
                    # START: 7
                    # LDclick: 8
                    # RDclick: 9
                    # --- Axis:
                    # DL x: 0
                    # DL y: 1
                    # LT/RT: 2
                    # DR y: 3
                    # DR x: 4

                    self.chipkick_button = 4
                    self.kick_button = 5
                    self.dribble_button = 3
                    # if there's no need for a straffe, we set it as -1
                    self.straffe_button = 9
                    self.normal_axis = 1   # y
                    self.aux_axis = 0      # x
                    self.power_axis = 2    # power
                    self.straffe_axis = 5  # s
                    self.angle_axis = 4    # a
                    self.turbo_button = None
                    # locks
                    self.lock_x = None
                    self.lock_y = None
                elif template == 'attack3':
                    # ATTACK3 Template:
                    # --- Buttons:
                    # 1: 0
                    # 2: 1
                    # 3: 2
                    # 4: 3
                    # 5: 4
                    # 6: 5
                    # 7: 6
                    # 8: 7
                    # 9: 8
                    # 10: 9
                    # 11: 10
                    # --- Axis:
                    # x: 0
                    # y: 1
                    # +: 2
                    self.chipkick_button = 3
                    self.kick_button = 0
                    self.dribble_button = 2
                    self.straffe_button = 1
                    self.normal_axis = 1
                    self.aux_axis = 0
                    self.power_axis = 2
                    self.turbo_button = None
                    # locks
                    self.lock_x = 7
                    self.lock_y = 8
                elif template == 'maxprint':
                    self.kick_button = 2
                    self.chipkick_button = 1
                    self.dribble_button = 0
                    self.straffe_button = 3
                    self.normal_axis = 1
                    self.aux_axis = 0
                    self.power_axis = 2
                    self.turbo_button = 10
                    # some custom ratios
                    self.angle_ratio = 30
                    self.speed_ratio = 0.5
                    self.power_ratio = 4.0
                    # locks
                    self.lock_x = None
                    self.lock_y = None
                else:
                    print 'ERROR: Unrecognized joystick template.'
                    # fallback to nojoysticks if we have no bindings
                    self.joystick_found = False
            else:
                print 'ERROR: Joystick index not found.'

    def set_speeds(self, speeds):
        if self.relative:
            self.robot.action.speeds = speeds
        else:
            self.robot.action.absolute_speeds = speeds

    def locked_x(self):
        if self.lock_x is None:
            return False
        return self.joystick.get_button(self.lock_x) == 1

    def locked_y(self):
        if self.lock_y is None:
            return False
        return self.joystick.get_button(self.lock_y) == 1

    def _step(self):
        if self.joystick_found:
            # User did something
            for event in pygame.event.get():
                pass
                ## Possible joystick actions: JOYAXISMOTION JOYBALLMOTION JOYBUTTONDOWN JOYBUTTONUP JOYHATMOTION
                #if event.type == pygame.JOYBUTTONDOWN:
                #    print("Joystick button pressed.")
                #if event.type == pygame.JOYBUTTONUP:
                #    print("Joystick button released.")

            # 'attack3-like' joysticks have straffe_button != -1
            if self.straffe_button != -1:
                x = self.joystick.get_axis(self.aux_axis)
                y = -self.joystick.get_axis(self.normal_axis)
                if self.locked_x():
                    y = 0
                if self.locked_y():
                    x = 0
                power = (1 - self.joystick.get_axis(self.power_axis)) / 2
                if self.joystick.get_button(self.kick_button):
                    self.robot.action.kick = power * self.power_ratio
                elif self.joystick.get_button(self.chipkick_button):
                    self.robot.action.chipkick = power * self.power_ratio
                elif self.joystick.get_button(self.dribble_button):
                    self.robot.action.dribble = power * self.power_ratio
                if self.turbo_button is not None:
                    if self.joystick.get_button(self.turbo_button):
                        x *= self.turbo_ratio
                        y *= self.turbo_ratio
                if self.joystick.get_button(self.straffe_button):
                    self.set_speeds((y * self.speed_ratio, -x * self.speed_ratio, 0.0))
                else:
                    self.set_speeds((y * self.speed_ratio, 0.0, -x * self.angle_ratio))

            # 'straffe-ready' joysticks have straffe_button == -1
            else:
                x = -self.joystick.get_axis(self.aux_axis)
                y = -self.joystick.get_axis(self.normal_axis)
                if self.locked_x():
                    y = 0
                if self.locked_y():
                    x = 0
                #s = self.joystick.get_axis(self.straffe_axis)
                a = -self.joystick.get_axis(self.angle_axis)
                power = abs(self.joystick.get_axis(self.power_axis))

                # implementation of the deadzone:
                x = x if abs(x) > self.deadzone else 0
                y = y if abs(y) > self.deadzone else 0
                a = a if abs(a) > self.deadzone else 0
                power = power if abs(power) > self.deadzone else 0

                self.set_speeds((y * self.speed_ratio, x * self.speed_ratio, a * 300))

                if self.joystick.get_button(self.kick_button):
                    self.robot.action.kick = power
                elif self.joystick.get_button(self.chipkick_button):
                    self.robot.action.chipkick = power
                elif self.joystick.get_button(self.dribble_button):
                    self.robot.action.dribble = 1
        else:
            speed = 0.6
            power = 0.8
            pressed = pygame.key.get_pressed()

            # Displacement through asdw
            if pressed[pygame.K_w] and pressed[pygame.K_a]:    # '\
                self.set_speeds((speed, speed, 0))
            elif pressed[pygame.K_s] and pressed[pygame.K_a]:  # ./
                self.set_speeds((-speed, speed, 0))
            elif pressed[pygame.K_s] and pressed[pygame.K_d]:  # \.
                self.set_speeds((-speed, -speed, 0))
            elif pressed[pygame.K_w] and pressed[pygame.K_d]:  # /'
                self.set_speeds((speed, -speed, 0))
            elif pressed[pygame.K_s]:              # v
                self.set_speeds((-speed, 0, 0))
            elif pressed[pygame.K_a]:              # <-o
                self.set_speeds((0, speed, 0))
            elif pressed[pygame.K_w]:              # ^
                self.set_speeds((speed, 0, 0))
            elif pressed[pygame.K_d]:              # o->
                self.set_speeds((0, -speed, 0))
            else:
                self.set_speeds((0, 0, 0))

            # Orientation through <- and ->
            if pressed[pygame.K_LEFT]:
                self.set_speeds(self.robot.action.speeds[:-1] + (self.angle_ratio,))
            elif pressed[pygame.K_RIGHT]:
                self.set_speeds(self.robot.action.speeds[:-1] + (-self.angle_ratio,))
            else:
                self.set_speeds(self.robot.action.speeds[:-1] + (0,))

            # Chipkicking, kicking and dribbling
            if pressed[pygame.K_SPACE] and pressed[pygame.K_LCTRL]:
                self.robot.action.chipkick = power
            elif pressed[pygame.K_SPACE]:
                self.robot.action.kick = power
            else:
                self.robot.action.chipkick = 0
                self.robot.action.kick = 0
            if pressed[pygame.K_LSHIFT]:
                self.robot.action.dribble = 1
            else:
                self.robot.action.dribble = 0

            # TODO: 'Esc' to quit!
            # if pressed[pygame.K_ESCAPE]:
               # self.robot.action.speeds = 0, 0, 0
               # pygame.quit()
<commit_msg>Small improvement on xbox joystick.<commit_after>#
# Copyright (C) 2013 RoboIME
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
import pygame
from numpy import abs

from .. import Skill

class Joystick(Skill):
    """
    This tactic enables joystick control for a robot.
    It uses the library 'pygame'.
    """

    speed_ratio = 2.0
    turbo_ratio = 2.0  # used to multiply speed_ratio
    angle_ratio = 180.0
    power_ratio = 1.0
    deadzone = 0.2

    def __init__(self, robot, relative=True, **kwargs):
        super(Joystick, self).__init__(robot, deterministic=True, **kwargs)
        pygame.init()
        pygame.joystick.init()
        self.check_joysticks()
        self.relative = relative

    def __del__(self):
        pygame.quit()

    def check_joysticks(self, index=0):
        if pygame.joystick.get_count() == 0:
            self.joystick_found = False
            # raise RuntimeError('No joysticks found.')
            print 'WARNING: No joysticks found.'
            print 'Setting keyboard as main controller.'
            pygame.joystick.quit()

            #key_control = pygame.display.set_mode((150, 150))
            pygame.display.set_mode((150, 150))
            pygame.display.set_caption('Keyboard Controller')

        else:
            if index < pygame.joystick.get_count():
                # Available joysticks: ADD NEW JOYSTICK TEMPLATES HERE!
                # name_of_joystick = (numaxes, numhats, numbuttons)
                available_templates = {
                    (5, 1, 10): 'xbox',
                    (6, 0, 15): 'xbox',
                    (3, 0, 11): 'attack3',
                    (5, 1, 12): 'maxprint',
                }

                # Yes, we do have a joystick.
                self.joystick_found = True
                # Assigns a joystick to the corresponding index and initializes it.
                self.joystick = pygame.joystick.Joystick(index)
                self.joystick.init()
                # Not yet in use:
                # self.fishingrod = FishingRod(self.robot, b_angle=0)

                # Acquires the actual template of the joystick
                joystick_topology = (self.joystick.get_numaxes(), self.joystick.get_numhats(), self.joystick.get_numbuttons())
                template = available_templates.get(joystick_topology)
                print 'template:', template

                # Available joysticks: ADD NEW JOYSTICK KEYMAPS HERE!
                if template == 'xbox':
                    # XBOX Template:
                    # --- Buttons:
                    # A: 0
                    # B: 1
                    # X: 2
                    # Y: 3
                    # LB: 4
                    # RB: 5
                    # BACK: 6
                    # START: 7
                    # LDclick: 8
                    # RDclick: 9
                    # --- Axis:
                    # DL x: 0
                    # DL y: 1
                    # LT/RT: 2
                    # DR y: 3
                    # DR x: 4

                    self.chipkick_button = 4
                    self.kick_button = 5
                    self.dribble_button = 3
                    # if there's no need for a straffe, we set it as -1
                    self.straffe_button = 9
                    self.normal_axis = 1   # y
                    self.aux_axis = 0      # x
                    self.power_axis = 2    # power
                    self.straffe_axis = 5  # s
                    self.angle_axis = 4    # a
                    self.turbo_button = None
                    # locks
                    self.lock_x = 1
                    self.lock_y = 0
                elif template == 'attack3':
                    # ATTACK3 Template:
                    # --- Buttons:
                    # 1: 0
                    # 2: 1
                    # 3: 2
                    # 4: 3
                    # 5: 4
                    # 6: 5
                    # 7: 6
                    # 8: 7
                    # 9: 8
                    # 10: 9
                    # 11: 10
                    # --- Axis:
                    # x: 0
                    # y: 1
                    # +: 2
                    self.chipkick_button = 3
                    self.kick_button = 0
                    self.dribble_button = 2
                    self.straffe_button = 1
                    self.normal_axis = 1
                    self.aux_axis = 0
                    self.power_axis = 2
                    self.turbo_button = None
                    # locks
                    self.lock_x = 7
                    self.lock_y = 8
                elif template == 'maxprint':
                    self.kick_button = 2
                    self.chipkick_button = 1
                    self.dribble_button = 0
                    self.straffe_button = 3
                    self.normal_axis = 1
                    self.aux_axis = 0
                    self.power_axis = 2
                    self.turbo_button = 10
                    # some custom ratios
                    self.angle_ratio = 30
                    self.speed_ratio = 0.5
                    self.power_ratio = 4.0
                    # locks
                    self.lock_x = None
                    self.lock_y = None
                else:
                    print 'ERROR: Unrecognized joystick template.'
                    # fallback to nojoysticks if we have no bindings
                    self.joystick_found = False
            else:
                print 'ERROR: Joystick index not found.'

    def set_speeds(self, speeds):
        if self.relative:
            self.robot.action.speeds = speeds
        else:
            self.robot.action.absolute_speeds = speeds

    def locked_x(self):
        if self.lock_x is None:
            return False
        return self.joystick.get_button(self.lock_x) == 1

    def locked_y(self):
        if self.lock_y is None:
            return False
        return self.joystick.get_button(self.lock_y) == 1

    def _step(self):
        if self.joystick_found:
            # User did something
            for event in pygame.event.get():
                pass
                ## Possible joystick actions: JOYAXISMOTION JOYBALLMOTION JOYBUTTONDOWN JOYBUTTONUP JOYHATMOTION
                #if event.type == pygame.JOYBUTTONDOWN:
                #    print("Joystick button pressed.")
                #if event.type == pygame.JOYBUTTONUP:
                #    print("Joystick button released.")

            # 'attack3-like' joysticks have straffe_button != -1
            if self.straffe_button != -1:
                x = self.joystick.get_axis(self.aux_axis)
                y = -self.joystick.get_axis(self.normal_axis)
                if self.locked_x():
                    y = 0
                if self.locked_y():
                    x = 0
                power = (1 - self.joystick.get_axis(self.power_axis)) / 2
                if self.joystick.get_button(self.kick_button):
                    self.robot.action.kick = power * self.power_ratio
                elif self.joystick.get_button(self.chipkick_button):
                    self.robot.action.chipkick = power * self.power_ratio
                elif self.joystick.get_button(self.dribble_button):
                    self.robot.action.dribble = power * self.power_ratio
                if self.turbo_button is not None:
                    if self.joystick.get_button(self.turbo_button):
                        x *= self.turbo_ratio
                        y *= self.turbo_ratio
                if self.joystick.get_button(self.straffe_button):
                    self.set_speeds((y * self.speed_ratio, -x * self.speed_ratio, 0.0))
                else:
                    self.set_speeds((y * self.speed_ratio, 0.0, -x * self.angle_ratio))

            # 'straffe-ready' joysticks have straffe_button == -1
            else:
                x = -self.joystick.get_axis(self.aux_axis)
                y = -self.joystick.get_axis(self.normal_axis)
                if self.locked_x():
                    y = 0
                if self.locked_y():
                    x = 0
                #s = self.joystick.get_axis(self.straffe_axis)
                a = -self.joystick.get_axis(self.angle_axis)
                power = abs(self.joystick.get_axis(self.power_axis))

                # implementation of the deadzone:
                x = x if abs(x) > self.deadzone else 0
                y = y if abs(y) > self.deadzone else 0
                a = a if abs(a) > self.deadzone else 0
                power = power if abs(power) > self.deadzone else 0

                self.set_speeds((y * self.speed_ratio, x * self.speed_ratio, a * 300))

                if self.joystick.get_button(self.kick_button):
                    self.robot.action.kick = power
                elif self.joystick.get_button(self.chipkick_button):
                    self.robot.action.chipkick = power
                elif self.joystick.get_button(self.dribble_button):
                    self.robot.action.dribble = 1
        else:
            speed = 0.6
            power = 0.8
            pressed = pygame.key.get_pressed()

            # Displacement through asdw
            if pressed[pygame.K_w] and pressed[pygame.K_a]:    # '\
                self.set_speeds((speed, speed, 0))
            elif pressed[pygame.K_s] and pressed[pygame.K_a]:  # ./
                self.set_speeds((-speed, speed, 0))
            elif pressed[pygame.K_s] and pressed[pygame.K_d]:  # \.
                self.set_speeds((-speed, -speed, 0))
            elif pressed[pygame.K_w] and pressed[pygame.K_d]:  # /'
                self.set_speeds((speed, -speed, 0))
            elif pressed[pygame.K_s]:              # v
                self.set_speeds((-speed, 0, 0))
            elif pressed[pygame.K_a]:              # <-o
                self.set_speeds((0, speed, 0))
            elif pressed[pygame.K_w]:              # ^
                self.set_speeds((speed, 0, 0))
            elif pressed[pygame.K_d]:              # o->
                self.set_speeds((0, -speed, 0))
            else:
                self.set_speeds((0, 0, 0))

            # Orientation through <- and ->
            if pressed[pygame.K_LEFT]:
                self.set_speeds(self.robot.action.speeds[:-1] + (self.angle_ratio,))
            elif pressed[pygame.K_RIGHT]:
                self.set_speeds(self.robot.action.speeds[:-1] + (-self.angle_ratio,))
            else:
                self.set_speeds(self.robot.action.speeds[:-1] + (0,))

            # Chipkicking, kicking and dribbling
            if pressed[pygame.K_SPACE] and pressed[pygame.K_LCTRL]:
                self.robot.action.chipkick = power
            elif pressed[pygame.K_SPACE]:
                self.robot.action.kick = power
            else:
                self.robot.action.chipkick = 0
                self.robot.action.kick = 0
            if pressed[pygame.K_LSHIFT]:
                self.robot.action.dribble = 1
            else:
                self.robot.action.dribble = 0

            # TODO: 'Esc' to quit!
            # if pressed[pygame.K_ESCAPE]:
               # self.robot.action.speeds = 0, 0, 0
               # pygame.quit()
<|endoftext|>


<commit_before>from urlparse import parse_qs

from moto.ec2.utils import camelcase_to_underscores, method_namess_from_class

from .amazon_dev_pay import AmazonDevPay
from .amis import AmisResponse
from .availability_zones_and_regions import AvailabilityZonesAndRegions
from .customer_gateways import CustomerGateways
from .dhcp_options import DHCPOptions
from .elastic_block_store import ElasticBlockStore
from .elastic_ip_addresses import ElasticIPAddresses
from .elastic_network_interfaces import ElasticNetworkInterfaces
from .general import General
from .internet_gateways import InternetGateways
from .ip_addresses import IPAddresses
from .key_pairs import KeyPairs
from .monitoring import Monitoring
from .network_acls import NetworkACLs
from .placement_groups import PlacementGroups
from .reserved_instances import ReservedInstances
from .route_tables import RouteTables
from .security_groups import SecurityGroups
from .spot_instances import SpotInstances
from .subnets import Subnets
from .virtual_private_gateways import VirtualPrivateGateways
from .vm_export import VMExport
from .vm_import import VMImport
from .vpcs import VPCs
from .vpn_connections import VPNConnections
from .windows import Windows

from .instances import InstanceResponse
from .tags import TagResponse

class EC2Response(object):

    sub_responses = [InstanceResponse, TagResponse, AmisResponse, AvailabilityZonesAndRegions, SecurityGroups]

    def dispatch(self, uri, body, headers):
        if body:
            querystring = parse_qs(body)
        else:
            querystring = parse_qs(headers)

        action = querystring['Action'][0]
        action = camelcase_to_underscores(action)

        for sub_response in self.sub_responses:
            method_names = method_namess_from_class(sub_response)
            if action in method_names:
                response = sub_response(querystring)
                method = getattr(response, action)
                return method()
        raise NotImplementedError("The {} action has not been implemented".format(action))
<commit_msg>clean up ec2 response init<commit_after>from urlparse import parse_qs

from moto.ec2.utils import camelcase_to_underscores, method_namess_from_class

from .amazon_dev_pay import AmazonDevPay
from .amis import AmisResponse
from .availability_zones_and_regions import AvailabilityZonesAndRegions
from .customer_gateways import CustomerGateways
from .dhcp_options import DHCPOptions
from .elastic_block_store import ElasticBlockStore
from .elastic_ip_addresses import ElasticIPAddresses
from .elastic_network_interfaces import ElasticNetworkInterfaces
from .general import General
from .instances import InstanceResponse
from .internet_gateways import InternetGateways
from .ip_addresses import IPAddresses
from .key_pairs import KeyPairs
from .monitoring import Monitoring
from .network_acls import NetworkACLs
from .placement_groups import PlacementGroups
from .reserved_instances import ReservedInstances
from .route_tables import RouteTables
from .security_groups import SecurityGroups
from .spot_instances import SpotInstances
from .subnets import Subnets
from .tags import TagResponse
from .virtual_private_gateways import VirtualPrivateGateways
from .vm_export import VMExport
from .vm_import import VMImport
from .vpcs import VPCs
from .vpn_connections import VPNConnections
from .windows import Windows

class EC2Response(object):

    sub_responses = [
        AmazonDevPay,
        AmisResponse,
        AvailabilityZonesAndRegions,
        CustomerGateways,
        DHCPOptions,
        ElasticBlockStore,
        ElasticIPAddresses,
        ElasticNetworkInterfaces,
        General,
        InstanceResponse,
        InternetGateways,
        IPAddresses,
        KeyPairs,
        Monitoring,
        NetworkACLs,
        PlacementGroups,
        ReservedInstances,
        RouteTables,
        SecurityGroups,
        SpotInstances,
        Subnets,
        TagResponse,
        VirtualPrivateGateways,
        VMExport,
        VMImport,
        VPCs,
        VPNConnections,
        Windows,
    ]

    def dispatch(self, uri, body, headers):
        if body:
            querystring = parse_qs(body)
        else:
            querystring = parse_qs(headers)

        action = querystring['Action'][0]
        action = camelcase_to_underscores(action)

        for sub_response in self.sub_responses:
            method_names = method_namess_from_class(sub_response)
            if action in method_names:
                response = sub_response(querystring)
                method = getattr(response, action)
                return method()
        raise NotImplementedError("The {} action has not been implemented".format(action))
<|endoftext|>


<commit_before>from django.db import models
from django.db.models.query import QuerySet
import datetime
import re
"""
    build_file_template will be the script
    that gets copied into the mock environment
    this will ultimately be what builds the package
    Getting the output logs out requires this to be
    the case unfortunately
"""
build_file_template = """#!/bin/bash
rm -rf /tmp/build
mkdir /tmp/build
`%s 2> /tmp/errors 1> /tmp/log`
"""
class MozillaPackage(models.Model):
    arch_type = models.CharField(max_length=128)
    output_type = models.CharField(max_length=128)
    build_package_name = models.CharField(max_length=128, blank=True, null=True)
    build_status = models.CharField(max_length=128, blank=True, null=True)
    rhel_version = models.CharField(max_length=128)
    input_type = models.CharField(max_length=128)
    """
        package has the label Remote Package
    """
    package = models.CharField(max_length=128)
    prefix_dir = models.CharField(max_length=128)
    """
        install_package_name has the label Package Name
        this very well will end up being redundant and deprecated
        The idea is that you can give the output package a new
        name other than what it is known as via the uploaded
        file or via pypi/gem
    """
    install_package_name = models.CharField(max_length=128)
    package_version = models.CharField(max_length=128)
    conflicts = models.CharField(max_length=128)
    provides = models.CharField(max_length=128)
    package_url = models.CharField(max_length=128)
    application_group = models.CharField(max_length=128)
    celery_id = models.CharField(max_length=128, blank=True, null=True)
    created_on = models.DateTimeField(blank=True, null=True)
    completed_on = models.DateTimeField(blank=True, null=True)
    updated_on = models.DateTimeField(blank=True, null=True)
    search_fields = (
            'arch_type',
            'package'
        )

    class Meta:
        db_table = 'mozilla_package'

    @property
    def os(self):
        return self.os

    @property
    def version(self):
        return self.version

    def add_log(self, log_type, log_message):
        MozillaPackageLog(mozilla_package = self,
                log_type = log_type,
                log_message = log_message).save()

    def generate_build_string(self):
        """
            build_string will be our computed
            fpm command that builds the package
        """
        dependency_string = ""
        version_string = "-v %s" % self.package_version if self.package_version != '' else ""
        if self.mozillapackagedependency_set:
            for dep in self.mozillapackagedependency_set.all():
                dependency_string = "%s -d \"%s\"" % (dependency_string, dep.name)

        if self.input_type == 'python' or self.input_type == 'gem':
            build_string = "setarch %s fpm -s %s -t %s %s %s %s" % (
                    self.arch_type,
                    self.input_type,
                    self.output_type,
                    version_string,
                    dependency_string,
                    self.install_package_name,
                    )
        """
            Clean up the build string a bit to remove double spaces
            that can get added if no version or dependencies
        """
        build_string = re.sub('\s+', ' ', build_string)
        return str(build_string)

    def generate_build_file_content(self):
        build_file_content = build_file_template % self.generate_build_string()
        return build_file_content

    def write_build_file(self):
        file_content = self.generate_build_file_content()
        pass

    def save(self, *args, **kwargs):
        if not self.id:
            self.created_on = datetime.datetime.now()
            self.build_status = 'Building'
        self.updated_on = datetime.datetime.now()
        super(MozillaPackage, self).save(*args, **kwargs)

class MozillaPackageLog(models.Model):
    mozilla_package = models.ForeignKey('MozillaPackage', null=False, blank=False)
    log_type = models.CharField(max_length=128)
    log_message = models.TextField()
    log_time = models.DateTimeField(blank = False, null=False, default=datetime.datetime.now)

    @property
    def log_message_to_html(self):
        """
            Here we want to take the log message and replace the
            newline characters with
 for html output
        """
        message_list = self.log_message.split("\n")
        return "
".join(message_list)

    class Meta:
        db_table = 'mozilla_package_log'

class MozillaPackageDependency(models.Model):
    mozilla_package = models.ForeignKey('MozillaPackage', null=False, blank=False)
    name = models.CharField(max_length=128)

    class Meta:
        db_table = 'mozilla_package_dependency'
<commit_msg>Removed vendor and os properties from model<commit_after>from django.db import models
from django.db.models.query import QuerySet
import datetime
import re
"""
    build_file_template will be the script
    that gets copied into the mock environment
    this will ultimately be what builds the package
    Getting the output logs out requires this to be
    the case unfortunately
"""
build_file_template = """#!/bin/bash
rm -rf /tmp/build
mkdir /tmp/build
`%s 2> /tmp/errors 1> /tmp/log`
"""
class MozillaPackage(models.Model):
    arch_type = models.CharField(max_length=128)
    output_type = models.CharField(max_length=128)
    build_package_name = models.CharField(max_length=128, blank=True, null=True)
    build_status = models.CharField(max_length=128, blank=True, null=True)
    rhel_version = models.CharField(max_length=128)
    input_type = models.CharField(max_length=128)
    """
        package has the label Remote Package
    """
    package = models.CharField(max_length=128)
    prefix_dir = models.CharField(max_length=128)
    """
        install_package_name has the label Package Name
        this very well will end up being redundant and deprecated
        The idea is that you can give the output package a new
        name other than what it is known as via the uploaded
        file or via pypi/gem
    """
    install_package_name = models.CharField(max_length=128)
    package_version = models.CharField(max_length=128)
    conflicts = models.CharField(max_length=128)
    provides = models.CharField(max_length=128)
    package_url = models.CharField(max_length=128)
    application_group = models.CharField(max_length=128)
    celery_id = models.CharField(max_length=128, blank=True, null=True)
    created_on = models.DateTimeField(blank=True, null=True)
    completed_on = models.DateTimeField(blank=True, null=True)
    updated_on = models.DateTimeField(blank=True, null=True)
    search_fields = (
            'arch_type',
            'package'
        )

    class Meta:
        db_table = 'mozilla_package'

    def add_log(self, log_type, log_message):
        MozillaPackageLog(mozilla_package = self,
                log_type = log_type,
                log_message = log_message).save()

    def generate_build_string(self):
        """
            build_string will be our computed
            fpm command that builds the package
        """
        dependency_string = ""
        version_string = "-v %s" % self.package_version if self.package_version != '' else ""
        if self.mozillapackagedependency_set:
            for dep in self.mozillapackagedependency_set.all():
                dependency_string = "%s -d \"%s\"" % (dependency_string, dep.name)

        if self.input_type == 'python' or self.input_type == 'gem':
            build_string = "setarch %s fpm -s %s -t %s %s %s %s" % (
                    self.arch_type,
                    self.input_type,
                    self.output_type,
                    version_string,
                    dependency_string,
                    self.install_package_name,
                    )
        """
            Clean up the build string a bit to remove double spaces
            that can get added if no version or dependencies
        """
        build_string = re.sub('\s+', ' ', build_string)
        return str(build_string)

    def generate_build_file_content(self):
        build_file_content = build_file_template % self.generate_build_string()
        return build_file_content

    def write_build_file(self):
        file_content = self.generate_build_file_content()
        pass

    def save(self, *args, **kwargs):
        if not self.id:
            self.created_on = datetime.datetime.now()
            self.build_status = 'Building'
        self.updated_on = datetime.datetime.now()
        super(MozillaPackage, self).save(*args, **kwargs)

class MozillaPackageLog(models.Model):
    mozilla_package = models.ForeignKey('MozillaPackage', null=False, blank=False)
    log_type = models.CharField(max_length=128)
    log_message = models.TextField()
    log_time = models.DateTimeField(blank = False, null=False, default=datetime.datetime.now)

    @property
    def log_message_to_html(self):
        """
            Here we want to take the log message and replace the
            newline characters with
 for html output
        """
        message_list = self.log_message.split("\n")
        return "
".join(message_list)

    class Meta:
        db_table = 'mozilla_package_log'

class MozillaPackageDependency(models.Model):
    mozilla_package = models.ForeignKey('MozillaPackage', null=False, blank=False)
    name = models.CharField(max_length=128)

    class Meta:
        db_table = 'mozilla_package_dependency'
<|endoftext|>


<commit_before>from plugins.contracts import ICartridgeAgentPlugin
from modules.util.log import LogFactory
from modules.topology.topologycontext import *
import mdsclient

class WkaMemberConfigurator(ICartridgeAgentPlugin):

    def __init__(self):
        self.log = None
        self.my_member_id = None

    def publish_metadata(self, properties_data):
        publish_data = mdsclient.MDSPutRequest()
        publish_data.properties = properties_data
        mdsclient.put(publish_data, app=True)

    def add_to_restart_queue(self, member_id):
        data = {"key":"restart","values":member_id}
        self.log.info("Publishing members to be restarted data=%s" % (data))
        self.publish_metadata(data)

    def publish_as_wka_member(self, ip):
        private_ip = ip
        data = {"key":"wka","values":private_ip}
        self.log.info("Publishing wka members data=%s " % (data))
        self.publish_metadata(data)

    def remove_me_from_queue(self):
        self.log.info("Removing me %s from restart queue" % self.my_member_id);
        mdsclient.delete_property_value("restart", self.my_member_id)

    def publish_wka_members(self, service_name, cluster_id):
        topology = TopologyContext.get_topology()
        service = topology.get_service(service_name)
        cluster = service.get_cluster(cluster_id)

        members = cluster.get_members()
        for member in members:
            if(member.member_id == self.my_member_id):
                self.log.info("My Ips %s , %s" % (member.member_default_private_ip, member.member_default_public_ip))
                self.publish_as_wka_member(member.member_default_private_ip)
            else:
                self.log.info("Other WKA members memberid=%s privateip=%s, public ip=%s " % (member.member_id, member.member_default_private_ip, member.member_default_public_ip))
                self.add_to_restart_queue(member.member_id)

        #configure me with other wka members
        # remove me from queue if i am there

        return None, None

    @staticmethod
    def isTrue(self, str):
        #should be an utility method
        return str.lower() in ("true", "True", "1" , "yes", "Yes")

    def fetch_wka_members(self):
        mds_response = mdsclient.get(app=True)
        wka_members= None
        if mds_response is not None:
            wka_members = mds_response.properties.get("wka")

        self.log.info("WKA members %s " % wka_members);

    def run_plugin(self, values):
        self.log = LogFactory().get_log(__name__)
        self.log.info("Starting Clustering Configuration")

        clusterId = values['CLUSTER_ID']
        self.log.info("CLUSTER_ID %s" % clusterId)

        service_name = values['SERVICE_NAME']
        self.log.info("SERVICE_NAME %s" % service_name)

        cluering_type = values['CLUSTERING_TYPE']
        self.log.info("CLUSTERING_TYPE %s" % cluering_type)

        is_wka_member = values['WKA_MEMBER']
        self.log.info("WKA_MEMBER %s" % is_wka_member)

        self.my_member_id = values['MEMBER_ID']
        self.log.info("MEMBER_ID %s" % self.my_member_id)

        if self.is_wka(WkaMemberConfigurator.isTrue(is_wka_member)):
            self.log.info("This is a WKA member")
            self.remove_me_from_queue()
            self.get_all_members(service_name, clusterId)
        else:
            self.log.info("This is not a WKA member")
            self.fetch_wka_members()
<commit_msg>expose wka members as  evn variables<commit_after>import os
from plugins.contracts import ICartridgeAgentPlugin
from modules.util.log import LogFactory
from modules.topology.topologycontext import *
import mdsclient

class WkaMemberConfigurator(ICartridgeAgentPlugin):

    def __init__(self):
        self.log = None
        self.my_member_id = None

    def publish_metadata(self, properties_data):
        publish_data = mdsclient.MDSPutRequest()
        publish_data.properties = properties_data
        mdsclient.put(publish_data, app=True)

    def add_to_restart_queue(self, member_id):
        data = {"key":"restart","values":member_id}
        self.log.info("Publishing members to be restarted data=%s" % (data))
        self.publish_metadata(data)

    def publish_as_wka_member(self, ip):
        private_ip = ip
        data = {"key":"wka","values":private_ip}
        self.log.info("Publishing wka members data=%s " % (data))
        self.publish_metadata(data)

    def remove_me_from_queue(self):
        self.log.info("Removing me %s from restart queue" % self.my_member_id);
        mdsclient.delete_property_value("restart", self.my_member_id)

    def publish_wka_members(self, service_name, cluster_id):
        topology = TopologyContext.get_topology()
        service = topology.get_service(service_name)
        cluster = service.get_cluster(cluster_id)

        members = cluster.get_members()
        wka_members=[]
        local_member_port=4000
        for member in members:
            if(member.member_id == self.my_member_id):
                self.log.info("My Ips %s , %s" % (member.member_default_private_ip, member.member_default_public_ip))
                self.publish_as_wka_member(member.member_default_private_ip)
            else:
                self.log.info("Other WKA members memberid=%s privateip=%s, public ip=%s " % (member.member_id, member.member_default_private_ip, member.member_default_public_ip))
                wka_members.append(member.member_default_private_ip+':'+str(local_member_port))
                self.add_to_restart_queue(member.member_id)

        #configure me with other wka members
        # remove me from queue if i am there
        local_members = ','.join(map(str, wka_members))
        local_members= "'{}'".format(local_members)
        self.log.info("*** local_members=%s " % (local_members))

        os.environ['STRATOS_MEMBERS'] = str(local_members)
        self.log.info("env local members=%s" % (os.environ.get('STRATOS_MEMBERS')))

        return None, None

    @staticmethod
    def isTrue(str):
        #should be an utility method
        return str.lower() in ("true", "True", "1" , "yes", "Yes")

    def fetch_wka_members(self):
        wka_members=[]
        local_member_port=4000
        mds_response = mdsclient.get(app=True)
        wka_members= None
        if mds_response is not None:
            wka_members = mds_response.properties.get("wka")

        self.log.info("WKA members %s " % wka_members);

    def execute_clustring_configurater(self):
        #invoke clustering confugurater
        pass

    def run_plugin(self, values):
        self.log = LogFactory().get_log(__name__)
        self.log.info("Starting Clustering Configuration")

        clusterId = values['CLUSTER_ID']
        self.log.info("CLUSTER_ID %s" % clusterId)

        service_name = values['SERVICE_NAME']
        self.log.info("SERVICE_NAME %s" % service_name)

        cluering_type = values['CLUSTERING_TYPE']
        self.log.info("CLUSTERING_TYPE %s" % cluering_type)

        is_wka_member = values['WKA_MEMBER']
        self.log.info("WKA_MEMBER %s" % is_wka_member)

        self.my_member_id = values['MEMBER_ID']
        self.log.info("MEMBER_ID %s" % self.my_member_id)

        sub_domain = values['SUB_DOMAIN']
        sub_domain= "'{}'".format(sub_domain)
        self.log.info("SUB_DOMAIN %s" % (sub_domain))

        os.environ['STRATOS_SUB_DOMAIN'] = str(sub_domain)
        self.log.info("env clustering  SUB_DOMAIN=%s" % (os.environ.get('SUB_DOMAIN','worker')))

        if WkaMemberConfigurator.isTrue(is_wka_member):
            self.log.info("This is a WKA member")
            self.remove_me_from_queue()
            self.publish_wka_members(service_name, clusterId)
        else:
            self.log.info("This is not a WKA member")
            self.fetch_wka_members()

        self.execute_clustring_configurater()
<|endoftext|>


<commit_before># This file is part of the MapProxy project.
# Copyright (C) 2013 Omniscale GmbH & Co. KG <http://omniscale.com>
#
# 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.

import time
import multiprocessing

from mp_renderd.queue import STOP
from mp_renderd.task import Task

from nose.tools import eq_
from mp_renderd.worker import BaseWorker
from mp_renderd.worker import SeedWorker

class SleepWorker(BaseWorker):
    def do_sleep(self, doc):
        time.sleep(doc['time'])
        return {}

class TestSleepWorker(object):
    def setup(self):
        self.in_queue = multiprocessing.Queue(2)
        self.out_queue = multiprocessing.Queue(2)
        self.worker = SleepWorker(in_queue=self.in_queue, out_queue=self.out_queue)

    def test_start_stop(self):
        self.worker.start()
        assert self.worker.is_alive()

        self.in_queue.put(STOP)
        time.sleep(0.1)
        assert not self.worker.is_alive()

    def test_dispatch(self):
        self.worker.dispatch(Task('foo', doc={'command': 'sleep', 'time': 0}))
        assert self.worker.handle_task_message()
        result = self.out_queue.get()
        eq_(result.doc, {'status': 'ok'})

class TestBaseWorker(object):
    def setup(self):
        self.in_queue = multiprocessing.Queue(2)
        self.out_queue = multiprocessing.Queue(2)
        self.worker = BaseWorker(in_queue=self.in_queue, out_queue=self.out_queue)

    def test_dispatch_unknown_command(self):
        self.worker.dispatch(Task('foo', doc={'command': 'foo'}))
        assert self.worker.handle_task_message()
        result = self.out_queue.get()
        eq_(result.doc, {
            'status': 'error',
            'error_message': 'unknown command: foo',
        })
    def test_stop(self):
        self.worker.in_queue.put(STOP)
        assert not self.worker.handle_task_message()

class ExceptionWorker(BaseWorker):
    def do_exception(self, doc):
        raise ValueError('foo')

class TestSleepWorker(object):
    def setup(self):
        self.in_queue = multiprocessing.Queue(2)
        self.out_queue = multiprocessing.Queue(2)
        self.worker = ExceptionWorker(in_queue=self.in_queue, out_queue=self.out_queue)

    def test_dispatch(self):
        self.worker.dispatch(Task('foo', doc={'command': 'exception'}))
        assert self.worker.handle_task_message()
        result = self.out_queue.get()
        eq_(result.doc['status'], 'error')
        eq_(result.doc['error_message'], "exception while processing 'exception': foo")
        assert 'raise ValueError' in result.doc['error_detail']

class DummyCache(object):
    def __init__(self):
        self.requested_tiles = []

    def load_tile_coords(self, tiles):
        self.requested_tiles.extend(tiles)

class TestSeedWorker(object):
    def setup(self):
        self.in_queue = multiprocessing.Queue(2)
        self.out_queue = multiprocessing.Queue(2)
        self.caches = {
            'test_cache': DummyCache()
        }
        self.worker = SeedWorker(
            caches=self.caches,
            base_config={},
            in_queue=self.in_queue,
            out_queue=self.out_queue,
        )

    def test_create_tile(self):
        self.worker.dispatch(Task('foo', doc={'command': 'tile', 'cache_identifier': 'test_cache', 'tiles': [[0, 0, 0]]}))
        assert self.worker.handle_task_message()
        result = self.out_queue.get()
        eq_(result.doc, {'status': 'ok'})
        eq_(self.caches['test_cache'].requested_tiles, [(0, 0, 0)])

    def test_create_tile_unknown_cache(self):
        self.worker.dispatch(Task('foo', doc={'command': 'tile', 'cache_identifier': 'unknown', 'tiles': [[0, 0, 0]]}))
        assert self.worker.handle_task_message()
        result = self.out_queue.get()
        eq_(result.doc, {'status': 'error', 'error_message': "unknown cache 'unknown'"})

        eq_(self.caches['test_cache'].requested_tiles, [])

    def test_start_create_tiles_stop(self):
        # use manager.list to receive requested_tiles from worker process
        manager = multiprocessing.Manager()
        self.caches['test_cache'].requested_tiles = manager.list()
        self.worker.start()

        self.worker.dispatch(Task('foo', doc={'command': 'tile', 'cache_identifier': 'test_cache', 'tiles': [[0, 0, 0]]}))
        self.worker.dispatch(Task('foo', doc={'command': 'tile', 'cache_identifier': 'test_cache', 'tiles': [[5, 0, 0], [5, 1, 0], [5, 2, 0]]}))
        self.worker.dispatch(Task('foo', doc={'command': 'tile', 'cache_identifier': 'test_cache', 'tiles': [[2, 3, 4]]}))
        self.in_queue.put(STOP)

        for _ in range(3):
            result = self.out_queue.get()
            eq_(result.doc, {'status': 'ok'})

        eq_(list(self.caches['test_cache'].requested_tiles), [(0, 0, 0), (5, 0, 0), (5, 1, 0), (5, 2, 0), (2, 3, 4)])<commit_msg>fixed class name of test<commit_after># This file is part of the MapProxy project.
# Copyright (C) 2013 Omniscale GmbH & Co. KG <http://omniscale.com>
#
# 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.

import time
import multiprocessing

from mp_renderd.queue import STOP
from mp_renderd.task import Task

from nose.tools import eq_
from mp_renderd.worker import BaseWorker
from mp_renderd.worker import SeedWorker

class SleepWorker(BaseWorker):
    def do_sleep(self, doc):
        time.sleep(doc['time'])
        return {}

class TestSleepWorker(object):
    def setup(self):
        self.in_queue = multiprocessing.Queue(2)
        self.out_queue = multiprocessing.Queue(2)
        self.worker = SleepWorker(in_queue=self.in_queue, out_queue=self.out_queue)

    def test_start_stop(self):
        self.worker.start()
        assert self.worker.is_alive()

        self.in_queue.put(STOP)
        time.sleep(0.1)
        assert not self.worker.is_alive()

    def test_dispatch(self):
        self.worker.dispatch(Task('foo', doc={'command': 'sleep', 'time': 0}))
        assert self.worker.handle_task_message()
        result = self.out_queue.get()
        eq_(result.doc, {'status': 'ok'})

class TestBaseWorker(object):
    def setup(self):
        self.in_queue = multiprocessing.Queue(2)
        self.out_queue = multiprocessing.Queue(2)
        self.worker = BaseWorker(in_queue=self.in_queue, out_queue=self.out_queue)

    def test_dispatch_unknown_command(self):
        self.worker.dispatch(Task('foo', doc={'command': 'foo'}))
        assert self.worker.handle_task_message()
        result = self.out_queue.get()
        eq_(result.doc, {
            'status': 'error',
            'error_message': 'unknown command: foo',
        })
    def test_stop(self):
        self.worker.in_queue.put(STOP)
        assert not self.worker.handle_task_message()

class ExceptionWorker(BaseWorker):
    def do_exception(self, doc):
        raise ValueError('foo')

class TestExceptionWorker(object):
    def setup(self):
        self.in_queue = multiprocessing.Queue(2)
        self.out_queue = multiprocessing.Queue(2)
        self.worker = ExceptionWorker(in_queue=self.in_queue, out_queue=self.out_queue)

    def test_dispatch(self):
        self.worker.dispatch(Task('foo', doc={'command': 'exception'}))
        assert self.worker.handle_task_message()
        result = self.out_queue.get()
        eq_(result.doc['status'], 'error')
        eq_(result.doc['error_message'], "exception while processing 'exception': foo")
        assert 'raise ValueError' in result.doc['error_detail']

class DummyCache(object):
    def __init__(self):
        self.requested_tiles = []

    def load_tile_coords(self, tiles):
        self.requested_tiles.extend(tiles)

class TestSeedWorker(object):
    def setup(self):
        self.in_queue = multiprocessing.Queue(2)
        self.out_queue = multiprocessing.Queue(2)
        self.caches = {
            'test_cache': DummyCache()
        }
        self.worker = SeedWorker(
            caches=self.caches,
            base_config={},
            in_queue=self.in_queue,
            out_queue=self.out_queue,
        )

    def test_create_tile(self):
        self.worker.dispatch(Task('foo', doc={'command': 'tile', 'cache_identifier': 'test_cache', 'tiles': [[0, 0, 0]]}))
        assert self.worker.handle_task_message()
        result = self.out_queue.get()
        eq_(result.doc, {'status': 'ok'})
        eq_(self.caches['test_cache'].requested_tiles, [(0, 0, 0)])

    def test_create_tile_unknown_cache(self):
        self.worker.dispatch(Task('foo', doc={'command': 'tile', 'cache_identifier': 'unknown', 'tiles': [[0, 0, 0]]}))
        assert self.worker.handle_task_message()
        result = self.out_queue.get()
        eq_(result.doc, {'status': 'error', 'error_message': "unknown cache 'unknown'"})

        eq_(self.caches['test_cache'].requested_tiles, [])

    def test_start_create_tiles_stop(self):
        # use manager.list to receive requested_tiles from worker process
        manager = multiprocessing.Manager()
        self.caches['test_cache'].requested_tiles = manager.list()
        self.worker.start()

        self.worker.dispatch(Task('foo', doc={'command': 'tile', 'cache_identifier': 'test_cache', 'tiles': [[0, 0, 0]]}))
        self.worker.dispatch(Task('foo', doc={'command': 'tile', 'cache_identifier': 'test_cache', 'tiles': [[5, 0, 0], [5, 1, 0], [5, 2, 0]]}))
        self.worker.dispatch(Task('foo', doc={'command': 'tile', 'cache_identifier': 'test_cache', 'tiles': [[2, 3, 4]]}))
        self.in_queue.put(STOP)

        for _ in range(3):
            result = self.out_queue.get()
            eq_(result.doc, {'status': 'ok'})

        eq_(list(self.caches['test_cache'].requested_tiles), [(0, 0, 0), (5, 0, 0), (5, 1, 0), (5, 2, 0), (2, 3, 4)])<|endoftext|>