<commit_before>import datetime
from collections import namedtuple

import pytz
from django.test import SimpleTestCase, TestCase
from mock import patch

from corehq.apps.export.filters import FormSubmittedByFilter, OwnerFilter
from corehq.apps.export.forms import (
    BaseFilterExportDownloadForm,
    EmwfFilterFormExport,
    LocationRestrictedMobileWorkerFilter,
    FilterCaseESExportDownloadForm
)
from corehq.apps.domain.models import Domain
from corehq.apps.reports.filters.case_list import CaseListFilter
from corehq.apps.groups.models import Group
from corehq.apps.locations.models import SQLLocation

@patch('corehq.apps.reports.util.get_first_form_submission_received', lambda x: datetime.datetime(2015, 1, 1))
class TestEmwfFilterFormExport(SimpleTestCase):

    def setUp(self):
        DomainObject = namedtuple('DomainObject', ['uses_locations', 'name'])
        self.project = DomainObject(False, "foo")

    def test_get_datespan_filter(self):
        form_data = {'date_range': '2015-06-25 to 2016-02-19'}
        form = EmwfFilterFormExport(self.project, pytz.utc, form_data)
        self.assertTrue(form.is_valid())
        datespan_filter = form._get_datespan_filter()
        self.assertEqual(datespan_filter.lt, datetime.datetime(2016, 2, 20, tzinfo=pytz.utc))
        self.assertEqual(datespan_filter.gte, datetime.datetime(2015, 6, 25, tzinfo=pytz.utc))
        self.assertEqual(datespan_filter.lte, None)
        self.assertEqual(datespan_filter.gt, None)

    def test_export_to_es_user_types_map(self):
        mapping = {'mobile': ['mobile'], 'demo_user': ['demo'], 'supply': ['supply'],
                   'unknown': ['unknown', 'system', 'web']}
        self.assertEqual(
            EmwfFilterFormExport._EXPORT_TO_ES_USER_TYPES_MAP,
            mapping
        )

    def test_skip_layout_default(self):
        self.assertFalse(BaseFilterExportDownloadForm.skip_layout)

@patch('corehq.apps.reports.util.get_first_form_submission_received', lambda x: datetime.datetime(2015, 1, 1))
class TestEmwfFilterExportMixin(TestCase):
    def setUp(self):
        self.domain = Domain(name="testapp", is_active=True)
        self.user_ids = ['e80c5e54ab552245457d2546d0cdbb03', 'e80c5e54ab552245457d2546d0cdbb04']
        self.user_ids_slug = ['u__' + user_id for user_id in self.user_ids]
        self.location_ids = ['e80c5e54ab552245457d2546d0cdbb03', 'e80c5e54ab552245457d2546d0cdbb04']
        self.location_ids_slug = ['l__' + location_id for location_id in self.location_ids]

    def test_get_user_ids(self):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)

        self.assertEqual(self.filter_export._get_user_ids(self.user_ids_slug), self.user_ids)

    def test_get_users_filter(self):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)
        user_filter = self.filter_export._get_users_filter(self.user_ids_slug)

        self.assertIsInstance(user_filter, FormSubmittedByFilter)
        self.assertEqual(user_filter.submitted_by, self.user_ids)
        expected_filter = {
            'terms': {'form.meta.userID': self.user_ids}
        }
        self.assertEqual(user_filter.to_es_filter(), expected_filter)

    def test_get_location_ids(self):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)

        self.assertEqual(self.filter_export._get_locations_ids(self.location_ids_slug), self.location_ids)

    @patch('corehq.apps.export.forms.user_ids_at_locations_and_descendants')
    def test_get_locations_filter(self, users_patch):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)
        users_patch.return_value = self.user_ids
        locations_filter = self.filter_export._get_locations_filter(self.location_ids_slug)

        self.assertIsInstance(locations_filter, FormSubmittedByFilter)
        self.assertEqual(locations_filter.submitted_by, self.user_ids)
        expected_filter = {
            'terms': {'form.meta.userID': self.user_ids}
        }
        self.assertEqual(locations_filter.to_es_filter(), expected_filter)

    def test_get_group_ids(self):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)
        group_ids = ['e80c5e54ab552245457d2546d0cdbb03', 'e80c5e54ab552245457d2546d0cdbb04']
        group_ids_slug = ['g__e80c5e54ab552245457d2546d0cdbb03', 'g__e80c5e54ab552245457d2546d0cdbb04']

        self.assertEqual(self.filter_export._get_group_ids(group_ids_slug), group_ids)

@patch('corehq.apps.reports.util.get_first_form_submission_received', lambda x: datetime.datetime(2015, 1, 1))
@patch.object(EmwfFilterFormExport, '_get_datespan_filter', lambda x: None)
@patch.object(EmwfFilterFormExport, '_get_group_filter')
@patch.object(EmwfFilterFormExport, '_get_user_type_filter')
@patch.object(EmwfFilterFormExport, '_get_users_filter')
@patch.object(EmwfFilterFormExport, '_get_locations_filter')
class TestEmwfFilterFormExportFilters(TestCase):
    def test_attributes(self, *patches):
        domain = Domain(name="testapp", is_active=True)
        export_filter = EmwfFilterFormExport(domain, pytz.utc)

        self.assertTrue(export_filter.skip_layout)
        self.assertEqual(export_filter.export_user_filter, FormSubmittedByFilter)
        self.assertEqual(export_filter.dynamic_filter_class, LocationRestrictedMobileWorkerFilter)

    def test_get_form_filter_for_all_locations_access(self, locations_patch, users_patch, user_type_patch,
                                                      group_patch):
        domain = Domain(name="testapp", is_active=True)
        group_ids_slug = ['g__e80c5e54ab552245457d2546d0cdbb03', 'g__e80c5e54ab552245457d2546d0cdbb04']
        EmwfFilterFormExport(domain, pytz.utc).get_form_filter(group_ids_slug, True)

        group_patch.assert_called_once_with(group_ids_slug)
        user_type_patch.assert_called_once_with(group_ids_slug)
        users_patch.assert_called_once_with(group_ids_slug)
        locations_patch.assert_called_once_with(group_ids_slug)

    def test_get_form_filter_for_restricted_locations_access(self, locations_patch, users_patch,
                                                             user_type_patch, group_patch):
        domain = Domain(name="testapp", is_active=True)
        group_ids_slug = ['g__e80c5e54ab552245457d2546d0cdbb03', 'g__e80c5e54ab552245457d2546d0cdbb04']
        EmwfFilterFormExport(domain, pytz.utc).get_form_filter(group_ids_slug, False)

        assert not group_patch.called, 'User Filter Called for restricted location access'
        assert not user_type_patch.called, 'User Type Filter Called for restricted location access'
        users_patch.assert_called_once_with(group_ids_slug)
        locations_patch.assert_called_once_with(group_ids_slug)

@patch('corehq.apps.reports.util.get_first_form_submission_received', lambda x: datetime.datetime(2015, 1, 1))
@patch.object(FilterCaseESExportDownloadForm, '_get_datespan_filter', lambda x: [])
@patch.object(FilterCaseESExportDownloadForm, '_get_group_ids')
@patch.object(Group, 'get_static_user_ids_for_groups')
@patch.object(Group, 'get_case_sharing_groups')
@patch.object(SQLLocation, 'get_case_sharing_locations_ids')
class TestFilterCaseESExportDownloadForm(TestCase):
    def setUp(self):
        self.domain = Domain(name="testapp", is_active=True)
        self.group_ids = ['e80c5e54ab552245457d2546d0cdbb03', 'e80c5e54ab552245457d2546d0cdbb04']
        self.group_ids_slug = ['g__e80c5e54ab552245457d2546d0cdbb03', 'g__e80c5e54ab552245457d2546d0cdbb04']

    @patch.object(FilterCaseESExportDownloadForm, '_get_group_independent_filters', lambda x, y, z: [])
    def test_attributes(self, *patches):
        self.export_filter = FilterCaseESExportDownloadForm(self.domain, pytz.utc)

        self.assertTrue(self.export_filter.skip_layout)
        self.assertEqual(self.export_filter.export_user_filter, OwnerFilter)
        self.assertEqual(self.export_filter.dynamic_filter_class, CaseListFilter)

    @patch.object(FilterCaseESExportDownloadForm, '_get_group_independent_filters', lambda x, y, z: [])
    def test_get_case_filter_for_all_locations_access(self, case_sharing_locations_ids_patch,
                                                      case_sharing_groups_patch, static_user_ids_for_group_patch,
                                                      group_ids_patch):
        group_ids_patch.return_value = self.group_ids
        self.export_filter = FilterCaseESExportDownloadForm(self.domain, pytz.utc)
        self.export_filter.get_case_filter(self.group_ids_slug, True)

        group_ids_patch.assert_called_once_with(self.group_ids_slug)
        static_user_ids_for_group_patch.assert_called_once_with(self.group_ids)
        assert not case_sharing_groups_patch.called
        assert not case_sharing_locations_ids_patch.called

    @patch.object(FilterCaseESExportDownloadForm, '_get_group_independent_filters', lambda x, y, z: [])
    def test_get_case_filter_for_restricted_locations_access(self, case_sharing_locations_ids_patch,
                                                             case_sharing_groups_patch,
                                                             static_user_ids_for_group_patch, group_ids_patch):
        self.export_filter = FilterCaseESExportDownloadForm(self.domain, pytz.utc)
        self.export_filter.get_case_filter(self.group_ids_slug, False)

        assert not group_ids_patch.called
        assert not static_user_ids_for_group_patch.called
        assert not case_sharing_groups_patch.called
        case_sharing_locations_ids_patch.assert_called_once_with(self.domain.name)

    @patch.object(FilterCaseESExportDownloadForm.dynamic_filter_class, 'selected_user_types')
    @patch.object(FilterCaseESExportDownloadForm, '_get_locations_filter')
    @patch.object(FilterCaseESExportDownloadForm, '_get_locations_ids')
    @patch.object(FilterCaseESExportDownloadForm, '_get_users_filter')
    def test_get_group_independent_filters_for_all_access(self, get_users_filter,
                                                          get_locations_ids, get_locations_filter,
                                                          get_es_user_types, *patches):
        self.export_filter = FilterCaseESExportDownloadForm(self.domain, pytz.utc)
        self.export_filter._get_group_independent_filters(self.group_ids_slug, True)
        get_es_user_types.assert_called_once_with(self.group_ids_slug)
        get_locations_filter.assert_called_once_with(self.group_ids_slug)
        get_locations_ids.assert_called_once_with(self.group_ids_slug)
        get_users_filter.assert_called_once_with(self.group_ids_slug)

    @patch.object(FilterCaseESExportDownloadForm.dynamic_filter_class, 'selected_user_types')
    @patch.object(FilterCaseESExportDownloadForm, '_get_locations_filter')
    @patch.object(FilterCaseESExportDownloadForm, '_get_locations_ids')
    @patch.object(FilterCaseESExportDownloadForm, '_get_users_filter')
    def test_get_group_independent_filters_for_restricted_access(self, get_users_filter,
                                                                 get_locations_ids, get_locations_filter,
                                                                 get_es_user_types, *patches):
        self.export_filter = FilterCaseESExportDownloadForm(self.domain, pytz.utc)
        self.export_filter._get_group_independent_filters(self.group_ids_slug, False)
        assert not get_es_user_types.called
        get_locations_filter.assert_called_once_with(self.group_ids_slug)
        get_locations_ids.assert_called_once_with(self.group_ids_slug)
        get_users_filter.assert_called_once_with(self.group_ids_slug)
<commit_msg>Add tests for _get_es_user_types and _get_group_filter<commit_after>import datetime
from collections import namedtuple

import pytz
from django.test import SimpleTestCase, TestCase
from mock import patch

from corehq.apps.export.filters import FormSubmittedByFilter, OwnerFilter, GroupFormSubmittedByFilter
from corehq.apps.export.forms import (
    BaseFilterExportDownloadForm,
    EmwfFilterFormExport,
    LocationRestrictedMobileWorkerFilter,
    FilterCaseESExportDownloadForm
)
from corehq.apps.domain.models import Domain
from corehq.apps.reports.filters.case_list import CaseListFilter
from corehq.apps.groups.models import Group
from corehq.apps.locations.models import SQLLocation

@patch('corehq.apps.reports.util.get_first_form_submission_received', lambda x: datetime.datetime(2015, 1, 1))
class TestEmwfFilterFormExport(SimpleTestCase):

    def setUp(self):
        DomainObject = namedtuple('DomainObject', ['uses_locations', 'name'])
        self.project = DomainObject(False, "foo")

    def test_get_datespan_filter(self):
        form_data = {'date_range': '2015-06-25 to 2016-02-19'}
        form = EmwfFilterFormExport(self.project, pytz.utc, form_data)
        self.assertTrue(form.is_valid())
        datespan_filter = form._get_datespan_filter()
        self.assertEqual(datespan_filter.lt, datetime.datetime(2016, 2, 20, tzinfo=pytz.utc))
        self.assertEqual(datespan_filter.gte, datetime.datetime(2015, 6, 25, tzinfo=pytz.utc))
        self.assertEqual(datespan_filter.lte, None)
        self.assertEqual(datespan_filter.gt, None)

    def test_export_to_es_user_types_map(self):
        mapping = {'mobile': ['mobile'], 'demo_user': ['demo'], 'supply': ['supply'],
                   'unknown': ['unknown', 'system', 'web']}
        self.assertEqual(
            EmwfFilterFormExport._EXPORT_TO_ES_USER_TYPES_MAP,
            mapping
        )

    def test_skip_layout_default(self):
        self.assertFalse(BaseFilterExportDownloadForm.skip_layout)

@patch('corehq.apps.reports.util.get_first_form_submission_received', lambda x: datetime.datetime(2015, 1, 1))
class TestEmwfFilterExportMixin(TestCase):
    def setUp(self):
        self.domain = Domain(name="testapp", is_active=True)
        self.user_ids = ['e80c5e54ab552245457d2546d0cdbb03', 'e80c5e54ab552245457d2546d0cdbb04']
        self.user_ids_slug = ['u__' + user_id for user_id in self.user_ids]
        self.location_ids = ['e80c5e54ab552245457d2546d0cdbb03', 'e80c5e54ab552245457d2546d0cdbb04']
        self.location_ids_slug = ['l__' + location_id for location_id in self.location_ids]

    def test_get_user_ids(self):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)
        self.assertEqual(self.filter_export._get_user_ids(self.user_ids_slug), self.user_ids)

    def test_get_users_filter(self):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)
        user_filter = self.filter_export._get_users_filter(self.user_ids_slug)

        self.assertIsInstance(user_filter, EmwfFilterFormExport.export_user_filter)
        self.assertEqual(user_filter.submitted_by, self.user_ids)
        expected_filter = {
            'terms': {'form.meta.userID': self.user_ids}
        }
        self.assertEqual(user_filter.to_es_filter(), expected_filter)

    def test_get_location_ids(self):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)
        self.assertEqual(self.filter_export._get_locations_ids(self.location_ids_slug), self.location_ids)

    @patch('corehq.apps.export.forms.user_ids_at_locations_and_descendants')
    def test_get_locations_filter(self, users_patch):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)
        users_patch.return_value = self.user_ids
        locations_filter = self.filter_export._get_locations_filter(self.location_ids_slug)

        self.assertIsInstance(locations_filter, EmwfFilterFormExport.export_user_filter)
        self.assertEqual(locations_filter.submitted_by, self.user_ids)
        expected_filter = {
            'terms': {'form.meta.userID': self.user_ids}
        }
        self.assertEqual(locations_filter.to_es_filter(), expected_filter)

    def test_get_group_ids(self):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)
        group_ids = ['e80c5e54ab552245457d2546d0cdbb03', 'e80c5e54ab552245457d2546d0cdbb04']
        group_ids_slug = ['g__e80c5e54ab552245457d2546d0cdbb03', 'g__e80c5e54ab552245457d2546d0cdbb04']

        self.assertEqual(self.filter_export._get_group_ids(group_ids_slug), group_ids)

    def test_get_es_user_types(self):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)
        self.assertEqual(self.filter_export._get_es_user_types(['t__0', 't__1']), [0, 1])

    @patch('corehq.apps.export.filters.get_groups_user_ids')
    def test_get_group_filter(self, patch_object):
        self.filter_export = EmwfFilterFormExport(self.domain, pytz.utc)
        group_ids = ['e80c5e54ab552245457d2546d0cdbb03', 'e80c5e54ab552245457d2546d0cdbb04']
        group_ids_slug = ['g__e80c5e54ab552245457d2546d0cdbb03', 'g__e80c5e54ab552245457d2546d0cdbb04']
        patch_object.return_value = group_ids
        group_filter = self.filter_export._get_group_filter(group_ids_slug)

        self.assertIsInstance(group_filter, GroupFormSubmittedByFilter)
        self.assertEqual(group_filter.group_ids, group_ids)
        expected_filter = {
            'terms': {'form.meta.userID': group_ids}
        }
        self.assertEqual(group_filter.to_es_filter(), expected_filter)

@patch('corehq.apps.reports.util.get_first_form_submission_received', lambda x: datetime.datetime(2015, 1, 1))
@patch.object(EmwfFilterFormExport, '_get_datespan_filter', lambda x: None)
@patch.object(EmwfFilterFormExport, '_get_group_filter')
@patch.object(EmwfFilterFormExport, '_get_user_type_filter')
@patch.object(EmwfFilterFormExport, '_get_users_filter')
@patch.object(EmwfFilterFormExport, '_get_locations_filter')
class TestEmwfFilterFormExportFilters(TestCase):
    def test_attributes(self, *patches):
        domain = Domain(name="testapp", is_active=True)
        export_filter = EmwfFilterFormExport(domain, pytz.utc)

        self.assertTrue(export_filter.skip_layout)
        self.assertEqual(export_filter.export_user_filter, FormSubmittedByFilter)
        self.assertEqual(export_filter.dynamic_filter_class, LocationRestrictedMobileWorkerFilter)

    def test_get_form_filter_for_all_locations_access(self, locations_patch, users_patch, user_type_patch,
                                                      group_patch):
        domain = Domain(name="testapp", is_active=True)
        group_ids_slug = ['g__e80c5e54ab552245457d2546d0cdbb03', 'g__e80c5e54ab552245457d2546d0cdbb04']
        EmwfFilterFormExport(domain, pytz.utc).get_form_filter(group_ids_slug, True)

        group_patch.assert_called_once_with(group_ids_slug)
        user_type_patch.assert_called_once_with(group_ids_slug)
        users_patch.assert_called_once_with(group_ids_slug)
        locations_patch.assert_called_once_with(group_ids_slug)

    def test_get_form_filter_for_restricted_locations_access(self, locations_patch, users_patch,
                                                             user_type_patch, group_patch):
        domain = Domain(name="testapp", is_active=True)
        group_ids_slug = ['g__e80c5e54ab552245457d2546d0cdbb03', 'g__e80c5e54ab552245457d2546d0cdbb04']
        EmwfFilterFormExport(domain, pytz.utc).get_form_filter(group_ids_slug, False)

        assert not group_patch.called, 'User Filter Called for restricted location access'
        assert not user_type_patch.called, 'User Type Filter Called for restricted location access'
        users_patch.assert_called_once_with(group_ids_slug)
        locations_patch.assert_called_once_with(group_ids_slug)

@patch('corehq.apps.reports.util.get_first_form_submission_received', lambda x: datetime.datetime(2015, 1, 1))
@patch.object(FilterCaseESExportDownloadForm, '_get_datespan_filter', lambda x: [])
@patch.object(FilterCaseESExportDownloadForm, '_get_group_ids')
@patch.object(Group, 'get_static_user_ids_for_groups')
@patch.object(Group, 'get_case_sharing_groups')
@patch.object(SQLLocation, 'get_case_sharing_locations_ids')
class TestFilterCaseESExportDownloadForm(TestCase):
    def setUp(self):
        self.domain = Domain(name="testapp", is_active=True)
        self.group_ids = ['e80c5e54ab552245457d2546d0cdbb03', 'e80c5e54ab552245457d2546d0cdbb04']
        self.group_ids_slug = ['g__e80c5e54ab552245457d2546d0cdbb03', 'g__e80c5e54ab552245457d2546d0cdbb04']

    @patch.object(FilterCaseESExportDownloadForm, '_get_group_independent_filters', lambda x, y, z: [])
    def test_attributes(self, *patches):
        self.export_filter = FilterCaseESExportDownloadForm(self.domain, pytz.utc)

        self.assertTrue(self.export_filter.skip_layout)
        self.assertEqual(self.export_filter.export_user_filter, OwnerFilter)
        self.assertEqual(self.export_filter.dynamic_filter_class, CaseListFilter)

    @patch.object(FilterCaseESExportDownloadForm, '_get_group_independent_filters', lambda x, y, z: [])
    def test_get_case_filter_for_all_locations_access(self, case_sharing_locations_ids_patch,
                                                      case_sharing_groups_patch, static_user_ids_for_group_patch,
                                                      group_ids_patch):
        group_ids_patch.return_value = self.group_ids
        self.export_filter = FilterCaseESExportDownloadForm(self.domain, pytz.utc)
        self.export_filter.get_case_filter(self.group_ids_slug, True)

        group_ids_patch.assert_called_once_with(self.group_ids_slug)
        static_user_ids_for_group_patch.assert_called_once_with(self.group_ids)
        assert not case_sharing_groups_patch.called
        assert not case_sharing_locations_ids_patch.called

    @patch.object(FilterCaseESExportDownloadForm, '_get_group_independent_filters', lambda x, y, z: [])
    def test_get_case_filter_for_restricted_locations_access(self, case_sharing_locations_ids_patch,
                                                             case_sharing_groups_patch,
                                                             static_user_ids_for_group_patch, group_ids_patch):
        self.export_filter = FilterCaseESExportDownloadForm(self.domain, pytz.utc)
        self.export_filter.get_case_filter(self.group_ids_slug, False)

        assert not group_ids_patch.called
        assert not static_user_ids_for_group_patch.called
        assert not case_sharing_groups_patch.called
        case_sharing_locations_ids_patch.assert_called_once_with(self.domain.name)

    @patch.object(FilterCaseESExportDownloadForm.dynamic_filter_class, 'selected_user_types')
    @patch.object(FilterCaseESExportDownloadForm, '_get_locations_filter')
    @patch.object(FilterCaseESExportDownloadForm, '_get_locations_ids')
    @patch.object(FilterCaseESExportDownloadForm, '_get_users_filter')
    def test_get_group_independent_filters_for_all_access(self, get_users_filter,
                                                          get_locations_ids, get_locations_filter,
                                                          get_es_user_types, *patches):
        self.export_filter = FilterCaseESExportDownloadForm(self.domain, pytz.utc)
        self.export_filter._get_group_independent_filters(self.group_ids_slug, True)
        get_es_user_types.assert_called_once_with(self.group_ids_slug)
        get_locations_filter.assert_called_once_with(self.group_ids_slug)
        get_locations_ids.assert_called_once_with(self.group_ids_slug)
        get_users_filter.assert_called_once_with(self.group_ids_slug)

    @patch.object(FilterCaseESExportDownloadForm.dynamic_filter_class, 'selected_user_types')
    @patch.object(FilterCaseESExportDownloadForm, '_get_locations_filter')
    @patch.object(FilterCaseESExportDownloadForm, '_get_locations_ids')
    @patch.object(FilterCaseESExportDownloadForm, '_get_users_filter')
    def test_get_group_independent_filters_for_restricted_access(self, get_users_filter,
                                                                 get_locations_ids, get_locations_filter,
                                                                 get_es_user_types, *patches):
        self.export_filter = FilterCaseESExportDownloadForm(self.domain, pytz.utc)
        self.export_filter._get_group_independent_filters(self.group_ids_slug, False)
        assert not get_es_user_types.called
        get_locations_filter.assert_called_once_with(self.group_ids_slug)
        get_locations_ids.assert_called_once_with(self.group_ids_slug)
        get_users_filter.assert_called_once_with(self.group_ids_slug)
<|endoftext|>


<commit_before># -*- coding: utf8 -*-
"""Functions to uncompact details inside field values.

.. note::

    :func:`re.compile` operations are inside of expand functions since
    module-level function bytecode is cached in python.

"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals, with_statement)

import re
import zhon.pinyin
import zhon.hanzi

from unihan_etl.constants import SPACE_DELIMITED_FIELDS

#: IRG G Sources from http://www.unicode.org/reports/tr38/#kIRG_GSource
IRG_G_SOURCES = {
    'G0': 'GB2312-80',
    'G1': 'GB12345-90',
    'G3': 'GB7589-87 unsimplified forms',
    'G5': 'GB7590-87 unsimplified forms',
    'G7': ('General Purpose Hanzi List for Modern Chinese Language, and ',
           'General List of Simplified Hanzi'),
    'GS': 'Singapore Characters',
    'G8': 'GB8565-88',
    'G9': 'GB18030-2000',
    'GE': 'GB16500-95',
    'G4K': 'Siku Quanshu (四庫全書)',
    'GBK': 'Chinese Encyclopedia (中國大百科全書)',
    'GCH': 'Ci Hai (辞海)',
    'GCY': 'Ci Yuan (辭源)',
    'GCYY ': (
        'Chinese Academy of Surveying and Mapping Ideographs '
        '(中国测绘科学院用字)'
    ),
    'GDZ': 'Geographic Publishing House Ideographs (地质出版社用字)',
    'GFZ': 'Founder Press System (方正排版系统)',
    'GGH': 'Gudai Hanyu Cidian (古代汉语词典)',
    'GH': 'GB/T 15564-1995',
    'GHC': 'Hanyu Dacidian (漢語大詞典)',
    'GHZ': 'Hanyu Dazidian ideographs (漢語大字典)',
    'GIDC': 'ID system of the Ministry of Public Security of China, 2009',
    'GJZ': 'Commercial Press Ideographs (商务印书馆用字)',
    'GK': 'GB 12052-89',
    'GKX': (
        'Kangxi Dictionary ideographs (康熙字典) 9th edition (1958) '
        'including the addendum (康熙字典)補遺'
    ),
    'GRM': 'People’s Daily Ideographs (人民日报用字)',
    'GWZ': (
        'Hanyu Dacidian Publishing House Ideographs (漢語大詞典出版社用字)'
    ),
    'GXC': 'Xiandai Hanyu Cidian (现代汉语词典)',
    'GXH': 'Xinhua Zidian (新华字典)',
    'GZFY': 'Hanyu Fangyan Dacidian (汉语方言大词典)',
    'GZH': 'ZhongHua ZiHai (中华字海)',
    'GZJW': 'Yinzhou Jinwen Jicheng Yinde (殷周金文集成引得)',
    'GFC': (
        'Modern Chinese Standard Dictionary (现代汉语规范词典第二版。主编:李'
        '行健。北京:外语 教学与研究出版社) 2010, ISBN:978-7-5600-9518-9'
    ),
    'GGFZ': 'Tongyong Guifan Hanzi Zidian (通用规范汉字字典)'
}

def expand_kDefinition(value):
    return [c.strip() for c in value.split(';')]

def expand_kMandarin(value):
    cn = value[0]
    if len(value) == 1:
        tw = value[0]
    else:
        tw = value[1]
    return {
        "zh-Hans": cn,
        "zh-Hant": tw
    }

def expand_kTotalStrokes(value):
    cn = value[0]
    if len(value) == 1:
        tw = value[0]
    else:
        tw = value[1]
    return {
        "zh-Hans": int(cn),
        "zh-Hant": int(tw)
    }

def expand_kHanYu(value):
    for i, v in enumerate(value):
        value[i] = {
            "volume": int(v[0]),
            "page": int(v[1:5]),
            "character": int(v[6:8]),
            "virtual": int(v[8])
        }
    return value

def expand_kHanyuPinyin(value):
    location_pattern = re.compile(r"""
        (?P<volume>[1-8])
        (?P<page>[0-9]{4})\.
        (?P<character>[0-3][0-9])
        (?P<virtual>[0-3])
    """, re.X)

    for i, v in enumerate(value):
        v = [s.split(',') for s in v.split(':')]
        value[i] = {
            "locations": v[0],
            "readings": v[1]
        }

        for n, loc in enumerate(value[i]['locations']):
            m = location_pattern.match(loc).groupdict()
            value[i]['locations'][n] = {
                "volume": int(m['volume']),
                "page": int(m['page']),
                "character": int(m['character']),
                "virtual": int(m['virtual'])
            }
    return value

def expand_kXHC1983(value):
    pattern = re.compile(r"""
        (?P<page>[0-9]{4})\.
        (?P<position>[0-9]{2})
        (?P<entry>[0-9]{1})
        (?P<substituted>\*?)
    """, re.X)

    for i, v in enumerate(value):
        vals = v.split(':')
        value[i] = {
            "locations": vals[0].split(','),
            "reading": vals[1],
        }

        for n, loc in enumerate(value[i]['locations']):
            m = pattern.match(loc).groupdict()
            value[i]['locations'][n] = {
                "page": int(m['page']),
                "position": int(m['position']),
                "entry": int(m['entry']),
                "substituted": m['substituted'] == "*"
            }
    return value

def expand_kCheungBauer(value):
    pattern = re.compile(r"""
        (?P<radical>[0-9]{3})\/(?P<strokes>[0-9]{2});
        (?P<cangjie>[A-Z]*);
        (?P<readings>[a-z1-6\[\]\/,]+)
    """, re.X)
    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "cangjie": m['cangjie'] or None,
            "readings": m['readings'].split(',')
        }
    return value

def expand_kRSAdobe_Japan1_6(value):
    pattern = re.compile(r"""
        (?P<type>[CV])\+
        (?P<cid>[0-9]{1,5})\+
        (?P<radical>[1-9][0-9]{0,2})\.
        (?P<strokes>[1-9][0-9]?)\.
        (?P<strokes_residue>[0-9]{1,2})
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()

        value[i] = {
            "type": m['type'],
            "cid": int(m['cid']),
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "strokes-residue": int(m['strokes_residue'])
        }
    return value

def expand_kCihaiT(value):
    for i, v in enumerate(value):
        vals = v.split('.')
        value[i] = {
            "page": int(vals[0]),
            "row": int(vals[1][0]),
            "position": int(vals[1][1:3]),
        }
    return value

def expand_kDaeJaweon(value):
    match = re.split(r'([0-9]{4})\.([0-9]{2})([01])', value)
    value = {
        "page": int(match[1]),
        "position": int(match[2]),
        "virtual": int(match[3]),
    }
    return value

def expand_kFenn(value):
    for i, v in enumerate(value):
        m = re.split(r'([0-9]+a?)([A-KP*])', value[i])
        value[i] = {
            "phonetic": int(m[1]),
            "frequency": m[2]
        }
    return value

def expand_kHanyuPinlu(value):
    pattern = re.compile(r"""
        (?P<phonetic>[a-z{}]+)
        \((?P<frequency>[0-9]+)\)
    """.format(zhon.pinyin.lowercase), re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "phonetic": m['phonetic'],
            "frequency": int(m['frequency'])
        }
    return value

def expand_kHDZRadBreak(value):
    rad, loc = value.split(':')

    pattern = re.compile(r"""
        (?P<radical>[{}]+)
        \[(?P<ucn>U\+2F[0-9A-D][0-9A-F])\]
    """.format(zhon.hanzi.radicals), re.X)
    m = pattern.match(rad).groupdict()

    return {
        "radical": m['radical'],
        "ucn": m['ucn'],
        "location": loc
    }

def expand_kSBGY(value):
    for i, v in enumerate(value):
        vals = v.split('.')
        value[i] = {
            "page": int(vals[0]),
            "character": int(vals[1])
        }
    return value

def expand_kRSUnicode(value):
    pattern = re.compile(r"""
        (?P<radical>[1-9][0-9]{0,2})
        (?P<simplified>\'?)\.
        (?P<strokes>-?[0-9]{1,2})
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "simplified": m['simplified'] == "'"
        }
    return value

def _expand_kRSGeneric(value):
    for i, v in enumerate(value):
        vals = v.split('.')
        value[i] = {
            "radical": int(vals[0]),
            "strokes": int(vals[1]),
        }
    return value

expand_kRSJapanese = _expand_kRSGeneric
expand_kRSKangXi = _expand_kRSGeneric
expand_kRSKanWa = _expand_kRSGeneric
expand_kRSKorean = _expand_kRSGeneric

def _expand_kIRG_GenericSource(value):
    v = value.split('-')
    return {
        "source": v[0],
        "location": v[1]
    }

expand_kIRG_GSource = _expand_kIRG_GenericSource
expand_kIRG_HSource = _expand_kIRG_GenericSource
expand_kIRG_JSource = _expand_kIRG_GenericSource
expand_kIRG_KPSource = _expand_kIRG_GenericSource
expand_kIRG_KSource = _expand_kIRG_GenericSource
expand_kIRG_MSource = _expand_kIRG_GenericSource
expand_kIRG_TSource = _expand_kIRG_GenericSource
expand_kIRG_USource = _expand_kIRG_GenericSource
expand_kIRG_VSource = _expand_kIRG_GenericSource

def expand_kGSR(value):
    for i, v in enumerate(value):
        m = re.split(r'([0-9]{4})([a-vx-z])(\')?', v)
        value[i] = {
            "set": int(m[1]),
            "letter": m[2],
            "apostrophe": m[3] == "'"
        }
    return value

def expand_field(field, fvalue):
    """Return structured value of information in UNIHAN field.

    :param field: field name
    :type field: str
    :param fvalue: value of field
    :type favalue: str
    :returns: list or dict of expanded field information per UNIHAN's
        documentation
    :rtype: list or dict
    """
    if field in SPACE_DELIMITED_FIELDS and fvalue:
        fvalue = fvalue.split(' ')

    try:
        expansion_func = eval('expand_%s' % field)
        return expansion_func(fvalue)
    except NameError:
        pass

    return fvalue
<commit_msg>used named group regex for kGSR<commit_after># -*- coding: utf8 -*-
"""Functions to uncompact details inside field values.

.. note::

    :func:`re.compile` operations are inside of expand functions since
    module-level function bytecode is cached in python.

"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals, with_statement)

import re
import zhon.pinyin
import zhon.hanzi

from unihan_etl.constants import SPACE_DELIMITED_FIELDS

#: IRG G Sources from http://www.unicode.org/reports/tr38/#kIRG_GSource
IRG_G_SOURCES = {
    'G0': 'GB2312-80',
    'G1': 'GB12345-90',
    'G3': 'GB7589-87 unsimplified forms',
    'G5': 'GB7590-87 unsimplified forms',
    'G7': ('General Purpose Hanzi List for Modern Chinese Language, and ',
           'General List of Simplified Hanzi'),
    'GS': 'Singapore Characters',
    'G8': 'GB8565-88',
    'G9': 'GB18030-2000',
    'GE': 'GB16500-95',
    'G4K': 'Siku Quanshu (四庫全書)',
    'GBK': 'Chinese Encyclopedia (中國大百科全書)',
    'GCH': 'Ci Hai (辞海)',
    'GCY': 'Ci Yuan (辭源)',
    'GCYY ': (
        'Chinese Academy of Surveying and Mapping Ideographs '
        '(中国测绘科学院用字)'
    ),
    'GDZ': 'Geographic Publishing House Ideographs (地质出版社用字)',
    'GFZ': 'Founder Press System (方正排版系统)',
    'GGH': 'Gudai Hanyu Cidian (古代汉语词典)',
    'GH': 'GB/T 15564-1995',
    'GHC': 'Hanyu Dacidian (漢語大詞典)',
    'GHZ': 'Hanyu Dazidian ideographs (漢語大字典)',
    'GIDC': 'ID system of the Ministry of Public Security of China, 2009',
    'GJZ': 'Commercial Press Ideographs (商务印书馆用字)',
    'GK': 'GB 12052-89',
    'GKX': (
        'Kangxi Dictionary ideographs (康熙字典) 9th edition (1958) '
        'including the addendum (康熙字典)補遺'
    ),
    'GRM': 'People’s Daily Ideographs (人民日报用字)',
    'GWZ': (
        'Hanyu Dacidian Publishing House Ideographs (漢語大詞典出版社用字)'
    ),
    'GXC': 'Xiandai Hanyu Cidian (现代汉语词典)',
    'GXH': 'Xinhua Zidian (新华字典)',
    'GZFY': 'Hanyu Fangyan Dacidian (汉语方言大词典)',
    'GZH': 'ZhongHua ZiHai (中华字海)',
    'GZJW': 'Yinzhou Jinwen Jicheng Yinde (殷周金文集成引得)',
    'GFC': (
        'Modern Chinese Standard Dictionary (现代汉语规范词典第二版。主编:李'
        '行健。北京:外语 教学与研究出版社) 2010, ISBN:978-7-5600-9518-9'
    ),
    'GGFZ': 'Tongyong Guifan Hanzi Zidian (通用规范汉字字典)'
}

def expand_kDefinition(value):
    return [c.strip() for c in value.split(';')]

def expand_kMandarin(value):
    cn = value[0]
    if len(value) == 1:
        tw = value[0]
    else:
        tw = value[1]
    return {
        "zh-Hans": cn,
        "zh-Hant": tw
    }

def expand_kTotalStrokes(value):
    cn = value[0]
    if len(value) == 1:
        tw = value[0]
    else:
        tw = value[1]
    return {
        "zh-Hans": int(cn),
        "zh-Hant": int(tw)
    }

def expand_kHanYu(value):
    for i, v in enumerate(value):
        value[i] = {
            "volume": int(v[0]),
            "page": int(v[1:5]),
            "character": int(v[6:8]),
            "virtual": int(v[8])
        }
    return value

def expand_kHanyuPinyin(value):
    location_pattern = re.compile(r"""
        (?P<volume>[1-8])
        (?P<page>[0-9]{4})\.
        (?P<character>[0-3][0-9])
        (?P<virtual>[0-3])
    """, re.X)

    for i, v in enumerate(value):
        v = [s.split(',') for s in v.split(':')]
        value[i] = {
            "locations": v[0],
            "readings": v[1]
        }

        for n, loc in enumerate(value[i]['locations']):
            m = location_pattern.match(loc).groupdict()
            value[i]['locations'][n] = {
                "volume": int(m['volume']),
                "page": int(m['page']),
                "character": int(m['character']),
                "virtual": int(m['virtual'])
            }
    return value

def expand_kXHC1983(value):
    pattern = re.compile(r"""
        (?P<page>[0-9]{4})\.
        (?P<position>[0-9]{2})
        (?P<entry>[0-9]{1})
        (?P<substituted>\*?)
    """, re.X)

    for i, v in enumerate(value):
        vals = v.split(':')
        value[i] = {
            "locations": vals[0].split(','),
            "reading": vals[1],
        }

        for n, loc in enumerate(value[i]['locations']):
            m = pattern.match(loc).groupdict()
            value[i]['locations'][n] = {
                "page": int(m['page']),
                "position": int(m['position']),
                "entry": int(m['entry']),
                "substituted": m['substituted'] == "*"
            }
    return value

def expand_kCheungBauer(value):
    pattern = re.compile(r"""
        (?P<radical>[0-9]{3})\/(?P<strokes>[0-9]{2});
        (?P<cangjie>[A-Z]*);
        (?P<readings>[a-z1-6\[\]\/,]+)
    """, re.X)
    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "cangjie": m['cangjie'] or None,
            "readings": m['readings'].split(',')
        }
    return value

def expand_kRSAdobe_Japan1_6(value):
    pattern = re.compile(r"""
        (?P<type>[CV])\+
        (?P<cid>[0-9]{1,5})\+
        (?P<radical>[1-9][0-9]{0,2})\.
        (?P<strokes>[1-9][0-9]?)\.
        (?P<strokes_residue>[0-9]{1,2})
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()

        value[i] = {
            "type": m['type'],
            "cid": int(m['cid']),
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "strokes-residue": int(m['strokes_residue'])
        }
    return value

def expand_kCihaiT(value):
    for i, v in enumerate(value):
        vals = v.split('.')
        value[i] = {
            "page": int(vals[0]),
            "row": int(vals[1][0]),
            "position": int(vals[1][1:3]),
        }
    return value

def expand_kDaeJaweon(value):
    pattern = re.compile(r"""
        (?P<page>[0-9]{4})\.
        (?P<position>[0-9]{2})
        (?P<virtual>[01])
    """, re.X)
    m = pattern.match(value).groupdict()

    value = {
        "page": int(m['page']),
        "position": int(m['position']),
        "virtual": int(m['virtual']),
    }
    return value

def expand_kFenn(value):
    pattern = re.compile(r"""
        (?P<phonetic>[0-9]+a?)
        (?P<frequency>[A-KP*])
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict(v)
        value[i] = {
            "phonetic": int(m['phonetic']),
            "frequency": m['frequency']
        }
    return value

def expand_kHanyuPinlu(value):
    pattern = re.compile(r"""
        (?P<phonetic>[a-z{}]+)
        \((?P<frequency>[0-9]+)\)
    """.format(zhon.pinyin.lowercase), re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "phonetic": m['phonetic'],
            "frequency": int(m['frequency'])
        }
    return value

def expand_kHDZRadBreak(value):
    rad, loc = value.split(':')

    pattern = re.compile(r"""
        (?P<radical>[{}]+)
        \[(?P<ucn>U\+2F[0-9A-D][0-9A-F])\]
    """.format(zhon.hanzi.radicals), re.X)
    m = pattern.match(rad).groupdict()

    return {
        "radical": m['radical'],
        "ucn": m['ucn'],
        "location": loc
    }

def expand_kSBGY(value):
    for i, v in enumerate(value):
        vals = v.split('.')
        value[i] = {
            "page": int(vals[0]),
            "character": int(vals[1])
        }
    return value

def expand_kRSUnicode(value):
    pattern = re.compile(r"""
        (?P<radical>[1-9][0-9]{0,2})
        (?P<simplified>\'?)\.
        (?P<strokes>-?[0-9]{1,2})
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "simplified": m['simplified'] == "'"
        }
    return value

def _expand_kRSGeneric(value):
    for i, v in enumerate(value):
        vals = v.split('.')
        value[i] = {
            "radical": int(vals[0]),
            "strokes": int(vals[1]),
        }
    return value

expand_kRSJapanese = _expand_kRSGeneric
expand_kRSKangXi = _expand_kRSGeneric
expand_kRSKanWa = _expand_kRSGeneric
expand_kRSKorean = _expand_kRSGeneric

def _expand_kIRG_GenericSource(value):
    v = value.split('-')
    return {
        "source": v[0],
        "location": v[1]
    }

expand_kIRG_GSource = _expand_kIRG_GenericSource
expand_kIRG_HSource = _expand_kIRG_GenericSource
expand_kIRG_JSource = _expand_kIRG_GenericSource
expand_kIRG_KPSource = _expand_kIRG_GenericSource
expand_kIRG_KSource = _expand_kIRG_GenericSource
expand_kIRG_MSource = _expand_kIRG_GenericSource
expand_kIRG_TSource = _expand_kIRG_GenericSource
expand_kIRG_USource = _expand_kIRG_GenericSource
expand_kIRG_VSource = _expand_kIRG_GenericSource

def expand_kGSR(value):
    pattern = re.compile(r"""
        (?P<set>[0-9]{4})
        (?P<letter>[a-vx-z])
        (?P<apostrophe>\')?
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "set": int(m['set']),
            "letter": m['letter'],
            "apostrophe": m['apostrophe'] == "'"
        }
    return value

def expand_field(field, fvalue):
    """Return structured value of information in UNIHAN field.

    :param field: field name
    :type field: str
    :param fvalue: value of field
    :type favalue: str
    :returns: list or dict of expanded field information per UNIHAN's
        documentation
    :rtype: list or dict
    """
    if field in SPACE_DELIMITED_FIELDS and fvalue:
        fvalue = fvalue.split(' ')

    try:
        expansion_func = eval('expand_%s' % field)
        return expansion_func(fvalue)
    except NameError:
        pass

    return fvalue
<|endoftext|>


<commit_before>import math
from sys import exit
import numpy as np
import uuid
from .onshape_api.client import Client
from .config import config, configFile
from colorama import Fore, Back, Style

# OnShape API client
client = Client(logging=False, creds=configFile)
client.useCollisionsConfigurations = config['useCollisionsConfigurations']

# If a versionId is provided, it will be used, else the main workspace is retrieved
if config['versionId'] == '':
    print("\n" + Style.BRIGHT + '* Retrieving workspace ID ...' + Style.RESET_ALL)
    document = client.get_document(config['documentId']).json()
    workspaceId = document['defaultWorkspace']['id']
    print(Fore.GREEN + "+ Using workspace id: " + workspaceId + Style.RESET_ALL)
else:
    print("\n" + Style.BRIGHT + '* Using configuration version ID ' +
          config['versionId']+' ...' + Style.RESET_ALL)

# Now, finding the assembly, according to given name in configuration, or else the first possible one
print("\n" + Style.BRIGHT +
      '* Retrieving elements in the document, searching for the assembly...' + Style.RESET_ALL)
if config['versionId'] == '':
    elements = client.list_elements(config['documentId'], workspaceId).json()
else:
    elements = client.list_elements(
        config['documentId'], config['versionId'], 'v').json()
assemblyId = None
assemblyName = ''
for element in elements:
    if element['type'] == 'Assembly' and \
            (config['assemblyName'] is False or element['name'] == config['assemblyName']):
        print(Fore.GREEN + "+ Found assembly, id: " +
              element['id']+', name: "'+element['name']+'"' + Style.RESET_ALL)
        assemblyName = element['name']
        assemblyId = element['id']

if assemblyId == None:
    print(Fore.RED + "ERROR: Unable to find assembly in this document" + Style.RESET_ALL)
    exit(1)

# Retrieving the assembly
print("\n" + Style.BRIGHT + '* Retrieving assembly "' +
      assemblyName+'" with id '+assemblyId + Style.RESET_ALL)
if config['versionId'] == '':
    assembly = client.get_assembly(
        config['documentId'], workspaceId, assemblyId)
else:
    assembly = client.get_assembly(
        config['documentId'], config['versionId'], assemblyId, 'v')

root = assembly['rootAssembly']

# Finds a (leaf) instance given the full path, typically A B C where A and B would be subassemblies and C
# the final part

def findInstance(path, instances=None):
    global assembly

    if instances is None:
        instances = assembly['rootAssembly']['instances']

    for instance in instances:
        if instance['id'] == path[0]:
            if len(path) == 1:
                # If the length of remaining path is 1, the part is in the current assembly/subassembly
                return instance
            else:
                # Else, we need to find the matching sub assembly to find the proper part (recursively)
                d = instance['documentId']
                m = instance['documentMicroversion']
                e = instance['elementId']
                for asm in assembly['subAssemblies']:
                    if asm['documentId'] == d and asm['documentMicroversion'] == m and asm['elementId'] == e:
                        return findInstance(path[1:], asm['instances'])

    print(Fore.RED + 'Could not find instance for ' + str(path) + Style.RESET_ALL)

# Collecting occurrences, the path is the assembly / sub assembly chain
occurrences = {}
for occurrence in root['occurrences']:
    occurrence['assignation'] = None
    occurrence['instance'] = findInstance(occurrence['path'])
    occurrence['transform'] = np.matrix(
        np.reshape(occurrence['transform'], (4, 4)))
    occurrence['linkName'] = None
    occurrences[tuple(occurrence['path'])] = occurrence

# Gets an occurrence given its full path

def getOccurrence(path):
    return occurrences[tuple(path)]

# Assignations are pieces that will be in the same link. Note that this is only for top-level
# item of the path (all sub assemblies and parts in assemblies are naturally in the same link as
# the parent), but other parts that can be connected with mates in top assemblies are then assigned to
# the link
assignations = {}

# Frames (mated with frame_ name) will be special links in the output file allowing to track some specific
# manually identified frames
frames = {}

def assignParts(root, parent):
    assignations[root] = parent
    for occurrence in occurrences.values():
        if occurrence['path'][0] == root:
            occurrence['assignation'] = parent

from .features import init as features_init, getLimits
features_init(client, config, root, workspaceId, assemblyId)

# First, features are scanned to find the DOFs. Links that they connects are then tagged
print("\n" + Style.BRIGHT +
      '* Getting assembly features, scanning for DOFs...' + Style.RESET_ALL)
trunk = None
relations = {}
features = root['features']
for feature in features:
    if feature['featureType'] == 'mateConnector':
        name = feature['featureData']['name']
        if name[0:5] == 'link_':
            name = name[5:]
            occurrences[(feature['featureData']['occurrence'][0],)
                        ]['linkName'] = name
    else:
        if feature['suppressed']:
            continue

        data = feature['featureData']

        if len(data['matedEntities']) != 2 or \
                len(data['matedEntities'][0]['matedOccurrence']) == 0 \
                or len(data['matedEntities'][1]['matedOccurrence']) == 0:
            continue

        child = data['matedEntities'][0]['matedOccurrence'][0]
        parent = data['matedEntities'][1]['matedOccurrence'][0]

        if data['name'][0:3] == 'dof':
            parts = data['name'].split('_')
            del parts[0]
            data['inverted'] = False
            if parts[-1] == 'inv' or parts[-1] == 'inverted':
                data['inverted'] = True
                del parts[-1]
            name = '_'.join(parts)
            if name == '':
                print(Fore.RED +
                      'ERROR: a DOF dones\'t have any name ("'+data['name']+'" should be "dof_...")' + Style.RESET_ALL)
                exit()

            if data['mateType'] == 'REVOLUTE' or data['mateType'] == 'CYLINDRICAL':
                jointType = 'revolute'
                limits = getLimits(jointType, data['name'])
            elif data['mateType'] == 'SLIDER':
                jointType = 'prismatic'
                limits = getLimits(jointType, data['name'])
            elif data['mateType'] == 'FASTENED':
                jointType = 'floating'
                limits = None
            else:
                print(Fore.RED + 'ERROR: "' + name +
                      '" is declared as a DOF but the mate type is '+data['mateType']+'')
                print(
                    '       Only REVOLUTE, CYLINDRICAL, SLIDER and FASTENED are supported' + Style.RESET_ALL)
                exit(1)

            # We compute the axis in the world frame
            matedEntity = data['matedEntities'][0]
            matedTransform = getOccurrence(
                matedEntity['matedOccurrence'])['transform']

            # jointToPart is the (rotation only) matrix from joint to the part
            # it is attached to
            jointToPart = np.eye(4)
            jointToPart[:3, :3] = np.stack((
                np.array(matedEntity['matedCS']['xAxis']),
                np.array(matedEntity['matedCS']['yAxis']),
                np.array(matedEntity['matedCS']['zAxis'])
            )).T

            if data['inverted']:
                if limits is not None:
                    limits = (-limits[1], -limits[0])

                # Flipping the joint around X axis
                flip = np.array([[1, 0, 0, 0],
                                [0, -1, 0, 0],
                                [0, 0, -1, 0],
                                [0, 0,  0,  1]])
                jointToPart = flip.dot(jointToPart)

            zAxis = np.array([0, 0, 1])

            origin = matedEntity['matedCS']['origin']
            translation = np.matrix(np.identity(4))
            translation[0, 3] += origin[0]
            translation[1, 3] += origin[1]
            translation[2, 3] += origin[2]
            worldAxisFrame = matedTransform * translation

            # Resulting frame of axis, always revolving around z
            worldAxisFrame = worldAxisFrame.dot(jointToPart)

            limitsStr = ''
            if limits is not None:
                limitsStr = '[' + str(round(limits[0], 3)) + \
                    ': ' + str(round(limits[1], 3)) + ']'
            print(Fore.GREEN + '+ Found DOF: '+name + ' ' + Style.DIM +
                  '('+jointType+')'+limitsStr + Style.RESET_ALL)

            if child in relations:
                print(Fore.RED)
                print('Error, the relation '+name +
                      ' is connected a child that is already connected')
                print('Be sure you ordered properly your relations, see:')
                print(
                    'https://onshape-to-robot.readthedocs.io/en/latest/design.html#specifying-degrees-of-freedom')
                print(Style.RESET_ALL)
                exit()

            relations[child] = {
                'parent': parent,
                'worldAxisFrame': worldAxisFrame,
                'zAxis': zAxis,
                'name': name,
                'type': jointType,
                'limits': limits
            }

            assignParts(child, child)
            assignParts(parent, parent)
            if child not in frames:
                frames[child] = []
            if parent not in frames:
                frames[parent] = []

print(Fore.GREEN + Style.BRIGHT + '* Found total ' +
      str(len(relations))+' DOFs' + Style.RESET_ALL)

# If we have no DOF
if len(relations) == 0:
    trunk = root['instances'][0]['id']
    assignParts(trunk, trunk)

def connectParts(child, parent):
    assignParts(child, parent)

# Spreading parts assignations, this parts mainly does two things:
# 1. Finds the parts of the top level assembly that are not directly in a sub assembly and try to assign them
#    to an existing link that was identified before
# 2. Among those parts, finds the ones that are frames (connected with a frame_* connector)
changed = True
while changed:
    changed = False
    for feature in features:
        if feature['featureType'] != 'mate' or feature['suppressed']:
            continue

        data = feature['featureData']

        if len(data['matedEntities']) != 2 \
                or len(data['matedEntities'][0]['matedOccurrence']) == 0 \
                or len(data['matedEntities'][1]['matedOccurrence']) == 0:
            continue

        occurrenceA = data['matedEntities'][0]['matedOccurrence'][0]
        occurrenceB = data['matedEntities'][1]['matedOccurrence'][0]

        if (occurrenceA not in assignations) != (occurrenceB not in assignations):
            if data['name'][0:5] == 'frame':
                name = '_'.join(data['name'].split('_')[1:])
                if occurrenceA in assignations:
                    frames[occurrenceA].append(
                        [name, data['matedEntities'][1]['matedOccurrence']])
                    assignParts(occurrenceB, {True: assignations[occurrenceA], False: 'frame'}[
                                config['drawFrames']])
                    changed = True
                else:
                    frames[occurrenceB].append(
                        [name, data['matedEntities'][0]['matedOccurrence']])
                    assignParts(occurrenceA, {True: assignations[occurrenceB], False: 'frame'}[
                                config['drawFrames']])
                    changed = True
            else:
                if occurrenceA in assignations:
                    connectParts(occurrenceB, assignations[occurrenceA])
                    changed = True
                else:
                    connectParts(occurrenceA, assignations[occurrenceB])
                    changed = True

# Building and checking robot tree, here we:
# 1. Search for robot trunk (which will be the top-level link)
# 2. Scan for orphaned parts (if you add something floating with no mate to anything)
#    that are then assigned to trunk by default
# 3. Collect all the pieces of the robot tree
print("\n" + Style.BRIGHT + '* Building robot tree' + Style.RESET_ALL)

for childId in relations:
    entry = relations[childId]
    if entry['parent'] not in relations:
        trunk = entry['parent']
        break
trunkOccurrence = getOccurrence([trunk])
print(Style.BRIGHT + '* Trunk is ' +
      trunkOccurrence['instance']['name'] + Style.RESET_ALL)

for occurrence in occurrences.values():
    if occurrence['assignation'] is None:
        print(Fore.YELLOW + 'WARNING: part (' +
              occurrence['instance']['name']+') has no assignation, connecting it with trunk' + Style.RESET_ALL)
        child = occurrence['path'][0]
        connectParts(child, trunk)

def collect(id):
    part = {}
    part['id'] = id
    part['children'] = []
    for childId in relations:
        entry = relations[childId]
        if entry['parent'] == id:
            child = collect(childId)
            child['axis_frame'] = entry['worldAxisFrame']
            child['z_axis'] = entry['zAxis']
            child['dof_name'] = entry['name']
            child['jointType'] = entry['type']
            child['jointLimits'] = entry['limits']
            part['children'].append(child)
    return part

tree = collect(trunk)
<commit_msg>Checking if matedEntities is present in data dict (see #45)<commit_after>import math
from sys import exit
import numpy as np
import uuid
from .onshape_api.client import Client
from .config import config, configFile
from colorama import Fore, Back, Style

# OnShape API client
client = Client(logging=False, creds=configFile)
client.useCollisionsConfigurations = config['useCollisionsConfigurations']

# If a versionId is provided, it will be used, else the main workspace is retrieved
if config['versionId'] == '':
    print("\n" + Style.BRIGHT + '* Retrieving workspace ID ...' + Style.RESET_ALL)
    document = client.get_document(config['documentId']).json()
    workspaceId = document['defaultWorkspace']['id']
    print(Fore.GREEN + "+ Using workspace id: " + workspaceId + Style.RESET_ALL)
else:
    print("\n" + Style.BRIGHT + '* Using configuration version ID ' +
          config['versionId']+' ...' + Style.RESET_ALL)

# Now, finding the assembly, according to given name in configuration, or else the first possible one
print("\n" + Style.BRIGHT +
      '* Retrieving elements in the document, searching for the assembly...' + Style.RESET_ALL)
if config['versionId'] == '':
    elements = client.list_elements(config['documentId'], workspaceId).json()
else:
    elements = client.list_elements(
        config['documentId'], config['versionId'], 'v').json()
assemblyId = None
assemblyName = ''
for element in elements:
    if element['type'] == 'Assembly' and \
            (config['assemblyName'] is False or element['name'] == config['assemblyName']):
        print(Fore.GREEN + "+ Found assembly, id: " +
              element['id']+', name: "'+element['name']+'"' + Style.RESET_ALL)
        assemblyName = element['name']
        assemblyId = element['id']

if assemblyId == None:
    print(Fore.RED + "ERROR: Unable to find assembly in this document" + Style.RESET_ALL)
    exit(1)

# Retrieving the assembly
print("\n" + Style.BRIGHT + '* Retrieving assembly "' +
      assemblyName+'" with id '+assemblyId + Style.RESET_ALL)
if config['versionId'] == '':
    assembly = client.get_assembly(
        config['documentId'], workspaceId, assemblyId)
else:
    assembly = client.get_assembly(
        config['documentId'], config['versionId'], assemblyId, 'v')

root = assembly['rootAssembly']

# Finds a (leaf) instance given the full path, typically A B C where A and B would be subassemblies and C
# the final part

def findInstance(path, instances=None):
    global assembly

    if instances is None:
        instances = assembly['rootAssembly']['instances']

    for instance in instances:
        if instance['id'] == path[0]:
            if len(path) == 1:
                # If the length of remaining path is 1, the part is in the current assembly/subassembly
                return instance
            else:
                # Else, we need to find the matching sub assembly to find the proper part (recursively)
                d = instance['documentId']
                m = instance['documentMicroversion']
                e = instance['elementId']
                for asm in assembly['subAssemblies']:
                    if asm['documentId'] == d and asm['documentMicroversion'] == m and asm['elementId'] == e:
                        return findInstance(path[1:], asm['instances'])

    print(Fore.RED + 'Could not find instance for ' + str(path) + Style.RESET_ALL)

# Collecting occurrences, the path is the assembly / sub assembly chain
occurrences = {}
for occurrence in root['occurrences']:
    occurrence['assignation'] = None
    occurrence['instance'] = findInstance(occurrence['path'])
    occurrence['transform'] = np.matrix(
        np.reshape(occurrence['transform'], (4, 4)))
    occurrence['linkName'] = None
    occurrences[tuple(occurrence['path'])] = occurrence

# Gets an occurrence given its full path

def getOccurrence(path):
    return occurrences[tuple(path)]

# Assignations are pieces that will be in the same link. Note that this is only for top-level
# item of the path (all sub assemblies and parts in assemblies are naturally in the same link as
# the parent), but other parts that can be connected with mates in top assemblies are then assigned to
# the link
assignations = {}

# Frames (mated with frame_ name) will be special links in the output file allowing to track some specific
# manually identified frames
frames = {}

def assignParts(root, parent):
    assignations[root] = parent
    for occurrence in occurrences.values():
        if occurrence['path'][0] == root:
            occurrence['assignation'] = parent

from .features import init as features_init, getLimits
features_init(client, config, root, workspaceId, assemblyId)

# First, features are scanned to find the DOFs. Links that they connects are then tagged
print("\n" + Style.BRIGHT +
      '* Getting assembly features, scanning for DOFs...' + Style.RESET_ALL)
trunk = None
relations = {}
features = root['features']
for feature in features:
    if feature['featureType'] == 'mateConnector':
        name = feature['featureData']['name']
        if name[0:5] == 'link_':
            name = name[5:]
            occurrences[(feature['featureData']['occurrence'][0],)
                        ]['linkName'] = name
    else:
        if feature['suppressed']:
            continue

        data = feature['featureData']

        if 'matedEntities' not in data or len(data['matedEntities']) != 2 or \
                len(data['matedEntities'][0]['matedOccurrence']) == 0 \
                or len(data['matedEntities'][1]['matedOccurrence']) == 0:
            continue

        child = data['matedEntities'][0]['matedOccurrence'][0]
        parent = data['matedEntities'][1]['matedOccurrence'][0]

        if data['name'][0:3] == 'dof':
            parts = data['name'].split('_')
            del parts[0]
            data['inverted'] = False
            if parts[-1] == 'inv' or parts[-1] == 'inverted':
                data['inverted'] = True
                del parts[-1]
            name = '_'.join(parts)
            if name == '':
                print(Fore.RED +
                      'ERROR: a DOF dones\'t have any name ("'+data['name']+'" should be "dof_...")' + Style.RESET_ALL)
                exit()

            if data['mateType'] == 'REVOLUTE' or data['mateType'] == 'CYLINDRICAL':
                jointType = 'revolute'
                limits = getLimits(jointType, data['name'])
            elif data['mateType'] == 'SLIDER':
                jointType = 'prismatic'
                limits = getLimits(jointType, data['name'])
            elif data['mateType'] == 'FASTENED':
                jointType = 'floating'
                limits = None
            else:
                print(Fore.RED + 'ERROR: "' + name +
                      '" is declared as a DOF but the mate type is '+data['mateType']+'')
                print(
                    '       Only REVOLUTE, CYLINDRICAL, SLIDER and FASTENED are supported' + Style.RESET_ALL)
                exit(1)

            # We compute the axis in the world frame
            matedEntity = data['matedEntities'][0]
            matedTransform = getOccurrence(
                matedEntity['matedOccurrence'])['transform']

            # jointToPart is the (rotation only) matrix from joint to the part
            # it is attached to
            jointToPart = np.eye(4)
            jointToPart[:3, :3] = np.stack((
                np.array(matedEntity['matedCS']['xAxis']),
                np.array(matedEntity['matedCS']['yAxis']),
                np.array(matedEntity['matedCS']['zAxis'])
            )).T

            if data['inverted']:
                if limits is not None:
                    limits = (-limits[1], -limits[0])

                # Flipping the joint around X axis
                flip = np.array([[1, 0, 0, 0],
                                [0, -1, 0, 0],
                                [0, 0, -1, 0],
                                [0, 0,  0,  1]])
                jointToPart = flip.dot(jointToPart)

            zAxis = np.array([0, 0, 1])

            origin = matedEntity['matedCS']['origin']
            translation = np.matrix(np.identity(4))
            translation[0, 3] += origin[0]
            translation[1, 3] += origin[1]
            translation[2, 3] += origin[2]
            worldAxisFrame = matedTransform * translation

            # Resulting frame of axis, always revolving around z
            worldAxisFrame = worldAxisFrame.dot(jointToPart)

            limitsStr = ''
            if limits is not None:
                limitsStr = '[' + str(round(limits[0], 3)) + \
                    ': ' + str(round(limits[1], 3)) + ']'
            print(Fore.GREEN + '+ Found DOF: '+name + ' ' + Style.DIM +
                  '('+jointType+')'+limitsStr + Style.RESET_ALL)

            if child in relations:
                print(Fore.RED)
                print('Error, the relation '+name +
                      ' is connected a child that is already connected')
                print('Be sure you ordered properly your relations, see:')
                print(
                    'https://onshape-to-robot.readthedocs.io/en/latest/design.html#specifying-degrees-of-freedom')
                print(Style.RESET_ALL)
                exit()

            relations[child] = {
                'parent': parent,
                'worldAxisFrame': worldAxisFrame,
                'zAxis': zAxis,
                'name': name,
                'type': jointType,
                'limits': limits
            }

            assignParts(child, child)
            assignParts(parent, parent)
            if child not in frames:
                frames[child] = []
            if parent not in frames:
                frames[parent] = []

print(Fore.GREEN + Style.BRIGHT + '* Found total ' +
      str(len(relations))+' DOFs' + Style.RESET_ALL)

# If we have no DOF
if len(relations) == 0:
    trunk = root['instances'][0]['id']
    assignParts(trunk, trunk)

def connectParts(child, parent):
    assignParts(child, parent)

# Spreading parts assignations, this parts mainly does two things:
# 1. Finds the parts of the top level assembly that are not directly in a sub assembly and try to assign them
#    to an existing link that was identified before
# 2. Among those parts, finds the ones that are frames (connected with a frame_* connector)
changed = True
while changed:
    changed = False
    for feature in features:
        if feature['featureType'] != 'mate' or feature['suppressed']:
            continue

        data = feature['featureData']

        if len(data['matedEntities']) != 2 \
                or len(data['matedEntities'][0]['matedOccurrence']) == 0 \
                or len(data['matedEntities'][1]['matedOccurrence']) == 0:
            continue

        occurrenceA = data['matedEntities'][0]['matedOccurrence'][0]
        occurrenceB = data['matedEntities'][1]['matedOccurrence'][0]

        if (occurrenceA not in assignations) != (occurrenceB not in assignations):
            if data['name'][0:5] == 'frame':
                name = '_'.join(data['name'].split('_')[1:])
                if occurrenceA in assignations:
                    frames[occurrenceA].append(
                        [name, data['matedEntities'][1]['matedOccurrence']])
                    assignParts(occurrenceB, {True: assignations[occurrenceA], False: 'frame'}[
                                config['drawFrames']])
                    changed = True
                else:
                    frames[occurrenceB].append(
                        [name, data['matedEntities'][0]['matedOccurrence']])
                    assignParts(occurrenceA, {True: assignations[occurrenceB], False: 'frame'}[
                                config['drawFrames']])
                    changed = True
            else:
                if occurrenceA in assignations:
                    connectParts(occurrenceB, assignations[occurrenceA])
                    changed = True
                else:
                    connectParts(occurrenceA, assignations[occurrenceB])
                    changed = True

# Building and checking robot tree, here we:
# 1. Search for robot trunk (which will be the top-level link)
# 2. Scan for orphaned parts (if you add something floating with no mate to anything)
#    that are then assigned to trunk by default
# 3. Collect all the pieces of the robot tree
print("\n" + Style.BRIGHT + '* Building robot tree' + Style.RESET_ALL)

for childId in relations:
    entry = relations[childId]
    if entry['parent'] not in relations:
        trunk = entry['parent']
        break
trunkOccurrence = getOccurrence([trunk])
print(Style.BRIGHT + '* Trunk is ' +
      trunkOccurrence['instance']['name'] + Style.RESET_ALL)

for occurrence in occurrences.values():
    if occurrence['assignation'] is None:
        print(Fore.YELLOW + 'WARNING: part (' +
              occurrence['instance']['name']+') has no assignation, connecting it with trunk' + Style.RESET_ALL)
        child = occurrence['path'][0]
        connectParts(child, trunk)

def collect(id):
    part = {}
    part['id'] = id
    part['children'] = []
    for childId in relations:
        entry = relations[childId]
        if entry['parent'] == id:
            child = collect(childId)
            child['axis_frame'] = entry['worldAxisFrame']
            child['z_axis'] = entry['zAxis']
            child['dof_name'] = entry['name']
            child['jointType'] = entry['type']
            child['jointLimits'] = entry['limits']
            part['children'].append(child)
    return part

tree = collect(trunk)
<|endoftext|>


<commit_before>from corehq.apps.commtrack.util import num_periods_late
from dimagi.utils.decorators.memoized import memoized
from corehq.apps.locations.models import Location
from corehq.apps.commtrack.models import Product, SupplyPointCase
from dimagi.utils.couch.loosechange import map_reduce
from corehq.apps.reports.api import ReportDataSource
from datetime import datetime, timedelta
from casexml.apps.stock.models import StockState, StockTransaction
from django.db.models import Sum, Avg
from corehq.apps.reports.commtrack.util import get_relevant_supply_point_ids, product_ids_filtered_by_program
from corehq.apps.reports.commtrack.const import STOCK_SECTION_TYPE
from casexml.apps.stock.utils import months_of_stock_remaining, stock_category

# TODO make settings
REPORTING_PERIOD = 'weekly'
REPORTING_PERIOD_ARGS = (1,)

def is_timely(case, limit=0):
    return num_periods_late(case, REPORTING_PERIOD, *REPORTING_PERIOD_ARGS) <= limit

def reporting_status(transaction, start_date, end_date):
    if transaction:
        last_reported = transaction.report.date.date()
    else:
        last_reported = None

    if last_reported and last_reported < start_date:
        return 'ontime'
    elif last_reported and start_date <= last_reported <= end_date:
        return 'late'
    else:
        return 'nonreporting'

class CommtrackDataSourceMixin(object):

    @property
    def domain(self):
        return self.config.get('domain')

    @property
    @memoized
    def active_location(self):
        loc_id = self.config.get('location_id')
        if loc_id:
            return Location.get(loc_id)

    @property
    @memoized
    def active_product(self):
        prod_id = self.config.get('product_id')
        if prod_id:
            return Product.get(prod_id)

    @property
    @memoized
    def program_id(self):
        prog_id = self.config.get('program_id')
        if prog_id != '':
            return prog_id

    @property
    def start_date(self):
        date = self.config.get('start_date')
        if date:
            return datetime.strptime(date, '%Y-%m-%d').date()
        else:
            return (datetime.now() - timedelta(30)).date()

    @property
    def end_date(self):
        date = self.config.get('end_date')
        if date:
            return datetime.strptime(date, '%Y-%m-%d').date()
        else:
            return datetime.now().date()

class StockStatusDataSource(ReportDataSource, CommtrackDataSourceMixin):
    """
    Config:
        domain: The domain to report on.
        location_id: ID of location to get data for. Omit for all locations.
        product_id: ID of product to get data for. Omit for all products.
        aggregate: True to aggregate the indicators by product for the current location.

    Data Slugs:
        product_name: Name of the product
        product_id: ID of the product
        location_id: The ID of the current location.
        location_lineage: The lineage of the current location.
        current_stock: The current stock level
        consumption: The current monthly consumption rate
        months_remaining: The number of months remaining until stock out
        category: The status category. See casexml.apps.stock.models.StockState.stock_category
        resupply_quantity_needed: Max amount - current amount

    """
    slug = 'agg_stock_status'

    SLUG_PRODUCT_NAME = 'product_name'
    SLUG_PRODUCT_ID = 'product_id'
    SLUG_MONTHS_REMAINING = 'months_remaining'
    SLUG_CONSUMPTION = 'consumption'
    SLUG_CURRENT_STOCK = 'current_stock'
    SLUG_LOCATION_ID = 'location_id'
    SLUG_LOCATION_LINEAGE = 'location_lineage'
    SLUG_STOCKOUT_SINCE = 'stockout_since'
    SLUG_STOCKOUT_DURATION = 'stockout_duration'
    SLUG_LAST_REPORTED = 'last_reported'
    SLUG_CATEGORY = 'category'
    SLUG_RESUPPLY_QUANTITY_NEEDED = 'resupply_quantity_needed'

    _slug_attrib_map = {
        SLUG_PRODUCT_NAME: lambda s: Product.get(s.product_id).name,
        SLUG_PRODUCT_ID: lambda s: s.product_id,
        SLUG_LOCATION_ID: lambda s: SupplyPointCase.get(s.case_id).location_[-1],
        # SLUG_LOCATION_LINEAGE: lambda p: list(reversed(p.location_[:-1])),
        SLUG_CURRENT_STOCK: 'stock_on_hand',
        SLUG_CONSUMPTION: lambda s: s.daily_consumption * 30 if s.daily_consumption else None,
        SLUG_MONTHS_REMAINING: 'months_remaining',
        SLUG_CATEGORY: 'stock_category',
        # SLUG_STOCKOUT_SINCE: 'stocked_out_since',
        # SLUG_STOCKOUT_DURATION: 'stockout_duration_in_months',
        SLUG_LAST_REPORTED: 'last_modified_date',
        SLUG_RESUPPLY_QUANTITY_NEEDED: 'resupply_quantity_needed',
    }

    def slugs(self):
        return self._slug_attrib_map.keys()

    def filter_by_program(self, stock_states):
        return stock_states.filter(
            product_id__in=product_ids_filtered_by_program(
                self.domain,
                self.program_id
            )
        )

    def get_data(self, slugs=None):
        sp_ids = get_relevant_supply_point_ids(self.domain, self.active_location)

        if len(sp_ids) == 1:
            stock_states = StockState.objects.filter(
                case_id=sp_ids[0],
                section_id=STOCK_SECTION_TYPE,
                last_modified_date__lte=self.end_date,
                last_modified_date__gte=self.start_date,
            )

            if self.program_id:
                stock_states = self.filter_by_program(stock_states)

            return self.leaf_node_data(stock_states)
        else:
            stock_states = StockState.objects.filter(
                case_id__in=sp_ids,
                section_id=STOCK_SECTION_TYPE,
                last_modified_date__lte=self.end_date,
                last_modified_date__gte=self.start_date,
            )

            if self.program_id:
                stock_states = self.filter_by_program(stock_states)

            if self.config.get('aggregate'):
                aggregates = stock_states.values('product_id').annotate(
                    avg_consumption=Avg('daily_consumption'),
                    total_stock=Sum('stock_on_hand')
                )
                return self.aggregated_data(aggregates)
            else:
                return self.raw_product_states(stock_states, slugs)

    def leaf_node_data(self, stock_states):
        for state in stock_states:
            product = Product.get(state.product_id)
            yield {
                'category': state.stock_category,
                'product_id': product._id,
                'consumption': state.daily_consumption * 30 if state.daily_consumption else None,
                'months_remaining': state.months_remaining,
                'location_id': state.case_id,
                'product_name': product.name,
                'current_stock': state.stock_on_hand,
                'location_lineage': None,
                'resupply_quantity_needed': state.resupply_quantity_needed
            }

    def aggregated_data(self, stock_states):
        for state in stock_states:
            product = Product.get(state['product_id'])
            yield {
                'category': stock_category(state['total_stock'], state['avg_consumption']),
                'product_id': product._id,
                'consumption': state['avg_consumption'],
                'months_remaining': months_of_stock_remaining(state['total_stock'], state['avg_consumption']),
                'location_id': None,
                'product_name': product.name,
                'current_stock': state['total_stock'],
                'location_lineage': None,
                'resupply_quantity_needed': None
            }

    def raw_product_states(self, stock_states, slugs):
        def _slug_attrib(slug, attrib, product, output):
            if not slugs or slug in slugs:
                if callable(attrib):
                    output[slug] = attrib(product)
                else:
                    output[slug] = getattr(product, attrib, '')

        for state in stock_states:
            out = {}
            for slug, attrib in self._slug_attrib_map.items():
                _slug_attrib(slug, attrib, state, out)

            yield out

class StockStatusBySupplyPointDataSource(StockStatusDataSource):
    def get_data(self):
        data = list(super(StockStatusBySupplyPointDataSource, self).get_data())

        products = dict((r['product_id'], r['product_name']) for r in data)
        product_ids = sorted(products.keys(), key=lambda e: products[e])

        by_supply_point = map_reduce(lambda e: [(e['location_id'],)], data=data, include_docs=True)
        locs = dict((loc._id, loc) for loc in Location.view(
                '_all_docs',
                keys=by_supply_point.keys(),
                include_docs=True))

        for loc_id, subcases in by_supply_point.iteritems():
            loc = locs[loc_id]
            by_product = dict((c['product_id'], c) for c in subcases)

            rec = {
                'name': loc.name,
                'type': loc.location_type,
                'geo': loc._geopoint,
            }
            for prod in product_ids:
                rec.update(dict(('%s-%s' % (prod, key), by_product.get(prod, {}).get(key)) for key in
                                ('current_stock', 'consumption', 'months_remaining', 'category')))
            yield rec

class ReportingStatusDataSource(ReportDataSource, CommtrackDataSourceMixin):
    """
    Config:
        domain: The domain to report on.
        location_id: ID of location to get data for. Omit for all locations.
    """

    def get_data(self):
        sp_ids = get_relevant_supply_point_ids(
            self.domain,
            self.active_location
        )

        products = Product.by_domain(self.domain)
        if self.program_id:
            products = filter(
                lambda product: product.program_id == self.program_id, products
            )

        for sp_id in sp_ids:
            loc = SupplyPointCase.get(sp_id).location
            transactions = StockTransaction.objects.filter(
                case_id=sp_id,
                section_id=STOCK_SECTION_TYPE,
            )

            if transactions:
                last_transaction = sorted(
                    transactions,
                    key=lambda trans: trans.report.date
                )[-1]
            else:
                last_transaction = None

            yield {
                'loc_id': loc._id,
                'loc_path': loc.path,
                'name': loc.name,
                'type': loc.location_type,
                'reporting_status': reporting_status(
                    last_transaction,
                    self.start_date,
                    self.end_date
                ),
                'geo': loc._geopoint,
            }
<commit_msg>Fix map for single location selections<commit_after>from corehq.apps.commtrack.util import num_periods_late
from dimagi.utils.decorators.memoized import memoized
from corehq.apps.locations.models import Location
from corehq.apps.commtrack.models import Product, SupplyPointCase
from dimagi.utils.couch.loosechange import map_reduce
from corehq.apps.reports.api import ReportDataSource
from datetime import datetime, timedelta
from casexml.apps.stock.models import StockState, StockTransaction
from django.db.models import Sum, Avg
from corehq.apps.reports.commtrack.util import get_relevant_supply_point_ids, product_ids_filtered_by_program
from corehq.apps.reports.commtrack.const import STOCK_SECTION_TYPE
from casexml.apps.stock.utils import months_of_stock_remaining, stock_category

# TODO make settings
REPORTING_PERIOD = 'weekly'
REPORTING_PERIOD_ARGS = (1,)

def is_timely(case, limit=0):
    return num_periods_late(case, REPORTING_PERIOD, *REPORTING_PERIOD_ARGS) <= limit

def reporting_status(transaction, start_date, end_date):
    if transaction:
        last_reported = transaction.report.date.date()
    else:
        last_reported = None

    if last_reported and last_reported < start_date:
        return 'ontime'
    elif last_reported and start_date <= last_reported <= end_date:
        return 'late'
    else:
        return 'nonreporting'

class CommtrackDataSourceMixin(object):

    @property
    def domain(self):
        return self.config.get('domain')

    @property
    @memoized
    def active_location(self):
        loc_id = self.config.get('location_id')
        if loc_id:
            return Location.get(loc_id)

    @property
    @memoized
    def active_product(self):
        prod_id = self.config.get('product_id')
        if prod_id:
            return Product.get(prod_id)

    @property
    @memoized
    def program_id(self):
        prog_id = self.config.get('program_id')
        if prog_id != '':
            return prog_id

    @property
    def start_date(self):
        date = self.config.get('start_date')
        if date:
            return datetime.strptime(date, '%Y-%m-%d').date()
        else:
            return (datetime.now() - timedelta(30)).date()

    @property
    def end_date(self):
        date = self.config.get('end_date')
        if date:
            return datetime.strptime(date, '%Y-%m-%d').date()
        else:
            return datetime.now().date()

class StockStatusDataSource(ReportDataSource, CommtrackDataSourceMixin):
    """
    Config:
        domain: The domain to report on.
        location_id: ID of location to get data for. Omit for all locations.
        product_id: ID of product to get data for. Omit for all products.
        aggregate: True to aggregate the indicators by product for the current location.

    Data Slugs:
        product_name: Name of the product
        product_id: ID of the product
        location_id: The ID of the current location.
        location_lineage: The lineage of the current location.
        current_stock: The current stock level
        consumption: The current monthly consumption rate
        months_remaining: The number of months remaining until stock out
        category: The status category. See casexml.apps.stock.models.StockState.stock_category
        resupply_quantity_needed: Max amount - current amount

    """
    slug = 'agg_stock_status'

    SLUG_PRODUCT_NAME = 'product_name'
    SLUG_PRODUCT_ID = 'product_id'
    SLUG_MONTHS_REMAINING = 'months_remaining'
    SLUG_CONSUMPTION = 'consumption'
    SLUG_CURRENT_STOCK = 'current_stock'
    SLUG_LOCATION_ID = 'location_id'
    SLUG_LOCATION_LINEAGE = 'location_lineage'
    SLUG_STOCKOUT_SINCE = 'stockout_since'
    SLUG_STOCKOUT_DURATION = 'stockout_duration'
    SLUG_LAST_REPORTED = 'last_reported'
    SLUG_CATEGORY = 'category'
    SLUG_RESUPPLY_QUANTITY_NEEDED = 'resupply_quantity_needed'

    _slug_attrib_map = {
        SLUG_PRODUCT_NAME: lambda s: Product.get(s.product_id).name,
        SLUG_PRODUCT_ID: lambda s: s.product_id,
        SLUG_LOCATION_ID: lambda s: SupplyPointCase.get(s.case_id).location_[-1],
        # SLUG_LOCATION_LINEAGE: lambda p: list(reversed(p.location_[:-1])),
        SLUG_CURRENT_STOCK: 'stock_on_hand',
        SLUG_CONSUMPTION: lambda s: s.daily_consumption * 30 if s.daily_consumption else None,
        SLUG_MONTHS_REMAINING: 'months_remaining',
        SLUG_CATEGORY: 'stock_category',
        # SLUG_STOCKOUT_SINCE: 'stocked_out_since',
        # SLUG_STOCKOUT_DURATION: 'stockout_duration_in_months',
        SLUG_LAST_REPORTED: 'last_modified_date',
        SLUG_RESUPPLY_QUANTITY_NEEDED: 'resupply_quantity_needed',
    }

    def slugs(self):
        return self._slug_attrib_map.keys()

    def filter_by_program(self, stock_states):
        return stock_states.filter(
            product_id__in=product_ids_filtered_by_program(
                self.domain,
                self.program_id
            )
        )

    def get_data(self, slugs=None):
        sp_ids = get_relevant_supply_point_ids(self.domain, self.active_location)

        if len(sp_ids) == 1:
            stock_states = StockState.objects.filter(
                case_id=sp_ids[0],
                section_id=STOCK_SECTION_TYPE,
                last_modified_date__lte=self.end_date,
                last_modified_date__gte=self.start_date,
            )

            if self.program_id:
                stock_states = self.filter_by_program(stock_states)

            return self.leaf_node_data(stock_states)
        else:
            stock_states = StockState.objects.filter(
                case_id__in=sp_ids,
                section_id=STOCK_SECTION_TYPE,
                last_modified_date__lte=self.end_date,
                last_modified_date__gte=self.start_date,
            )

            if self.program_id:
                stock_states = self.filter_by_program(stock_states)

            if self.config.get('aggregate'):
                aggregates = stock_states.values('product_id').annotate(
                    avg_consumption=Avg('daily_consumption'),
                    total_stock=Sum('stock_on_hand')
                )
                return self.aggregated_data(aggregates)
            else:
                return self.raw_product_states(stock_states, slugs)

    def leaf_node_data(self, stock_states):
        for state in stock_states:
            product = Product.get(state.product_id)
            yield {
                'category': state.stock_category,
                'product_id': product._id,
                'consumption': state.daily_consumption * 30 if state.daily_consumption else None,
                'months_remaining': state.months_remaining,
                'location_id': SupplyPointCase.get(state.case_id).location_id,
                'product_name': product.name,
                'current_stock': state.stock_on_hand,
                'location_lineage': None,
                'resupply_quantity_needed': state.resupply_quantity_needed
            }

    def aggregated_data(self, stock_states):
        for state in stock_states:
            product = Product.get(state['product_id'])
            yield {
                'category': stock_category(state['total_stock'], state['avg_consumption']),
                'product_id': product._id,
                'consumption': state['avg_consumption'],
                'months_remaining': months_of_stock_remaining(state['total_stock'], state['avg_consumption']),
                'location_id': None,
                'product_name': product.name,
                'current_stock': state['total_stock'],
                'location_lineage': None,
                'resupply_quantity_needed': None
            }

    def raw_product_states(self, stock_states, slugs):
        def _slug_attrib(slug, attrib, product, output):
            if not slugs or slug in slugs:
                if callable(attrib):
                    output[slug] = attrib(product)
                else:
                    output[slug] = getattr(product, attrib, '')

        for state in stock_states:
            out = {}
            for slug, attrib in self._slug_attrib_map.items():
                _slug_attrib(slug, attrib, state, out)

            yield out

class StockStatusBySupplyPointDataSource(StockStatusDataSource):
    def get_data(self):
        data = list(super(StockStatusBySupplyPointDataSource, self).get_data())

        products = dict((r['product_id'], r['product_name']) for r in data)
        product_ids = sorted(products.keys(), key=lambda e: products[e])

        by_supply_point = map_reduce(lambda e: [(e['location_id'],)], data=data, include_docs=True)
        locs = dict((loc._id, loc) for loc in Location.view(
                '_all_docs',
                keys=by_supply_point.keys(),
                include_docs=True))

        for loc_id, subcases in by_supply_point.iteritems():
            loc = locs[loc_id]
            by_product = dict((c['product_id'], c) for c in subcases)

            rec = {
                'name': loc.name,
                'type': loc.location_type,
                'geo': loc._geopoint,
            }
            for prod in product_ids:
                rec.update(dict(('%s-%s' % (prod, key), by_product.get(prod, {}).get(key)) for key in
                                ('current_stock', 'consumption', 'months_remaining', 'category')))
            yield rec

class ReportingStatusDataSource(ReportDataSource, CommtrackDataSourceMixin):
    """
    Config:
        domain: The domain to report on.
        location_id: ID of location to get data for. Omit for all locations.
    """

    def get_data(self):
        sp_ids = get_relevant_supply_point_ids(
            self.domain,
            self.active_location
        )

        products = Product.by_domain(self.domain)
        if self.program_id:
            products = filter(
                lambda product: product.program_id == self.program_id, products
            )

        for sp_id in sp_ids:
            loc = SupplyPointCase.get(sp_id).location
            transactions = StockTransaction.objects.filter(
                case_id=sp_id,
                section_id=STOCK_SECTION_TYPE,
            )

            if transactions:
                last_transaction = sorted(
                    transactions,
                    key=lambda trans: trans.report.date
                )[-1]
            else:
                last_transaction = None

            yield {
                'loc_id': loc._id,
                'loc_path': loc.path,
                'name': loc.name,
                'type': loc.location_type,
                'reporting_status': reporting_status(
                    last_transaction,
                    self.start_date,
                    self.end_date
                ),
                'geo': loc._geopoint,
            }
<|endoftext|>