<commit_before>import requests
import pickle
import json

with open('sites.json', 'r') as f:
    sites = json.load(f)

for url, data in sites.items():
    with requests.Session() as session:
        result = session.post(url, data=data)
        if result.ok:
            print("logged in to {}".format(url))

with open('cookies.jar', 'wb') as f:
    pickle.dump(session.cookies, f)
<commit_msg>make this file safe for test collection<commit_after>import requests
import pickle
import json

if __name__ == "__main__":
    with open('sites.json', 'r') as f:
        sites = json.load(f)

    for url, data in sites.items():
        with requests.Session() as session:
            result = session.post(url, data=data)
            if result.ok:
                print("logged in to {}".format(url))

    with open('cookies.jar', 'wb') as f:
        pickle.dump(session.cookies, f)
<|endoftext|>


<commit_before>import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils.rnn as rnn_utils

class LSTMModule(nn.Module):
    """An LSTM-based input module"""

    def __init__(
        self,
        embed_size,
        hidden_size,
        vocab_size=None,
        embeddings=None,
        lstm_reduction="max",
        freeze=False,
        bidirectional=True,
        verbose=True,
        **lstm_kwargs,
    ):
        """
        Args:
            embed_size: The (integer) size of the input at each time
                step; usually this is the size of the embeddings
            hidden_size: The size of the hidden layer in the LSTM
            vocab_size: The size of the vocabulary of the embeddings
                If embeddings=None, this helps to set the size of the randomly
                    initilialized embeddings
                If embeddings!=None, this is used to double check that the
                    provided embeddings have the intended size
            embeddings: An optional embedding Tensor
            lstm_reduction: One of ['mean', 'max', 'last', 'attention']
                denoting what to return as the output of the LSTMLayer
            freeze: If False, allow the embeddings to be updated
        """
        super().__init__()
        self.lstm_reduction = lstm_reduction
        self.output_dim = hidden_size * 2 if bidirectional else hidden_size
        self.verbose = verbose

        # Load provided embeddings or randomly initialize new ones
        if embeddings is None:
            self.embeddings = nn.Embedding(vocab_size, embed_size)
            if self.verbose:
                print(f"Using randomly initialized embeddings.")
        else:
            self.embeddings = self._load_pretrained(embeddings)
            if self.verbose:
                print(f"Using pretrained embeddings.")

        # Freeze or not
        self.embeddings.weight.requires_grad = not freeze

        if self.verbose:
            print(
                f"Embeddings shape = ({self.embeddings.num_embeddings}, "
                f"{self.embeddings.embedding_dim})"
            )
            print(f"The embeddings are {'' if freeze else 'NOT '}FROZEN")
            print(f"Using lstm_reduction = '{lstm_reduction}'")

        # Create lstm core
        self.lstm = nn.LSTM(
            embed_size,
            hidden_size,
            batch_first=True,
            bidirectional=bidirectional,
            **lstm_kwargs,
        )
        if lstm_reduction == "attention":
            att_size = hidden_size * (self.lstm.bidirectional + 1)
            att_param = nn.Parameter(torch.FloatTensor(att_size, 1))
            nn.init.xavier_normal_(att_param)
            self.attention_param = att_param

    def _load_pretrained(self, pretrained):
        if not pretrained.dim() == 2:
            msg = (
                f"Provided embeddings have shape {pretrained.shape}. "
                "Expected a 2-dimensional tesnor."
            )
            raise ValueError(msg)
        rows, cols = pretrained.shape
        embedding = nn.Embedding(num_embeddings=rows, embedding_dim=cols)
        embedding.weight.data.copy_(pretrained)
        return embedding

    def _attention(self, output):
        # output is of shape (seq_length, hidden_size)
        score = torch.matmul(output, self.attention_param).squeeze()
        score = F.softmax(score, dim=0).view(output.size(0), 1)
        scored_output = output * score
        condensed_output = torch.sum(scored_output, dim=0)
        return condensed_output

    def reset_parameters(self):
        # Note: Classifier.reset() calls reset_parameters() recursively on all
        # children, so this method need not reset children modules such as
        # nn.lstm or nn.Embedding
        pass

    def _reduce_output(self, outputs, seq_lengths):
        """Reduces the output of an LSTM step

        Args:
            outputs: (torch.FloatTensor) the hidden state outputs from the
                lstm, with shape [batch_size, max_seq_length, hidden_size]
        """
        batch_size = outputs.shape[0]
        reduced = []
        # Necessary to iterate over batch because of different sequence lengths
        for i in range(batch_size):
            if self.lstm_reduction == "mean":
                # Average over all non-padding reduced
                # Use dim=0 because first dimension disappears after indexing
                reduced.append(outputs[i, : seq_lengths[i], :].mean(dim=0))
            elif self.lstm_reduction == "max":
                # Max-pool over all non-padding reduced
                # Use dim=0 because first dimension disappears after indexing
                reduced.append(outputs[i, : seq_lengths[i], :].max(dim=0)[0])
            elif self.lstm_reduction == "last":
                # Take the last output of the sequence (before padding starts)
                # NOTE: maybe better to take first and last?
                reduced.append(outputs[i, seq_lengths[i] - 1, :])
            elif self.lstm_reduction == "attention":
                reduced.append(self._attention(outputs[i, : seq_lengths[i], :]))
            else:
                msg = (
                    f"Did not recognize lstm kwarg 'lstm_reduction' == "
                    f"{self.lstm_reduction}"
                )
                raise ValueError(msg)
        return torch.stack(reduced, dim=0)

    def forward(self, X):
        """Applies one step of an lstm (plus reduction) to the input X

        Args:
            X: (torch.LongTensor) of shape [batch_size, max_seq_length].
                The indices of the embeddings to look up for each item in the
                batch.
        """
        # Identify the first non-zero integer from the right (i.e., the length
        # of the sequence before padding starts).
        batch_size, max_seq = X.shape
        seq_lengths = torch.zeros(batch_size, dtype=torch.long)
        for i in range(batch_size):
            for j in range(max_seq - 1, -1, -1):
                if X[i, j] != 0:
                    seq_lengths[i] = j + 1
                    break
        # Sort by length because pack_padded_sequence requires it
        # Save original order to restore before returning
        seq_lengths, perm_idx = seq_lengths.sort(0, descending=True)
        X = X[perm_idx, :]
        inv_perm_idx = torch.tensor(
            [i for i, _ in sorted(enumerate(perm_idx), key=lambda idx: idx[1])],
            dtype=torch.long,
        )

        X_encoded = self.embeddings(X)
        X_packed = rnn_utils.pack_padded_sequence(
            X_encoded, seq_lengths, batch_first=True
        )

        outputs, (h_t, c_t) = self.lstm(X_packed)

        outputs_unpacked, _ = rnn_utils.pad_packed_sequence(
            outputs, batch_first=True
        )
        reduced = self._reduce_output(outputs_unpacked, seq_lengths)
        return reduced[inv_perm_idx, :]
<commit_msg>Add option to pass features directly into LSTM and skip embeddings<commit_after>import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils.rnn as rnn_utils

class LSTMModule(nn.Module):
    """An LSTM-based input module"""

    def __init__(
        self,
        embed_size,
        hidden_size,
        vocab_size=None,
        embeddings=None,
        lstm_reduction="max",
        freeze=False,
        skip_embeddings=False,
        bidirectional=True,
        verbose=True,
        **lstm_kwargs,
    ):
        """
        Args:
            embed_size: The (integer) size of the input at each time
                step; usually this is the size of the embeddings
            hidden_size: The size of the hidden layer in the LSTM
            vocab_size: The size of the vocabulary of the embeddings
                If embeddings=None, this helps to set the size of the randomly
                    initilialized embeddings
                If embeddings!=None, this is used to double check that the
                    provided embeddings have the intended size
            embeddings: An optional embedding Tensor
            lstm_reduction: One of ['mean', 'max', 'last', 'attention']
                denoting what to return as the output of the LSTMLayer
            freeze: If False, allow the embeddings to be updated
            skip_embeddings: If True, directly accept X without using embeddings
        """
        super().__init__()
        self.lstm_reduction = lstm_reduction
        self.output_dim = hidden_size * 2 if bidirectional else hidden_size
        self.verbose = verbose
        self.skip_embeddings = skip_embeddings

        # Load provided embeddings or randomly initialize new ones
        if embeddings is None:
            self.embeddings = nn.Embedding(vocab_size, embed_size)
            if self.verbose:
                print(f"Using randomly initialized embeddings.")
        else:
            self.embeddings = self._load_pretrained(embeddings)
            if self.verbose:
                print(f"Using pretrained embeddings.")

        # Freeze or not
        self.embeddings.weight.requires_grad = not freeze

        if self.verbose:
            if self.skip_embeddings:
                print("Skipping embeddings and using direct input.")
            else:
                print(
                    f"Embeddings shape = ({self.embeddings.num_embeddings}, "
                    f"{self.embeddings.embedding_dim})"
                )
                print(f"The embeddings are {'' if freeze else 'NOT '}FROZEN")
                print(f"Using lstm_reduction = '{lstm_reduction}'")

        # Create lstm core
        self.lstm = nn.LSTM(
            embed_size,
            hidden_size,
            batch_first=True,
            bidirectional=bidirectional,
            **lstm_kwargs,
        )
        if lstm_reduction == "attention":
            att_size = hidden_size * (self.lstm.bidirectional + 1)
            att_param = nn.Parameter(torch.FloatTensor(att_size, 1))
            nn.init.xavier_normal_(att_param)
            self.attention_param = att_param

    def _load_pretrained(self, pretrained):
        if not pretrained.dim() == 2:
            msg = (
                f"Provided embeddings have shape {pretrained.shape}. "
                "Expected a 2-dimensional tesnor."
            )
            raise ValueError(msg)
        rows, cols = pretrained.shape
        embedding = nn.Embedding(num_embeddings=rows, embedding_dim=cols)
        embedding.weight.data.copy_(pretrained)
        return embedding

    def _attention(self, output):
        # output is of shape (seq_length, hidden_size)
        score = torch.matmul(output, self.attention_param).squeeze()
        score = F.softmax(score, dim=0).view(output.size(0), 1)
        scored_output = output * score
        condensed_output = torch.sum(scored_output, dim=0)
        return condensed_output

    def reset_parameters(self):
        # Note: Classifier.reset() calls reset_parameters() recursively on all
        # children, so this method need not reset children modules such as
        # nn.lstm or nn.Embedding
        pass

    def _reduce_output(self, outputs, seq_lengths):
        """Reduces the output of an LSTM step

        Args:
            outputs: (torch.FloatTensor) the hidden state outputs from the
                lstm, with shape [batch_size, max_seq_length, hidden_size]
        """
        batch_size = outputs.shape[0]
        reduced = []
        # Necessary to iterate over batch because of different sequence lengths
        for i in range(batch_size):
            if self.lstm_reduction == "mean":
                # Average over all non-padding reduced
                # Use dim=0 because first dimension disappears after indexing
                reduced.append(outputs[i, : seq_lengths[i], :].mean(dim=0))
            elif self.lstm_reduction == "max":
                # Max-pool over all non-padding reduced
                # Use dim=0 because first dimension disappears after indexing
                reduced.append(outputs[i, : seq_lengths[i], :].max(dim=0)[0])
            elif self.lstm_reduction == "last":
                # Take the last output of the sequence (before padding starts)
                # NOTE: maybe better to take first and last?
                reduced.append(outputs[i, seq_lengths[i] - 1, :])
            elif self.lstm_reduction == "attention":
                reduced.append(self._attention(outputs[i, : seq_lengths[i], :]))
            else:
                msg = (
                    f"Did not recognize lstm kwarg 'lstm_reduction' == "
                    f"{self.lstm_reduction}"
                )
                raise ValueError(msg)
        return torch.stack(reduced, dim=0)

    def forward(self, X):
        """Applies one step of an lstm (plus reduction) to the input X

        Args:
            X: (torch.LongTensor) of shape [batch_size, max_seq_length].
                The indices of the embeddings to look up for each item in the
                batch.
        """
        # Identify the first non-zero integer from the right (i.e., the length
        # of the sequence before padding starts).
        batch_size, max_seq = X.shape
        seq_lengths = torch.zeros(batch_size, dtype=torch.long)
        for i in range(batch_size):
            for j in range(max_seq - 1, -1, -1):
                if X[i, j] != 0:
                    seq_lengths[i] = j + 1
                    break
        # Sort by length because pack_padded_sequence requires it
        # Save original order to restore before returning
        seq_lengths, perm_idx = seq_lengths.sort(0, descending=True)
        X = X[perm_idx, :]
        inv_perm_idx = torch.tensor(
            [i for i, _ in sorted(enumerate(perm_idx), key=lambda idx: idx[1])],
            dtype=torch.long,
        )

        X_encoded = X if self.skip_embeddings else self.embeddings(X)
        X_packed = rnn_utils.pack_padded_sequence(
            X_encoded, seq_lengths, batch_first=True
        )

        outputs, (h_t, c_t) = self.lstm(X_packed)

        outputs_unpacked, _ = rnn_utils.pad_packed_sequence(
            outputs, batch_first=True
        )
        reduced = self._reduce_output(outputs_unpacked, seq_lengths)
        return reduced[inv_perm_idx, :]
<|endoftext|>


<commit_before># Copyright (c) 2010-2012, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

from openquake.calculators.hazard.uhs.ath import completed_task_count
from openquake.calculators.hazard.uhs.ath import remaining_tasks_in_block
from openquake.utils import stats

from tests.calculators.hazard.uhs.core_unittest import UHSBaseTestCase

class UHSTaskHandlerTestCase(UHSBaseTestCase):
    """Tests for functionality related to the asynchronous task handler code,
    which is used by the mini-framework
    :function:`openquake.utils.tasks.distribute`.
    """

    def test_completed_task_count_no_stats(self):
        # Test `complete_task_count` with no counters set;
        # it should just return 0.
        self.assertEqual(0, completed_task_count(self.job_id))

    def test_complete_task_count_success(self):
        stats.incr_counter(self.job_id, 'h', 'compute_uhs_task')
        self.assertEqual(1, completed_task_count(self.job_id))

    def test_complete_task_count_success(self):
        stats.incr_counter(self.job_id, 'h', 'compute_uhs_task-failures')
        self.assertEqual(1, completed_task_count(self.job_id))

    def test_complete_task_count_success_and_fail(self):
        # Test `complete_task_count` with success and fail counters:
        stats.incr_counter(self.job_id, 'h', 'compute_uhs_task')
        stats.incr_counter(self.job_id, 'h', 'compute_uhs_task-failures')
        self.assertEqual(2, completed_task_count(self.job_id))

    def test_remaining_tasks_in_block(self):
        # Tasks should be submitted to works for one block (of sites) at a
        # time. For each block, we want to look at Redis counters to determine
        # when the block is finished calculating.
        # `remaining_tasks_in_block` is a generator that yields the remaining
        # number of tasks in a block. When there are no more tasks left in the
        # block, a `StopIteration` is raised.
        gen = remaining_tasks_in_block(self.job_id, 4, 0)

        incr_count = lambda: stats.incr_counter(
            self.job_id, 'h', 'compute_uhs_task')

        self.assertEqual(4, gen.next())
        incr_count()
        self.assertEqual(3, gen.next())
        incr_count()
        incr_count()
        self.assertEqual(1, gen.next())
        incr_count()
        self.assertRaises(StopIteration, gen.next)

    def test_remaining_tasks_in_block_nonzero_start_count(self):
        # Same as the above test, except test with the start_count
        # set to something > 0 (to simulate a mid-calculation block).

        incr_count = lambda: stats.incr_counter(
            self.job_id, 'h', 'compute_uhs_task')

        # Just for variety, set 5 successful and 5 failed task counters:
        for _ in xrange(5):
            stats.incr_counter(self.job_id, 'h', 'compute_uhs_task')
        for _ in xrange(5):
            stats.incr_counter(self.job_id, 'h', 'compute_uhs_task-failures')

        # count starts at 10
        gen = remaining_tasks_in_block(self.job_id, 4, 10)

        self.assertEqual(4, gen.next())
        incr_count()
        self.assertEqual(3, gen.next())
        incr_count()
        incr_count()
        self.assertEqual(1, gen.next())
        incr_count()
        self.assertRaises(StopIteration, gen.next)
<commit_msg>Updated an import<commit_after># Copyright (c) 2010-2012, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

from openquake.calculators.hazard.uhs.ath import completed_task_count
from openquake.calculators.hazard.uhs.ath import remaining_tasks_in_block
from openquake.utils import stats

from tests.calculators.hazard.uhs.core_test import UHSBaseTestCase

class UHSTaskHandlerTestCase(UHSBaseTestCase):
    """Tests for functionality related to the asynchronous task handler code,
    which is used by the mini-framework
    :function:`openquake.utils.tasks.distribute`.
    """

    def test_completed_task_count_no_stats(self):
        # Test `complete_task_count` with no counters set;
        # it should just return 0.
        self.assertEqual(0, completed_task_count(self.job_id))

    def test_complete_task_count_success(self):
        stats.incr_counter(self.job_id, 'h', 'compute_uhs_task')
        self.assertEqual(1, completed_task_count(self.job_id))

    def test_complete_task_count_success(self):
        stats.incr_counter(self.job_id, 'h', 'compute_uhs_task-failures')
        self.assertEqual(1, completed_task_count(self.job_id))

    def test_complete_task_count_success_and_fail(self):
        # Test `complete_task_count` with success and fail counters:
        stats.incr_counter(self.job_id, 'h', 'compute_uhs_task')
        stats.incr_counter(self.job_id, 'h', 'compute_uhs_task-failures')
        self.assertEqual(2, completed_task_count(self.job_id))

    def test_remaining_tasks_in_block(self):
        # Tasks should be submitted to works for one block (of sites) at a
        # time. For each block, we want to look at Redis counters to determine
        # when the block is finished calculating.
        # `remaining_tasks_in_block` is a generator that yields the remaining
        # number of tasks in a block. When there are no more tasks left in the
        # block, a `StopIteration` is raised.
        gen = remaining_tasks_in_block(self.job_id, 4, 0)

        incr_count = lambda: stats.incr_counter(
            self.job_id, 'h', 'compute_uhs_task')

        self.assertEqual(4, gen.next())
        incr_count()
        self.assertEqual(3, gen.next())
        incr_count()
        incr_count()
        self.assertEqual(1, gen.next())
        incr_count()
        self.assertRaises(StopIteration, gen.next)

    def test_remaining_tasks_in_block_nonzero_start_count(self):
        # Same as the above test, except test with the start_count
        # set to something > 0 (to simulate a mid-calculation block).

        incr_count = lambda: stats.incr_counter(
            self.job_id, 'h', 'compute_uhs_task')

        # Just for variety, set 5 successful and 5 failed task counters:
        for _ in xrange(5):
            stats.incr_counter(self.job_id, 'h', 'compute_uhs_task')
        for _ in xrange(5):
            stats.incr_counter(self.job_id, 'h', 'compute_uhs_task-failures')

        # count starts at 10
        gen = remaining_tasks_in_block(self.job_id, 4, 10)

        self.assertEqual(4, gen.next())
        incr_count()
        self.assertEqual(3, gen.next())
        incr_count()
        incr_count()
        self.assertEqual(1, gen.next())
        incr_count()
        self.assertRaises(StopIteration, gen.next)
<|endoftext|>


<commit_before>from __future__ import absolute_import
from abc import ABCMeta, abstractmethod
from django.test import SimpleTestCase
import mock
from functools import partial

from corehq.apps.locations.tests.util import LocationHierarchyTestCase
from corehq.apps.es.fake.groups_fake import GroupESFake
from corehq.apps.es.fake.users_fake import UserESFake
from corehq.apps.groups.models import Group
from corehq.apps.locations.tests.util import delete_all_locations
from corehq.apps.users.dbaccessors.all_commcare_users import delete_all_users
from corehq.apps.reports_core.filters import Choice
from corehq.apps.userreports.models import ReportConfiguration
from corehq.apps.userreports.reports.filters.choice_providers import (
    ChoiceQueryContext, LocationChoiceProvider, UserChoiceProvider, GroupChoiceProvider,
    OwnerChoiceProvider, StaticChoiceProvider, SearchableChoice)
from corehq.apps.users.models import CommCareUser, WebUser, DomainMembership
from corehq.apps.users.util import normalize_username
import six

class StaticChoiceProviderTest(SimpleTestCase):

    @classmethod
    def setUpClass(cls):
        super(StaticChoiceProviderTest, cls).setUpClass()
        cls.choice_provider = StaticChoiceProvider([Choice('1', 'One'), Choice('2', 'Two'), Choice('3', 'Three')])

    def test_query_no_search(self):
        self.assertEqual(self.choice_provider.query(ChoiceQueryContext('', 2, page=0)),
                         [Choice('1', 'One'), Choice('2', 'Two')])
        self.assertEqual(self.choice_provider.query(ChoiceQueryContext('', 2, page=1)),
                         [Choice('3', 'Three')])

    def test_get_choices_for_values(self):
        self.assertEqual(
            set(self.choice_provider.get_choices_for_values(['2', '4', '6'], None)),
            {Choice('2', 'Two'), Choice('4', '4'), Choice('6', '6')}
        )

class ChoiceProviderTestMixin(six.with_metaclass(ABCMeta, object)):
    """
    A mixin for a creating uniform tests for different ChoiceProvider subclasses.

    ChoiceProviderTestMixin creates a simple framework in which the real choice provider
    being tested gets compared to a static choice provider (one that simply filters,
    sorts, and slices Choice objects in memory), and verifies they produce the same result.

    Classes that want to use this framework must

    1. subclass TestCase/SimpleTestCase and ChoiceProviderTestMixin
    2. Initialize `choice_provider` and the data it requires in setUpClass
    3. Initialize `static_choice_provider` with Choices (or SearchableChoices)
       that correspond to the data set up in part (2) in setUpClass
    4. Implement the abstract test methods according to the suggestions in the docstrings

    """
    choice_provider = None
    static_choice_provider = None
    choice_query_context = ChoiceQueryContext

    def _test_query(self, query_context):
        self.assertEqual(
            self.choice_provider.query(query_context),
            self.static_choice_provider.query(query_context))

    def _test_get_choices_for_values(self, values, user):
        self.assertEqual(
            self.choice_provider.get_choices_for_values(values, user),
            self.static_choice_provider.get_choices_for_values(values, user),
        )

    def test_query_no_search_all(self):
        self._test_query(self.choice_query_context('', limit=20, page=0))

    def test_query_no_search_first_short_page(self):
        self._test_query(self.choice_query_context('', 2, page=0))

    def test_query_no_search_second_short_page(self):
        self._test_query(self.choice_query_context('', 2, page=1))

    @abstractmethod
    def test_query_search(self):
        """
        Suggested implementation:

            self._test_query(ChoiceQueryContext('relevant_search_term', 2, page=1))

        """
        pass

    @abstractmethod
    def test_get_choices_for_values(self):
        """
        Suggested implementation:

            self._test_get_choices_for_values(
                [irrelevant_value, relevant_value, relevant_value],
                web_user,
            )
        """
        pass

class LocationChoiceProviderTest(ChoiceProviderTestMixin, LocationHierarchyTestCase):
    domain = 'location-choice-provider'
    location_type_names = ['state', 'county', 'city']
    location_structure = [
        ('Massachusetts', [
            ('Middlesex', [
                ('Cambridge', []),
                ('Somerville', []),
            ]),
            ('Suffolke', [      # Make all locations contain the letter 'e'
                ('Bostone', []),
            ])
        ])
    ]

    @classmethod
    def setUpClass(cls):
        delete_all_locations()
        delete_all_users()
        super(LocationChoiceProviderTest, cls).setUpClass()
        report = ReportConfiguration(domain=cls.domain)
        choice_tuples = [
            (location.name, SearchableChoice(
                location.location_id,
                location.get_path_display(),
                searchable_text=[location.site_code, location.name]
            ))
            for location in six.itervalues(cls.locations)
        ]
        choice_tuples.sort()
        choices = [choice for name, choice in choice_tuples]
        cls.web_user = WebUser.create(cls.domain, 'blah', 'password')
        cls.choice_provider = LocationChoiceProvider(report, None)
        cls.choice_provider.configure({
            "include_descendants": False,
            "show_full_path": True,
        })
        cls.static_choice_provider = StaticChoiceProvider(choices)
        cls.choice_query_context = partial(ChoiceQueryContext, user=cls.web_user)

    @classmethod
    def tearDownClass(cls):
        cls.domain_obj.delete()
        delete_all_locations()

    def test_query_search(self):
        # Searching for something common to all locations gets you all locations
        self._test_query(self.choice_query_context('e', page=0))
        self._test_query(self.choice_query_context('e', page=1))

    def test_get_choices_for_values(self):
        self._test_get_choices_for_values(
            ['made-up', self.locations['Cambridge'].location_id, self.locations['Middlesex'].location_id],
            self.web_user,
        )

    def test_scoped_to_location_search(self):
        self.web_user.set_location(self.domain, self.locations['Middlesex'])
        self.restrict_user_to_assigned_locations(self.web_user)
        scoped_choices = [
            SearchableChoice(
                location.location_id,
                location.get_path_display(),
                searchable_text=[location.site_code, location.name]
            )
            for location in [
                self.locations['Cambridge'],
                self.locations['Middlesex'],
                self.locations['Somerville'],
            ]
        ]
        self.static_choice_provider = StaticChoiceProvider(scoped_choices)

        # When an empty query is given, the user receives all the choices they can access
        self._test_query(self.choice_query_context('', page=0))
        # Searching for something common to all locations give only the accessible locations
        self._test_query(self.choice_query_context('e', page=0))
        # When a user queries for something they can access, it gets returned
        self._test_query(self.choice_query_context('Somerville', page=0))
        # When a user searches for something they can't access, it isn't returned
        self._test_query(self.choice_query_context('Boston', page=0))

    def test_cant_circumvent_restrictions(self):
        # Boston should be inaccessible to this user, if they edit the URL to
        # search for it anyways, it shouldn't be returned
        self.web_user.set_location(self.domain, self.locations['Middlesex'])
        self.restrict_user_to_assigned_locations(self.web_user)
        loc_ids_by_name = {name: loc.location_id for name, loc in self.locations.items()}
        choices = self.choice_provider.get_choices_for_known_values([
            loc_ids_by_name['Cambridge'], loc_ids_by_name['Somerville'], loc_ids_by_name['Bostone'],
        ], self.web_user)
        self.assertItemsEqual(
            [choice.value for choice in choices],
            [loc_ids_by_name["Cambridge"], loc_ids_by_name["Somerville"]]
        )
        self.web_user.set_role(self.domain, 'none')
        self.web_user.unset_location(self.domain)
        self.web_user.reset_locations(self.domain, [])

@mock.patch('corehq.apps.users.analytics.UserES', UserESFake)
@mock.patch('corehq.apps.userreports.reports.filters.choice_providers.UserES', UserESFake)
class UserChoiceProviderTest(SimpleTestCase, ChoiceProviderTestMixin):
    domain = 'user-choice-provider'

    @classmethod
    def make_mobile_worker(cls, username, domain=None):
        domain = domain or cls.domain
        user = CommCareUser(username=normalize_username(username, domain),
                            domain=domain)
        user.domain_membership = DomainMembership(domain=domain)
        doc = user._doc
        doc['username.exact'] = doc['username']
        doc['base_username'] = username
        UserESFake.save_doc(doc)
        return user

    @classmethod
    def make_web_user(cls, email, domain=None):
        domain = domain or cls.domain
        domains = [domain]
        user = WebUser(username=email, domains=domains)
        user.domain_memberships = [DomainMembership(domain=domain)]
        doc = user._doc
        doc['username.exact'] = doc['username']
        doc['base_username'] = email
        UserESFake.save_doc(doc)
        return user

    @classmethod
    def setUpClass(cls):
        report = ReportConfiguration(domain=cls.domain)

        cls.web_user = cls.make_web_user('candice@example.com')
        cls.users = [
            cls.make_mobile_worker('bernice'),
            cls.web_user,
            cls.make_mobile_worker('dennis'),
            cls.make_mobile_worker('elizabeth'),
            cls.make_mobile_worker('albert'),
            # test that docs in other domains are filtered out
            cls.make_mobile_worker('aaa', domain='some-other-domain'),
        ]
        choices = [
            SearchableChoice(
                user.get_id, user.raw_username,
                searchable_text=[user.username, user.last_name, user.first_name])
            for user in cls.users if user.is_member_of(cls.domain)
        ]
        choices.sort(key=lambda choice: choice.display)
        cls.choice_provider = UserChoiceProvider(report, None)
        cls.static_choice_provider = StaticChoiceProvider(choices)

    @classmethod
    def tearDownClass(cls):
        UserESFake.reset_docs()

    def test_query_search(self):
        self._test_query(ChoiceQueryContext(query='ni', limit=10, page=0))

    def test_get_choices_for_values(self):
        self._test_get_choices_for_values(
            ['unknown-user'] + [user._id for user in self.users],
            self.web_user,
        )

@mock.patch('corehq.apps.userreports.reports.filters.choice_providers.GroupES', GroupESFake)
class GroupChoiceProviderTest(SimpleTestCase, ChoiceProviderTestMixin):
    domain = 'group-choice-provider'

    @classmethod
    def make_group(cls, name, domain=None):
        domain = domain or cls.domain
        group = Group(name=name, domain=domain, case_sharing=True)
        GroupESFake.save_doc(group._doc)
        return group

    @classmethod
    def setUpClass(cls):
        report = ReportConfiguration(domain=cls.domain)

        cls.groups = [
            cls.make_group('Team B no'),
            cls.make_group('Team C no'),
            cls.make_group('Team D', domain='not-this-domain'),
            cls.make_group('Team E yes'),
            cls.make_group('Team A yes'),
        ]
        choices = [
            SearchableChoice(group.get_id, group.name, searchable_text=[group.name])
            for group in cls.groups if group.domain == cls.domain]
        choices.sort(key=lambda choice: choice.display)
        cls.choice_provider = GroupChoiceProvider(report, None)
        cls.static_choice_provider = StaticChoiceProvider(choices)
        cls.web_user = UserChoiceProviderTest.make_web_user('web-user@example.com', domain=cls.domain)

    @classmethod
    def tearDownClass(cls):
        GroupESFake.reset_docs()

    def test_query_search(self):
        self._test_query(ChoiceQueryContext('yes', limit=10, page=0))

    def test_get_choices_for_values(self):
        self._test_get_choices_for_values(
            ['unknown-group'] + [group.get_id for group in self.groups],
            self.web_user,
        )

@mock.patch('corehq.apps.users.analytics.UserES', UserESFake)
@mock.patch('corehq.apps.userreports.reports.filters.choice_providers.UserES', UserESFake)
@mock.patch('corehq.apps.userreports.reports.filters.choice_providers.GroupES', GroupESFake)
class OwnerChoiceProviderTest(LocationHierarchyTestCase, ChoiceProviderTestMixin):
    domain = 'owner-choice-provider'
    location_type_names = ['state']
    location_structure = [('Massachusetts', [])]

    @classmethod
    def setUpClass(cls):
        super(OwnerChoiceProviderTest, cls).setUpClass()
        report = ReportConfiguration(domain=cls.domain)
        cls.group = GroupChoiceProviderTest.make_group('Group', domain=cls.domain)
        cls.mobile_worker = UserChoiceProviderTest.make_mobile_worker('mobile-worker', domain=cls.domain)
        cls.web_user = UserChoiceProviderTest.make_web_user('web-user@example.com', domain=cls.domain)
        cls.location = cls.locations['Massachusetts']
        cls.docs = [cls.group, cls.mobile_worker, cls.web_user, cls.location]
        cls.choices = [
            SearchableChoice(cls.group.get_id, cls.group.name, [cls.group.name]),
            SearchableChoice(cls.mobile_worker.get_id, cls.mobile_worker.raw_username,
                             [cls.mobile_worker.username]),
            SearchableChoice(cls.web_user.get_id, cls.web_user.username,
                             [cls.web_user.username]),
            SearchableChoice(cls.location.location_id, cls.location.display_name,
                             [cls.location.name, cls.location.site_code]),
        ]
        cls.choice_provider = OwnerChoiceProvider(report, None)
        cls.static_choice_provider = StaticChoiceProvider(cls.choices)
        cls.choice_query_context = partial(ChoiceQueryContext, user=cls.mobile_worker)

    @classmethod
    def tearDownClass(cls):
        UserESFake.reset_docs()
        GroupESFake.reset_docs()
        cls.domain_obj.delete()
        delete_all_locations()

    def test_query_search(self):
        self._test_query(self.choice_query_context('o', limit=10, offset=0))
        self._test_query(self.choice_query_context('l', limit=10, offset=0))
        self._test_query(self.choice_query_context('no-match', limit=10, offset=0))
        self._test_query(self.choice_query_context('o', limit=3, offset=2))

    def test_get_choices_for_values(self):
        self._test_get_choices_for_values(
            ['unknown', self.group._id, self.web_user._id, self.location.location_id,
             self.mobile_worker._id],
            self.web_user,
        )
<commit_msg>Don't cleanup before test<commit_after>from __future__ import absolute_import
from abc import ABCMeta, abstractmethod
from django.test import SimpleTestCase
import mock
from functools import partial

from corehq.apps.locations.tests.util import LocationHierarchyTestCase
from corehq.apps.es.fake.groups_fake import GroupESFake
from corehq.apps.es.fake.users_fake import UserESFake
from corehq.apps.groups.models import Group
from corehq.apps.locations.tests.util import delete_all_locations
from corehq.apps.users.dbaccessors.all_commcare_users import delete_all_users
from corehq.apps.reports_core.filters import Choice
from corehq.apps.userreports.models import ReportConfiguration
from corehq.apps.userreports.reports.filters.choice_providers import (
    ChoiceQueryContext, LocationChoiceProvider, UserChoiceProvider, GroupChoiceProvider,
    OwnerChoiceProvider, StaticChoiceProvider, SearchableChoice)
from corehq.apps.users.models import CommCareUser, WebUser, DomainMembership
from corehq.apps.users.util import normalize_username
import six

class StaticChoiceProviderTest(SimpleTestCase):

    @classmethod
    def setUpClass(cls):
        super(StaticChoiceProviderTest, cls).setUpClass()
        cls.choice_provider = StaticChoiceProvider([Choice('1', 'One'), Choice('2', 'Two'), Choice('3', 'Three')])

    def test_query_no_search(self):
        self.assertEqual(self.choice_provider.query(ChoiceQueryContext('', 2, page=0)),
                         [Choice('1', 'One'), Choice('2', 'Two')])
        self.assertEqual(self.choice_provider.query(ChoiceQueryContext('', 2, page=1)),
                         [Choice('3', 'Three')])

    def test_get_choices_for_values(self):
        self.assertEqual(
            set(self.choice_provider.get_choices_for_values(['2', '4', '6'], None)),
            {Choice('2', 'Two'), Choice('4', '4'), Choice('6', '6')}
        )

class ChoiceProviderTestMixin(six.with_metaclass(ABCMeta, object)):
    """
    A mixin for a creating uniform tests for different ChoiceProvider subclasses.

    ChoiceProviderTestMixin creates a simple framework in which the real choice provider
    being tested gets compared to a static choice provider (one that simply filters,
    sorts, and slices Choice objects in memory), and verifies they produce the same result.

    Classes that want to use this framework must

    1. subclass TestCase/SimpleTestCase and ChoiceProviderTestMixin
    2. Initialize `choice_provider` and the data it requires in setUpClass
    3. Initialize `static_choice_provider` with Choices (or SearchableChoices)
       that correspond to the data set up in part (2) in setUpClass
    4. Implement the abstract test methods according to the suggestions in the docstrings

    """
    choice_provider = None
    static_choice_provider = None
    choice_query_context = ChoiceQueryContext

    def _test_query(self, query_context):
        self.assertEqual(
            self.choice_provider.query(query_context),
            self.static_choice_provider.query(query_context))

    def _test_get_choices_for_values(self, values, user):
        self.assertEqual(
            self.choice_provider.get_choices_for_values(values, user),
            self.static_choice_provider.get_choices_for_values(values, user),
        )

    def test_query_no_search_all(self):
        self._test_query(self.choice_query_context('', limit=20, page=0))

    def test_query_no_search_first_short_page(self):
        self._test_query(self.choice_query_context('', 2, page=0))

    def test_query_no_search_second_short_page(self):
        self._test_query(self.choice_query_context('', 2, page=1))

    @abstractmethod
    def test_query_search(self):
        """
        Suggested implementation:

            self._test_query(ChoiceQueryContext('relevant_search_term', 2, page=1))

        """
        pass

    @abstractmethod
    def test_get_choices_for_values(self):
        """
        Suggested implementation:

            self._test_get_choices_for_values(
                [irrelevant_value, relevant_value, relevant_value],
                web_user,
            )
        """
        pass

class LocationChoiceProviderTest(ChoiceProviderTestMixin, LocationHierarchyTestCase):
    domain = 'location-choice-provider'
    location_type_names = ['state', 'county', 'city']
    location_structure = [
        ('Massachusetts', [
            ('Middlesex', [
                ('Cambridge', []),
                ('Somerville', []),
            ]),
            ('Suffolke', [      # Make all locations contain the letter 'e'
                ('Bostone', []),
            ])
        ])
    ]

    @classmethod
    def setUpClass(cls):
        super(LocationChoiceProviderTest, cls).setUpClass()
        report = ReportConfiguration(domain=cls.domain)
        choice_tuples = [
            (location.name, SearchableChoice(
                location.location_id,
                location.get_path_display(),
                searchable_text=[location.site_code, location.name]
            ))
            for location in six.itervalues(cls.locations)
        ]
        choice_tuples.sort()
        choices = [choice for name, choice in choice_tuples]
        cls.web_user = WebUser.create(cls.domain, 'blah', 'password')
        cls.choice_provider = LocationChoiceProvider(report, None)
        cls.choice_provider.configure({
            "include_descendants": False,
            "show_full_path": True,
        })
        cls.static_choice_provider = StaticChoiceProvider(choices)
        cls.choice_query_context = partial(ChoiceQueryContext, user=cls.web_user)

    @classmethod
    def tearDownClass(cls):
        cls.domain_obj.delete()
        delete_all_locations()
        delete_all_users()

    def test_query_search(self):
        # Searching for something common to all locations gets you all locations
        self._test_query(self.choice_query_context('e', page=0))
        self._test_query(self.choice_query_context('e', page=1))

    def test_get_choices_for_values(self):
        self._test_get_choices_for_values(
            ['made-up', self.locations['Cambridge'].location_id, self.locations['Middlesex'].location_id],
            self.web_user,
        )

    def test_scoped_to_location_search(self):
        self.web_user.set_location(self.domain, self.locations['Middlesex'])
        self.restrict_user_to_assigned_locations(self.web_user)
        scoped_choices = [
            SearchableChoice(
                location.location_id,
                location.get_path_display(),
                searchable_text=[location.site_code, location.name]
            )
            for location in [
                self.locations['Cambridge'],
                self.locations['Middlesex'],
                self.locations['Somerville'],
            ]
        ]
        self.static_choice_provider = StaticChoiceProvider(scoped_choices)

        # When an empty query is given, the user receives all the choices they can access
        self._test_query(self.choice_query_context('', page=0))
        # Searching for something common to all locations give only the accessible locations
        self._test_query(self.choice_query_context('e', page=0))
        # When a user queries for something they can access, it gets returned
        self._test_query(self.choice_query_context('Somerville', page=0))
        # When a user searches for something they can't access, it isn't returned
        self._test_query(self.choice_query_context('Boston', page=0))

    def test_cant_circumvent_restrictions(self):
        # Boston should be inaccessible to this user, if they edit the URL to
        # search for it anyways, it shouldn't be returned
        self.web_user.set_location(self.domain, self.locations['Middlesex'])
        self.restrict_user_to_assigned_locations(self.web_user)
        loc_ids_by_name = {name: loc.location_id for name, loc in self.locations.items()}
        choices = self.choice_provider.get_choices_for_known_values([
            loc_ids_by_name['Cambridge'], loc_ids_by_name['Somerville'], loc_ids_by_name['Bostone'],
        ], self.web_user)
        self.assertItemsEqual(
            [choice.value for choice in choices],
            [loc_ids_by_name["Cambridge"], loc_ids_by_name["Somerville"]]
        )
        self.web_user.set_role(self.domain, 'none')
        self.web_user.unset_location(self.domain)
        self.web_user.reset_locations(self.domain, [])

@mock.patch('corehq.apps.users.analytics.UserES', UserESFake)
@mock.patch('corehq.apps.userreports.reports.filters.choice_providers.UserES', UserESFake)
class UserChoiceProviderTest(SimpleTestCase, ChoiceProviderTestMixin):
    domain = 'user-choice-provider'

    @classmethod
    def make_mobile_worker(cls, username, domain=None):
        domain = domain or cls.domain
        user = CommCareUser(username=normalize_username(username, domain),
                            domain=domain)
        user.domain_membership = DomainMembership(domain=domain)
        doc = user._doc
        doc['username.exact'] = doc['username']
        doc['base_username'] = username
        UserESFake.save_doc(doc)
        return user

    @classmethod
    def make_web_user(cls, email, domain=None):
        domain = domain or cls.domain
        domains = [domain]
        user = WebUser(username=email, domains=domains)
        user.domain_memberships = [DomainMembership(domain=domain)]
        doc = user._doc
        doc['username.exact'] = doc['username']
        doc['base_username'] = email
        UserESFake.save_doc(doc)
        return user

    @classmethod
    def setUpClass(cls):
        report = ReportConfiguration(domain=cls.domain)

        cls.web_user = cls.make_web_user('candice@example.com')
        cls.users = [
            cls.make_mobile_worker('bernice'),
            cls.web_user,
            cls.make_mobile_worker('dennis'),
            cls.make_mobile_worker('elizabeth'),
            cls.make_mobile_worker('albert'),
            # test that docs in other domains are filtered out
            cls.make_mobile_worker('aaa', domain='some-other-domain'),
        ]
        choices = [
            SearchableChoice(
                user.get_id, user.raw_username,
                searchable_text=[user.username, user.last_name, user.first_name])
            for user in cls.users if user.is_member_of(cls.domain)
        ]
        choices.sort(key=lambda choice: choice.display)
        cls.choice_provider = UserChoiceProvider(report, None)
        cls.static_choice_provider = StaticChoiceProvider(choices)

    @classmethod
    def tearDownClass(cls):
        UserESFake.reset_docs()

    def test_query_search(self):
        self._test_query(ChoiceQueryContext(query='ni', limit=10, page=0))

    def test_get_choices_for_values(self):
        self._test_get_choices_for_values(
            ['unknown-user'] + [user._id for user in self.users],
            self.web_user,
        )

@mock.patch('corehq.apps.userreports.reports.filters.choice_providers.GroupES', GroupESFake)
class GroupChoiceProviderTest(SimpleTestCase, ChoiceProviderTestMixin):
    domain = 'group-choice-provider'

    @classmethod
    def make_group(cls, name, domain=None):
        domain = domain or cls.domain
        group = Group(name=name, domain=domain, case_sharing=True)
        GroupESFake.save_doc(group._doc)
        return group

    @classmethod
    def setUpClass(cls):
        report = ReportConfiguration(domain=cls.domain)

        cls.groups = [
            cls.make_group('Team B no'),
            cls.make_group('Team C no'),
            cls.make_group('Team D', domain='not-this-domain'),
            cls.make_group('Team E yes'),
            cls.make_group('Team A yes'),
        ]
        choices = [
            SearchableChoice(group.get_id, group.name, searchable_text=[group.name])
            for group in cls.groups if group.domain == cls.domain]
        choices.sort(key=lambda choice: choice.display)
        cls.choice_provider = GroupChoiceProvider(report, None)
        cls.static_choice_provider = StaticChoiceProvider(choices)
        cls.web_user = UserChoiceProviderTest.make_web_user('web-user@example.com', domain=cls.domain)

    @classmethod
    def tearDownClass(cls):
        GroupESFake.reset_docs()

    def test_query_search(self):
        self._test_query(ChoiceQueryContext('yes', limit=10, page=0))

    def test_get_choices_for_values(self):
        self._test_get_choices_for_values(
            ['unknown-group'] + [group.get_id for group in self.groups],
            self.web_user,
        )

@mock.patch('corehq.apps.users.analytics.UserES', UserESFake)
@mock.patch('corehq.apps.userreports.reports.filters.choice_providers.UserES', UserESFake)
@mock.patch('corehq.apps.userreports.reports.filters.choice_providers.GroupES', GroupESFake)
class OwnerChoiceProviderTest(LocationHierarchyTestCase, ChoiceProviderTestMixin):
    domain = 'owner-choice-provider'
    location_type_names = ['state']
    location_structure = [('Massachusetts', [])]

    @classmethod
    def setUpClass(cls):
        super(OwnerChoiceProviderTest, cls).setUpClass()
        report = ReportConfiguration(domain=cls.domain)
        cls.group = GroupChoiceProviderTest.make_group('Group', domain=cls.domain)
        cls.mobile_worker = UserChoiceProviderTest.make_mobile_worker('mobile-worker', domain=cls.domain)
        cls.web_user = UserChoiceProviderTest.make_web_user('web-user@example.com', domain=cls.domain)
        cls.location = cls.locations['Massachusetts']
        cls.docs = [cls.group, cls.mobile_worker, cls.web_user, cls.location]
        cls.choices = [
            SearchableChoice(cls.group.get_id, cls.group.name, [cls.group.name]),
            SearchableChoice(cls.mobile_worker.get_id, cls.mobile_worker.raw_username,
                             [cls.mobile_worker.username]),
            SearchableChoice(cls.web_user.get_id, cls.web_user.username,
                             [cls.web_user.username]),
            SearchableChoice(cls.location.location_id, cls.location.display_name,
                             [cls.location.name, cls.location.site_code]),
        ]
        cls.choice_provider = OwnerChoiceProvider(report, None)
        cls.static_choice_provider = StaticChoiceProvider(cls.choices)
        cls.choice_query_context = partial(ChoiceQueryContext, user=cls.mobile_worker)

    @classmethod
    def tearDownClass(cls):
        UserESFake.reset_docs()
        GroupESFake.reset_docs()
        cls.domain_obj.delete()
        delete_all_locations()

    def test_query_search(self):
        self._test_query(self.choice_query_context('o', limit=10, offset=0))
        self._test_query(self.choice_query_context('l', limit=10, offset=0))
        self._test_query(self.choice_query_context('no-match', limit=10, offset=0))
        self._test_query(self.choice_query_context('o', limit=3, offset=2))

    def test_get_choices_for_values(self):
        self._test_get_choices_for_values(
            ['unknown', self.group._id, self.web_user._id, self.location.location_id,
             self.mobile_worker._id],
            self.web_user,
        )
<|endoftext|>


<commit_before>from corehq.util.elastic import ensure_production_index_deleted
from pillowtop.es_utils import set_index_reindex_settings, \
    set_index_normal_settings, get_index_info_from_pillow, initialize_mapping_if_necessary
from pillowtop.pillow.interface import PillowRuntimeContext

class PillowReindexer(object):

    def __init__(self, pillow, change_provider):
        self.pillow = pillow
        self.change_provider = change_provider

    def reindex(self, start_from=None):
        reindexer_context = PillowRuntimeContext(do_set_checkpoint=False)
        for change in self.change_provider.iter_changes(start_from=start_from):
            self.pillow.processor(change, reindexer_context)

class ElasticPillowReindexer(PillowReindexer):

    def __init__(self, pillow, change_provider, elasticsearch, index_info):
        super(ElasticPillowReindexer, self).__init__(pillow, change_provider)
        self.es = elasticsearch
        self.index_info = index_info

    def reindex(self, start_from=None):
        if not start_from:
            # when not resuming force delete and create the index
            self._delete_and_prepare_index_for_reindex()

        super(ElasticPillowReindexer, self).reindex(start_from)

        self._prepare_index_for_usage()

    def _delete_and_prepare_index_for_reindex(self):
        ensure_production_index_deleted(self.index_info.index)
        self.es.indices.create(index=self.index_info.index, body=self.index_info.meta)
        initialize_mapping_if_necessary(self.es, self.index_info)
        set_index_reindex_settings(self.es, self.index_info.index)

    def _prepare_index_for_usage(self):
        set_index_normal_settings(self.es, self.index_info.index)
        self.es.indices.refresh(self.index_info.index)

def get_default_reindexer_for_elastic_pillow(pillow, change_provider):
    return ElasticPillowReindexer(
        pillow=pillow,
        change_provider=change_provider,
        elasticsearch=pillow.get_es_new(),
        index_info=get_index_info_from_pillow(pillow),
    )
<commit_msg>don't delete the index during reindexing<commit_after>from pillowtop.es_utils import set_index_reindex_settings, \
    set_index_normal_settings, get_index_info_from_pillow, initialize_mapping_if_necessary
from pillowtop.pillow.interface import PillowRuntimeContext

class PillowReindexer(object):

    def __init__(self, pillow, change_provider):
        self.pillow = pillow
        self.change_provider = change_provider

    def reindex(self, start_from=None):
        reindexer_context = PillowRuntimeContext(do_set_checkpoint=False)
        for change in self.change_provider.iter_changes(start_from=start_from):
            self.pillow.processor(change, reindexer_context)

class ElasticPillowReindexer(PillowReindexer):

    def __init__(self, pillow, change_provider, elasticsearch, index_info):
        super(ElasticPillowReindexer, self).__init__(pillow, change_provider)
        self.es = elasticsearch
        self.index_info = index_info

    def reindex(self, start_from=None):
        if not start_from:
            # when not resuming force delete and create the index
            self._prepare_index_for_reindex()

        super(ElasticPillowReindexer, self).reindex(start_from)

        self._prepare_index_for_usage()

    def _prepare_index_for_reindex(self):
        if not self.es.indices.exists(self.index_info.index):
            self.es.indices.create(index=self.index_info.index, body=self.index_info.meta)
        initialize_mapping_if_necessary(self.es, self.index_info)
        set_index_reindex_settings(self.es, self.index_info.index)

    def _prepare_index_for_usage(self):
        set_index_normal_settings(self.es, self.index_info.index)
        self.es.indices.refresh(self.index_info.index)

def get_default_reindexer_for_elastic_pillow(pillow, change_provider):
    return ElasticPillowReindexer(
        pillow=pillow,
        change_provider=change_provider,
        elasticsearch=pillow.get_es_new(),
        index_info=get_index_info_from_pillow(pillow),
    )
<|endoftext|>


<commit_before>"""
The upgraded CMS online admin area.

This is an enhanced version of the Django admin area, providing a more
user-friendly appearance and providing additional functionality over the
standard implementation.
"""

from __future__ import with_statement
import threading

from django.core.exceptions import PermissionDenied
from django.contrib.contenttypes.models import ContentType
from django.http import Http404, HttpResponseRedirect
from django.shortcuts import render, redirect
from django import forms

import reversion

from cms.core.admin import PageBaseAdmin, site, PAGE_FROM_KEY, PAGE_FROM_SITEMAP_VALUE
from cms.core.db import locked
from cms.apps.historylinks.models import HistoryLink
from cms.apps.pages.models import Page, get_registered_content

# The GET parameter used to indicate content type.
PAGE_TYPE_PARAMETER = "type"

# HACK: Used to provide thread safety for the patch inlines hack (below).
page_inlines_lock = threading.RLock()

class PageAdmin(PageBaseAdmin):

    """Admin settings for Page models."""

    fieldsets = (
        (None, {
            "fields": ("title", "url_title", "parent",),
        },),
        ("Publication", {
            "fields": ("publication_date", "expiry_date", "is_online",),
            "classes": ("collapse",),
        }),
        ("Navigation", {
            "fields": ("short_title", "permalink", "in_navigation",),
            "classes": ("collapse",),
        }),
        ("Search engine optimization", {
            "fields": ("browser_title", "meta_keywords", "meta_description", "robots_index", "robots_follow", "robots_archive", "sitemap_priority", "sitemap_changefreq",),
            "classes": ("collapse",),
        }),
    )

    def __init__(self, *args, **kwargs):
        """Initialzies the PageAdmin."""
        super(PageAdmin, self).__init__(*args, **kwargs)
        # Register all content classes with reversion.
        self.content_inline_instances = []
        for content_cls in get_registered_content():
            self._autoregister(content_cls, follow=["page"])

    def register_content_inline(self, content_cls, inline, inline_admin):
        """Registers an inline model with the page admin."""
        # Register the admin class.
        self.inlines = list(self.inlines)
        self.inlines.append(inline_admin)
        inline_instance = inline_admin(self.model, self.admin_site)
        self.content_inline_instances.append((content_cls, inline_instance))
        self.inline_instances.append(inline_instance)
        # Register with reversion.
        self._autoregister(inline, follow=["page"])

    def queryset(self, request):
        """Only allows editing of pages in this site."""
        return request.pages.all

    # Reversion

    def get_revision_form_data(self, request, obj, version):
        """
        Returns a dictionary of data to set in the admin form in order to revert
        to the given revision.
        """
        data = super(PageAdmin, self).get_revision_form_data(request, obj, version)
        content_version = version.revision.version_set.all().get(content_type=obj.content_type_id)
        data.update(content_version.field_dict)
        return data

    # Plugable content types.

    def get_page_content_cls(self, request, obj=None):
        """Retrieves the page content type slug."""
        if PAGE_TYPE_PARAMETER in request.GET:
            return ContentType.objects.get_for_id(request.GET[PAGE_TYPE_PARAMETER]).model_class()
        if obj and obj.content_type:
            return obj.content_type.model_class()
        raise Http404, "You must specify a page content type."

    def get_fieldsets(self, request, obj=None):
        """Generates the custom content fieldsets."""
        content_cls = self.get_page_content_cls(request, obj)
        content_fields = [field.name for field in content_cls._meta.fields if field.name != "page"]
        fieldsets = super(PageAdmin, self).get_fieldsets(request, obj)
        if content_fields:
            content_fieldsets = content_cls.fieldsets or (
                ("Page content", {
                    "fields": content_fields,
                }),
            )
            fieldsets = tuple(fieldsets[0:1]) + content_fieldsets + tuple(fieldsets[1:])
        return fieldsets

    def get_form(self, request, obj=None, **kwargs):
        """Adds the template area fields to the form."""
        content_cls = self.get_page_content_cls(request, obj)
        form_attrs = {}
        for field in content_cls._meta.fields:
            if field.name == "page":
                continue
            form_field = self.formfield_for_dbfield(field, request=request)
            # HACK: add in the initial value. No other way to pass this on.
            if obj:
                try:
                    form_field.initial = getattr(obj.content, field.name, "")
                except content_cls.DoesNotExist:
                    pass  # This means that we're in a reversion recovery, or something weird has happened to the databae.
            # Store the field.
            form_attrs[field.name] = form_field
        ContentForm = type("%sForm" % self.__class__.__name__, (forms.ModelForm,), form_attrs)
        defaults = {"form": ContentForm}
        defaults.update(kwargs)
        PageForm = super(PageAdmin, self).get_form(request, obj, **defaults)
        # HACK: Need to limit parents field based on object. This should be done in
        # formfield_for_foreignkey, but that method does not know about the object instance.
        if obj:
            invalid_parents = set(child.id for child in obj.all_children)
            invalid_parents.add(obj.id)
        else:
            invalid_parents = frozenset()
        homepage = request.pages.homepage
        if homepage:
            parent_choices = []
            for page in [homepage] + homepage.all_children:
                if not page.id in invalid_parents:
                    parent_choices.append((page.id, u" \u203a ".join(unicode(breadcrumb) for breadcrumb in page.breadcrumbs)))
        else:
            parent_choices = []
        if not parent_choices:
            parent_choices = (("", "---------"),)
        PageForm.base_fields["parent"].choices = parent_choices
        # Return the completed form.
        return PageForm

    def save_model(self, request, obj, form, change):
        """Saves the model and adds its content fields."""
        content_cls = self.get_page_content_cls(request, obj)
        with locked(Page, content_cls, ContentType, HistoryLink):
            content_cls_type = ContentType.objects.get_for_model(content_cls)
            # Delete the old page content, if it's expired.
            if change and ContentType.objects.get_for_id(obj.content_type_id) != content_cls_type:
                obj.content.delete()
            # Create the page content.
            obj.content_type = content_cls_type
            if change:
                try:
                    content_obj = content_cls_type.model_class().objects.get(page=obj)
                except content_cls.DoesNotExist:
                    content_obj = content_cls()  # We're either in a reversion recovery, or something has gone very wrong with the database...
            else:
                content_obj = content_cls()
            # Modify the page content.
            for field in content_obj._meta.fields:
                if field.name == "page":
                    continue
                setattr(content_obj, field.name, form.cleaned_data[field.name])
            # Get the page order.
            if not obj.order:
                try:
                    obj.order = self.model.objects.order_by("-order").values_list("order", flat=True)[0] + 1
                except IndexError:
                    obj.order = 1
            # Save the model.
            super(PageAdmin, self).save_model(request, obj, form, change)
            # Save the page content.
            content_obj.page = obj
            content_obj.save()

    # Permissions.

    def has_add_content_permission(self, request, model):
        """Checks whether the given user can edit the given content model."""
        opts = model._meta
        return request.user.has_perm("{0}.{1}".format(opts.app_label, opts.get_add_permission()))

    def has_add_permission(self, request):
        """Checks whether the user can edits pages and at least one content model."""
        if not super(PageAdmin, self).has_add_permission(request):
            return False
        for content_model in get_registered_content():
            if self.has_add_content_permission(request, content_model):
                return True
        return False

    def has_change_permission(self, request, obj=None):
        """Checks whether the user can edit the page and associated content model."""
        if not super(PageAdmin, self).has_change_permission(request, obj):
            return False
        if obj:
            content_model = ContentType.objects.get_for_id(obj.content_type_id).model_class()
            content_opts = content_model._meta
            return request.user.has_perm("{0}.{1}".format(
                content_opts.app_label,
                content_opts.get_change_permission(),
            ))
        return True

    def has_delete_permission(self, request, obj=None):
        """Checks whether the user can delete the page and associated content model."""
        if not super(PageAdmin, self).has_delete_permission(request, obj):
            return False
        if obj:
            content_model = ContentType.objects.get_for_id(obj.content_type_id).model_class()
            content_opts = content_model._meta
            return request.user.has_perm("{0}.{1}".format(
                content_opts.app_label,
                content_opts.get_delete_permission(),
            ))
        return True

    # Custom views.

    def patch_response_location(self, request, response):
        """Perpetuates the 'from' key in all redirect responses."""
        if isinstance(response, HttpResponseRedirect):
            if PAGE_FROM_KEY in request.GET:
                response["Location"] += "?%s=%s" % (PAGE_FROM_KEY, request.GET[PAGE_FROM_KEY])
        return response

    def patch_inlines(self, request, obj, func):
        """Updates the admin inlines to only display those relevant to the content class."""
        # HACK: Not thread-safe, only suitable for process-based hosting.
        def do_patch_inlines(*args, **kwargs):
            with page_inlines_lock:
                inlines = self.inlines
                inline_instances = self.inline_instances
                try:
                    content_cls = self.get_page_content_cls(request, obj)
                    content_inline_instances = [
                        content_inline_instance
                        for content_inline_cls, content_inline_instance
                        in self.content_inline_instances
                        if content_inline_cls == content_cls
                    ]
                    # Trim current inline instances.
                    self.inlines = inlines[:-len(self.content_inline_instances)]
                    self.inline_instances = inline_instances[:-len(self.content_inline_instances)]
                    # Add in new ones.
                    for content_inline_instance in content_inline_instances:
                        self.inlines.append(content_inline_instance.__class__)
                        self.inline_instances.append(content_inline_instance)
                    # Run the func.
                    return func(*args, **kwargs)
                finally:
                    self.inlines = inlines
                    self.inline_instances = inline_instances
        return do_patch_inlines

    def changelist_view(self, request, *args, **kwargs):
        """Redirects to the sitemap, if appropriate."""
        if PAGE_FROM_KEY in request.GET:
            redirect_slug = request.GET[PAGE_FROM_KEY]
            if redirect_slug == PAGE_FROM_SITEMAP_VALUE:
                return redirect("admin:index")
        return super(PageAdmin, self).changelist_view(request, *args, **kwargs)

    def add_view(self, request, *args, **kwargs):
        """Ensures that a valid content type is chosen."""
        if not PAGE_TYPE_PARAMETER in request.GET:
            # Generate the available content items.
            content_items = get_registered_content()
            content_items.sort(lambda a, b: cmp(a.classifier, b.classifier) or cmp(a._meta.verbose_name.lower(), b._meta.verbose_name.lower()))
            content_types = []
            for content_type in content_items:
                if self.has_add_content_permission(request, content_type):
                    # If we get this far, then we have permisison to add a page of this type.
                    get_params = request.GET.copy()
                    get_params[PAGE_TYPE_PARAMETER] = ContentType.objects.get_for_model(content_type).id
                    query_string = get_params.urlencode()
                    url = request.path + "?" + query_string
                    content_type_context = {
                        "name": content_type._meta.verbose_name,
                        "icon": content_type.icon,
                        "url": url,
                        "classifier": content_type.classifier
                    }
                    content_types.append(content_type_context)
            # Shortcut for when there is a single content type.
            if len(content_types) == 1:
                return redirect(content_types[0]["url"])
            # Render the select page template.
            context = {
                "title": "Select page type",
                "content_types": content_types
            }
            return render(request, "admin/pages/page/select_page_type.html", context)
        else:
            if not self.has_add_content_permission(request, ContentType.objects.get_for_id(request.GET[PAGE_TYPE_PARAMETER]).model_class()):
                raise PermissionDenied, "You are not allowed to add pages of that content type."
        return self.patch_inlines(request, None, super(PageAdmin, self).add_view)(request, *args, **kwargs)

    def change_view(self, request, object_id, *args, **kwargs):
        """Uses only the correct inlines for the page."""
        page = Page.objects.get(id=object_id)
        return self.patch_inlines(request, page, super(PageAdmin, self).change_view)(request, object_id, *args, **kwargs)

    def response_add(self, request, *args, **kwargs):
        """Redirects to the sitemap if appropriate."""
        response = super(PageAdmin, self).response_add(request, *args, **kwargs)
        return self.patch_response_location(request, response)

    def response_change(self, request, *args, **kwargs):
        """Redirects to the sitemap if appropriate."""
        response = super(PageAdmin, self).response_change(request, *args, **kwargs)
        return self.patch_response_location(request, response)

    def delete_view(self, request, *args, **kwargs):
        """Redirects to the sitemap if appropriate."""
        response = super(PageAdmin, self).delete_view(request, *args, **kwargs)
        return self.patch_response_location(request, response)

site.register(Page, PageAdmin)
site.register_link_list(Page)

page_admin = site._registry[Page]<commit_msg>Added in fix for page recover view.<commit_after>"""
The upgraded CMS online admin area.

This is an enhanced version of the Django admin area, providing a more
user-friendly appearance and providing additional functionality over the
standard implementation.
"""

from __future__ import with_statement
import threading

from django.core.exceptions import PermissionDenied
from django.contrib.contenttypes.models import ContentType
from django.http import Http404, HttpResponseRedirect
from django.shortcuts import render, redirect
from django import forms

import reversion

from cms.core.admin import PageBaseAdmin, site, PAGE_FROM_KEY, PAGE_FROM_SITEMAP_VALUE
from cms.core.db import locked
from cms.apps.historylinks.models import HistoryLink
from cms.apps.pages.models import Page, get_registered_content

# The GET parameter used to indicate content type.
PAGE_TYPE_PARAMETER = "type"

# HACK: Used to provide thread safety for the patch inlines hack (below).
page_inlines_lock = threading.RLock()

class PageAdmin(PageBaseAdmin):

    """Admin settings for Page models."""

    fieldsets = (
        (None, {
            "fields": ("title", "url_title", "parent",),
        },),
        ("Publication", {
            "fields": ("publication_date", "expiry_date", "is_online",),
            "classes": ("collapse",),
        }),
        ("Navigation", {
            "fields": ("short_title", "permalink", "in_navigation",),
            "classes": ("collapse",),
        }),
        ("Search engine optimization", {
            "fields": ("browser_title", "meta_keywords", "meta_description", "robots_index", "robots_follow", "robots_archive", "sitemap_priority", "sitemap_changefreq",),
            "classes": ("collapse",),
        }),
    )

    def __init__(self, *args, **kwargs):
        """Initialzies the PageAdmin."""
        super(PageAdmin, self).__init__(*args, **kwargs)
        # Register all content classes with reversion.
        self.content_inline_instances = []
        for content_cls in get_registered_content():
            self._autoregister(content_cls, follow=["page"])

    def register_content_inline(self, content_cls, inline, inline_admin):
        """Registers an inline model with the page admin."""
        # Register the admin class.
        self.inlines = list(self.inlines)
        self.inlines.append(inline_admin)
        inline_instance = inline_admin(self.model, self.admin_site)
        self.content_inline_instances.append((content_cls, inline_instance))
        self.inline_instances.append(inline_instance)
        # Register with reversion.
        self._autoregister(inline, follow=["page"])

    def queryset(self, request):
        """Only allows editing of pages in this site."""
        return request.pages.all

    # Reversion

    def get_revision_form_data(self, request, obj, version):
        """
        Returns a dictionary of data to set in the admin form in order to revert
        to the given revision.
        """
        data = super(PageAdmin, self).get_revision_form_data(request, obj, version)
        content_version = version.revision.version_set.all().get(
            content_type = obj.content_type_id,
            object_id_int = obj.pk,
        )
        data.update(content_version.field_dict)
        return data

    # Plugable content types.

    def get_page_content_cls(self, request, obj=None):
        """Retrieves the page content type slug."""
        if PAGE_TYPE_PARAMETER in request.GET:
            return ContentType.objects.get_for_id(request.GET[PAGE_TYPE_PARAMETER]).model_class()
        if obj and obj.content_type:
            return obj.content_type.model_class()
        raise Http404, "You must specify a page content type."

    def get_fieldsets(self, request, obj=None):
        """Generates the custom content fieldsets."""
        content_cls = self.get_page_content_cls(request, obj)
        content_fields = [field.name for field in content_cls._meta.fields if field.name != "page"]
        fieldsets = super(PageAdmin, self).get_fieldsets(request, obj)
        if content_fields:
            content_fieldsets = content_cls.fieldsets or (
                ("Page content", {
                    "fields": content_fields,
                }),
            )
            fieldsets = tuple(fieldsets[0:1]) + content_fieldsets + tuple(fieldsets[1:])
        return fieldsets

    def get_form(self, request, obj=None, **kwargs):
        """Adds the template area fields to the form."""
        content_cls = self.get_page_content_cls(request, obj)
        form_attrs = {}
        for field in content_cls._meta.fields:
            if field.name == "page":
                continue
            form_field = self.formfield_for_dbfield(field, request=request)
            # HACK: add in the initial value. No other way to pass this on.
            if obj:
                try:
                    form_field.initial = getattr(obj.content, field.name, "")
                except content_cls.DoesNotExist:
                    pass  # This means that we're in a reversion recovery, or something weird has happened to the databae.
            # Store the field.
            form_attrs[field.name] = form_field
        ContentForm = type("%sForm" % self.__class__.__name__, (forms.ModelForm,), form_attrs)
        defaults = {"form": ContentForm}
        defaults.update(kwargs)
        PageForm = super(PageAdmin, self).get_form(request, obj, **defaults)
        # HACK: Need to limit parents field based on object. This should be done in
        # formfield_for_foreignkey, but that method does not know about the object instance.
        if obj:
            invalid_parents = set(child.id for child in obj.all_children)
            invalid_parents.add(obj.id)
        else:
            invalid_parents = frozenset()
        homepage = request.pages.homepage
        if homepage:
            parent_choices = []
            for page in [homepage] + homepage.all_children:
                if not page.id in invalid_parents:
                    parent_choices.append((page.id, u" \u203a ".join(unicode(breadcrumb) for breadcrumb in page.breadcrumbs)))
        else:
            parent_choices = []
        if not parent_choices:
            parent_choices = (("", "---------"),)
        PageForm.base_fields["parent"].choices = parent_choices
        # Return the completed form.
        return PageForm

    def save_model(self, request, obj, form, change):
        """Saves the model and adds its content fields."""
        content_cls = self.get_page_content_cls(request, obj)
        with locked(Page, content_cls, ContentType, HistoryLink):
            content_cls_type = ContentType.objects.get_for_model(content_cls)
            # Delete the old page content, if it's expired.
            if change and ContentType.objects.get_for_id(obj.content_type_id) != content_cls_type:
                obj.content.delete()
            # Create the page content.
            obj.content_type = content_cls_type
            if change:
                try:
                    content_obj = content_cls_type.model_class().objects.get(page=obj)
                except content_cls.DoesNotExist:
                    content_obj = content_cls()  # We're either in a reversion recovery, or something has gone very wrong with the database...
            else:
                content_obj = content_cls()
            # Modify the page content.
            for field in content_obj._meta.fields:
                if field.name == "page":
                    continue
                setattr(content_obj, field.name, form.cleaned_data[field.name])
            # Get the page order.
            if not obj.order:
                try:
                    obj.order = self.model.objects.order_by("-order").values_list("order", flat=True)[0] + 1
                except IndexError:
                    obj.order = 1
            # Save the model.
            super(PageAdmin, self).save_model(request, obj, form, change)
            # Save the page content.
            content_obj.page = obj
            content_obj.save()

    # Permissions.

    def has_add_content_permission(self, request, model):
        """Checks whether the given user can edit the given content model."""
        opts = model._meta
        return request.user.has_perm("{0}.{1}".format(opts.app_label, opts.get_add_permission()))

    def has_add_permission(self, request):
        """Checks whether the user can edits pages and at least one content model."""
        if not super(PageAdmin, self).has_add_permission(request):
            return False
        for content_model in get_registered_content():
            if self.has_add_content_permission(request, content_model):
                return True
        return False

    def has_change_permission(self, request, obj=None):
        """Checks whether the user can edit the page and associated content model."""
        if not super(PageAdmin, self).has_change_permission(request, obj):
            return False
        if obj:
            content_model = ContentType.objects.get_for_id(obj.content_type_id).model_class()
            content_opts = content_model._meta
            return request.user.has_perm("{0}.{1}".format(
                content_opts.app_label,
                content_opts.get_change_permission(),
            ))
        return True

    def has_delete_permission(self, request, obj=None):
        """Checks whether the user can delete the page and associated content model."""
        if not super(PageAdmin, self).has_delete_permission(request, obj):
            return False
        if obj:
            content_model = ContentType.objects.get_for_id(obj.content_type_id).model_class()
            content_opts = content_model._meta
            return request.user.has_perm("{0}.{1}".format(
                content_opts.app_label,
                content_opts.get_delete_permission(),
            ))
        return True

    # Custom views.

    def patch_response_location(self, request, response):
        """Perpetuates the 'from' key in all redirect responses."""
        if isinstance(response, HttpResponseRedirect):
            if PAGE_FROM_KEY in request.GET:
                response["Location"] += "?%s=%s" % (PAGE_FROM_KEY, request.GET[PAGE_FROM_KEY])
        return response

    def patch_inlines(self, request, obj, func):
        """Updates the admin inlines to only display those relevant to the content class."""
        # HACK: Not thread-safe, only suitable for process-based hosting.
        def do_patch_inlines(*args, **kwargs):
            with page_inlines_lock:
                inlines = self.inlines
                inline_instances = self.inline_instances
                try:
                    content_cls = self.get_page_content_cls(request, obj)
                    content_inline_instances = [
                        content_inline_instance
                        for content_inline_cls, content_inline_instance
                        in self.content_inline_instances
                        if content_inline_cls == content_cls
                    ]
                    # Trim current inline instances.
                    self.inlines = inlines[:-len(self.content_inline_instances)]
                    self.inline_instances = inline_instances[:-len(self.content_inline_instances)]
                    # Add in new ones.
                    for content_inline_instance in content_inline_instances:
                        self.inlines.append(content_inline_instance.__class__)
                        self.inline_instances.append(content_inline_instance)
                    # Run the func.
                    return func(*args, **kwargs)
                finally:
                    self.inlines = inlines
                    self.inline_instances = inline_instances
        return do_patch_inlines

    def changelist_view(self, request, *args, **kwargs):
        """Redirects to the sitemap, if appropriate."""
        if PAGE_FROM_KEY in request.GET:
            redirect_slug = request.GET[PAGE_FROM_KEY]
            if redirect_slug == PAGE_FROM_SITEMAP_VALUE:
                return redirect("admin:index")
        return super(PageAdmin, self).changelist_view(request, *args, **kwargs)

    def add_view(self, request, *args, **kwargs):
        """Ensures that a valid content type is chosen."""
        if not PAGE_TYPE_PARAMETER in request.GET:
            # Generate the available content items.
            content_items = get_registered_content()
            content_items.sort(lambda a, b: cmp(a.classifier, b.classifier) or cmp(a._meta.verbose_name.lower(), b._meta.verbose_name.lower()))
            content_types = []
            for content_type in content_items:
                if self.has_add_content_permission(request, content_type):
                    # If we get this far, then we have permisison to add a page of this type.
                    get_params = request.GET.copy()
                    get_params[PAGE_TYPE_PARAMETER] = ContentType.objects.get_for_model(content_type).id
                    query_string = get_params.urlencode()
                    url = request.path + "?" + query_string
                    content_type_context = {
                        "name": content_type._meta.verbose_name,
                        "icon": content_type.icon,
                        "url": url,
                        "classifier": content_type.classifier
                    }
                    content_types.append(content_type_context)
            # Shortcut for when there is a single content type.
            if len(content_types) == 1:
                return redirect(content_types[0]["url"])
            # Render the select page template.
            context = {
                "title": "Select page type",
                "content_types": content_types
            }
            return render(request, "admin/pages/page/select_page_type.html", context)
        else:
            if not self.has_add_content_permission(request, ContentType.objects.get_for_id(request.GET[PAGE_TYPE_PARAMETER]).model_class()):
                raise PermissionDenied, "You are not allowed to add pages of that content type."
        return self.patch_inlines(request, None, super(PageAdmin, self).add_view)(request, *args, **kwargs)

    def change_view(self, request, object_id, *args, **kwargs):
        """Uses only the correct inlines for the page."""
        page = Page.objects.get(id=object_id)
        return self.patch_inlines(request, page, super(PageAdmin, self).change_view)(request, object_id, *args, **kwargs)

    def response_add(self, request, *args, **kwargs):
        """Redirects to the sitemap if appropriate."""
        response = super(PageAdmin, self).response_add(request, *args, **kwargs)
        return self.patch_response_location(request, response)

    def response_change(self, request, *args, **kwargs):
        """Redirects to the sitemap if appropriate."""
        response = super(PageAdmin, self).response_change(request, *args, **kwargs)
        return self.patch_response_location(request, response)

    def delete_view(self, request, *args, **kwargs):
        """Redirects to the sitemap if appropriate."""
        response = super(PageAdmin, self).delete_view(request, *args, **kwargs)
        return self.patch_response_location(request, response)

site.register(Page, PageAdmin)
site.register_link_list(Page)

page_admin = site._registry[Page]<|endoftext|>


<commit_before>"""Core views used by the CMS."""

from django import template
from django.conf import settings
from django.contrib.contenttypes.models import ContentType
from django.http import Http404, HttpResponse, HttpResponseRedirect, HttpResponseNotFound, HttpResponseServerError
from django.shortcuts import render_to_response

from cms.apps.pages.models import Page

def render_page(request, path_info=""):
    """Dispatches the request to the site pages."""
    # Attempt to retrieve the homepage.
    try:
        homepage = Page.objects.get_homepage()
    except Page.DoesNotExist:
        raise Http404, "The site does not have a homepage."
    # Check publication state.
    if not homepage.is_published:
        raise Http404, "The site homepage has not been published."
    # Dispatch the request!
    request.breadcrumbs = []
    return homepage.content.dispatch(request, path_info)

def permalink_redirect(request, content_type_id, object_id):
    """Redirects to the object encoded in the permalink."""
    # Attempt to access the encoded content type.
    try:
        content_type = ContentType.objects.get_for_id(content_type_id)
    except ContentType.DoesNotExist, ex:
        raise Http404, str(ex)
    # Attempt to access the encoded object.
    try:
        obj = content_type.get_object_for_this_type(pk=object_id)
    except content_type.model_class().DoesNotExist, ex:
        raise Http404, str(ex)
    # Attempt to redirect to the object's absolute URL.
    try:
        redirect_url = obj.get_absolute_url()
    except AttributeError:
        raise Http404, "%s objects do not publish an absolute URL." % content_type.name.title()
    return HttpResponseRedirect(redirect_url)

def render_template(request, path, base_path=""):
    """
    Serves static template files based on the given path.

    If supplied, `base_path` will be prepended onto the path.
    """
    template_name = base_path + path
    if not template_name or template_name.endswith("/"):
        template_name += "base.html"
    try:
        return render_to_response(template_name, {}, template.RequestContext(request))
    except template.TemplateDoesNotExist:
        raise Http404, "The template '%s' does not exist." % template_name

def tinymce_init(request):
    """Renders the TinyMCE initialization script."""
    context = {"TINYMCE_CONTENT_CSS": settings.TINYMCE_CONTENT_CSS}
    return render_to_response("admin/tinymce_init.js", context, template.RequestContext(request), mimetype="text/javascript")

def reorder_pages(request):
    """Swaps the ordering of two pages."""
    # Get the POST variables.
    page_ids = request.POST.getlist("pages")
    # Get the page objects.
    pages = Page.objects.filter(id__in=page_ids)
    page_0_order = pages[0].order
    page_1_order = pages[1].order
    # Blank their order fields.
    for page in pages:
        page.order = None
        page.save()
    # Swap their order fields.
    pages[0].order = page_1_order
    pages[1].order = page_0_order
    pages[0].save()
    pages[1].save()
    # Send a positive response.
    return HttpResponse("Swapped page '%s' with page '%s'." % (pages[0], pages[1]))

def handler500(self, request):
    """Renders the server error document."""
    content = template.loader.render_to_string("500.html", {}, template.RequestContext(request))
    return HttpResponseServerError(content)

<commit_msg>Homepage publication takes place in render method, not view.<commit_after>"""Core views used by the CMS."""

from django import template
from django.conf import settings
from django.contrib.contenttypes.models import ContentType
from django.http import Http404, HttpResponse, HttpResponseRedirect, HttpResponseNotFound, HttpResponseServerError
from django.shortcuts import render_to_response

from cms.apps.pages.models import Page

def render_page(request, path_info=""):
    """Dispatches the request to the site pages."""
    # Attempt to retrieve the homepage.
    try:
        homepage = Page.objects.get_homepage()
    except Page.DoesNotExist:
        raise Http404, "The site does not have a homepage."
    # Dispatch the request!
    request.breadcrumbs = []
    return homepage.content.dispatch(request, path_info)

def permalink_redirect(request, content_type_id, object_id):
    """Redirects to the object encoded in the permalink."""
    # Attempt to access the encoded content type.
    try:
        content_type = ContentType.objects.get_for_id(content_type_id)
    except ContentType.DoesNotExist, ex:
        raise Http404, str(ex)
    # Attempt to access the encoded object.
    try:
        obj = content_type.get_object_for_this_type(pk=object_id)
    except content_type.model_class().DoesNotExist, ex:
        raise Http404, str(ex)
    # Attempt to redirect to the object's absolute URL.
    try:
        redirect_url = obj.get_absolute_url()
    except AttributeError:
        raise Http404, "%s objects do not publish an absolute URL." % content_type.name.title()
    return HttpResponseRedirect(redirect_url)

def render_template(request, path, base_path=""):
    """
    Serves static template files based on the given path.

    If supplied, `base_path` will be prepended onto the path.
    """
    template_name = base_path + path
    if not template_name or template_name.endswith("/"):
        template_name += "base.html"
    try:
        return render_to_response(template_name, {}, template.RequestContext(request))
    except template.TemplateDoesNotExist:
        raise Http404, "The template '%s' does not exist." % template_name

def tinymce_init(request):
    """Renders the TinyMCE initialization script."""
    context = {"TINYMCE_CONTENT_CSS": settings.TINYMCE_CONTENT_CSS}
    return render_to_response("admin/tinymce_init.js", context, template.RequestContext(request), mimetype="text/javascript")

def reorder_pages(request):
    """Swaps the ordering of two pages."""
    # Get the POST variables.
    page_ids = request.POST.getlist("pages")
    # Get the page objects.
    pages = Page.objects.filter(id__in=page_ids)
    page_0_order = pages[0].order
    page_1_order = pages[1].order
    # Blank their order fields.
    for page in pages:
        page.order = None
        page.save()
    # Swap their order fields.
    pages[0].order = page_1_order
    pages[1].order = page_0_order
    pages[0].save()
    pages[1].save()
    # Send a positive response.
    return HttpResponse("Swapped page '%s' with page '%s'." % (pages[0], pages[1]))

def handler500(self, request):
    """Renders the server error document."""
    content = template.loader.render_to_string("500.html", {}, template.RequestContext(request))
    return HttpResponseServerError(content)

<|endoftext|>