<commit_before>

# TODO: no need to test codelist fields separately; instead test the whole serializer in once along with the code and vocabulary fields. Or is testing the fields separately preferable?

from django.test import TestCase # Runs each test in a transaction and flushes database
from unittest import skip
import datetime

from django.test import RequestFactory, Client
from rest_framework.test import APIClient
from iati.factory import iati_factory
from iati.transaction import factories as transaction_factory
from iati_codelists.factory import codelist_factory
from iati_vocabulary.factory import vocabulary_factory
from api.activity import serializers
from iati import models as iati_models
from iati.transaction import models as transaction_models
from django.core.exceptions import ObjectDoesNotExist

from decimal import Decimal

from iati.models import Activity
from iati.factory.utils import _create_test_activity
from api.activity.serializers import ActivitySerializer

class ActivitySaveTestCase(TestCase):
    request_dummy = RequestFactory().get('/')
    c = APIClient()

    def setUp(self):
        self.activity1 = _create_test_activity(id="0001", iati_identifier="0001")
        self.activity2 = _create_test_activity(id="0002", iati_identifier="0002")
        self.activity3 = _create_test_activity(id="0003", iati_identifier="0003")

    def test_prefetch_reporting_organisations(self):
        """
        Test if the prefetches are applied correctly
        Here we expect 3 queries:
        1. Fetch Activity objects
        2. Fetch ReportingOrganisation objects
        3. Fetch corresponding narratives
        """

        with self.assertNumQueries(3):
            queryset = Activity.objects.all().prefetch_reporting_organisations()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('reporting_organisations',))

            list(serializer.data)

    def test_prefetch_title(self):
        """
        Test if the prefetches are applied correctly
        Here we expect 3 queries:
        1. Fetch Activity objects
        2. Fetch title narratives
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_title()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('title',))

            list(serializer.data)

    def test_prefetch_descriptions(self):
        """
        Test if the prefetches are applied correctly
        Here we expect 3 queries:
        1. Fetch Activity objects
        2. Fetch Description objects
        3. Fetch corresponding narratives
        """

        with self.assertNumQueries(3):
            queryset = Activity.objects.all().prefetch_descriptions()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('descriptions',))

            list(serializer.data)

    def test_prefetch_participating_organisations(self):
        """
        Test if the prefetches are applied correctly
        Here we expect 3 queries:
        1. Fetch Activity objects
        2. Fetch ParticipatingOrganisation objects
        3. Fetch corresponding narratives objects
        """

        with self.assertNumQueries(3):
            queryset = Activity.objects.all().prefetch_participating_organisations()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('participating_organisations',))

            list(serializer.data)

    def test_prefetch_recipient_countries(self):
        """
        Test if the prefetches are applied correctly
        Here we expect 3 queries:
        1. Fetch Activity objects
        2. Fetch ActivityRecipientCountry objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_recipient_countries()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('recipient_countries',))

            list(serializer.data)

    def test_prefetch_recipient_regions(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 8 queries:
        1. Fetch Activity objects
        2. Fetch ActivityRecipientRegion objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_recipient_regions()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('recipient_regions',))

            list(serializer.data)

    def test_prefetch_sectors(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch ActivitySector objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_sectors()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('sectors',))

            list(serializer.data)

    def test_prefetch_activity_dates(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch ActivityDate objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_activity_dates()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('activity_dates',))

            list(serializer.data)

    def test_prefetch_policy_markers(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch ActivityPolicyMarker objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_policy_markers()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('policy_markers',))

            list(serializer.data)

    def test_prefetch_budgets(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch Budget objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_budgets()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('budgets',))

            list(serializer.data)

    def test_prefetch_document_links(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 7 queries:
        1. Fetch Activity objects.
        2. Fetch DocumentLink objects.
        3. Fetch DocumentLinkLanguage objects.
        4. Fetch DocumentLinkCategory objects.
        5. Fetch Narrative objects.
        6. Fetch DocumentLinkCategory objects.
        7. Fetch DocumentLinkLanguage objects.
        TODO: Verify if the queries 6 and 7 can be deleted.
        """

        with self.assertNumQueries(7):
            queryset = Activity.objects.all().prefetch_document_links()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('document_links',))

            list(serializer.data)

    def test_prefetch_results(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch Result objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_results()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('results',))

            list(serializer.data)

    def test_prefetch_locations(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch Location objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_locations()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('locations',))

            list(serializer.data)

    def test_prefetch_related_activities(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch RelatedActivity objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_related_activities()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('related_activities',))

            list(serializer.data)

    def test_prefetch_aggregations(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch ActivityAggregation objects
        3. Fetch ChildAggregation objects
        4. Fetch  ActivityPlusChildAggregationobjects
        """

        with self.assertNumQueries(4):
            queryset = Activity.objects.all().prefetch_aggregations()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('related_aggregations',))

            list(serializer.data)

    def test_prefetch_default_aid_type(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        """

        with self.assertNumQueries(1):
            queryset = Activity.objects.all().prefetch_default_aid_type()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('default_aid_type',))

            list(serializer.data)

    def test_prefetch_default_finance_type(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        """

        with self.assertNumQueries(1):
            queryset = Activity.objects.all().prefetch_default_finance_type()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('default_finance_type',))

            list(serializer.data)

    def test_prefetch_other_identifiers(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch otheridentifier objects
        3. Fetch Narrative objects
        """

        with self.assertNumQueries(3):
            queryset = Activity.objects.all().prefetch_other_identifiers()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('other_identifiers',))

            list(serializer.data)

    def test_prefetch_contact_info(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch ContactInfo objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_contact_info()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('contact_info',))

            list(serializer.data)

    def test_prefetch_legacy_data(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch LegacyData objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_legacy_data()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('legacy_data',))

            list(serializer.data)

    def test_prefetch_conditions(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch conditions objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_conditions()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('conditions',))

            list(serializer.data)

    def test_prefetch_crs_add(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch CrsAdd objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_crs_add()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('crs_add',))

            list(serializer.data)

    def test_prefetch_fss(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch Fss objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_fss()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('fss',))

            list(serializer.data)

    def test_prefetch_planned_disbursement(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch PlannedDisbursement objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_planned_disbursement()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('planned_disbursement',))

            list(serializer.data)

    def test_prefetch_humanitarian_scope(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch HumanitarianScope objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_humanitarian_scope()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('humanitarian_scope',))

            list(serializer.data)

    def test_prefetch_country_budget_items(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch CountryBudgetItem objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_country_budget_items()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('country_budget_items',))

            list(serializer.data)

<commit_msg>Oredering Prefetch Tests methods<commit_after>

# TODO: no need to test codelist fields separately; instead test the whole serializer in once along with the code and vocabulary fields. Or is testing the fields separately preferable?

from django.test import TestCase # Runs each test in a transaction and flushes database
from unittest import skip
import datetime

from django.test import RequestFactory, Client
from rest_framework.test import APIClient
from iati.factory import iati_factory
from iati.transaction import factories as transaction_factory
from iati_codelists.factory import codelist_factory
from iati_vocabulary.factory import vocabulary_factory
from api.activity import serializers
from iati import models as iati_models
from iati.transaction import models as transaction_models
from django.core.exceptions import ObjectDoesNotExist

from decimal import Decimal

from iati.models import Activity
from iati.factory.utils import _create_test_activity
from api.activity.serializers import ActivitySerializer

class ActivitySaveTestCase(TestCase):
    request_dummy = RequestFactory().get('/')
    c = APIClient()

    def setUp(self):
        self.activity1 = _create_test_activity(id="0001", iati_identifier="0001")
        self.activity2 = _create_test_activity(id="0002", iati_identifier="0002")
        self.activity3 = _create_test_activity(id="0003", iati_identifier="0003")

    def test_prefetch_reporting_organisations(self):
        """
        Test if the prefetches are applied correctly
        Here we expect 3 queries:
        1. Fetch Activity objects
        2. Fetch ReportingOrganisation objects
        3. Fetch corresponding narratives
        """

        with self.assertNumQueries(3):
            queryset = Activity.objects.all().prefetch_reporting_organisations()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('reporting_organisations',))

            list(serializer.data)

    def test_prefetch_title(self):
        """
        Test if the prefetches are applied correctly
        Here we expect 3 queries:
        1. Fetch Activity objects
        2. Fetch title narratives
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_title()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('title',))

            list(serializer.data)

    def test_prefetch_descriptions(self):
        """
        Test if the prefetches are applied correctly
        Here we expect 3 queries:
        1. Fetch Activity objects
        2. Fetch Description objects
        3. Fetch corresponding narratives
        """

        with self.assertNumQueries(3):
            queryset = Activity.objects.all().prefetch_descriptions()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('descriptions',))

            list(serializer.data)

    def test_prefetch_participating_organisations(self):
        """
        Test if the prefetches are applied correctly
        Here we expect 3 queries:
        1. Fetch Activity objects
        2. Fetch ParticipatingOrganisation objects
        3. Fetch corresponding narratives objects
        """

        with self.assertNumQueries(3):
            queryset = Activity.objects.all().prefetch_participating_organisations()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('participating_organisations',))

            list(serializer.data)

    def test_prefetch_other_identifiers(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch otheridentifier objects
        3. Fetch Narrative objects
        """

        with self.assertNumQueries(3):
            queryset = Activity.objects.all().prefetch_other_identifiers()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('other_identifiers',))

            list(serializer.data)

    def test_prefetch_activity_dates(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch ActivityDate objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_activity_dates()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('activity_dates',))

            list(serializer.data)

    def test_prefetch_contact_info(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch ContactInfo objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_contact_info()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('contact_info',))

            list(serializer.data)

    def test_prefetch_recipient_countries(self):
        """
        Test if the prefetches are applied correctly
        Here we expect 3 queries:
        1. Fetch Activity objects
        2. Fetch ActivityRecipientCountry objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_recipient_countries()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('recipient_countries',))

            list(serializer.data)

    def test_prefetch_recipient_regions(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 8 queries:
        1. Fetch Activity objects
        2. Fetch ActivityRecipientRegion objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_recipient_regions()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('recipient_regions',))

            list(serializer.data)

    def test_prefetch_locations(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch Location objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_locations()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('locations',))

            list(serializer.data)

    def test_prefetch_sectors(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch ActivitySector objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_sectors()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('sectors',))

            list(serializer.data)

    def test_prefetch_country_budget_items(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch CountryBudgetItem objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_country_budget_items()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('country_budget_items',))

            list(serializer.data)

    def test_prefetch_humanitarian_scope(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch HumanitarianScope objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_humanitarian_scope()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('humanitarian_scope',))

            list(serializer.data)

    def test_prefetch_policy_markers(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch ActivityPolicyMarker objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_policy_markers()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('policy_markers',))

            list(serializer.data)

    def test_prefetch_budgets(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch Budget objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_budgets()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('budgets',))

            list(serializer.data)

    def test_prefetch_planned_disbursement(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch PlannedDisbursement objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_planned_disbursement()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('planned_disbursement',))

            list(serializer.data)

    def test_prefetch_document_links(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 7 queries:
        1. Fetch Activity objects.
        2. Fetch DocumentLink objects.
        3. Fetch DocumentLinkLanguage objects.
        4. Fetch DocumentLinkCategory objects.
        5. Fetch Narrative objects.
        6. Fetch DocumentLinkCategory objects.
        7. Fetch DocumentLinkLanguage objects.
        TODO: Verify if the queries 6 and 7 can be deleted.
        """

        with self.assertNumQueries(7):
            queryset = Activity.objects.all().prefetch_document_links()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('document_links',))

            list(serializer.data)

    def test_prefetch_related_activities(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch RelatedActivity objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_related_activities()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('related_activities',))

            list(serializer.data)

    def test_prefetch_legacy_data(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch LegacyData objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_legacy_data()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('legacy_data',))

            list(serializer.data)

    def test_prefetch_conditions(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch conditions objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_conditions()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('conditions',))

            list(serializer.data)

    def test_prefetch_results(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch Result objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_results()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('results',))

            list(serializer.data)

    def test_prefetch_crs_add(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch CrsAdd objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_crs_add()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('crs_add',))

            list(serializer.data)

    def test_prefetch_fss(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch Fss objects
        """

        with self.assertNumQueries(2):
            queryset = Activity.objects.all().prefetch_fss()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('fss',))

            list(serializer.data)

    def test_prefetch_default_aid_type(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        """

        with self.assertNumQueries(1):
            queryset = Activity.objects.all().prefetch_default_aid_type()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('default_aid_type',))

            list(serializer.data)

    def test_prefetch_default_finance_type(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        """

        with self.assertNumQueries(1):
            queryset = Activity.objects.all().prefetch_default_finance_type()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('default_finance_type',))

            list(serializer.data)

    def test_prefetch_aggregations(self):

        """
        Test if the prefetches are applied correctly
        Here we expect 2 queries:
        1. Fetch Activity objects
        2. Fetch ActivityAggregation objects
        3. Fetch ChildAggregation objects
        4. Fetch  ActivityPlusChildAggregationobjects
        """

        with self.assertNumQueries(4):
            queryset = Activity.objects.all().prefetch_aggregations()
            serializer = ActivitySerializer(
                    queryset,
                    many=True,
                    context={'request': self.request_dummy},
                    fields=('related_aggregations',))

            list(serializer.data)

<|endoftext|>


<commit_before>from . import about
from . import ev
from . import gs
from collections import namedtuple
import os
import re
import string
import subprocess
import time

try:
	STARTUPINFO = subprocess.STARTUPINFO()
	STARTUPINFO.dwFlags |= subprocess.STARTF_USESHOWWINDOW
	STARTUPINFO.wShowWindow = subprocess.SW_HIDE
except (AttributeError):
	STARTUPINFO = None

Proc = namedtuple('Proc', 'p input orig_cmd cmd_lst env wd ok exc')
Result = namedtuple('Result', 'out cmd_lst err ok exc')

class _command(object):
	def __init__(self):
		self.input = None
		self.stdin = subprocess.PIPE
		self.stdout = subprocess.PIPE
		self.stderr = subprocess.PIPE
		self.startupinfo = STARTUPINFO
		self.wd = None
		self.env = {}

	def proc(self):
		if self.wd:
			wd = self.wd
			try:
				os.makedirs(wd)
			except Exception:
				pass
		else:
			wd = None

		if self.input:
			input = gs.astr(self.input)
		else:
			input = None

		try:
			setsid = os.setsid
		except Exception:
			setsid = None

		out = ''
		err = ''
		exc = None

		nv = env(self.env)
		cmd_lst = self.cmd(nv)
		orig_cmd = cmd_lst[0]
		cmd_lst[0] = _which(orig_cmd, nv.get('PATH'))

		try:
			if not cmd_lst[0]:
				raise Exception('Cannot find command `%s`' % orig_cmd)

			p = subprocess.Popen(
				cmd_lst,
				stdout=self.stdout,
				stderr=self.stderr,
				stdin=self.stdin,
				startupinfo=self.startupinfo,
				shell=False,
				env=nv,
				cwd=wd,
				preexec_fn=setsid,
				bufsize=0
			)
		except Exception as e:
			exc = e
			p = None

		return Proc(
			p=p,
			input=input,
			orig_cmd=orig_cmd,
			cmd_lst=cmd_lst,
			env=nv,
			wd=wd,
			ok=(not exc),
			exc=exc
		)

	def run(self):
		out = ''
		err = ''
		exc = None

		pr = self.proc()
		if pr.ok:
			ev.debug('sh.run', pr)

			try:
				out, err = pr.p.communicate(input=pr.input)
			except Exception as e:
				exc = e
		else:
			exc = pr.exc

		return Result(
			out=gs.ustr(out),
			err=gs.ustr(err),
			cmd_lst=pr.cmd_lst,
			ok=(not exc),
			exc=exc
		)

class ShellCommand(_command):
	def __init__(self, cmd_str):
		_command.__init__(self)
		self.cmd_str = gs.astr(cmd_str)

	def cmd(self, e):
		return _cmd(self.cmd_str, e)

class Command(_command):
	def __init__(self, cmd_lst):
		_command.__init__(self)
		self.cmd_lst = [gs.astr(s) for s in cmd_lst]

	def cmd(self, e):
		return self.cmd_lst

def shl(m={}):
	return _shl(env(m))

def _shl(e):
	l = gs.setting('shell', [])
	if not l:
		fn = e.get('SHELL') or e.get('COMSPEC')
		if fn:
			name, _ = os.path.splitext(os.path.basename(fn))
			f = globals().get('_shl_%s' % name)
			if f:
				l = f(fn)

	if not l:
		if gs.os_is_windows():
			l = _shl_cmd('cmd')
		else:
			l = _shl_sh('sh')

	return l

def _shl_cmd(fn):
	return [fn, '/C', '${CMD}']

def _shl_sh(fn):
	return [fn, '-l', '-c', '${CMD}']

_shl_fish = _shl_sh
_shl_bash = _shl_sh
_shl_zsh = _shl_sh
_shl_rc = _shl_sh

def cmd(cmd_str, m={}):
	return _cmd(cmd_str, env(m))

def _cmd(cmd_str, e):
	cmdm = {'CMD': cmd_str}
	cmdl = []
	for s in _shl(e):
		s = string.Template(s).safe_substitute(cmdm)
		s = gs.astr(s)
		if s:
			cmdl.append(s)

	return cmdl

def gs_init(_={}):
	global _env_ext
	global GO_VERSION

	start = time.time()

	vars = [
		'PATH',
		'GOBIN',
		'GOPATH',
		'GOROOT',
	]

	cmdl = []
	for k in vars:
		cmdl.append('[[[$'+k+']]'+k+'[[%'+k+'%]]]')
	cmd_str = 'echo "%s"' % ' '.join(cmdl)

	cr = ShellCommand(cmd_str).run()
	if cr.exc:
		_print('error loading env vars: %s' % cr.exc)

	out = cr.out + cr.err

	mats = re.findall(r'\[\[\[(.*?)\]\](%s)\[\[(.*?)\]\]\]' % '|'.join(vars), out)
	for m in mats:
		a, k, b = m
		v = ''
		if a != '$'+k:
			v = a
		elif b != '%'+k+'%':
			v = b

		if v:
			_env_ext[k] = v

	cr_go = ShellCommand('go version').run()
	m = about.GO_VERSION_OUTPUT_PAT.search(cr_go.out + cr_go.err)
	if m:
		GO_VERSION = '-'.join(s for s in m.groups() if s)

	dur = (time.time() - start)

	ev.debug('sh.init', {
		'cr.init': cr,
		'cr.go': cr_go,
		'go_version': GO_VERSION,
		'env': _env_ext,
		'dur': dur,
	})

	cmd_lst = []
	for v in cr.cmd_lst:
		v = v.replace(cmd_str, 'echo "..."')
		cmd_lst.append(v)

	_print('load env vars %s: go version %s = %s: %0.3fs' % (
		cmd_lst,
		cr_go.cmd_lst,
		(GO_VERSION if GO_VERSION != about.DEFAULT_GO_VERSION else cr_go),
		dur,
	))

def _print(s):
	print('GoSblime %s sh: %s' % (about.VERSION, s))

def _shell_pathsep():
	return gs.setting('shell_pathsep') or os.pathsep

def _sj_path(p):
	return _shell_pathsep().join(p.split(os.pathsep))

def getenv(name, default='', m={}):
	return env(m).get(name, default)

def env(m={}):
	"""
	Assemble environment information needed for correct operation. In particular,
	ensure that directories containing binaries are included in PATH.
	"""
	e = os.environ.copy()

	# the system's env may be compatible with the shell
	# so try to fix the env vars that depend on shell_pathsep
	e['PATH'] = _sj_path(e.get('PATH', ''))
	e['GOPATH'] = _sj_path(e.get('GOPATH', ''))

	e.update(_env_ext)
	e.update(m)

	roots = gs.lst(e.get('GOPATH', '').split(os.pathsep), e.get('GOROOT', ''))
	lfn = gs.attr('last_active_go_fn', '')

	comps = lfn.split(os.sep)
	gs_gopath = []
	for i, s in enumerate(comps):
		if s.lower() == "src":
			p = os.sep.join(comps[:i])
			if p not in roots:
				gs_gopath.append(p)
	gs_gopath.reverse()
	e['GS_GOPATH'] = os.pathsep.join(gs_gopath)

	uenv = gs.setting('env', {})
	for k in uenv:
		try:
			uenv[k] = string.Template(uenv[k]).safe_substitute(e)
		except Exception as ex:
			gs.println('%s: Cannot expand env var `%s`: %s' % (NAME, k, ex))

	e.update(uenv)
	e.update(m)

	# For custom values of GOPATH, installed binaries via go install
	# will go into the "bin" dir of the corresponding GOPATH path.
	# Therefore, make sure these paths are included in PATH.

	add_path = [gs.home_dir_path('bin')]

	for s in gs.lst(e.get('GOROOT', ''), e.get('GOPATH', '').split(os.pathsep)):
		if s:
			s = os.path.join(s, 'bin')
			if s not in add_path:
				add_path.append(s)

	gobin = e.get('GOBIN', '')
	if gobin and gobin not in add_path:
		add_path.append(gobin)

	if gs.os_is_windows():
		l = [
			'~\\bin',
			'~\\go\\bin',
			'C:\\Go\\bin',
		]
	else:
		l = [
			'~/bin',
			'~/go/bin',
			'/usr/local/go/bin',
			'/usr/local/opt/go/bin',
			'/usr/local/bin',
			'/usr/bin',
		]

	for s in l:
		s = os.path.expanduser(s)
		if s not in add_path:
			add_path.append(s)

	psep = _shell_pathsep()

	for s in e.get('PATH', '').split(psep):
		if s and s not in add_path:
			add_path.append(s)

	e['PATH'] = psep.join(add_path)

	# Ensure no unicode objects leak through. The reason is twofold:
	# 	* On Windows, Python 2.6 (used by Sublime Text) subprocess.Popen
	# 	  can only take bytestrings as environment variables in the
	#	  "env"	parameter. Reference:
	# 	  https://github.com/DisposaBoy/GoSublime/issues/112
	# 	  http://stackoverflow.com/q/12253014/1670
	#   * Avoids issues with networking too.
	clean_env = {}
	for k, v in e.items():
		try:
			clean_env[gs.astr(k)] = gs.astr(v)
		except Exception as ex:
			gs.println('%s: Bad env: %s' % (NAME, ex))

	return clean_env

def which_ok(fn):
	try:
		return os.path.isfile(fn) and os.access(fn, os.X_OK)
	except Exception:
		return False

def which(cmd):
	return _which(cmd, getenv('PATH', ''))

def _which(cmd, env_path):
	if os.path.isabs(cmd):
		return cmd if which_ok(cmd) else ''

	# not supporting PATHEXT. period.
	if gs.os_is_windows() and not cmd.endswith('.exe'):
		cmd = '%s.exe' % cmd

	seen = {}
	for p in env_path.split(_shell_pathsep()):
		p = os.path.join(p, cmd)
		if p not in seen and which_ok(p):
			return p

		seen[p] = True

	return ''

def go(subcmd_str):
	cr = ShellCommand('go '+subcmd_str).run()
	out = cr.out.strip() + '\n' + cr.err.strip()
	return out.strip()

GO_VERSION = about.DEFAULT_GO_VERSION
_env_ext = {}
<commit_msg>* if GOROOT is not set in the shell env, get it from the go command<commit_after>from . import about
from . import ev
from . import gs
from collections import namedtuple
import os
import re
import string
import subprocess
import time

try:
	STARTUPINFO = subprocess.STARTUPINFO()
	STARTUPINFO.dwFlags |= subprocess.STARTF_USESHOWWINDOW
	STARTUPINFO.wShowWindow = subprocess.SW_HIDE
except (AttributeError):
	STARTUPINFO = None

Proc = namedtuple('Proc', 'p input orig_cmd cmd_lst env wd ok exc')
Result = namedtuple('Result', 'out cmd_lst err ok exc')

class _command(object):
	def __init__(self):
		self.input = None
		self.stdin = subprocess.PIPE
		self.stdout = subprocess.PIPE
		self.stderr = subprocess.PIPE
		self.startupinfo = STARTUPINFO
		self.wd = None
		self.env = {}

	def proc(self):
		if self.wd:
			wd = self.wd
			try:
				os.makedirs(wd)
			except Exception:
				pass
		else:
			wd = None

		if self.input:
			input = gs.astr(self.input)
		else:
			input = None

		try:
			setsid = os.setsid
		except Exception:
			setsid = None

		out = ''
		err = ''
		exc = None

		nv = env(self.env)
		cmd_lst = self.cmd(nv)
		orig_cmd = cmd_lst[0]
		cmd_lst[0] = _which(orig_cmd, nv.get('PATH'))

		try:
			if not cmd_lst[0]:
				raise Exception('Cannot find command `%s`' % orig_cmd)

			p = subprocess.Popen(
				cmd_lst,
				stdout=self.stdout,
				stderr=self.stderr,
				stdin=self.stdin,
				startupinfo=self.startupinfo,
				shell=False,
				env=nv,
				cwd=wd,
				preexec_fn=setsid,
				bufsize=0
			)
		except Exception as e:
			exc = e
			p = None

		return Proc(
			p=p,
			input=input,
			orig_cmd=orig_cmd,
			cmd_lst=cmd_lst,
			env=nv,
			wd=wd,
			ok=(not exc),
			exc=exc
		)

	def run(self):
		out = ''
		err = ''
		exc = None

		pr = self.proc()
		if pr.ok:
			ev.debug('sh.run', pr)

			try:
				out, err = pr.p.communicate(input=pr.input)
			except Exception as e:
				exc = e
		else:
			exc = pr.exc

		return Result(
			out=gs.ustr(out),
			err=gs.ustr(err),
			cmd_lst=pr.cmd_lst,
			ok=(not exc),
			exc=exc
		)

class ShellCommand(_command):
	def __init__(self, cmd_str):
		_command.__init__(self)
		self.cmd_str = gs.astr(cmd_str)

	def cmd(self, e):
		return _cmd(self.cmd_str, e)

class Command(_command):
	def __init__(self, cmd_lst):
		_command.__init__(self)
		self.cmd_lst = [gs.astr(s) for s in cmd_lst]

	def cmd(self, e):
		return self.cmd_lst

def shl(m={}):
	return _shl(env(m))

def _shl(e):
	l = gs.setting('shell', [])
	if not l:
		fn = e.get('SHELL') or e.get('COMSPEC')
		if fn:
			name, _ = os.path.splitext(os.path.basename(fn))
			f = globals().get('_shl_%s' % name)
			if f:
				l = f(fn)

	if not l:
		if gs.os_is_windows():
			l = _shl_cmd('cmd')
		else:
			l = _shl_sh('sh')

	return l

def _shl_cmd(fn):
	return [fn, '/C', '${CMD}']

def _shl_sh(fn):
	return [fn, '-l', '-c', '${CMD}']

_shl_fish = _shl_sh
_shl_bash = _shl_sh
_shl_zsh = _shl_sh
_shl_rc = _shl_sh

def cmd(cmd_str, m={}):
	return _cmd(cmd_str, env(m))

def _cmd(cmd_str, e):
	cmdm = {'CMD': cmd_str}
	cmdl = []
	for s in _shl(e):
		s = string.Template(s).safe_substitute(cmdm)
		s = gs.astr(s)
		if s:
			cmdl.append(s)

	return cmdl

def gs_init(_={}):
	global _env_ext
	global GO_VERSION

	start = time.time()

	vars = [
		'PATH',
		'GOBIN',
		'GOPATH',
		'GOROOT',
	]

	cmdl = []
	for k in vars:
		cmdl.append('[[[$'+k+']]'+k+'[[%'+k+'%]]]')
	cmd_str = 'echo "%s"' % ' '.join(cmdl)

	cr = ShellCommand(cmd_str).run()
	if cr.exc:
		_print('error loading env vars: %s' % cr.exc)

	out = cr.out + cr.err

	mats = re.findall(r'\[\[\[(.*?)\]\](%s)\[\[(.*?)\]\]\]' % '|'.join(vars), out)
	for m in mats:
		a, k, b = m
		v = ''
		if a != '$'+k:
			v = a
		elif b != '%'+k+'%':
			v = b

		if v:
			_env_ext[k] = v

	if not _env_ext.get('GOROOT'):
		_env_ext['GOROOT'] = go('env GOROOT')

	cr_go = ShellCommand('go version').run()
	m = about.GO_VERSION_OUTPUT_PAT.search(cr_go.out + cr_go.err)
	if m:
		GO_VERSION = '-'.join(s for s in m.groups() if s)

	dur = (time.time() - start)

	ev.debug('sh.init', {
		'cr.init': cr,
		'cr.go': cr_go,
		'go_version': GO_VERSION,
		'env': _env_ext,
		'dur': dur,
	})

	cmd_lst = []
	for v in cr.cmd_lst:
		v = v.replace(cmd_str, 'echo "..."')
		cmd_lst.append(v)

	_print('load env vars %s: go version %s = %s: %0.3fs' % (
		cmd_lst,
		cr_go.cmd_lst,
		(GO_VERSION if GO_VERSION != about.DEFAULT_GO_VERSION else cr_go),
		dur,
	))

def _print(s):
	print('GoSblime %s sh: %s' % (about.VERSION, s))

def _shell_pathsep():
	return gs.setting('shell_pathsep') or os.pathsep

def _sj_path(p):
	return _shell_pathsep().join(p.split(os.pathsep))

def getenv(name, default='', m={}):
	return env(m).get(name, default)

def env(m={}):
	"""
	Assemble environment information needed for correct operation. In particular,
	ensure that directories containing binaries are included in PATH.
	"""
	e = os.environ.copy()

	# the system's env may be compatible with the shell
	# so try to fix the env vars that depend on shell_pathsep
	e['PATH'] = _sj_path(e.get('PATH', ''))
	e['GOPATH'] = _sj_path(e.get('GOPATH', ''))

	e.update(_env_ext)
	e.update(m)

	roots = gs.lst(e.get('GOPATH', '').split(os.pathsep), e.get('GOROOT', ''))
	lfn = gs.attr('last_active_go_fn', '')

	comps = lfn.split(os.sep)
	gs_gopath = []
	for i, s in enumerate(comps):
		if s.lower() == "src":
			p = os.sep.join(comps[:i])
			if p not in roots:
				gs_gopath.append(p)
	gs_gopath.reverse()
	e['GS_GOPATH'] = os.pathsep.join(gs_gopath)

	uenv = gs.setting('env', {})
	for k in uenv:
		try:
			uenv[k] = string.Template(uenv[k]).safe_substitute(e)
		except Exception as ex:
			gs.println('%s: Cannot expand env var `%s`: %s' % (NAME, k, ex))

	e.update(uenv)
	e.update(m)

	# For custom values of GOPATH, installed binaries via go install
	# will go into the "bin" dir of the corresponding GOPATH path.
	# Therefore, make sure these paths are included in PATH.

	add_path = [gs.home_dir_path('bin')]

	for s in gs.lst(e.get('GOROOT', ''), e.get('GOPATH', '').split(os.pathsep)):
		if s:
			s = os.path.join(s, 'bin')
			if s not in add_path:
				add_path.append(s)

	gobin = e.get('GOBIN', '')
	if gobin and gobin not in add_path:
		add_path.append(gobin)

	if gs.os_is_windows():
		l = [
			'~\\bin',
			'~\\go\\bin',
			'C:\\Go\\bin',
		]
	else:
		l = [
			'~/bin',
			'~/go/bin',
			'/usr/local/go/bin',
			'/usr/local/opt/go/bin',
			'/usr/local/bin',
			'/usr/bin',
		]

	for s in l:
		s = os.path.expanduser(s)
		if s not in add_path:
			add_path.append(s)

	psep = _shell_pathsep()

	for s in e.get('PATH', '').split(psep):
		if s and s not in add_path:
			add_path.append(s)

	e['PATH'] = psep.join(add_path)

	# Ensure no unicode objects leak through. The reason is twofold:
	# 	* On Windows, Python 2.6 (used by Sublime Text) subprocess.Popen
	# 	  can only take bytestrings as environment variables in the
	#	  "env"	parameter. Reference:
	# 	  https://github.com/DisposaBoy/GoSublime/issues/112
	# 	  http://stackoverflow.com/q/12253014/1670
	#   * Avoids issues with networking too.
	clean_env = {}
	for k, v in e.items():
		try:
			clean_env[gs.astr(k)] = gs.astr(v)
		except Exception as ex:
			gs.println('%s: Bad env: %s' % (NAME, ex))

	return clean_env

def which_ok(fn):
	try:
		return os.path.isfile(fn) and os.access(fn, os.X_OK)
	except Exception:
		return False

def which(cmd):
	return _which(cmd, getenv('PATH', ''))

def _which(cmd, env_path):
	if os.path.isabs(cmd):
		return cmd if which_ok(cmd) else ''

	# not supporting PATHEXT. period.
	if gs.os_is_windows() and not cmd.endswith('.exe'):
		cmd = '%s.exe' % cmd

	seen = {}
	for p in env_path.split(_shell_pathsep()):
		p = os.path.join(p, cmd)
		if p not in seen and which_ok(p):
			return p

		seen[p] = True

	return ''

def go(subcmd_str):
	cr = ShellCommand('go '+subcmd_str).run()
	out = cr.out.strip() + '\n' + cr.err.strip()
	return out.strip()

GO_VERSION = about.DEFAULT_GO_VERSION
_env_ext = {}
<|endoftext|>


<commit_before>from astropy.io import fits
from scipy import optimize
from matplotlib import pyplot as plt
import numpy as np

data = fits.getdata("../ferengi_data.fits",1)

def fzeta_exp(p,x):

    y = p[0] * np.exp(-1 * (x-p[1])/p[2])

    return y

def fzeta_lin(p,x):

    y = p[0] + p[1] * x

    return y

def errfunc(p,x,y):

    err = y - fzeta_exp(p,x)

    return err

def common_labels(fig,xlabel=None,ylabel=None,xfontsize=12,yfontsize=12):

    # Set common labels
    cax = fig.add_subplot(111)    # The big subplot
    cax.set_axis_bgcolor('none')
    cax.spines['top'].set_color('none')
    cax.spines['bottom'].set_color('none')
    cax.spines['left'].set_color('none')
    cax.spines['right'].set_color('none')
    cax.tick_params(labelcolor='none', top='off', bottom='off', left='off', right='off')
    cax.set_xlabel(xlabel,fontsize=xfontsize)
    cax.set_ylabel(ylabel,fontsize=yfontsize)

    return cax

unique_galaxies = set(data['sdss_id'])
p_guess = np.array([0.5,0.3,1.])

# Set up plot
fig,axarr = plt.subplots(nrows=2,ncols=2,sharex=True,sharey=True)
bigax = common_labels(fig,'Redshift',r'$f_{features}$',14,20)
zarr = np.linspace(0,1,50)

for ax in axarr.ravel():

    slen = 0
    # Make sure there are enough points in the data to fit a function
    while slen < (len(p_guess)+1):
        ind = (data['sdss_id'] == unique_galaxies.pop())
        slen = sum(ind)

    galaxy1 = data[ind & (data['sim_evolution'] == 0)]

    z_gal = galaxy1['sim_redshift']
    f_gal = galaxy1['t00_smooth_or_features_a1_features_frac_weighted_2']

    # Data must be explicitly cast as double-type precision for optimization to work. Incredibly frustrating.
    # Fix: http://stackoverflow.com/questions/12473406/scipy-optimize-leastsq-returns-best-guess-parameters-not-new-best-fit

    p, cov, infodict, mesg, ier = optimize.leastsq(errfunc,p_guess,args=(z_gal.astype(np.float64),f_gal.astype(np.float64)),full_output=1)
    ax.plot(z_gal,f_gal,lw=2)
    ax.plot(zarr,fzeta_exp(p,zarr),'--',lw=1)
    ax.set_xlim(0,1)
    ax.set_ylim(0,1)

plt.show()
<commit_msg>Refactored to make plotting its own module<commit_after>from astropy.io import fits
from scipy import optimize
from matplotlib import pyplot as plt
import numpy as np

def get_data():
    data = fits.getdata("../ferengi_data.fits",1)
    unique_galaxies = set(data['sdss_id'])

    return data,unique_galaxies

def fzeta_exp(p,x):

    y = p[0] * np.exp(-1 * (x-p[1])/p[2])

    return y

def fzeta_lin(p,x):

    y = p[0] + p[1] * x

    return y

def errfunc(p,x,y):

    err = y - fzeta_exp(p,x)

    return err

def common_labels(fig,xlabel=None,ylabel=None,xfontsize=12,yfontsize=12):

    # Set common labels
    cax = fig.add_subplot(111)    # The big subplot
    cax.set_axis_bgcolor('none')
    cax.spines['top'].set_color('none')
    cax.spines['bottom'].set_color('none')
    cax.spines['left'].set_color('none')
    cax.spines['right'].set_color('none')
    cax.tick_params(labelcolor='none', top='off', bottom='off', left='off', right='off')
    cax.set_xlabel(xlabel,fontsize=xfontsize)
    cax.set_ylabel(ylabel,fontsize=yfontsize)

    return cax

def plot_examples(data,unique_galaxies):

    p_guess = np.array([0.5,0.3,1.])
    # Set up plot
    fig,axarr = plt.subplots(nrows=2,ncols=2,sharex=True,sharey=True)
    bigax = common_labels(fig,'Redshift',r'$f_{features}$',14,20)
    zarr = np.linspace(0,1,50)

    for ax in axarr.ravel():

        slen = 0
        # Make sure there are enough points in the data to fit a function
        while slen < (len(p_guess)+1):
            ind = (data['sdss_id'] == unique_galaxies.pop())
            slen = sum(ind)

        galaxy1 = data[ind & (data['sim_evolution'] == 0)]

        z_gal = galaxy1['sim_redshift']
        f_gal = galaxy1['p_features']

        # Data must be explicitly cast as double-type precision for optimization to work. Incredibly frustrating.
        # Fix: http://stackoverflow.com/questions/12473406/scipy-optimize-leastsq-returns-best-guess-parameters-not-new-best-fit

        p, cov, infodict, mesg, ier = optimize.leastsq(errfunc,p_guess,args=(z_gal.astype(np.float64),f_gal.astype(np.float64)),full_output=1)
        ax.plot(z_gal,f_gal,lw=2)
        ax.plot(zarr,fzeta_exp(p,zarr),'--',lw=1)
        ax.set_xlim(0,1)
        ax.set_ylim(0,1)

    plt.show()

    return None

def plot_zeta_sb(data):

if __name__ == "__main__":

    data,unique_galaxies = get_data()
    plot_examples(data,unique_galaxies)
<|endoftext|>


<commit_before># Copyright (c) 2015 Nicolas JOUANIN
#
# See the file license.txt for copying permission.
import logging
import asyncio
from asyncio import futures
from hbmqtt.mqtt.packet import MQTTFixedHeader
from hbmqtt.mqtt import packet_class
from hbmqtt.errors import NoDataException
from hbmqtt.mqtt.packet import PacketType
from hbmqtt.mqtt.connack import ConnackPacket
from hbmqtt.mqtt.connect import ConnectPacket
from hbmqtt.mqtt.pingresp import PingRespPacket
from hbmqtt.mqtt.pingreq import PingReqPacket
from hbmqtt.mqtt.publish import PublishPacket
from hbmqtt.mqtt.pubrel import PubrelPacket
from hbmqtt.mqtt.puback import PubackPacket
from hbmqtt.mqtt.pubrec import PubrecPacket
from hbmqtt.mqtt.pubcomp import PubcompPacket
from hbmqtt.mqtt.suback import SubackPacket
from hbmqtt.mqtt.unsuback import UnsubackPacket
from hbmqtt.mqtt.disconnect import DisconnectPacket
from hbmqtt.session import Session
from transitions import Machine

class InFlightMessage:
    states = ['new', 'published', 'acknowledged', 'received', 'released', 'completed']

    def __init__(self, packet, qos):
        self.packet = packet
        self.qos = qos
        self.puback = None
        self.pubrec = None
        self.pubcomp = None
        self.pubrel = None
        self._init_states()

    def _init_states(self):
        self.machine = Machine(model=self, states=InFlightMessage.states, initial='new')
        self.machine.add_transition(trigger='publish', source='new', dest='published')
        if self.qos == 0x01:
            self.machine.add_transition(trigger='acknowledge', source='published', dest='acknowledged')
        if self.qos == 0x02:
            self.machine.add_transition(trigger='receive', source='published', dest='received')
            self.machine.add_transition(trigger='release', source='received', dest='released')
            self.machine.add_transition(trigger='complete', source='released', dest='completed')

class ProtocolHandler:
    """
    Class implementing the MQTT communication protocol using asyncio features
    """

    def __init__(self, session: Session, loop=None):
        self.logger = logging.getLogger(__name__)
        self.session = session
        if loop is None:
            self._loop = asyncio.get_event_loop()
        else:
            self._loop = loop
        self._reader_task = None
        self._writer_task = None
        self._reader_ready = asyncio.Event(loop=self._loop)
        self._writer_ready = asyncio.Event(loop=self._loop)

        self._running = False

        extra_info = self.session.writer.get_extra_info('sockname')
        self.session.local_address = extra_info[0]
        self.session.local_port = extra_info[1]

        self.incoming_queues = dict()
        self.application_messages = asyncio.Queue()
        for p in PacketType:
            self.incoming_queues[p] = asyncio.Queue()
        self.outgoing_queue = asyncio.Queue()
        self._puback_waiters = dict()
        self._pubrec_waiters = dict()
        self._pubrel_waiters = dict()
        self._pubcomp_waiters = dict()
        self.delivered_message = asyncio.Queue()

    @asyncio.coroutine
    def start(self):
        self._running = True
        self._reader_task = asyncio.async(self._reader_coro(), loop=self._loop)
        self._writer_task = asyncio.async(self._writer_coro(), loop=self._loop)
        yield from asyncio.wait(
            [self._reader_ready.wait(), self._writer_ready.wait()], loop=self._loop)
        self.logger.debug("Handler tasks started")

    @asyncio.coroutine
    def mqtt_publish(self, topic, message, packet_id, dup, qos, retain):
        if packet_id in self.session.inflight_out:
            self.logger.warn("A message with the same packet ID is already in flight")
        packet = PublishPacket.build(topic, message, packet_id, dup, qos, retain)
        yield from self.outgoing_queue.put(packet)
        inflight_message = InFlightMessage(packet, qos)
        self.session.inflight_out[packet.variable_header.packet_id] = inflight_message

        inflight_message.publish()
        if qos == 0x01:
            waiter = futures.Future(loop=self._loop)
            self._puback_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.puback = waiter.result()
            inflight_message.acknowledge()
            del self._puback_waiters[packet_id]
        if qos == 0x02:
            # Wait for PUBREC
            waiter = futures.Future(loop=self._loop)
            self._pubrec_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.pubrec = waiter.result()
            del self._pubrec_waiters[packet_id]
            inflight_message.receive()

            # Send pubrel
            pubrel = PubrelPacket.build(packet_id)
            yield from self.outgoing_queue.put(pubrel)
            inflight_message.pubrel = pubrel
            inflight_message.release()

            # Wait for pubcomp
            waiter = futures.Future(loop=self._loop)
            self._pubcomp_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.pubcomp = waiter.result()
            del self._pubcomp_waiters[packet_id]
            inflight_message.complete()

        del self.session.inflight_out[packet_id]
        return inflight_message

    @asyncio.coroutine
    def stop(self):
        self._running = False
        self.session.reader.feed_eof()
        yield from asyncio.wait([self._writer_task, self._reader_task], loop=self._loop)

    @asyncio.coroutine
    def _reader_coro(self):
        self.logger.debug("Starting reader coro")
        while self._running:
            try:
                self._reader_ready.set()
                fixed_header = yield from asyncio.wait_for(MQTTFixedHeader.from_stream(self.session.reader), 5)
                if fixed_header:
                    cls = packet_class(fixed_header)
                    packet = yield from cls.from_stream(self.session.reader, fixed_header=fixed_header)
                    self.logger.debug(" <-in-- " + repr(packet))

                    if packet.fixed_header.packet_type == PacketType.CONNACK:
                        asyncio.Task(self.handle_connack(packet))
                    elif packet.fixed_header.packet_type == PacketType.SUBACK:
                        asyncio.Task(self.handle_suback(packet))
                    elif packet.fixed_header.packet_type == PacketType.UNSUBACK:
                        asyncio.Task(self.handle_unsuback(packet))
                    elif packet.fixed_header.packet_type == PacketType.PUBACK:
                        asyncio.Task(self.handle_puback(packet))
                    elif packet.fixed_header.packet_type == PacketType.PUBREC:
                        asyncio.Task(self.handle_pubrec(packet))
                    elif packet.fixed_header.packet_type == PacketType.PUBREL:
                        asyncio.Task(self.handle_pubrel(packet))
                    elif packet.fixed_header.packet_type == PacketType.PUBCOMP:
                        asyncio.Task(self.handle_pubcomp(packet))
                    elif packet.fixed_header.packet_type == PacketType.PINGREQ:
                        asyncio.Task(self.handle_pingreq(packet))
                    elif packet.fixed_header.packet_type == PacketType.PINGRESP:
                        asyncio.Task(self.handle_pingresp(packet))
                    elif packet.fixed_header.packet_type == PacketType.PUBLISH:
                        asyncio.Task(self.handle_publish(packet))
                    elif packet.fixed_header.packet_type == PacketType.DISCONNECT:
                        asyncio.Task(self.handle_disconnect(packet))
                    elif packet.fixed_header.packet_type == PacketType.CONNECT:
                        asyncio.Task(self.handle_connect(packet))
                    else:
                        self.logger.warn("Unhandled packet type: %s" % packet.fixed_header.packet_type)
                else:
                    self.logger.debug("No more data, stopping reader coro")
                    break
            except asyncio.TimeoutError:
                self.logger.debug("Input stream read timeout")
            except NoDataException as nde:
                self.logger.debug("No data available")
            except Exception as e:
                self.logger.warn("Unhandled exception in reader coro: %s" % e)
                break
        self.logger.debug("Reader coro stopped")

    @asyncio.coroutine
    def _writer_coro(self):
        self.logger.debug("Starting writer coro")
        while self._running:
            try:
                self._writer_ready.set()
                keepalive_timeout = self.session.keep_alive
                if keepalive_timeout <= 0:
                    keepalive_timeout = None
                packet = yield from asyncio.wait_for(self.outgoing_queue.get(), keepalive_timeout)
                yield from packet.to_stream(self.session.writer)
                self.logger.debug(" -out-> " + repr(packet))
                yield from self.session.writer.drain()
                #self.outgoing_queue.task_done() # to be used with Python 3.5
            except asyncio.TimeoutError as ce:
                self.logger.debug("Output queue get timeout")
                if self._running:
                    self.logger.debug("PING for keepalive")
                    self.handle_keepalive()
            except Exception as e:
                self.logger.warn("Unhandled exception in writer coro: %s" % e)
                break
        self.logger.debug("Writer coro stopping")
        # Flush queue before stopping
        if not self.outgoing_queue.empty():
            while True:
                try:
                    packet = self.outgoing_queue.get_nowait()
                    yield from packet.to_stream(self.session.writer)
                    self.logger.debug(" -out-> " + repr(packet))
                except asyncio.QueueEmpty:
                    break
                except Exception as e:
                    self.logger.warn("Unhandled exception in writer coro: %s" % e)
        self.logger.debug("Writer coro stopped")

    @asyncio.coroutine
    def mqtt_deliver_next_message(self):
        inflight_message = yield from self.delivered_message.get()
        return inflight_message

    def handle_keepalive(self):
        pass

    @asyncio.coroutine
    def handle_connack(self, connack: ConnackPacket):
        pass

    @asyncio.coroutine
    def handle_connect(self, connect: ConnectPacket):
        pass

    @asyncio.coroutine
    def handle_suback(self, suback: SubackPacket):
        pass

    @asyncio.coroutine
    def handle_unsuback(self, unsuback: UnsubackPacket):
        pass

    @asyncio.coroutine
    def handle_pingresp(self, pingresp: PingRespPacket):
        pass

    @asyncio.coroutine
    def handle_pingreq(self, pingreq: PingReqPacket):
        pass

    @asyncio.coroutine
    def handle_disconnect(self, disconnect: DisconnectPacket):
        pass

    @asyncio.coroutine
    def handle_puback(self, puback: PubackPacket):
        packet_id = puback.variable_header.packet_id
        try:
            waiter = self._puback_waiters[packet_id]
            waiter.set_result(puback)
        except KeyError as ke:
            self.logger.warn("Received PUBACK for unknown pending subscription with Id: %s" % packet_id)

    @asyncio.coroutine
    def handle_pubrec(self, pubrec: PubrecPacket):
        packet_id = pubrec.variable_header.packet_id
        try:
            waiter = self._pubrec_waiters[packet_id]
            waiter.set_result(pubrec)
        except KeyError as ke:
            self.logger.warn("Received PUBREC for unknown pending subscription with Id: %s" % packet_id)

    @asyncio.coroutine
    def handle_pubcomp(self, pubcomp: PubcompPacket):
        packet_id = pubcomp.variable_header.packet_id
        try:
            waiter = self._pubcomp_waiters[packet_id]
            waiter.set_result(pubcomp)
        except KeyError as ke:
            self.logger.warn("Received PUBCOMP for unknown pending subscription with Id: %s" % packet_id)

    @asyncio.coroutine
    def handle_pubrel(self, pubrel: PubrecPacket):
        packet_id = pubrel.variable_header.packet_id
        try:
            waiter = self._pubrel_waiters[packet_id]
            waiter.set_result(pubrel)
        except KeyError as ke:
            self.logger.warn("Received PUBREL for unknown pending subscription with Id: %s" % packet_id)

    @asyncio.coroutine
    def handle_publish(self, publish : PublishPacket):
        inflight_message = None
        packet_id = publish.variable_header.packet_id
        qos = (publish.fixed_header.flags >> 1) & 0x03
        if packet_id in self.session.inflight_in:
            inflight_message = self.session.inflight_in[packet_id]
        else:
            inflight_message = InFlightMessage(publish, qos)
            self.session.inflight_in[packet_id] = inflight_message
            inflight_message.publish()

        if qos == 1:
            puback = PubackPacket.build(packet_id)
            yield from self.outgoing_queue.put(puback)
            inflight_message.acknowledge()
        if qos == 2:
            pubrec = PubrecPacket.build(packet_id)
            yield from self.outgoing_queue.put(pubrec)
            inflight_message.receive()
            waiter = futures.Future(loop=self._loop)
            self._pubrel_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.pubrel = waiter.result()
            del self._pubrel_waiters[packet_id]
            inflight_message.release()
            pubcomp = PubcompPacket.build(packet_id)
            yield from self.outgoing_queue.put(pubcomp)
            inflight_message.complete()
        yield from self.delivered_message.put(inflight_message)
        del self.session.inflight_in[packet_id]
<commit_msg>Add logging for unhandled messages<commit_after># Copyright (c) 2015 Nicolas JOUANIN
#
# See the file license.txt for copying permission.
import logging
import asyncio
from asyncio import futures
from hbmqtt.mqtt.packet import MQTTFixedHeader
from hbmqtt.mqtt import packet_class
from hbmqtt.errors import NoDataException
from hbmqtt.mqtt.packet import PacketType
from hbmqtt.mqtt.connack import ConnackPacket
from hbmqtt.mqtt.connect import ConnectPacket
from hbmqtt.mqtt.pingresp import PingRespPacket
from hbmqtt.mqtt.pingreq import PingReqPacket
from hbmqtt.mqtt.publish import PublishPacket
from hbmqtt.mqtt.pubrel import PubrelPacket
from hbmqtt.mqtt.puback import PubackPacket
from hbmqtt.mqtt.pubrec import PubrecPacket
from hbmqtt.mqtt.pubcomp import PubcompPacket
from hbmqtt.mqtt.suback import SubackPacket
from hbmqtt.mqtt.subscribe import SubscribePacket
from hbmqtt.mqtt.unsuback import UnsubackPacket
from hbmqtt.mqtt.disconnect import DisconnectPacket
from hbmqtt.session import Session
from transitions import Machine

class InFlightMessage:
    states = ['new', 'published', 'acknowledged', 'received', 'released', 'completed']

    def __init__(self, packet, qos):
        self.packet = packet
        self.qos = qos
        self.puback = None
        self.pubrec = None
        self.pubcomp = None
        self.pubrel = None
        self._init_states()

    def _init_states(self):
        self.machine = Machine(model=self, states=InFlightMessage.states, initial='new')
        self.machine.add_transition(trigger='publish', source='new', dest='published')
        if self.qos == 0x01:
            self.machine.add_transition(trigger='acknowledge', source='published', dest='acknowledged')
        if self.qos == 0x02:
            self.machine.add_transition(trigger='receive', source='published', dest='received')
            self.machine.add_transition(trigger='release', source='received', dest='released')
            self.machine.add_transition(trigger='complete', source='released', dest='completed')

class ProtocolHandler:
    """
    Class implementing the MQTT communication protocol using asyncio features
    """

    def __init__(self, session: Session, loop=None):
        self.logger = logging.getLogger(__name__)
        self.session = session
        if loop is None:
            self._loop = asyncio.get_event_loop()
        else:
            self._loop = loop
        self._reader_task = None
        self._writer_task = None
        self._reader_ready = asyncio.Event(loop=self._loop)
        self._writer_ready = asyncio.Event(loop=self._loop)

        self._running = False

        extra_info = self.session.writer.get_extra_info('sockname')
        self.session.local_address = extra_info[0]
        self.session.local_port = extra_info[1]

        self.incoming_queues = dict()
        self.application_messages = asyncio.Queue()
        for p in PacketType:
            self.incoming_queues[p] = asyncio.Queue()
        self.outgoing_queue = asyncio.Queue()
        self._puback_waiters = dict()
        self._pubrec_waiters = dict()
        self._pubrel_waiters = dict()
        self._pubcomp_waiters = dict()
        self.delivered_message = asyncio.Queue()

    @asyncio.coroutine
    def start(self):
        self._running = True
        self._reader_task = asyncio.async(self._reader_coro(), loop=self._loop)
        self._writer_task = asyncio.async(self._writer_coro(), loop=self._loop)
        yield from asyncio.wait(
            [self._reader_ready.wait(), self._writer_ready.wait()], loop=self._loop)
        self.logger.debug("Handler tasks started")

    @asyncio.coroutine
    def mqtt_publish(self, topic, message, packet_id, dup, qos, retain):
        if packet_id in self.session.inflight_out:
            self.logger.warn("A message with the same packet ID is already in flight")
        packet = PublishPacket.build(topic, message, packet_id, dup, qos, retain)
        yield from self.outgoing_queue.put(packet)
        inflight_message = InFlightMessage(packet, qos)
        self.session.inflight_out[packet.variable_header.packet_id] = inflight_message

        inflight_message.publish()
        if qos == 0x01:
            waiter = futures.Future(loop=self._loop)
            self._puback_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.puback = waiter.result()
            inflight_message.acknowledge()
            del self._puback_waiters[packet_id]
        if qos == 0x02:
            # Wait for PUBREC
            waiter = futures.Future(loop=self._loop)
            self._pubrec_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.pubrec = waiter.result()
            del self._pubrec_waiters[packet_id]
            inflight_message.receive()

            # Send pubrel
            pubrel = PubrelPacket.build(packet_id)
            yield from self.outgoing_queue.put(pubrel)
            inflight_message.pubrel = pubrel
            inflight_message.release()

            # Wait for pubcomp
            waiter = futures.Future(loop=self._loop)
            self._pubcomp_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.pubcomp = waiter.result()
            del self._pubcomp_waiters[packet_id]
            inflight_message.complete()

        del self.session.inflight_out[packet_id]
        return inflight_message

    @asyncio.coroutine
    def stop(self):
        self._running = False
        self.session.reader.feed_eof()
        yield from asyncio.wait([self._writer_task, self._reader_task], loop=self._loop)

    @asyncio.coroutine
    def _reader_coro(self):
        self.logger.debug("Starting reader coro")
        while self._running:
            try:
                self._reader_ready.set()
                fixed_header = yield from asyncio.wait_for(MQTTFixedHeader.from_stream(self.session.reader), 5)
                if fixed_header:
                    cls = packet_class(fixed_header)
                    packet = yield from cls.from_stream(self.session.reader, fixed_header=fixed_header)
                    self.logger.debug(" <-in-- " + repr(packet))

                    if packet.fixed_header.packet_type == PacketType.CONNACK:
                        asyncio.Task(self.handle_connack(packet))
                    elif packet.fixed_header.packet_type == PacketType.SUBSCRIBE:
                        asyncio.Task(self.handle_subscribe(packet))
                    elif packet.fixed_header.packet_type == PacketType.SUBACK:
                        asyncio.Task(self.handle_suback(packet))
                    elif packet.fixed_header.packet_type == PacketType.UNSUBACK:
                        asyncio.Task(self.handle_unsuback(packet))
                    elif packet.fixed_header.packet_type == PacketType.PUBACK:
                        asyncio.Task(self.handle_puback(packet))
                    elif packet.fixed_header.packet_type == PacketType.PUBREC:
                        asyncio.Task(self.handle_pubrec(packet))
                    elif packet.fixed_header.packet_type == PacketType.PUBREL:
                        asyncio.Task(self.handle_pubrel(packet))
                    elif packet.fixed_header.packet_type == PacketType.PUBCOMP:
                        asyncio.Task(self.handle_pubcomp(packet))
                    elif packet.fixed_header.packet_type == PacketType.PINGREQ:
                        asyncio.Task(self.handle_pingreq(packet))
                    elif packet.fixed_header.packet_type == PacketType.PINGRESP:
                        asyncio.Task(self.handle_pingresp(packet))
                    elif packet.fixed_header.packet_type == PacketType.PUBLISH:
                        asyncio.Task(self.handle_publish(packet))
                    elif packet.fixed_header.packet_type == PacketType.DISCONNECT:
                        asyncio.Task(self.handle_disconnect(packet))
                    elif packet.fixed_header.packet_type == PacketType.CONNECT:
                        asyncio.Task(self.handle_connect(packet))
                    else:
                        self.logger.warn("Unhandled packet type: %s" % packet.fixed_header.packet_type)
                else:
                    self.logger.debug("No more data, stopping reader coro")
                    break
            except asyncio.TimeoutError:
                self.logger.debug("Input stream read timeout")
            except NoDataException as nde:
                self.logger.debug("No data available")
            except Exception as e:
                self.logger.warn("Unhandled exception in reader coro: %s" % e)
                break
        self.logger.debug("Reader coro stopped")

    @asyncio.coroutine
    def _writer_coro(self):
        self.logger.debug("Starting writer coro")
        while self._running:
            try:
                self._writer_ready.set()
                keepalive_timeout = self.session.keep_alive
                if keepalive_timeout <= 0:
                    keepalive_timeout = None
                packet = yield from asyncio.wait_for(self.outgoing_queue.get(), keepalive_timeout)
                yield from packet.to_stream(self.session.writer)
                self.logger.debug(" -out-> " + repr(packet))
                yield from self.session.writer.drain()
                #self.outgoing_queue.task_done() # to be used with Python 3.5
            except asyncio.TimeoutError as ce:
                self.logger.debug("Output queue get timeout")
                if self._running:
                    self.logger.debug("PING for keepalive")
                    self.handle_keepalive()
            except Exception as e:
                self.logger.warn("Unhandled exception in writer coro: %s" % e)
                break
        self.logger.debug("Writer coro stopping")
        # Flush queue before stopping
        if not self.outgoing_queue.empty():
            while True:
                try:
                    packet = self.outgoing_queue.get_nowait()
                    yield from packet.to_stream(self.session.writer)
                    self.logger.debug(" -out-> " + repr(packet))
                except asyncio.QueueEmpty:
                    break
                except Exception as e:
                    self.logger.warn("Unhandled exception in writer coro: %s" % e)
        self.logger.debug("Writer coro stopped")

    @asyncio.coroutine
    def mqtt_deliver_next_message(self):
        inflight_message = yield from self.delivered_message.get()
        return inflight_message

    def handle_keepalive(self):
        self.logger.warn('keepalive unhandled')

    @asyncio.coroutine
    def handle_connack(self, connack: ConnackPacket):
        self.logger.warn('CONNACK unhandled')

    @asyncio.coroutine
    def handle_connect(self, connect: ConnectPacket):
        self.logger.warn('CONNECT unhandled')

    @asyncio.coroutine
    def handle_subscribe(self, subscribe: SubscribePacket):
        self.logger.warn('SUBSCRIBE unhandled')

    @asyncio.coroutine
    def handle_suback(self, suback: SubackPacket):
        self.logger.warn('SUBACK unhandled')

    @asyncio.coroutine
    def handle_unsuback(self, unsuback: UnsubackPacket):
        self.logger.warn('UNSUBACK unhandled')

    @asyncio.coroutine
    def handle_pingresp(self, pingresp: PingRespPacket):
        self.logger.warn('PINGRESP unhandled')

    @asyncio.coroutine
    def handle_pingreq(self, pingreq: PingReqPacket):
        self.logger.warn('PINGREQ unhandled')

    @asyncio.coroutine
    def handle_disconnect(self, disconnect: DisconnectPacket):
        self.logger.warn('DISCONNECT unhandled')

    @asyncio.coroutine
    def handle_puback(self, puback: PubackPacket):
        packet_id = puback.variable_header.packet_id
        try:
            waiter = self._puback_waiters[packet_id]
            waiter.set_result(puback)
        except KeyError as ke:
            self.logger.warn("Received PUBACK for unknown pending subscription with Id: %s" % packet_id)

    @asyncio.coroutine
    def handle_pubrec(self, pubrec: PubrecPacket):
        packet_id = pubrec.variable_header.packet_id
        try:
            waiter = self._pubrec_waiters[packet_id]
            waiter.set_result(pubrec)
        except KeyError as ke:
            self.logger.warn("Received PUBREC for unknown pending subscription with Id: %s" % packet_id)

    @asyncio.coroutine
    def handle_pubcomp(self, pubcomp: PubcompPacket):
        packet_id = pubcomp.variable_header.packet_id
        try:
            waiter = self._pubcomp_waiters[packet_id]
            waiter.set_result(pubcomp)
        except KeyError as ke:
            self.logger.warn("Received PUBCOMP for unknown pending subscription with Id: %s" % packet_id)

    @asyncio.coroutine
    def handle_pubrel(self, pubrel: PubrecPacket):
        packet_id = pubrel.variable_header.packet_id
        try:
            waiter = self._pubrel_waiters[packet_id]
            waiter.set_result(pubrel)
        except KeyError as ke:
            self.logger.warn("Received PUBREL for unknown pending subscription with Id: %s" % packet_id)

    @asyncio.coroutine
    def handle_publish(self, publish : PublishPacket):
        inflight_message = None
        packet_id = publish.variable_header.packet_id
        qos = (publish.fixed_header.flags >> 1) & 0x03
        if packet_id in self.session.inflight_in:
            inflight_message = self.session.inflight_in[packet_id]
        else:
            inflight_message = InFlightMessage(publish, qos)
            self.session.inflight_in[packet_id] = inflight_message
            inflight_message.publish()

        if qos == 1:
            puback = PubackPacket.build(packet_id)
            yield from self.outgoing_queue.put(puback)
            inflight_message.acknowledge()
        if qos == 2:
            pubrec = PubrecPacket.build(packet_id)
            yield from self.outgoing_queue.put(pubrec)
            inflight_message.receive()
            waiter = futures.Future(loop=self._loop)
            self._pubrel_waiters[packet_id] = waiter
            yield from waiter
            inflight_message.pubrel = waiter.result()
            del self._pubrel_waiters[packet_id]
            inflight_message.release()
            pubcomp = PubcompPacket.build(packet_id)
            yield from self.outgoing_queue.put(pubcomp)
            inflight_message.complete()
        yield from self.delivered_message.put(inflight_message)
        del self.session.inflight_in[packet_id]
<|endoftext|>