<commit_before>from copy import deepcopy
from cms.extensions.toolbar import ExtensionToolbar
from cms.toolbar_pool import toolbar_pool
from cms.utils.urlutils import admin_reverse
from django.contrib.auth.models import Permission
from django.contrib.sites.models import Site

from cms.api import create_page
from cms.constants import PUBLISHER_STATE_DIRTY
from cms.extensions import extension_pool
from cms.extensions import TitleExtension
from cms.extensions import PageExtension
from cms.models import Page
from cms.test_utils.project.extensionapp.models import (MyPageExtension,
                                                        MyTitleExtension)
from cms.test_utils.testcases import SettingsOverrideTestCase as TestCase
from cms.tests import AdminTestsBase
from cms.utils.compat.dj import get_user_model, force_unicode
from django.template.response import TemplateResponse

class ExtensionsTestCase(TestCase):
    def test_register_extension(self):
        initial_extension_count = len(extension_pool.page_extensions)
        # --- None extension registering -----------------------------
        from cms.exceptions import SubClassNeededError
        none_extension = self.get_none_extension_class()
        self.assertRaises(SubClassNeededError, extension_pool.register, none_extension)
        self.assertEqual(len(extension_pool.page_extensions), initial_extension_count)
        self.assertEqual(len(extension_pool.title_extensions), initial_extension_count)

        # --- Page registering ---------------------------------------
        page_extension = self.get_page_extension_class()

        # register first time
        extension_pool.register(page_extension)
        self.assertEqual(len(extension_pool.page_extensions), initial_extension_count+1)

        # register second time
        extension_pool.register(page_extension)
        self.assertEqual(len(extension_pool.page_extensions), initial_extension_count+1)

        self.assertIs(extension_pool.signaling_activated, True)

        # --- Title registering --------------------------------------
        title_extension = self.get_title_extension_class()

        # register first time
        extension_pool.register(title_extension)
        self.assertEqual(len(extension_pool.title_extensions), initial_extension_count+1)

        # register second time
        extension_pool.register(title_extension)
        self.assertEqual(len(extension_pool.title_extensions), initial_extension_count+1)

        self.assertIs(extension_pool.signaling_activated, True)

        # --- Unregister ---------------------------------------------
        extension_pool.unregister(page_extension)
        self.assertEqual(len(extension_pool.page_extensions), initial_extension_count)

        extension_pool.unregister(title_extension)
        self.assertEqual(len(extension_pool.title_extensions), initial_extension_count)

        # Unregister an object that is not registered yet
        extension_pool.unregister(page_extension)
        extension_pool.unregister(title_extension)

        try:
            from django.apps import apps
            del apps.all_models['cms']['testpageextension']
            del apps.all_models['cms']['testtitleextension']
        except ImportError:
            pass

    def get_page_extension_class(self):
        from django.db import models

        class TestPageExtension(PageExtension):
            content = models.CharField('Content', max_length=50)

        return TestPageExtension

    def get_title_extension_class(self):
        from django.db import models

        class TestTitleExtension(TitleExtension):
            content = models.CharField('Content', max_length=50)

        return TestTitleExtension

    def get_none_extension_class(self):
        class TestNoneExtension(object):
            pass

        return TestNoneExtension

    def test_publish_page_extension(self):
        page = create_page('Test Page Extension', "nav_playground.html", "en")
        page_extension = MyPageExtension(extended_object=page, extra='page extension 1')
        page_extension.save()
        page.mypageextension = page_extension

        # publish first time
        page.publish('en')
        self.assertEqual(page_extension.extra, page.publisher_public.mypageextension.extra)
        self.assertEqual(page.get_publisher_state('en'), 0)
        # change and publish again
        page = Page.objects.get(pk=page.pk)
        page_extension = page.mypageextension
        page_extension.extra = 'page extension 1 - changed'
        page_extension.save()
        self.assertEqual(page.get_publisher_state('en', True), PUBLISHER_STATE_DIRTY)
        page.publish('en')
        self.assertEqual(page.get_publisher_state('en', True), 0)
        # delete
        page_extension.delete()
        self.assertFalse(MyPageExtension.objects.filter(pk=page_extension.pk).exists())
        self.assertEqual(page.get_publisher_state('en', True), PUBLISHER_STATE_DIRTY)

    def test_publish_title_extension(self):
        page = create_page('Test Title Extension', "nav_playground.html", "en")
        title = page.get_title_obj()
        title_extension = MyTitleExtension(extended_object=title, extra_title='title extension 1')
        title_extension.save()
        page.mytitleextension = title_extension

        # publish first time
        page.publish('en')
        # import ipdb; ipdb.set_trace()
        self.assertEqual(page.get_publisher_state('en'), 0)
        self.assertEqual(title_extension.extra_title, page.publisher_public.get_title_obj().mytitleextension.extra_title)

        # change and publish again
        page = Page.objects.get(pk=page.pk)
        title = page.get_title_obj()
        title_extension = title.mytitleextension
        title_extension.extra_title = 'title extension 1 - changed'
        title_extension.save()
        self.assertEqual(page.get_publisher_state('en', True), PUBLISHER_STATE_DIRTY)
        page.publish('en')
        self.assertEqual(page.get_publisher_state('en', True), 0)

        # delete
        title_extension.delete()
        self.assertFalse(MyTitleExtension.objects.filter(pk=title_extension.pk).exists())

class ExtensionAdminTestCase(AdminTestsBase):
    def setUp(self):
        User = get_user_model()

        self.admin, self.normal_guy = self._get_guys()

        if get_user_model().USERNAME_FIELD == 'email':
            self.no_page_permission_user = User.objects.create_user('no_page_permission', 'test2@test.com', 'test2@test.com')
        else:
            self.no_page_permission_user = User.objects.create_user('no_page_permission', 'test2@test.com', 'no_page_permission')

        self.no_page_permission_user.is_staff = True
        self.no_page_permission_user.is_active = True
        self.no_page_permission_user.save()
        [self.no_page_permission_user.user_permissions.add(p) for p in Permission.objects.filter(
            codename__in=[
                'change_mypageextension', 'change_mytitleextension',
                'add_mypageextension', 'add_mytitleextension',
                'delete_mypageextension', 'delete_mytitleextension',
            ]
        )]
        self.site = Site.objects.get(pk=1)
        self.page = create_page(
            'My Extension Page', 'nav_playground.html', 'en',
            site=self.site, created_by=self.admin)
        self.page_title = self.page.get_title_obj()
        self.page_extension = MyPageExtension.objects.create(
            extended_object=self.page,
            extra="page extension text")
        self.title_extension = MyTitleExtension.objects.create(
            extended_object=self.page.get_title_obj(),
            extra_title="title extension text")

        self.page_without_extension = create_page(
            'A Page', 'nav_playground.html', 'en',
            site=self.site, created_by=self.admin)
        self.page_title_without_extension = self.page_without_extension.get_title_obj()

    def test_admin_page_extension(self):
        with self.login_user_context(self.admin):
            # add a new extension
            response = self.client.get(
                admin_reverse('extensionapp_mypageextension_add') + '?extended_object=%s' % self.page_without_extension.pk
            )
            self.assertEqual(response.status_code, 200)
            # make sure there is no extension yet
            self.assertFalse(MyPageExtension.objects.filter(extended_object=self.page_without_extension).exists())
            post_data = {
                'extra': 'my extra'
            }
            response = self.client.post(
                admin_reverse('extensionapp_mypageextension_add') + '?extended_object=%s' % self.page_without_extension.pk,
                post_data, follow=True
            )
            created_page_extension = MyPageExtension.objects.get(extended_object=self.page_without_extension)

            # can delete extension
            response = self.client.post(
                admin_reverse('extensionapp_mypageextension_delete', args=(created_page_extension.pk,)),
                {'post': 'yes'}, follow=True
            )
            self.assertFalse(MyPageExtension.objects.filter(extended_object=self.page_without_extension).exists())

            # accessing the add view on a page that already has an extension should redirect
            response = self.client.get(
                admin_reverse('extensionapp_mypageextension_add') + '?extended_object=%s' % self.page.pk
            )
            self.assertRedirects(response, admin_reverse('extensionapp_mypageextension_change', args=(self.page_extension.pk,)))

            # saving an extension should work without the GET parameter
            post_data = {
                'extra': 'my extra text'
            }
            self.client.post(
                admin_reverse('extensionapp_mypageextension_change', args=(self.page_extension.pk,)),
                post_data, follow=True
            )
            self.assertTrue(MyPageExtension.objects.filter(extra='my extra text', pk=self.page_extension.pk).exists())

        with self.login_user_context(self.no_page_permission_user):
            # can't save if user does not have permissions to change the page
            post_data = {
                'extra': 'try to change extra text'
            }
            response = self.client.post(
                admin_reverse('extensionapp_mypageextension_change', args=(self.page_extension.pk,)),
                post_data, follow=True
            )
            self.assertEqual(response.status_code, 403)

            # can't delete without page permission
            response = self.client.post(
                admin_reverse('extensionapp_mypageextension_delete', args=(self.page_extension.pk,)),
                {'post': 'yes'}, follow=True
            )
            self.assertEqual(response.status_code, 403)
            self.assertTrue(MyPageExtension.objects.filter(extended_object=self.page).exists())

    def test_toolbar_page_extension(self):
        old_toolbars = deepcopy(toolbar_pool.toolbars)
        class SampleExtension(ExtensionToolbar):
            model = MyPageExtension  # The PageExtension / TitleExtension you are working with

            def populate(self):
                current_page_menu = self._setup_extension_toolbar()
                if current_page_menu:
                    position = 0
                    page_extension, url = self.get_page_extension_admin()
                    if url:
                        current_page_menu.add_modal_item('TestItem', url=url,
                                                         disabled=not self.toolbar.edit_mode,
                                                         position=position)
        toolbar_pool.register(SampleExtension)
        with self.login_user_context(self.admin):
            response = self.client.get('/en/?edit')
            self.assertIn("TestItem", response.rendered_content)
        toolbar_pool.toolbars = old_toolbars

    def test_toolbar_title_extension(self):
        old_toolbars = deepcopy(toolbar_pool.toolbars)

        class SampleExtension(ExtensionToolbar):
            model = MyTitleExtension

            def populate(self):
                current_page_menu = self._setup_extension_toolbar()
                if current_page_menu:
                    position = 0
                    urls = self.get_title_extension_admin()
                    for title_extension, url in urls:
                        current_page_menu.add_modal_item('TestItem', url=url,
                                                         disabled=not self.toolbar.edit_mode,
                                                         position=position)
        toolbar_pool.register(SampleExtension)
        with self.login_user_context(self.admin):
            response = self.client.get('/en/?edit')
            self.assertIn("TestItem", response.rendered_content)
        toolbar_pool.toolbars = old_toolbars

    def test_admin_title_extension(self):
        with self.login_user_context(self.admin):
            # add a new extension
            response = self.client.get(
                admin_reverse('extensionapp_mytitleextension_add') + '?extended_object=%s' % self.page_title_without_extension.pk
            )
            self.assertEqual(response.status_code, 200)
            # make sure there is no extension yet
            self.assertFalse(MyTitleExtension.objects.filter(extended_object=self.page_title_without_extension).exists())
            post_data = {
                'extra_title': 'my extra title'
            }
            self.client.post(
                admin_reverse('extensionapp_mytitleextension_add') + '?extended_object=%s' % self.page_title_without_extension.pk,
                post_data, follow=True
            )
            created_title_extension = MyTitleExtension.objects.get(extended_object=self.page_title_without_extension)

            # can delete extension
            self.client.post(
                admin_reverse('extensionapp_mytitleextension_delete', args=(created_title_extension.pk,)),
                {'post': 'yes'}, follow=True
            )
            self.assertFalse(MyTitleExtension.objects.filter(extended_object=self.page_title_without_extension).exists())

            # accessing the add view on a page that already has an extension should redirect
            response = self.client.get(
                admin_reverse('extensionapp_mytitleextension_add') + '?extended_object=%s' % self.page_title.pk
            )
            self.assertRedirects(response, admin_reverse('extensionapp_mytitleextension_change', args=(self.title_extension.pk,)))

            # saving an extension should work without the GET parameter
            post_data = {
                'extra_title': 'my extra text'
            }
            self.client.post(
                admin_reverse('extensionapp_mytitleextension_change', args=(self.title_extension.pk,)),
                post_data, follow=True
            )
            self.assertTrue(MyTitleExtension.objects.filter(extra_title='my extra text', pk=self.title_extension.pk).exists())

        with self.login_user_context(self.no_page_permission_user):
            # can't save if user does not have permissions to change the page
            post_data = {
                'extra_title': 'try to change extra text'
            }
            response = self.client.post(
                admin_reverse('extensionapp_mytitleextension_change', args=(self.title_extension.pk,)),
                post_data, follow=True
            )
            self.assertEqual(response.status_code, 403)

            # can't delete without page permission
            response = self.client.post(
                admin_reverse('extensionapp_mytitleextension_delete', args=(self.title_extension.pk,)),
                {'post': 'yes'}, follow=True
            )
            self.assertEqual(response.status_code, 403)
            self.assertTrue(MyTitleExtension.objects.filter(extended_object=self.page_title).exists())
<commit_msg>fix extension test imports<commit_after>from copy import deepcopy
from cms.extensions.toolbar import ExtensionToolbar
from cms.toolbar_pool import toolbar_pool
from cms.utils.urlutils import admin_reverse
from django.contrib.auth.models import Permission
from django.contrib.sites.models import Site

from cms.api import create_page
from cms.constants import PUBLISHER_STATE_DIRTY
from cms.extensions import extension_pool
from cms.extensions import TitleExtension
from cms.extensions import PageExtension
from cms.models import Page
from cms.test_utils.project.extensionapp.models import (MyPageExtension,
                                                        MyTitleExtension)
from cms.test_utils.testcases import SettingsOverrideTestCase as TestCase
from cms.tests import AdminTestsBase
from cms.utils.compat.dj import get_user_model

class ExtensionsTestCase(TestCase):
    def test_register_extension(self):
        initial_extension_count = len(extension_pool.page_extensions)
        # --- None extension registering -----------------------------
        from cms.exceptions import SubClassNeededError
        none_extension = self.get_none_extension_class()
        self.assertRaises(SubClassNeededError, extension_pool.register, none_extension)
        self.assertEqual(len(extension_pool.page_extensions), initial_extension_count)
        self.assertEqual(len(extension_pool.title_extensions), initial_extension_count)

        # --- Page registering ---------------------------------------
        page_extension = self.get_page_extension_class()

        # register first time
        extension_pool.register(page_extension)
        self.assertEqual(len(extension_pool.page_extensions), initial_extension_count+1)

        # register second time
        extension_pool.register(page_extension)
        self.assertEqual(len(extension_pool.page_extensions), initial_extension_count+1)

        self.assertIs(extension_pool.signaling_activated, True)

        # --- Title registering --------------------------------------
        title_extension = self.get_title_extension_class()

        # register first time
        extension_pool.register(title_extension)
        self.assertEqual(len(extension_pool.title_extensions), initial_extension_count+1)

        # register second time
        extension_pool.register(title_extension)
        self.assertEqual(len(extension_pool.title_extensions), initial_extension_count+1)

        self.assertIs(extension_pool.signaling_activated, True)

        # --- Unregister ---------------------------------------------
        extension_pool.unregister(page_extension)
        self.assertEqual(len(extension_pool.page_extensions), initial_extension_count)

        extension_pool.unregister(title_extension)
        self.assertEqual(len(extension_pool.title_extensions), initial_extension_count)

        # Unregister an object that is not registered yet
        extension_pool.unregister(page_extension)
        extension_pool.unregister(title_extension)

        try:
            from django.apps import apps
            del apps.all_models['cms']['testpageextension']
            del apps.all_models['cms']['testtitleextension']
        except ImportError:
            pass

    def get_page_extension_class(self):
        from django.db import models

        class TestPageExtension(PageExtension):
            content = models.CharField('Content', max_length=50)

        return TestPageExtension

    def get_title_extension_class(self):
        from django.db import models

        class TestTitleExtension(TitleExtension):
            content = models.CharField('Content', max_length=50)

        return TestTitleExtension

    def get_none_extension_class(self):
        class TestNoneExtension(object):
            pass

        return TestNoneExtension

    def test_publish_page_extension(self):
        page = create_page('Test Page Extension', "nav_playground.html", "en")
        page_extension = MyPageExtension(extended_object=page, extra='page extension 1')
        page_extension.save()
        page.mypageextension = page_extension

        # publish first time
        page.publish('en')
        self.assertEqual(page_extension.extra, page.publisher_public.mypageextension.extra)
        self.assertEqual(page.get_publisher_state('en'), 0)
        # change and publish again
        page = Page.objects.get(pk=page.pk)
        page_extension = page.mypageextension
        page_extension.extra = 'page extension 1 - changed'
        page_extension.save()
        self.assertEqual(page.get_publisher_state('en', True), PUBLISHER_STATE_DIRTY)
        page.publish('en')
        self.assertEqual(page.get_publisher_state('en', True), 0)
        # delete
        page_extension.delete()
        self.assertFalse(MyPageExtension.objects.filter(pk=page_extension.pk).exists())
        self.assertEqual(page.get_publisher_state('en', True), PUBLISHER_STATE_DIRTY)

    def test_publish_title_extension(self):
        page = create_page('Test Title Extension', "nav_playground.html", "en")
        title = page.get_title_obj()
        title_extension = MyTitleExtension(extended_object=title, extra_title='title extension 1')
        title_extension.save()
        page.mytitleextension = title_extension

        # publish first time
        page.publish('en')
        # import ipdb; ipdb.set_trace()
        self.assertEqual(page.get_publisher_state('en'), 0)
        self.assertEqual(title_extension.extra_title, page.publisher_public.get_title_obj().mytitleextension.extra_title)

        # change and publish again
        page = Page.objects.get(pk=page.pk)
        title = page.get_title_obj()
        title_extension = title.mytitleextension
        title_extension.extra_title = 'title extension 1 - changed'
        title_extension.save()
        self.assertEqual(page.get_publisher_state('en', True), PUBLISHER_STATE_DIRTY)
        page.publish('en')
        self.assertEqual(page.get_publisher_state('en', True), 0)

        # delete
        title_extension.delete()
        self.assertFalse(MyTitleExtension.objects.filter(pk=title_extension.pk).exists())

class ExtensionAdminTestCase(AdminTestsBase):
    def setUp(self):
        User = get_user_model()

        self.admin, self.normal_guy = self._get_guys()

        if get_user_model().USERNAME_FIELD == 'email':
            self.no_page_permission_user = User.objects.create_user('no_page_permission', 'test2@test.com', 'test2@test.com')
        else:
            self.no_page_permission_user = User.objects.create_user('no_page_permission', 'test2@test.com', 'no_page_permission')

        self.no_page_permission_user.is_staff = True
        self.no_page_permission_user.is_active = True
        self.no_page_permission_user.save()
        [self.no_page_permission_user.user_permissions.add(p) for p in Permission.objects.filter(
            codename__in=[
                'change_mypageextension', 'change_mytitleextension',
                'add_mypageextension', 'add_mytitleextension',
                'delete_mypageextension', 'delete_mytitleextension',
            ]
        )]
        self.site = Site.objects.get(pk=1)
        self.page = create_page(
            'My Extension Page', 'nav_playground.html', 'en',
            site=self.site, created_by=self.admin)
        self.page_title = self.page.get_title_obj()
        self.page_extension = MyPageExtension.objects.create(
            extended_object=self.page,
            extra="page extension text")
        self.title_extension = MyTitleExtension.objects.create(
            extended_object=self.page.get_title_obj(),
            extra_title="title extension text")

        self.page_without_extension = create_page(
            'A Page', 'nav_playground.html', 'en',
            site=self.site, created_by=self.admin)
        self.page_title_without_extension = self.page_without_extension.get_title_obj()

    def test_admin_page_extension(self):
        with self.login_user_context(self.admin):
            # add a new extension
            response = self.client.get(
                admin_reverse('extensionapp_mypageextension_add') + '?extended_object=%s' % self.page_without_extension.pk
            )
            self.assertEqual(response.status_code, 200)
            # make sure there is no extension yet
            self.assertFalse(MyPageExtension.objects.filter(extended_object=self.page_without_extension).exists())
            post_data = {
                'extra': 'my extra'
            }
            response = self.client.post(
                admin_reverse('extensionapp_mypageextension_add') + '?extended_object=%s' % self.page_without_extension.pk,
                post_data, follow=True
            )
            created_page_extension = MyPageExtension.objects.get(extended_object=self.page_without_extension)

            # can delete extension
            response = self.client.post(
                admin_reverse('extensionapp_mypageextension_delete', args=(created_page_extension.pk,)),
                {'post': 'yes'}, follow=True
            )
            self.assertFalse(MyPageExtension.objects.filter(extended_object=self.page_without_extension).exists())

            # accessing the add view on a page that already has an extension should redirect
            response = self.client.get(
                admin_reverse('extensionapp_mypageextension_add') + '?extended_object=%s' % self.page.pk
            )
            self.assertRedirects(response, admin_reverse('extensionapp_mypageextension_change', args=(self.page_extension.pk,)))

            # saving an extension should work without the GET parameter
            post_data = {
                'extra': 'my extra text'
            }
            self.client.post(
                admin_reverse('extensionapp_mypageextension_change', args=(self.page_extension.pk,)),
                post_data, follow=True
            )
            self.assertTrue(MyPageExtension.objects.filter(extra='my extra text', pk=self.page_extension.pk).exists())

        with self.login_user_context(self.no_page_permission_user):
            # can't save if user does not have permissions to change the page
            post_data = {
                'extra': 'try to change extra text'
            }
            response = self.client.post(
                admin_reverse('extensionapp_mypageextension_change', args=(self.page_extension.pk,)),
                post_data, follow=True
            )
            self.assertEqual(response.status_code, 403)

            # can't delete without page permission
            response = self.client.post(
                admin_reverse('extensionapp_mypageextension_delete', args=(self.page_extension.pk,)),
                {'post': 'yes'}, follow=True
            )
            self.assertEqual(response.status_code, 403)
            self.assertTrue(MyPageExtension.objects.filter(extended_object=self.page).exists())

    def test_toolbar_page_extension(self):
        old_toolbars = deepcopy(toolbar_pool.toolbars)
        class SampleExtension(ExtensionToolbar):
            model = MyPageExtension  # The PageExtension / TitleExtension you are working with

            def populate(self):
                current_page_menu = self._setup_extension_toolbar()
                if current_page_menu:
                    position = 0
                    page_extension, url = self.get_page_extension_admin()
                    if url:
                        current_page_menu.add_modal_item('TestItem', url=url,
                                                         disabled=not self.toolbar.edit_mode,
                                                         position=position)
        toolbar_pool.register(SampleExtension)
        with self.login_user_context(self.admin):
            response = self.client.get('/en/?edit')
            self.assertIn("TestItem", response.rendered_content)
        toolbar_pool.toolbars = old_toolbars

    def test_toolbar_title_extension(self):
        old_toolbars = deepcopy(toolbar_pool.toolbars)

        class SampleExtension(ExtensionToolbar):
            model = MyTitleExtension

            def populate(self):
                current_page_menu = self._setup_extension_toolbar()
                if current_page_menu:
                    position = 0
                    urls = self.get_title_extension_admin()
                    for title_extension, url in urls:
                        current_page_menu.add_modal_item('TestItem', url=url,
                                                         disabled=not self.toolbar.edit_mode,
                                                         position=position)
        toolbar_pool.register(SampleExtension)
        with self.login_user_context(self.admin):
            response = self.client.get('/en/?edit')
            self.assertIn("TestItem", response.rendered_content)
        toolbar_pool.toolbars = old_toolbars

    def test_admin_title_extension(self):
        with self.login_user_context(self.admin):
            # add a new extension
            response = self.client.get(
                admin_reverse('extensionapp_mytitleextension_add') + '?extended_object=%s' % self.page_title_without_extension.pk
            )
            self.assertEqual(response.status_code, 200)
            # make sure there is no extension yet
            self.assertFalse(MyTitleExtension.objects.filter(extended_object=self.page_title_without_extension).exists())
            post_data = {
                'extra_title': 'my extra title'
            }
            self.client.post(
                admin_reverse('extensionapp_mytitleextension_add') + '?extended_object=%s' % self.page_title_without_extension.pk,
                post_data, follow=True
            )
            created_title_extension = MyTitleExtension.objects.get(extended_object=self.page_title_without_extension)

            # can delete extension
            self.client.post(
                admin_reverse('extensionapp_mytitleextension_delete', args=(created_title_extension.pk,)),
                {'post': 'yes'}, follow=True
            )
            self.assertFalse(MyTitleExtension.objects.filter(extended_object=self.page_title_without_extension).exists())

            # accessing the add view on a page that already has an extension should redirect
            response = self.client.get(
                admin_reverse('extensionapp_mytitleextension_add') + '?extended_object=%s' % self.page_title.pk
            )
            self.assertRedirects(response, admin_reverse('extensionapp_mytitleextension_change', args=(self.title_extension.pk,)))

            # saving an extension should work without the GET parameter
            post_data = {
                'extra_title': 'my extra text'
            }
            self.client.post(
                admin_reverse('extensionapp_mytitleextension_change', args=(self.title_extension.pk,)),
                post_data, follow=True
            )
            self.assertTrue(MyTitleExtension.objects.filter(extra_title='my extra text', pk=self.title_extension.pk).exists())

        with self.login_user_context(self.no_page_permission_user):
            # can't save if user does not have permissions to change the page
            post_data = {
                'extra_title': 'try to change extra text'
            }
            response = self.client.post(
                admin_reverse('extensionapp_mytitleextension_change', args=(self.title_extension.pk,)),
                post_data, follow=True
            )
            self.assertEqual(response.status_code, 403)

            # can't delete without page permission
            response = self.client.post(
                admin_reverse('extensionapp_mytitleextension_delete', args=(self.title_extension.pk,)),
                {'post': 'yes'}, follow=True
            )
            self.assertEqual(response.status_code, 403)
            self.assertTrue(MyTitleExtension.objects.filter(extended_object=self.page_title).exists())
<|endoftext|>


<commit_before>import argparse
import os

import pytest

import parsl
from parsl.app.app import bash_app
import parsl.app.errors as pe

from parsl.tests.configs.local_threads import config

local_config = config

@bash_app
def command_not_found(stderr='std.err', stdout='std.out'):
    cmd_line = 'catdogcat'
    return cmd_line

@bash_app
def bash_misuse(stderr='std.err', stdout='std.out'):
    cmd_line = 'exit 15'
    return cmd_line

@bash_app
def div_0(stderr='std.err', stdout='std.out'):
    cmd_line = '$((5/0))'
    return cmd_line

@bash_app
def invalid_exit(stderr='std.err', stdout='std.out'):
    cmd_line = 'exit 3.141'
    return cmd_line

@bash_app
def not_executable(stderr='std.err', stdout='std.out'):
    cmd_line = '/dev/null'
    return cmd_line

@bash_app
def bad_format(stderr='std.err', stdout='std.out'):
    cmd_line = 'echo {0}'
    return cmd_line

test_matrix = {
    div_0: {
        'exit_code': 1
    },
    bash_misuse: {
        'exit_code': 15
    },
    command_not_found: {
        'exit_code': 127
    },
    invalid_exit: {
        'exit_code': 128
    },
    not_executable: {
        'exit_code': 126
    }
}

whitelist = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'configs', '*threads*')

# @pytest.mark.whitelist(whitelist, reason='broken in IPP')
@pytest.mark.skip("Broke somewhere between PR #525 and PR #652")
def test_bash_formatting():

    f = bad_format()
    try:
        f.result()
    except Exception as e:
        print("Caught exception", e)
        assert isinstance(
            e, parsl.app.errors.AppBadFormatting), "Expected AppBadFormatting got : {0}".format(e)
    return True

# @pytest.mark.whitelist(whitelist, reason='broken in IPP')
@pytest.mark.skip("Broke somewhere between PR #525 and PR #652")
def test_div_0(test_fn=div_0):
    err_code = test_matrix[test_fn]['exit_code']
    f = test_fn()
    try:
        f.result()
    except Exception as e:
        print("Caught exception", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)
    print(os.listdir('.'))
    os.remove('std.err')
    os.remove('std.out')
    return True

@pytest.mark.issue363
def test_bash_misuse(test_fn=bash_misuse):
    err_code = test_matrix[test_fn]['exit_code']
    f = test_fn()
    try:
        f.result()
    except pe.BashExitFailure as e:
        print("Caught expected BashExitFailure", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)
    os.remove('std.err')
    os.remove('std.out')

@pytest.mark.issue363
def test_command_not_found(test_fn=command_not_found):
    err_code = test_matrix[test_fn]['exit_code']
    f = test_fn()
    try:
        f.result()
    except pe.BashExitFailure as e:
        print("Caught exception", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)

    os.remove('std.err')
    os.remove('std.out')
    return True

@pytest.mark.skip('broken')
def test_invalid_exit(test_fn=invalid_exit):
    err_code = test_matrix[test_fn]['exit_code']
    f = test_fn()
    try:
        f.result()
    except Exception as e:
        print("Caught exception", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)
    os.remove('std.err')
    os.remove('std.out')
    return True

# @pytest.mark.whitelist(whitelist, reason='broken in IPP')
@pytest.mark.skip("Broke somewhere between PR #525 and PR #652")
def test_not_executable(test_fn=not_executable):
    err_code = test_matrix[test_fn]['exit_code']
    f = test_fn()
    try:
        f.result()
    except Exception as e:
        print("Caught exception", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)
    os.remove('std.err')
    os.remove('std.out')
    return True

def run_app(test_fn, err_code):
    f = test_fn()
    print(f)
    try:
        f.result()
    except Exception as e:
        print("Caught exception", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)
    os.remove('std.err')
    os.remove('std.out')
    return True

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--debug", action='store_true',
                        help="Count of apps to launch")
    args = parser.parse_args()

    print(test_bash_formatting())

    exit(0)

    if args.debug:
        parsl.set_stream_logger()

    for test in test_matrix:
        try:
            run_app(test, test_matrix[test]['exit_code'])
        except AssertionError as e:
            print("Test {0:15} [FAILED]".format(test.__name__))
            print("Caught error : {0}".format(e))
        else:
            print("Test {0:15} [SUCCESS]".format(test.__name__))
<commit_msg>Re-enable test which appears to work now (#1934)<commit_after>import argparse
import os

import pytest

import parsl
from parsl.app.app import bash_app
import parsl.app.errors as pe

from parsl.tests.configs.local_threads import config

local_config = config

@bash_app
def command_not_found(stderr='std.err', stdout='std.out'):
    cmd_line = 'catdogcat'
    return cmd_line

@bash_app
def bash_misuse(stderr='std.err', stdout='std.out'):
    cmd_line = 'exit 15'
    return cmd_line

@bash_app
def div_0(stderr='std.err', stdout='std.out'):
    cmd_line = '$((5/0))'
    return cmd_line

@bash_app
def invalid_exit(stderr='std.err', stdout='std.out'):
    cmd_line = 'exit 3.141'
    return cmd_line

@bash_app
def not_executable(stderr='std.err', stdout='std.out'):
    cmd_line = '/dev/null'
    return cmd_line

@bash_app
def bad_format(stderr='std.err', stdout='std.out'):
    cmd_line = 'echo {0}'
    return cmd_line

test_matrix = {
    div_0: {
        'exit_code': 1
    },
    bash_misuse: {
        'exit_code': 15
    },
    command_not_found: {
        'exit_code': 127
    },
    invalid_exit: {
        'exit_code': 128
    },
    not_executable: {
        'exit_code': 126
    }
}

whitelist = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'configs', '*threads*')

# @pytest.mark.whitelist(whitelist, reason='broken in IPP')
@pytest.mark.skip("Broke somewhere between PR #525 and PR #652")
def test_bash_formatting():

    f = bad_format()
    try:
        f.result()
    except Exception as e:
        print("Caught exception", e)
        assert isinstance(
            e, parsl.app.errors.AppBadFormatting), "Expected AppBadFormatting got : {0}".format(e)
    return True

# @pytest.mark.whitelist(whitelist, reason='broken in IPP')
@pytest.mark.skip("Broke somewhere between PR #525 and PR #652")
def test_div_0(test_fn=div_0):
    err_code = test_matrix[test_fn]['exit_code']
    f = test_fn()
    try:
        f.result()
    except Exception as e:
        print("Caught exception", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)
    print(os.listdir('.'))
    os.remove('std.err')
    os.remove('std.out')
    return True

@pytest.mark.issue363
def test_bash_misuse(test_fn=bash_misuse):
    err_code = test_matrix[test_fn]['exit_code']
    f = test_fn()
    try:
        f.result()
    except pe.BashExitFailure as e:
        print("Caught expected BashExitFailure", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)
    os.remove('std.err')
    os.remove('std.out')

@pytest.mark.issue363
def test_command_not_found(test_fn=command_not_found):
    err_code = test_matrix[test_fn]['exit_code']
    f = test_fn()
    try:
        f.result()
    except pe.BashExitFailure as e:
        print("Caught exception", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)

    os.remove('std.err')
    os.remove('std.out')
    return True

@pytest.mark.skip('broken')
def test_invalid_exit(test_fn=invalid_exit):
    err_code = test_matrix[test_fn]['exit_code']
    f = test_fn()
    try:
        f.result()
    except Exception as e:
        print("Caught exception", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)
    os.remove('std.err')
    os.remove('std.out')
    return True

def test_not_executable(test_fn=not_executable):
    err_code = test_matrix[test_fn]['exit_code']
    f = test_fn()
    try:
        f.result()
    except Exception as e:
        print("Caught exception", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)
    os.remove('std.err')
    os.remove('std.out')
    return True

def run_app(test_fn, err_code):
    f = test_fn()
    print(f)
    try:
        f.result()
    except Exception as e:
        print("Caught exception", e)
        assert e.exitcode == err_code, "{0} expected err_code:{1} but got {2}".format(test_fn.__name__,
                                                                                err_code,
                                                                                e.exitcode)
    os.remove('std.err')
    os.remove('std.out')
    return True

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--debug", action='store_true',
                        help="Count of apps to launch")
    args = parser.parse_args()

    print(test_bash_formatting())

    exit(0)

    if args.debug:
        parsl.set_stream_logger()

    for test in test_matrix:
        try:
            run_app(test, test_matrix[test]['exit_code'])
        except AssertionError as e:
            print("Test {0:15} [FAILED]".format(test.__name__))
            print("Caught error : {0}".format(e))
        else:
            print("Test {0:15} [SUCCESS]".format(test.__name__))
<|endoftext|>


<commit_before>from django.conf.urls import patterns, url, include
from django.views.generic import TemplateView

from profiles.forms import SignupFormExtra

from comicsite.admin import projectadminurls

urlpatterns = patterns('',

    url(r'^test/send_email/$','comicsite.views.send_email'),
    url(r'^test/throw_exception/$','comicsite.views.throw_exception'),
    url(r'^test/throw_http404/$','comicsite.views.throw_http404'),

    url(r'^test/test_logging/$','comicsite.views.test_logging'),

    url(r'^(?P<site_short_name>[\w-]+)/$','comicsite.views.site'),

    # Include an admin url for each project in database. This stretches the django
    # Assumptions of urls being fixed a bit, but it is the only way to reuse much
    # of the automatic admin functionality whithout rewriting the whole interface
    # see issue #181
    url(r'^', include(projectadminurls.allurls),name='projectadmin'),

    url(r'^(?P<site_short_name>[\w-]+)/robots\.txt$', TemplateView.as_view(template_name='robots.html'),name="comicsite_robots_txt"),

    # these registration and account views are viewed in the context of a
    # project
    url(r'^(?P<site_short_name>[\w-]+)/accounts/signin/$','comicsite.views.signin',name="comicsite_signin"),
    url(r'^(?P<site_short_name>[\w-]+)/accounts/signup/$','comicsite.views.signup',{'signup_form':SignupFormExtra},name="comicsite_signup"),
    url(r'^(?P<site_short_name>[\w-]+)/accounts/signup_complete/$','comicsite.views.signup_complete',name="comicsite_signup_complete"),

    url(r'^(?P<site_short_name>[\w-]+)/files/$','comicmodels.views.upload_handler'),

    url(r'^(?P<project_name>[\w-]+)/serve/(?P<path>.+)/$','filetransfers.views.serve',name="project_serve_file"),

    #_register should be removed, moving to _request_participation. See #240 keeping it now for any links going here directly
    url(r'^(?P<site_short_name>[\w-]+)/_register/$','comicsite.views._register'),
    url(r'^(?P<site_short_name>[\w-]+)/_request_participation/$','comicsite.views._register'),

    url(r'^(?P<site_short_name>[\w-]+)/source/(?P<page_title>[\w-]+)/$','comicsite.views.pagesource'),

    url(r'^(?P<site_short_name>[\w-]+)/(?P<page_title>[\w-]+)/db/(?P<dropboxname>[\w-]+)/(?P<dropboxpath>.+)/$','comicsite.views.dropboxpage'),

    url(r'^(?P<site_short_name>[\w-]+)/(?P<page_title>[\w-]+)/insert/(?P<dropboxpath>.+)/$','comicsite.views.insertedpage'),

    url(r'^(?P<site_short_name>[\w-]+)/(?P<page_title>[\w-]+)/$','comicsite.views.page'),
)
<commit_msg>fixes image upload issue: ckeditor urls are now also available under /site/projectname/<commit_after>from django.conf.urls import patterns, url, include
from django.views.generic import TemplateView

from profiles.forms import SignupFormExtra

from comicsite.admin import projectadminurls

urlpatterns = patterns('',

    url(r'^test/send_email/$','comicsite.views.send_email'),
    url(r'^test/throw_exception/$','comicsite.views.throw_exception'),
    url(r'^test/throw_http404/$','comicsite.views.throw_http404'),

    url(r'^test/test_logging/$','comicsite.views.test_logging'),

    url(r'^(?P<site_short_name>[\w-]+)/$','comicsite.views.site'),

    # Include an admin url for each project in database. This stretches the django
    # Assumptions of urls being fixed a bit, but it is the only way to reuse much
    # of the automatic admin functionality whithout rewriting the whole interface
    # see issue #181
    url(r'^', include(projectadminurls.allurls),name='projectadmin'),

    url(r'^(?P<site_short_name>[\w-]+)/robots\.txt$', TemplateView.as_view(template_name='robots.html'),name="comicsite_robots_txt"),

    #url(r'^(?P<site_short_name>[\w-]+)/ckeditor/', include('ckeditor.urls')),

    # these registration and account views are viewed in the context of a
    # project
    url(r'^(?P<site_short_name>[\w-]+)/accounts/signin/$','comicsite.views.signin',name="comicsite_signin"),
    url(r'^(?P<site_short_name>[\w-]+)/accounts/signup/$','comicsite.views.signup',{'signup_form':SignupFormExtra},name="comicsite_signup"),
    url(r'^(?P<site_short_name>[\w-]+)/accounts/signup_complete/$','comicsite.views.signup_complete',name="comicsite_signup_complete"),

    url(r'^(?P<site_short_name>[\w-]+)/files/$','comicmodels.views.upload_handler'),

    url(r'^(?P<project_name>[\w-]+)/serve/(?P<path>.+)/$','filetransfers.views.serve',name="project_serve_file"),

    #_register should be removed, moving to _request_participation. See #240 keeping it now for any links going here directly
    url(r'^(?P<site_short_name>[\w-]+)/_register/$','comicsite.views._register'),
    url(r'^(?P<site_short_name>[\w-]+)/_request_participation/$','comicsite.views._register'),

    url(r'^(?P<site_short_name>[\w-]+)/source/(?P<page_title>[\w-]+)/$','comicsite.views.pagesource'),

    url(r'^(?P<site_short_name>[\w-]+)/(?P<page_title>[\w-]+)/db/(?P<dropboxname>[\w-]+)/(?P<dropboxpath>.+)/$','comicsite.views.dropboxpage'),

    url(r'^(?P<site_short_name>[\w-]+)/(?P<page_title>[\w-]+)/insert/(?P<dropboxpath>.+)/$','comicsite.views.insertedpage'),

    url(r'^(?P<site_short_name>[\w-]+)/(?P<page_title>[\w-]+)/$','comicsite.views.page'),
)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright 2016 KMEE - Luis Felipe Miléo <mileo@kmee.com.br>
# Copyright 2016 KMEE - Hendrix Costa <hendrix.costa@kmee.com.br>
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl).

import logging
from odoo import api, fields, models, _
from datetime import datetime, timedelta
from odoo import tools

_logger = logging.getLogger(__name__)

try:
    from pybrasil.feriado.constantes import (
        TIPO_FERIADO, ABRANGENCIA_FERIADO,
    )
except ImportError:
    _logger.warning('Cannot import pybrasil')

class ResourceCalendar(models.Model):

    _inherit = 'resource.calendar'

    def _compute_recursive_leaves(self, calendar):
        res = self.env['resource.calendar.leaves']
        res |= self.env['resource.calendar.leaves'].search([
            ('calendar_id', '=', calendar.id)
        ])
        if calendar.parent_id:
            res |= self._compute_recursive_leaves(calendar.parent_id)
        return res

    @api.multi
    @api.depends('parent_id')
    def _compute_leave_ids(self):
        for calendar in self:
            calendar.leave_ids = self._compute_recursive_leaves(calendar)

    parent_id = fields.Many2one(
        'resource.calendar',
        string='Parent Calendar',
        ondelete='restrict',
        index=True)
    child_ids = fields.One2many(
        'resource.calendar', 'parent_id',
        string='Child Calendar')
    _parent_store = True

    parent_left = fields.Integer(index=True)
    parent_right = fields.Integer(index=True)

    country_id = fields.Many2one('res.country', u'País')
    state_id = fields.Many2one(
        'res.country.state', u'Estado',
        domain="[('country_id','=',country_id)]")
    l10n_br_city_id = fields.Many2one(
        'l10n_br_base.city', u'Municipio',
        domain="[('state_id','=',state_id)]")
    leave_ids = fields.Many2many(
        comodel_name='resource.calendar.leaves',
        compute='_compute_leave_ids'
    )

    @api.constrains('parent_id')
    def _check_hierarchy(self):
        if not self._check_recursion():
            raise models.ValidationError(_(
                'Error! You cannot create recursive calendars.'))

    @api.multi
    def get_leave_intervals(self, resource_id=None, start_datetime=None,
                            end_datetime=None):
        """Get the leaves of the calendar. Leaves can be filtered on the
        resource, the start datetime or the end datetime.

        :param int resource_id: the id of the resource to take into account
                                when computing the leaves. If not set,
                                only general leaves are computed.
                                If set, generic and specific leaves
                                are computed.
        :param datetime start_datetime: if provided, do not take into
                                        account leaves ending before
                                        this date.
        :param datetime end_datetime: if provided, do not take into
                                        account leaves beginning
                                        after this date.

        :return list leaves: list of tuples (start_datetime, end_datetime) of
                             leave intervals
        """
        leaves = []
        for leave in self.leave_ids:
            if leave.resource_id and resource_id:
                if leave.resource_id and not resource_id == leave.\
                        resource_id.id:
                    continue
            elif leave.resource_id and not resource_id:
                continue
            date_from = datetime.strptime(
                leave.date_from, tools.DEFAULT_SERVER_DATETIME_FORMAT
            )
            if end_datetime and date_from > end_datetime:
                continue
            date_to = datetime.strptime(
                leave.date_to, tools.DEFAULT_SERVER_DATETIME_FORMAT
            )
            if start_datetime and date_to < start_datetime:
                continue
            leaves.append(leave)
        return leaves

    @api.multi
    def data_eh_feriado(self, data_referencia=datetime.now()):
        """Verificar se uma data é feriado.
        :param datetime data_referencia: Se nenhuma data referencia for passada
                                    verifique se hoje eh feriado no calendario
                                    corrente.
                                    Se a data referencia for passada, verifique
                                    se a data esta dentro de algum leave do
                                    calendario corrente
                                    date_start <= data_referencia <= data_end

        :return boolean True se a data referencia for feriado
                        False se a data referencia nao for feriado
        """
        for leave in self.leave_ids:
            if leave.date_from <= data_referencia.strftime(
                    "%Y-%m-%d %H:%M:%S"):
                if leave.date_to >= data_referencia.\
                        strftime("%Y-%m-%d %H:%M:%S"):
                    if leave.leave_kind == 'F':
                        return True
        return False

    @api.multi
    def data_eh_dia_util(self, data=datetime.now()):
        """Verificar se data é dia util.
        :param datetime data: Se nenhuma data referencia for passada
                              verifique o dia de hoje.
        :return boolean True: Se for dia útil
                        False: Se Não for dia útil
        """
        return not self.data_eh_feriado(data) and data.weekday() <= 4 or False

    @api.multi
    def quantidade_dias_uteis(
            self, data_inicio=datetime.now(), data_fim=datetime.now()):
        """Calcular a quantidade de dias úteis em determinado período.
        :param datetime data_inicio: Se nenhuma data referencia for passada
                                   verifique o dia de hoje.
        :param datetime data_fim: Se nenhuma data referencia for passada
                                   verifique o dia de hoje.
        :return int: Quantidade de dias úteis
        """
        dias_uteis = 0
        while data_inicio <= data_fim:
            if self.data_eh_dia_util(data_inicio):
                dias_uteis += 1
            data_inicio += timedelta(days=1)

        return dias_uteis

    @api.multi
    def proximo_dia_util(self, data_referencia=datetime.now()):
        """Retornar o próximo dia util.
        :param datetime data_referencia: Se nenhuma data referencia for passada
                                   verifique se amanha é dia útil.
        :return datetime Proximo dia util apartir da data referencia
        """
        data_referencia += timedelta(days=1)
        while data_referencia:
            if self.data_eh_dia_util(data_referencia):
                return data_referencia
            data_referencia += timedelta(days=1)

    @api.multi
    def get_dias_base(self, data_from=datetime.now(), data_to=datetime.now()):
        """Calcular a quantidade de dias que devem ser remunerados em
        determinado intervalo de tempo.
        :param datetime data_from: Data inicial do intervalo de tempo.
               datetime data_end: Data final do intervalo
        :return int : quantidade de dias que devem ser remunerada
       """
        quantidade_dias = (data_to - data_from).days + 1
        if quantidade_dias > 30:
            return 30
        else:
            return quantidade_dias

    @api.multi
    def data_eh_feriado_bancario(self, data_referencia=datetime.now()):
        """Verificar se uma data é feriado bancário.
        :param datetime data_referencia: Se nenhuma data referencia for passada
                                    verifique se hoje é feriado bancário.
                                    Se a data referencia for passada, verifique
                                    se a data esta dentro de algum leave
                                    date_start <= data_referencia <= data_end

        :return int leaves_count: +1 se for feriado bancário
                                   0 se a data nao for feriado bancário
        """
        domain = [
            ('date_from', '<=', data_referencia.strftime("%Y-%m-%d %H:%M:%S")),
            ('date_to', '>=', data_referencia.strftime("%Y-%m-%d %H:%M:%S")),
            ('leave_kind', 'in', ['F', 'B']),
        ]
        leaves_count = \
            self.env['resource.calendar.leaves'].search_count(domain)
        return leaves_count

    @api.multi
    def data_eh_feriado_emendado(self, data_referencia=datetime.now()):
        """Verificar se uma data é feriado emendado.
        :param datetime data_referencia: Se nenhuma data referencia for passada
                                   verifique se hoje é feriado emendado.
                                   Se a data referencia for passada, verifique
                                   se a data esta dentro de algum leave
                                   date_start <= data_referencia <= data_end

        :return retorna True ou False
        """
        dia_antes = data_referencia - timedelta(days=2)
        dia_depois = data_referencia + timedelta(days=2)

        dia_antes_eh_util = \
            True if dia_antes.weekday() > 5 or self.data_eh_feriado(
                dia_antes) else False
        dia_depois_eh_util = \
            True if dia_depois.weekday() > 4 or self.data_eh_feriado(
                dia_depois) else False

        return dia_antes_eh_util or dia_depois_eh_util

class ResourceCalendarLeave(models.Model):

    _inherit = 'resource.calendar.leaves'

    country_id = fields.Many2one(
        'res.country', string=u'País',
        related='calendar_id.country_id',
    )
    state_id = fields.Many2one(
        'res.country.state', u'Estado',
        related='calendar_id.state_id',
        domain="[('country_id','=',country_id)]",
        readonly=True
    )
    l10n_br_city_id = fields.Many2one(
        'l10n_br_base.city', u'Municipio',
        related='calendar_id.l10n_br_city_id',
        domain="[('state_id','=',state_id)]",
        readonly=True
    )
    leave_kind = fields.Selection(
        string=u'Leave Kind',
        selection=list(TIPO_FERIADO.iteritems()),
    )
    leave_scope = fields.Selection(
        string=u'Leave Scope',
        selection=list(ABRANGENCIA_FERIADO.iteritems()),
    )
<commit_msg>[FIX] Correções no calculo dos feriados<commit_after># -*- coding: utf-8 -*-
# Copyright 2016 KMEE - Luis Felipe Miléo <mileo@kmee.com.br>
# Copyright 2016 KMEE - Hendrix Costa <hendrix.costa@kmee.com.br>
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl).

import logging
from odoo import api, fields, models, _
from odoo.exceptions import AccessError
from datetime import datetime, timedelta
from odoo import tools

_logger = logging.getLogger(__name__)

try:
    from pybrasil.feriado.constantes import (
        TIPO_FERIADO, ABRANGENCIA_FERIADO,
    )
except ImportError:
    _logger.warning('Cannot import pybrasil')

class ResourceCalendar(models.Model):
    _inherit = 'resource.calendar'

    def _compute_recursive_leaves(self, calendar):
        res = self.env['resource.calendar.leaves']
        res |= self.env['resource.calendar.leaves'].search([
            ('calendar_id', '=', calendar.id)
        ])
        if calendar.parent_id:
            res |= self._compute_recursive_leaves(calendar.parent_id)
        return res

    @api.multi
    @api.depends('parent_id')
    def _compute_leave_ids(self):
        for calendar in self:
            calendar.leave_ids = self._compute_recursive_leaves(calendar)

    parent_id = fields.Many2one(
        'resource.calendar',
        string='Parent Calendar',
        ondelete='restrict',
        index=True)
    child_ids = fields.One2many(
        'resource.calendar', 'parent_id',
        string='Child Calendar')
    _parent_store = True

    parent_left = fields.Integer(index=True)
    parent_right = fields.Integer(index=True)

    country_id = fields.Many2one('res.country', u'País')
    state_id = fields.Many2one(
        'res.country.state', u'Estado',
        domain="[('country_id','=',country_id)]")
    l10n_br_city_id = fields.Many2one(
        'l10n_br_base.city', u'Municipio',
        domain="[('state_id','=',state_id)]")
    leave_ids = fields.Many2many(
        comodel_name='resource.calendar.leaves',
        compute='_compute_leave_ids'
    )

    @api.constrains('parent_id')
    def _check_hierarchy(self):
        if not self._check_recursion():
            raise models.ValidationError(_(
                'Error! You cannot create recursive calendars.'))

    @api.multi
    def get_leave_intervals(self, resource_id=None, start_datetime=None,
                            end_datetime=None):
        """Get the leaves of the calendar. Leaves can be filtered on the
        resource, the start datetime or the end datetime.

        :param int resource_id: the id of the resource to take into account
                                when computing the leaves. If not set,
                                only general leaves are computed.
                                If set, generic and specific leaves
                                are computed.
        :param datetime start_datetime: if provided, do not take into
                                        account leaves ending before
                                        this date.
        :param datetime end_datetime: if provided, do not take into
                                        account leaves beginning
                                        after this date.

        :return list leaves: list of tuples (start_datetime, end_datetime) of
                             leave intervals
        """
        leaves = []
        for leave in self.leave_ids:
            if leave.resource_id and resource_id:
                if leave.resource_id and not resource_id == leave. \
                        resource_id.id:
                    continue
            elif leave.resource_id and not resource_id:
                continue
            date_from = datetime.strptime(
                leave.date_from, tools.DEFAULT_SERVER_DATETIME_FORMAT
            )
            if end_datetime and date_from > end_datetime:
                continue
            date_to = datetime.strptime(
                leave.date_to, tools.DEFAULT_SERVER_DATETIME_FORMAT
            )
            if start_datetime and date_to < start_datetime:
                continue
            leaves.append(leave)
        return leaves

    @api.multi
    def data_eh_feriado(self, data_referencia=datetime.now()):
        """Verificar se uma data é feriado.
        :param datetime data_referencia: Se nenhuma data referencia for passada
                                    verifique se hoje eh feriado no calendario
                                    corrente.
                                    Se a data referencia for passada, verifique
                                    se a data esta dentro de algum leave do
                                    calendario corrente
                                    date_start <= data_referencia <= data_end

        :return boolean True se a data referencia for feriado
                        False se a data referencia nao for feriado
        """
        if not self.leave_ids:
            raise AccessError(
                "Você esta tentando calcular dadas e/ou ausencias de recursos "
                "sem especificar um calendário de referência!")
        for leave in self.leave_ids:
            if leave.date_from <= data_referencia.strftime(
                    "%Y-%m-%d %H:%M:%S"):
                if leave.date_to >= data_referencia. \
                        strftime("%Y-%m-%d %H:%M:%S"):
                    if leave.leave_kind in ('F', 'B'):
                        return True
        return False

    @api.multi
    def data_eh_dia_util(self, data=datetime.now()):
        """Verificar se data é dia util.
        :param datetime data: Se nenhuma data referencia for passada
                              verifique o dia de hoje.
        :return boolean True: Se for dia útil
                        False: Se Não for dia útil
        """
        return not self.data_eh_feriado(data) and data.weekday() <= 4 or False

    @api.multi
    def data_eh_dia_util_bancario(self, data=datetime.now()):
        """Verificar se data é dia util.
        :param datetime data: Se nenhuma data referencia for passada
                              verifique o dia de hoje.
        :return boolean True: Se for dia útil
                        False: Se Não for dia útil
        """
        if data.weekday() > 4:
            return False
        elif self.data_eh_feriado_bancario(data):
            return False
        return True

    @api.multi
    def quantidade_dias_uteis(
            self, data_inicio=datetime.now(), data_fim=datetime.now()):
        """Calcular a quantidade de dias úteis em determinado período.
        :param datetime data_inicio: Se nenhuma data referencia for passada
                                   verifique o dia de hoje.
        :param datetime data_fim: Se nenhuma data referencia for passada
                                   verifique o dia de hoje.
        :return int: Quantidade de dias úteis
        """
        dias_uteis = 0
        while data_inicio <= data_fim:
            if self.data_eh_dia_util(data_inicio):
                dias_uteis += 1
            data_inicio += timedelta(days=1)

        return dias_uteis

    @api.multi
    def proximo_dia_util(self, data_referencia=datetime.now()):
        """Retornar o próximo dia util.
        :param datetime data_referencia: Se nenhuma data referencia for passada
                                   verifique se amanha é dia útil.
        :return datetime Proximo dia util apartir da data referencia
        """
        while data_referencia:
            if self.data_eh_dia_util(data_referencia):
                return data_referencia
            data_referencia += timedelta(days=1)

    @api.multi
    def proximo_dia_util_bancario(self, data_referencia=datetime.now()):
        """Retornar o próximo dia util.
        :param datetime data_referencia: Se nenhuma data referencia for passada
                                   verifique se amanha é dia útil.
        :return datetime Proximo dia util apartir da data referencia
        """
        if self.data_eh_dia_util_bancario(data_referencia):
            return data_referencia
        data_referencia += timedelta(days=1)
        return self.proximo_dia_util_bancario(data_referencia)

    @api.multi
    def get_dias_base(self, data_from=datetime.now(), data_to=datetime.now()):
        """Calcular a quantidade de dias que devem ser remunerados em
        determinado intervalo de tempo.
        :param datetime data_from: Data inicial do intervalo de tempo.
               datetime data_end: Data final do intervalo
        :return int : quantidade de dias que devem ser remunerada
       """
        quantidade_dias = (data_to - data_from).days + 1
        if quantidade_dias > 30:
            return 30
        else:
            return quantidade_dias

    @api.multi
    def data_eh_feriado_bancario(self, data_referencia=datetime.now()):
        """Verificar se uma data é feriado bancário.
        :param datetime data_referencia: Se nenhuma data referencia for passada
                                    verifique se hoje é feriado bancário.
                                    Se a data referencia for passada, verifique
                                    se a data esta dentro de algum leave
                                    date_start <= data_referencia <= data_end

        :return int leaves_count: +1 se for feriado bancário
                                   0 se a data nao for feriado bancário
        """
        if not self.leave_ids:
            raise AccessError(
                "Você esta tentando calcular dadas e/ou ausencias de recursos "
                "sem especificar um calendário de referência!")
        if self.leave_ids.filtered(
                lambda record: data_referencia.strftime("%Y-%m-%d 00:00:00") <=
                record.date_from <=
                data_referencia.strftime("%Y-%m-%d 23:59:59") and
                record.leave_kind == 'B'):
            return True
        return False

    @api.multi
    def data_eh_feriado_emendado(self, data_referencia=datetime.now()):
        """Verificar se uma data é feriado emendado.
        :param datetime data_referencia: Se nenhuma data referencia for passada
                                   verifique se hoje é feriado emendado.
                                   Se a data referencia for passada, verifique
                                   se a data esta dentro de algum leave
                                   date_start <= data_referencia <= data_end

        :return retorna True ou False
        """
        dia_antes = data_referencia - timedelta(days=2)
        dia_depois = data_referencia + timedelta(days=2)

        dia_antes_eh_util = \
            True if dia_antes.weekday() > 5 or self.data_eh_feriado(
                dia_antes) else False
        dia_depois_eh_util = \
            True if dia_depois.weekday() > 4 or self.data_eh_feriado(
                dia_depois) else False

        return dia_antes_eh_util or dia_depois_eh_util

class ResourceCalendarLeave(models.Model):
    _inherit = 'resource.calendar.leaves'

    country_id = fields.Many2one(
        'res.country', string=u'País',
        related='calendar_id.country_id',
    )
    state_id = fields.Many2one(
        'res.country.state', u'Estado',
        related='calendar_id.state_id',
        domain="[('country_id','=',country_id)]",
        readonly=True
    )
    l10n_br_city_id = fields.Many2one(
        'l10n_br_base.city', u'Municipio',
        related='calendar_id.l10n_br_city_id',
        domain="[('state_id','=',state_id)]",
        readonly=True
    )
    leave_kind = fields.Selection(
        string=u'Leave Kind',
        selection=list(TIPO_FERIADO.iteritems()),
    )
    leave_scope = fields.Selection(
        string=u'Leave Scope',
        selection=list(ABRANGENCIA_FERIADO.iteritems()),
    )
<|endoftext|>


<commit_before>#!/usr/bin/env python3
"""Contains the different formats of CPDs used in PGM"""

import numpy as np
from pgmpy.Factor import Factor
from pgmpy import Exceptions

class TabularCPD(Factor):
    """
    Public Methods
    --------------
    marginalize([variables_list])
    reduce([values_list])
    normalize()
    """
    def __init__(self, event, event_card, values,
                 evidence=None, evidence_card=None):
        """
        Defines the conditional probability distribution table (cpd table)

        Parameters
        ----------
        event: string
            event whose cpd table is defined
        event_card: integer
            cardinality of event
        values: 2d array, 2d list
            values of the cpd table

            example
            for a distribution of P(grade|intel, diff)

            diff:       easy                 hard
            intel: dumb   avg   smart    dumb  avg   smart
            gradeA: 0.1    0.1    0.1     0.1  0.1    0.1
            gradeB: 0.1    0.1    0.1     0.1  0.1    0.1
            gradeC: 0.8    0.8    0.8     0.8  0.8    0.8

            values should be
            [[0.1,0.1,0.1,0.1,0.1,0.1],
            [0.1,0.1,0.1,0.1,0.1,0.1],
            [0.8,0.8,0.8,0.8,0.8,0.8]]

        evidence: string, list-type
            evidences(if any) w.r.t. which cpd is defined
        evidence_card: integer, list-type
            cardinality of evidences

        """
        if not isinstance(event, str):
            raise TypeError("Event must be a string")
        self.event = event
        variables = [event]
        if not isinstance(event_card, int):
            raise TypeError("Event cardinality must be an integer")
        self.event_card = event_card
        cardinality = [event_card]
        if evidence_card:
            if not isinstance(evidence_card, (list, set, tuple)):
                evidence_card = [evidence_card]
            cardinality.extend(evidence_card)
        self.evidence_card = evidence_card
        if evidence:
            if not isinstance(evidence, (list, set, tuple)):
                evidence = [evidence]
            variables.extend(evidence)
            if not len(evidence_card) == len(evidence):
                raise Exceptions.CardinalityError("Cardinality of all "
                                                  "evidences not specified")
        self.evidence = evidence
        if len(np.array(values).shape) is not 2:
            raise TypeError("Values must be a 2d list/array")
        self.cpd = np.array(values)
        Factor.__init__(self, variables, cardinality, self.cpd.flatten('F'))

    def marginalize(self, variables):
        """
        Modifies the cpd table with marginalized values.

        Paramters
        ---------
        variables: string, list-type
            name of variable to be marginalized

        Examples
        --------
        >>> from pgmpy.Factor.CPD import TabularCPD
        >>> cpd_table = TabularCPD('grade', 2,
        ...                        [[0.7, 0.6, 0.6, 0.2],[0.3, 0.4, 0.4, 0.8]],
        ...                        ['intel', 'diff'], [2, 2])
        >>> cpd_table.marginalize('diff')
        >>> cpd_table.get_cpd()
        array([[ 1.3,  0.8],
               [ 0.7,  1.2]])
        """
        if self.event in variables:
            self.event_card = 1
        Factor.marginalize(self, variables)
        self.cpd = self.values.reshape((self.event_card,
                                        np.product(self.cardinality)/self.event_card),
                                       order='F')
        self.evidence = [var for var in self.variables
                         if var is not self.event]
        self.evidence_card = [self.get_cardinality(variable)
                              for variable in self.evidence]

    def normalize(self):
        """
        Normalizes the cpd table

        Examples
        --------
        >>> from pgmpy.Factor.CPD import TabularCPD
        >>> cpd_table = TabularCPD('grade', 2,
        ...                        [[0.7, 0.2, 0.6, 0.2],[0.4, 0.4, 0.4, 0.8]],
        ...                        ['intel', 'diff'], [2, 2])
        >>> cpd_table.normalize()
        >>> cpd_table.get_cpd()
        array([[ 0.63636364,  0.33333333,  0.6       ,  0.2       ],
               [ 0.36363636,  0.66666667,  0.4       ,  0.8       ]])
        """
        self.cpd = self.cpd / self.cpd.sum(axis=0)

    def reduce(self, values):
        """
        Reduces the cpd table to the context of given variable values.

        Parameters
        ----------
        values: string, list-type
            name of the variable values

        Examples
        --------
        >>> from pgmpy.Factor.CPD import TabularCPD
        >>> cpd_table = TabularCPD('grade', 2,
        ...                        [[0.7, 0.6, 0.6, 0.2],[0.3, 0.4, 0.4, 0.8]],
        ...                        ['intel', 'diff'], [2, 2])
        >>> cpd_table.reduce('diff_0')
        >>> cpd_table.get_cpd()
        array([[ 0.7,  0.6],
               [ 0.3,  0.4]])
        """
        if not isinstance(values, (list, set, tuple)):
            values = [values]
        if any(self.event in value for value in values):
            self.event_card = 1
        Factor.reduce(self, values)
        self.cpd = self.values.reshape((self.event_card,
                                        np.product(self.cardinality)/self.event_card),
                                       order='F')
        self.evidence = [var for var in self.variables
                         if var is not self.event]
        self.evidence_card = [self.get_cardinality(variable)
                              for variable in self.evidence]

    def get_cpd(self):
        """
        Returns the cpd
        """
        return self.cpd
<commit_msg>fixed TabularCPD docstring<commit_after>#!/usr/bin/env python3
"""Contains the different formats of CPDs used in PGM"""

import numpy as np
from pgmpy.Factor import Factor
from pgmpy import Exceptions

class TabularCPD(Factor):
    """
    Defines the conditional probability distribution table (cpd table)

    Example
    -------
    For a distribution of P(grade|intel, diff)

    +-------+--------------------+------------------+
    |diff   |      easy          |    hard          |
    +-------+-----+------+-------+------+----+------+
    |intel  |dumb |  avg | smart | dumb |avg |smart |
    +-------+-----+------+-------+------+----+------+
    |gradeA |0.1  |  0.1 |  0.1  | 0.1  |0.1 | 0.1  |
    +-------+-----+------+-------+------+----+------+
    |gradeB |0.1  |  0.1 |  0.1  | 0.1  |0.1 | 0.1  |
    +-------+-----+------+-------+------+----+------+
    |gradeC |0.8  |  0.8 |  0.8  | 0.8  |0.8 | 0.8  |
    +-------+-----+------+-------+------+----+------+

    values should be
    [[0.1,0.1,0.1,0.1,0.1,0.1],
    [0.1,0.1,0.1,0.1,0.1,0.1],
    [0.8,0.8,0.8,0.8,0.8,0.8]]

    Parameters
    ----------
    event: string
        event whose cpd table is defined
    event_card: integer
        cardinality of event
    values: 2d array, 2d list
        values of the cpd table
    evidence: string, list-type
        evidences(if any) w.r.t. which cpd is defined
    evidence_card: integer, list-type
        cardinality of evidences

    Public Methods
    --------------
    marginalize([variables_list])
    reduce([values_list])
    normalize()
    """
    def __init__(self, event, event_card, values,
                 evidence=None, evidence_card=None):
        if not isinstance(event, str):
            raise TypeError("Event must be a string")
        self.event = event
        variables = [event]
        if not isinstance(event_card, int):
            raise TypeError("Event cardinality must be an integer")
        self.event_card = event_card
        cardinality = [event_card]
        if evidence_card:
            if not isinstance(evidence_card, (list, set, tuple)):
                evidence_card = [evidence_card]
            cardinality.extend(evidence_card)
        self.evidence_card = evidence_card
        if evidence:
            if not isinstance(evidence, (list, set, tuple)):
                evidence = [evidence]
            variables.extend(evidence)
            if not len(evidence_card) == len(evidence):
                raise Exceptions.CardinalityError("Cardinality of all "
                                                  "evidences not specified")
        self.evidence = evidence
        if len(np.array(values).shape) is not 2:
            raise TypeError("Values must be a 2d list/array")
        self.cpd = np.array(values)
        Factor.__init__(self, variables, cardinality, self.cpd.flatten('F'))

    def marginalize(self, variables):
        """
        Modifies the cpd table with marginalized values.

        Paramters
        ---------
        variables: string, list-type
            name of variable to be marginalized

        Examples
        --------
        >>> from pgmpy.Factor.CPD import TabularCPD
        >>> cpd_table = TabularCPD('grade', 2,
        ...                        [[0.7, 0.6, 0.6, 0.2],[0.3, 0.4, 0.4, 0.8]],
        ...                        ['intel', 'diff'], [2, 2])
        >>> cpd_table.marginalize('diff')
        >>> cpd_table.get_cpd()
        array([[ 1.3,  0.8],
               [ 0.7,  1.2]])
        """
        if self.event in variables:
            self.event_card = 1
        Factor.marginalize(self, variables)
        self.cpd = self.values.reshape((self.event_card,
                                        np.product(self.cardinality)/self.event_card),
                                       order='F')
        self.evidence = [var for var in self.variables
                         if var is not self.event]
        self.evidence_card = [self.get_cardinality(variable)
                              for variable in self.evidence]

    def normalize(self):
        """
        Normalizes the cpd table

        Examples
        --------
        >>> from pgmpy.Factor.CPD import TabularCPD
        >>> cpd_table = TabularCPD('grade', 2,
        ...                        [[0.7, 0.2, 0.6, 0.2],[0.4, 0.4, 0.4, 0.8]],
        ...                        ['intel', 'diff'], [2, 2])
        >>> cpd_table.normalize()
        >>> cpd_table.get_cpd()
        array([[ 0.63636364,  0.33333333,  0.6       ,  0.2       ],
               [ 0.36363636,  0.66666667,  0.4       ,  0.8       ]])
        """
        self.cpd = self.cpd / self.cpd.sum(axis=0)

    def reduce(self, values):
        """
        Reduces the cpd table to the context of given variable values.

        Parameters
        ----------
        values: string, list-type
            name of the variable values

        Examples
        --------
        >>> from pgmpy.Factor.CPD import TabularCPD
        >>> cpd_table = TabularCPD('grade', 2,
        ...                        [[0.7, 0.6, 0.6, 0.2],[0.3, 0.4, 0.4, 0.8]],
        ...                        ['intel', 'diff'], [2, 2])
        >>> cpd_table.reduce('diff_0')
        >>> cpd_table.get_cpd()
        array([[ 0.7,  0.6],
               [ 0.3,  0.4]])
        """
        if not isinstance(values, (list, set, tuple)):
            values = [values]
        if any(self.event in value for value in values):
            self.event_card = 1
        Factor.reduce(self, values)
        self.cpd = self.values.reshape((self.event_card,
                                        np.product(self.cardinality)/self.event_card),
                                       order='F')
        self.evidence = [var for var in self.variables
                         if var is not self.event]
        self.evidence_card = [self.get_cardinality(variable)
                              for variable in self.evidence]

    def get_cpd(self):
        """
        Returns the cpd
        """
        return self.cpd
<|endoftext|>


<commit_before>import re
from collections import defaultdict, namedtuple

import xlrd
from django.core.management.base import BaseCommand, CommandError
from django.utils.text import slugify
from mptindicators.scorecard.models import (
    Country, Section, Subsection, Indicator, IndicatorScore)

SECTION_RE = re.compile('(?:Section )?(\d).(\d)? (.*)')

#
# type definitions
#

CountryRecord = namedtuple('CountryRecord',
    ('name', 'column'))
SectionRecord = namedtuple('SectionRecord',
    ('section', 'subsection', 'title', 'row'))
QuestionRecord = namedtuple('QuestionRecord',
    ('number', 'text', 'description', 'comment', 'row', 'col', 'type'))
ScoreRecord = namedtuple('ScoreRecord',
    ('country', 'question', 'score', 'comment', 'sources', 'row', 'col'))

#
# data methods
#

def countries_iter(sheet):
    for col, cell in enumerate(sheet.row(1)):
        if col > 3:
            name = cell.value.split('-')[0].strip()
            yield CountryRecord(name, col)

def questions_iter(sheet):
    col = 1
    for row, cell in enumerate(sheet.col(col)):
        try:
            number = int(cell.value)
            text = sheet.cell(row, col + 1).value.strip()
            description = sheet.cell(row + 1, col + 1).value.strip()
            comment = sheet.cell(row + 2, col + 1).value.strip()

            if text.startswith('In law'):
                type = 1
            elif text.startswith('In practice'):
                type = 2
            elif text.startswith('Open Question'):
                type = 3
            else:
                type = 0

            yield QuestionRecord(number, text, description, comment, row, col, type)
        except:
            pass

def sections_iter(sheet):
    col = 0
    for row, cell in enumerate(sheet.col(col)):
        match = SECTION_RE.match(cell.value)
        if match:
            (sec, subsec, title) = match.groups()
            title = title.strip()
            yield SectionRecord(int(sec), int(subsec) if subsec else None, title, row)

def get_score(sheet, country, question):

    col = country.column
    row = question.row

    score = sheet.cell(row, col).value
    comment = sheet.cell(row + 1, col).value.strip().encode('utf-8')
    sources = sheet.cell(row + 2, col).value.strip().encode('utf-8')

    return ScoreRecord(country, question, score, comment, sources, row, col)

class Command(BaseCommand):

    help = 'Import scorecard data from master spreadsheet'

    def add_arguments(self, parser):
        parser.add_argument('path')

    def handle(self, *args, **options):

        book = xlrd.open_workbook(options['path'])
        sheet = book.sheets()[0]

        countries = list(countries_iter(sheet))
        questions = list(questions_iter(sheet))
        sections = list(sections_iter(sheet))

        def get_section(question):
            qsection = None
            for section in sections:
                if section.row > question.row:
                    break
                qsection = section
            return qsection

        #
        # load countries
        #

        print "loading countries..."

        Country.objects.all().delete()

        for record in countries:
            Country.objects.create(name=record.name,
                                   slug=slugify(record.name))

        #
        # load sections
        #

        print "loading sections and subsections..."

        Section.objects.all().delete()

        for record in sections:

            try:
                section = Section.objects.get(number=record.section)
            except Section.DoesNotExist:
                section = Section.objects.create(number=record.section)

            Subsection.objects.create(section=section,
                                      number=record.subsection,
                                      name=record.title)

        #
        # load indicators
        #

        print "loading indicators..."

        Indicator.objects.all().delete()

        for record in questions:

            ss = get_section(record)

            subsection = Subsection.objects.get(section__number=ss.section,
                                                number=ss.subsection)

            Indicator.objects.create(subsection=subsection,
                                     number=record.number,
                                     name=record.text,
                                     description=record.description,
                                     comment=record.comment,
                                     references='',
                                     type=record.type)

        #
        # load scores
        #

        print "loading indicator scores..."

        IndicatorScore.objects.all().delete()

        for country in countries:
            print '  {}'.format(country.name)
            for question in questions:
                score = get_score(sheet, country, question)

                value = None if score.score == '' else score.score

                iobj = Indicator.objects.get(number=question.number)
                cobj = Country.objects.get(name=country.name)

                IndicatorScore.objects.create(indicator=iobj,
                                              country=cobj,
                                              score=value,
                                              comment=score.comment,
                                              sources=score.sources)

<commit_msg>fix section names<commit_after>import re
from collections import defaultdict, namedtuple

import xlrd
from django.core.management.base import BaseCommand, CommandError
from django.utils.text import slugify
from mptindicators.scorecard.models import (
    Country, Section, Subsection, Indicator, IndicatorScore)

SECTION_RE = re.compile('(?:Section )?(\d).(\d)? (.*)')
SECTIONS = {
    1: 'Direct and Indirect Public Funding',
    2: 'Contribution and Expenditure Restrictions',
    3: 'Reporting and Public Disclosure',
    4: 'Third Party Actors',
    5: 'Monitoring and Enforcement',
}

#
# type definitions
#

CountryRecord = namedtuple('CountryRecord',
    ('name', 'column'))
SectionRecord = namedtuple('SectionRecord',
    ('section', 'subsection', 'title', 'row'))
QuestionRecord = namedtuple('QuestionRecord',
    ('number', 'text', 'description', 'comment', 'row', 'col', 'type'))
ScoreRecord = namedtuple('ScoreRecord',
    ('country', 'question', 'score', 'comment', 'sources', 'row', 'col'))

#
# data methods
#

def countries_iter(sheet):
    for col, cell in enumerate(sheet.row(1)):
        if col > 3:
            name = cell.value.split('-')[0].strip()
            yield CountryRecord(name, col)

def questions_iter(sheet):
    col = 1
    for row, cell in enumerate(sheet.col(col)):
        try:
            number = int(cell.value)
            text = sheet.cell(row, col + 1).value.strip()
            description = sheet.cell(row + 1, col + 1).value.strip()
            comment = sheet.cell(row + 2, col + 1).value.strip()

            if text.startswith('In law'):
                type = 1
            elif text.startswith('In practice'):
                type = 2
            elif text.startswith('Open Question'):
                type = 3
            else:
                type = 0

            yield QuestionRecord(number, text, description, comment, row, col, type)
        except:
            pass

def sections_iter(sheet):
    col = 0
    for row, cell in enumerate(sheet.col(col)):
        match = SECTION_RE.match(cell.value)
        if match:
            (sec, subsec, title) = match.groups()
            title = title.strip()
            yield SectionRecord(int(sec), int(subsec) if subsec else None, title, row)

def get_score(sheet, country, question):

    col = country.column
    row = question.row

    score = sheet.cell(row, col).value
    comment = sheet.cell(row + 1, col).value.strip().encode('utf-8')
    sources = sheet.cell(row + 2, col).value.strip().encode('utf-8')

    return ScoreRecord(country, question, score, comment, sources, row, col)

class Command(BaseCommand):

    help = 'Import scorecard data from master spreadsheet'

    def add_arguments(self, parser):
        parser.add_argument('path')

    def handle(self, *args, **options):

        book = xlrd.open_workbook(options['path'])
        sheet = book.sheets()[0]

        countries = list(countries_iter(sheet))
        questions = list(questions_iter(sheet))
        sections = list(sections_iter(sheet))

        def get_section(question):
            qsection = None
            for section in sections:
                if section.row > question.row:
                    break
                qsection = section
            return qsection

        #
        # load countries
        #

        print "loading countries..."

        Country.objects.all().delete()

        for record in countries:
            Country.objects.create(name=record.name,
                                   slug=slugify(record.name))

        #
        # load sections
        #

        print "loading sections and subsections..."

        Section.objects.all().delete()

        for record in sections:

            try:
                section = Section.objects.get(number=record.section)
            except Section.DoesNotExist:
                section = Section.objects.create(number=record.section,
                                                 name=SECTIONS[record.section])

            Subsection.objects.create(section=section,
                                      number=record.subsection,
                                      name=record.title)

        #
        # load indicators
        #

        print "loading indicators..."

        Indicator.objects.all().delete()

        for record in questions:

            ss = get_section(record)

            subsection = Subsection.objects.get(section__number=ss.section,
                                                number=ss.subsection)

            Indicator.objects.create(subsection=subsection,
                                     number=record.number,
                                     name=record.text,
                                     description=record.description,
                                     comment=record.comment,
                                     references='',
                                     type=record.type)

        #
        # load scores
        #

        print "loading indicator scores..."

        IndicatorScore.objects.all().delete()

        for country in countries:
            print '  {}'.format(country.name)
            for question in questions:
                score = get_score(sheet, country, question)

                value = None if score.score == '' else score.score

                iobj = Indicator.objects.get(number=question.number)
                cobj = Country.objects.get(name=country.name)

                IndicatorScore.objects.create(indicator=iobj,
                                              country=cobj,
                                              score=value,
                                              comment=score.comment,
                                              sources=score.sources)
<|endoftext|>