<commit_before>from typing import List, Union, Text, Optional, Any, Tuple, Dict

import logging
import scipy.sparse
import numpy as np

from rasa.utils.tensorflow.constants import SEQUENCE, BALANCED
from rasa.utils.tensorflow.model_data import RasaModelData, Data, FeatureArray
import tensorflow as tf

logger = logging.getLogger(__name__)

class RasaDataGenerator(tf.keras.utils.Sequence):
    """Abstract data generator."""

    def __init__(
        self,
        model_data: RasaModelData,
        batch_size: Union[int, List[int]],
        batch_strategy: Text = SEQUENCE,
        shuffle: bool = True,
    ):
        """Initializes the data generator.

        Args:
            model_data: The model data to use.
            batch_size: The batch size(s).
            batch_strategy: The batch strategy.
            shuffle: If 'True', data should be shuffled.
        """
        self.model_data = model_data
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.batch_strategy = batch_strategy

    def __len__(self) -> int:
        """Number of batches in the Sequence.

        Returns:
            The number of batches in the Sequence.
        """
        raise NotImplementedError

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """Gets batch at position `index`.

        Arguments:
            index: position of the batch in the Sequence.

        Returns:
            A batch (tuple of input data and target data).
        """
        raise NotImplementedError

    def on_epoch_end(self) -> None:
        """Update the data after every epoch."""
        raise NotImplementedError

    def _shuffle_and_balance(self, batch_size: int) -> Data:
        data = self.model_data.data

        if self.shuffle:
            data = self.model_data.shuffled_data(data)

        if self.batch_strategy == BALANCED:
            data = self.model_data.balanced_data(data, batch_size, self.shuffle)

        # do not override self.model_data.data, because we need original data for
        # balancing on the next epoch
        return data

    @staticmethod
    def prepare_batch(
        data: Data,
        start: Optional[int] = None,
        end: Optional[int] = None,
        tuple_sizes: Optional[Dict[Text, int]] = None,
    ) -> Tuple[Optional[np.ndarray]]:
        """Slices model data into batch using given start and end value.

        Args:
            data: The data to prepare.
            start: The start index of the batch
            end: The end index of the batch
            tuple_sizes: In case the feature is not present we propagate the batch with
              None. Tuple sizes contains the number of how many None values to add for
              what kind of feature.

        Returns:
            The features of the batch.
        """
        batch_data = []

        for key, attribute_data in data.items():
            for sub_key, f_data in attribute_data.items():
                # add None for not present values during processing
                if not f_data:
                    if tuple_sizes:
                        batch_data += [None] * tuple_sizes[key]
                    else:
                        batch_data.append(None)
                    continue

                for v in f_data:
                    if start is not None and end is not None:
                        _data = v[start:end]
                    elif start is not None:
                        _data = v[start:]
                    elif end is not None:
                        _data = v[:end]
                    else:
                        _data = v[:]

                    if _data.is_sparse:
                        batch_data.extend(
                            RasaDataGenerator._scipy_matrix_to_values(_data)
                        )
                    else:
                        batch_data.append(RasaDataGenerator._pad_dense_data(_data))

        # len of batch_data is equal to the number of keys in model data
        return tuple(batch_data)

    @staticmethod
    def _pad_dense_data(array_of_dense: FeatureArray) -> np.ndarray:
        """Pad data of different lengths.

        Sequential data is padded with zeros. Zeros are added to the end of data.

        Args:
            array_of_dense: The array to pad.

        Returns:
            The padded array.
        """
        if array_of_dense.number_of_dimensions == 4:
            return RasaDataGenerator._pad_4d_dense_data(array_of_dense)

        if array_of_dense[0].ndim < 2:
            # data doesn't contain a sequence
            return array_of_dense.astype(np.float32)

        data_size = len(array_of_dense)
        max_seq_len = max([x.shape[0] for x in array_of_dense])

        data_padded = np.zeros(
            [data_size, max_seq_len, array_of_dense[0].shape[-1]],
            dtype=array_of_dense[0].dtype,
        )
        for i in range(data_size):
            data_padded[i, : array_of_dense[i].shape[0], :] = array_of_dense[i]

        return data_padded.astype(np.float32)

    @staticmethod
    def _pad_4d_dense_data(array_of_array_of_dense: FeatureArray) -> np.ndarray:
        # in case of dialogue data we may have 4 dimensions
        # batch size x dialogue history length x sequence length x number of features

        # as transformers cannot handle 4D tensors pad and reshape the data
        # so that the resulting tensor is 3D
        # the shape is (sum of dialogue history length for all tensors in the
        # batch x max sequence length x number of features)
        # the original shape and the original dialogue length is passed on to the model
        # it can be used to transform the 3D tensor back into 4D

        # in order to create 4d tensor inputs, we created "fake" zero features
        # for nonexistent inputs. To save calculation we filter this features before
        # input to tf methods.
        number_of_features = array_of_array_of_dense[0][0].shape[-1]
        array_of_array_of_dense = RasaDataGenerator._filter_out_fake_inputs(
            array_of_array_of_dense
        )
        if not array_of_array_of_dense:
            # return empty 3d array with appropriate last dims
            return np.zeros((0, 0, number_of_features), dtype=np.float32)

        combined_dialogue_len = sum(
            len(array_of_dense) for array_of_dense in array_of_array_of_dense
        )
        max_seq_len = max(
            [
                x.shape[0]
                for array_of_dense in array_of_array_of_dense
                for x in array_of_dense
            ]
        )

        data_padded = np.zeros(
            [combined_dialogue_len, max_seq_len, number_of_features],
            dtype=array_of_array_of_dense[0][0].dtype,
        )

        current_sum_dialogue_len = 0
        for i, array_of_dense in enumerate(array_of_array_of_dense):
            for j, dense in enumerate(array_of_dense):
                data_padded[current_sum_dialogue_len + j, : dense.shape[0], :] = dense
            current_sum_dialogue_len += len(array_of_dense)

        return data_padded.astype(np.float32)

    @staticmethod
    def _scipy_matrix_to_values(array_of_sparse: FeatureArray) -> List[np.ndarray]:
        """Convert a scipy matrix into indices, data, and shape.

        Args:
            array_of_sparse: The sparse data array.

        Returns:
            A list of dense numpy arrays representing the sparse data.
        """
        if array_of_sparse.number_of_dimensions == 4:
            return RasaDataGenerator._4d_scipy_matrix_to_values(array_of_sparse)

        # we need to make sure that the matrices are coo_matrices otherwise the
        # transformation does not work (e.g. you cannot access x.row, x.col)
        if not isinstance(array_of_sparse[0], scipy.sparse.coo_matrix):
            array_of_sparse = [x.tocoo() for x in array_of_sparse]

        max_seq_len = max([x.shape[0] for x in array_of_sparse])

        # get the indices of values
        indices = np.hstack(
            [
                np.vstack([i * np.ones_like(x.row), x.row, x.col])
                for i, x in enumerate(array_of_sparse)
            ]
        ).T

        data = np.hstack([x.data for x in array_of_sparse])

        number_of_features = array_of_sparse[0].shape[-1]
        shape = np.array((len(array_of_sparse), max_seq_len, number_of_features))

        return [
            indices.astype(np.int64),
            data.astype(np.float32),
            shape.astype(np.int64),
        ]

    @staticmethod
    def _4d_scipy_matrix_to_values(
        array_of_array_of_sparse: FeatureArray,
    ) -> List[np.ndarray]:
        # in case of dialogue data we may have 4 dimensions
        # batch size x dialogue history length x sequence length x number of features

        # transformers cannot handle 4D tensors, therefore pad and reshape the data
        # so that the resulting tensor is 3D
        # the shape is (sum of dialogue history length for all tensors in the
        # batch x max sequence length x number of features)
        # the original shape and the original dialogue length is passed on to the model
        # it can be used to transform the 3D tensor back into 4D

        # in order to create 4d tensor inputs, we created "fake" zero features
        # for nonexistent inputs. To save calculation we filter this features before
        # input to tf methods.
        number_of_features = array_of_array_of_sparse[0][0].shape[-1]
        array_of_array_of_sparse = RasaDataGenerator._filter_out_fake_inputs(
            array_of_array_of_sparse
        )
        if not array_of_array_of_sparse:
            # create empty array with appropriate last dims
            return [
                np.empty((0, 3), dtype=np.int64),
                np.array([], dtype=np.float32),
                np.array([0, 0, number_of_features], dtype=np.int64),
            ]

        # we need to make sure that the matrices are coo_matrices otherwise the
        # transformation does not work (e.g. you cannot access x.row, x.col)
        if not isinstance(array_of_array_of_sparse[0][0], scipy.sparse.coo_matrix):
            array_of_array_of_sparse = [
                [x.tocoo() for x in array_of_sparse]
                for array_of_sparse in array_of_array_of_sparse
            ]

        dialogue_len = [
            len(array_of_sparse) for array_of_sparse in array_of_array_of_sparse
        ]
        combined_dialogue_len = sum(dialogue_len)
        max_seq_len = max(
            [
                x.shape[0]
                for array_of_sparse in array_of_array_of_sparse
                for x in array_of_sparse
            ]
        )
        # get the indices of values
        indices = np.hstack(
            [
                np.vstack(
                    [sum(dialogue_len[:i]) + j * np.ones_like(x.row), x.row, x.col]
                )
                for i, array_of_sparse in enumerate(array_of_array_of_sparse)
                for j, x in enumerate(array_of_sparse)
            ]
        ).T

        data = np.hstack(
            [
                x.data
                for array_of_sparse in array_of_array_of_sparse
                for x in array_of_sparse
            ]
        )

        shape = np.array((combined_dialogue_len, max_seq_len, number_of_features))

        return [
            indices.astype(np.int64),
            data.astype(np.float32),
            shape.astype(np.int64),
        ]

    @staticmethod
    def _filter_out_fake_inputs(
        array_of_array_of_features: FeatureArray,
    ) -> Union[List[List[np.ndarray]], List[List[scipy.sparse.spmatrix]]]:
        return list(
            filter(
                # filter empty lists created by another filter
                lambda x: len(x) > 0,
                [
                    # filter all the "fake" inputs, we know the input is "fake",
                    # when sequence dimension is `0`
                    list(filter(lambda x: x.shape[0] > 0, array_of_features))
                    for array_of_features in array_of_array_of_features
                ],
            )
        )

class RasaBatchDataGenerator(RasaDataGenerator):
    """Data generator with an optional increasing batch size."""

    def __init__(
        self,
        model_data: RasaModelData,
        batch_size: Union[List[int], int],
        epochs: int = 1,
        batch_strategy: Text = SEQUENCE,
        shuffle: bool = True,
    ):
        """Initializes the increasing batch size data generator.

        Args:
            model_data: The model data to use.
            batch_size: The batch size.
            epochs: The total number of epochs.
            batch_strategy: The batch strategy.
            shuffle: If 'True', data will be shuffled.
        """
        super().__init__(model_data, batch_size, batch_strategy, shuffle)

        if isinstance(batch_size, list):
            logger.debug(
                "The provided batch size is a list, this data generator will use a "
                "linear increasing batch size."
            )

        self._epochs = epochs
        # we use `on_epoch_end` method to prepare data for the next epoch
        # set current epoch to `-1`, so that `on_epoch_end` will increase it to `0`
        self._current_epoch = -1
        # actual batch size will be set inside `on_epoch_end`
        self._current_batch_size = None
        # create separate data variable that will store modified data for each batch
        self._data = None
        self.on_epoch_end()

    def __len__(self) -> int:
        """Number of batches in the Sequence.

        Returns:
            The number of batches in the Sequence.
        """
        # data was rebalanced, so need to recalculate number of examples
        num_examples = self.model_data.number_of_examples(self._data)
        batch_size = self._current_batch_size
        return num_examples // batch_size + int(num_examples % batch_size > 0)

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """Gets batch at position `index`.

        Arguments:
            index: position of the batch in the Sequence.

        Returns:
            A batch (tuple of input data and target data).
        """
        start = index * self._current_batch_size
        end = start + self._current_batch_size

        # return input and target data, as our target data is inside the input
        # data return None for the target data
        return self.prepare_batch(self._data, start, end), None

    def on_epoch_end(self) -> None:
        """Update the data after every epoch."""
        self._current_epoch += 1
        self._current_batch_size = self._linearly_increasing_batch_size()
        self._data = self._shuffle_and_balance(self._current_batch_size)

    def _linearly_increasing_batch_size(self) -> int:
        """Linearly increase batch size with every epoch.

        The idea comes from https://arxiv.org/abs/1711.00489.

        Returns:
            The batch size to use in this epoch.
        """
        if not isinstance(self.batch_size, list):
            return int(self.batch_size)

        if self._epochs > 1:
            return int(
                self.batch_size[0]
                + self._current_epoch
                * (self.batch_size[1] - self.batch_size[0])
                / (self._epochs - 1)
            )
        else:
            return int(self.batch_size[0])
<commit_msg>Fix import order so diff shows none<commit_after>from typing import List, Union, Text, Optional, Any, Tuple, Dict

import logging
import scipy.sparse
import numpy as np
import tensorflow as tf

from rasa.utils.tensorflow.constants import SEQUENCE, BALANCED
from rasa.utils.tensorflow.model_data import RasaModelData, Data, FeatureArray

logger = logging.getLogger(__name__)

class RasaDataGenerator(tf.keras.utils.Sequence):
    """Abstract data generator."""

    def __init__(
        self,
        model_data: RasaModelData,
        batch_size: Union[int, List[int]],
        batch_strategy: Text = SEQUENCE,
        shuffle: bool = True,
    ):
        """Initializes the data generator.

        Args:
            model_data: The model data to use.
            batch_size: The batch size(s).
            batch_strategy: The batch strategy.
            shuffle: If 'True', data should be shuffled.
        """
        self.model_data = model_data
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.batch_strategy = batch_strategy

    def __len__(self) -> int:
        """Number of batches in the Sequence.

        Returns:
            The number of batches in the Sequence.
        """
        raise NotImplementedError

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """Gets batch at position `index`.

        Arguments:
            index: position of the batch in the Sequence.

        Returns:
            A batch (tuple of input data and target data).
        """
        raise NotImplementedError

    def on_epoch_end(self) -> None:
        """Update the data after every epoch."""
        raise NotImplementedError

    def _shuffle_and_balance(self, batch_size: int) -> Data:
        data = self.model_data.data

        if self.shuffle:
            data = self.model_data.shuffled_data(data)

        if self.batch_strategy == BALANCED:
            data = self.model_data.balanced_data(data, batch_size, self.shuffle)

        # do not override self.model_data.data, because we need original data for
        # balancing on the next epoch
        return data

    @staticmethod
    def prepare_batch(
        data: Data,
        start: Optional[int] = None,
        end: Optional[int] = None,
        tuple_sizes: Optional[Dict[Text, int]] = None,
    ) -> Tuple[Optional[np.ndarray]]:
        """Slices model data into batch using given start and end value.

        Args:
            data: The data to prepare.
            start: The start index of the batch
            end: The end index of the batch
            tuple_sizes: In case the feature is not present we propagate the batch with
              None. Tuple sizes contains the number of how many None values to add for
              what kind of feature.

        Returns:
            The features of the batch.
        """
        batch_data = []

        for key, attribute_data in data.items():
            for sub_key, f_data in attribute_data.items():
                # add None for not present values during processing
                if not f_data:
                    if tuple_sizes:
                        batch_data += [None] * tuple_sizes[key]
                    else:
                        batch_data.append(None)
                    continue

                for v in f_data:
                    if start is not None and end is not None:
                        _data = v[start:end]
                    elif start is not None:
                        _data = v[start:]
                    elif end is not None:
                        _data = v[:end]
                    else:
                        _data = v[:]

                    if _data.is_sparse:
                        batch_data.extend(
                            RasaDataGenerator._scipy_matrix_to_values(_data)
                        )
                    else:
                        batch_data.append(RasaDataGenerator._pad_dense_data(_data))

        # len of batch_data is equal to the number of keys in model data
        return tuple(batch_data)

    @staticmethod
    def _pad_dense_data(array_of_dense: FeatureArray) -> np.ndarray:
        """Pad data of different lengths.

        Sequential data is padded with zeros. Zeros are added to the end of data.

        Args:
            array_of_dense: The array to pad.

        Returns:
            The padded array.
        """
        if array_of_dense.number_of_dimensions == 4:
            return RasaDataGenerator._pad_4d_dense_data(array_of_dense)

        if array_of_dense[0].ndim < 2:
            # data doesn't contain a sequence
            return array_of_dense.astype(np.float32)

        data_size = len(array_of_dense)
        max_seq_len = max([x.shape[0] for x in array_of_dense])

        data_padded = np.zeros(
            [data_size, max_seq_len, array_of_dense[0].shape[-1]],
            dtype=array_of_dense[0].dtype,
        )
        for i in range(data_size):
            data_padded[i, : array_of_dense[i].shape[0], :] = array_of_dense[i]

        return data_padded.astype(np.float32)

    @staticmethod
    def _pad_4d_dense_data(array_of_array_of_dense: FeatureArray) -> np.ndarray:
        # in case of dialogue data we may have 4 dimensions
        # batch size x dialogue history length x sequence length x number of features

        # as transformers cannot handle 4D tensors pad and reshape the data
        # so that the resulting tensor is 3D
        # the shape is (sum of dialogue history length for all tensors in the
        # batch x max sequence length x number of features)
        # the original shape and the original dialogue length is passed on to the model
        # it can be used to transform the 3D tensor back into 4D

        # in order to create 4d tensor inputs, we created "fake" zero features
        # for nonexistent inputs. To save calculation we filter this features before
        # input to tf methods.
        number_of_features = array_of_array_of_dense[0][0].shape[-1]
        array_of_array_of_dense = RasaDataGenerator._filter_out_fake_inputs(
            array_of_array_of_dense
        )
        if not array_of_array_of_dense:
            # return empty 3d array with appropriate last dims
            return np.zeros((0, 0, number_of_features), dtype=np.float32)

        combined_dialogue_len = sum(
            len(array_of_dense) for array_of_dense in array_of_array_of_dense
        )
        max_seq_len = max(
            [
                x.shape[0]
                for array_of_dense in array_of_array_of_dense
                for x in array_of_dense
            ]
        )

        data_padded = np.zeros(
            [combined_dialogue_len, max_seq_len, number_of_features],
            dtype=array_of_array_of_dense[0][0].dtype,
        )

        current_sum_dialogue_len = 0
        for i, array_of_dense in enumerate(array_of_array_of_dense):
            for j, dense in enumerate(array_of_dense):
                data_padded[current_sum_dialogue_len + j, : dense.shape[0], :] = dense
            current_sum_dialogue_len += len(array_of_dense)

        return data_padded.astype(np.float32)

    @staticmethod
    def _scipy_matrix_to_values(array_of_sparse: FeatureArray) -> List[np.ndarray]:
        """Convert a scipy matrix into indices, data, and shape.

        Args:
            array_of_sparse: The sparse data array.

        Returns:
            A list of dense numpy arrays representing the sparse data.
        """
        if array_of_sparse.number_of_dimensions == 4:
            return RasaDataGenerator._4d_scipy_matrix_to_values(array_of_sparse)

        # we need to make sure that the matrices are coo_matrices otherwise the
        # transformation does not work (e.g. you cannot access x.row, x.col)
        if not isinstance(array_of_sparse[0], scipy.sparse.coo_matrix):
            array_of_sparse = [x.tocoo() for x in array_of_sparse]

        max_seq_len = max([x.shape[0] for x in array_of_sparse])

        # get the indices of values
        indices = np.hstack(
            [
                np.vstack([i * np.ones_like(x.row), x.row, x.col])
                for i, x in enumerate(array_of_sparse)
            ]
        ).T

        data = np.hstack([x.data for x in array_of_sparse])

        number_of_features = array_of_sparse[0].shape[-1]
        shape = np.array((len(array_of_sparse), max_seq_len, number_of_features))

        return [
            indices.astype(np.int64),
            data.astype(np.float32),
            shape.astype(np.int64),
        ]

    @staticmethod
    def _4d_scipy_matrix_to_values(
        array_of_array_of_sparse: FeatureArray,
    ) -> List[np.ndarray]:
        # in case of dialogue data we may have 4 dimensions
        # batch size x dialogue history length x sequence length x number of features

        # transformers cannot handle 4D tensors, therefore pad and reshape the data
        # so that the resulting tensor is 3D
        # the shape is (sum of dialogue history length for all tensors in the
        # batch x max sequence length x number of features)
        # the original shape and the original dialogue length is passed on to the model
        # it can be used to transform the 3D tensor back into 4D

        # in order to create 4d tensor inputs, we created "fake" zero features
        # for nonexistent inputs. To save calculation we filter this features before
        # input to tf methods.
        number_of_features = array_of_array_of_sparse[0][0].shape[-1]
        array_of_array_of_sparse = RasaDataGenerator._filter_out_fake_inputs(
            array_of_array_of_sparse
        )
        if not array_of_array_of_sparse:
            # create empty array with appropriate last dims
            return [
                np.empty((0, 3), dtype=np.int64),
                np.array([], dtype=np.float32),
                np.array([0, 0, number_of_features], dtype=np.int64),
            ]

        # we need to make sure that the matrices are coo_matrices otherwise the
        # transformation does not work (e.g. you cannot access x.row, x.col)
        if not isinstance(array_of_array_of_sparse[0][0], scipy.sparse.coo_matrix):
            array_of_array_of_sparse = [
                [x.tocoo() for x in array_of_sparse]
                for array_of_sparse in array_of_array_of_sparse
            ]

        dialogue_len = [
            len(array_of_sparse) for array_of_sparse in array_of_array_of_sparse
        ]
        combined_dialogue_len = sum(dialogue_len)
        max_seq_len = max(
            [
                x.shape[0]
                for array_of_sparse in array_of_array_of_sparse
                for x in array_of_sparse
            ]
        )
        # get the indices of values
        indices = np.hstack(
            [
                np.vstack(
                    [sum(dialogue_len[:i]) + j * np.ones_like(x.row), x.row, x.col]
                )
                for i, array_of_sparse in enumerate(array_of_array_of_sparse)
                for j, x in enumerate(array_of_sparse)
            ]
        ).T

        data = np.hstack(
            [
                x.data
                for array_of_sparse in array_of_array_of_sparse
                for x in array_of_sparse
            ]
        )

        shape = np.array((combined_dialogue_len, max_seq_len, number_of_features))

        return [
            indices.astype(np.int64),
            data.astype(np.float32),
            shape.astype(np.int64),
        ]

    @staticmethod
    def _filter_out_fake_inputs(
        array_of_array_of_features: FeatureArray,
    ) -> Union[List[List[np.ndarray]], List[List[scipy.sparse.spmatrix]]]:
        return list(
            filter(
                # filter empty lists created by another filter
                lambda x: len(x) > 0,
                [
                    # filter all the "fake" inputs, we know the input is "fake",
                    # when sequence dimension is `0`
                    list(filter(lambda x: x.shape[0] > 0, array_of_features))
                    for array_of_features in array_of_array_of_features
                ],
            )
        )

class RasaBatchDataGenerator(RasaDataGenerator):
    """Data generator with an optional increasing batch size."""

    def __init__(
        self,
        model_data: RasaModelData,
        batch_size: Union[List[int], int],
        epochs: int = 1,
        batch_strategy: Text = SEQUENCE,
        shuffle: bool = True,
    ):
        """Initializes the increasing batch size data generator.

        Args:
            model_data: The model data to use.
            batch_size: The batch size.
            epochs: The total number of epochs.
            batch_strategy: The batch strategy.
            shuffle: If 'True', data will be shuffled.
        """
        super().__init__(model_data, batch_size, batch_strategy, shuffle)

        if isinstance(batch_size, list):
            logger.debug(
                "The provided batch size is a list, this data generator will use a "
                "linear increasing batch size."
            )

        self._epochs = epochs
        # we use `on_epoch_end` method to prepare data for the next epoch
        # set current epoch to `-1`, so that `on_epoch_end` will increase it to `0`
        self._current_epoch = -1
        # actual batch size will be set inside `on_epoch_end`
        self._current_batch_size = None
        # create separate data variable that will store modified data for each batch
        self._data = None
        self.on_epoch_end()

    def __len__(self) -> int:
        """Number of batches in the Sequence.

        Returns:
            The number of batches in the Sequence.
        """
        # data was rebalanced, so need to recalculate number of examples
        num_examples = self.model_data.number_of_examples(self._data)
        batch_size = self._current_batch_size
        return num_examples // batch_size + int(num_examples % batch_size > 0)

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """Gets batch at position `index`.

        Arguments:
            index: position of the batch in the Sequence.

        Returns:
            A batch (tuple of input data and target data).
        """
        start = index * self._current_batch_size
        end = start + self._current_batch_size

        # return input and target data, as our target data is inside the input
        # data return None for the target data
        return self.prepare_batch(self._data, start, end), None

    def on_epoch_end(self) -> None:
        """Update the data after every epoch."""
        self._current_epoch += 1
        self._current_batch_size = self._linearly_increasing_batch_size()
        self._data = self._shuffle_and_balance(self._current_batch_size)

    def _linearly_increasing_batch_size(self) -> int:
        """Linearly increase batch size with every epoch.

        The idea comes from https://arxiv.org/abs/1711.00489.

        Returns:
            The batch size to use in this epoch.
        """
        if not isinstance(self.batch_size, list):
            return int(self.batch_size)

        if self._epochs > 1:
            return int(
                self.batch_size[0]
                + self._current_epoch
                * (self.batch_size[1] - self.batch_size[0])
                / (self._epochs - 1)
            )
        else:
            return int(self.batch_size[0])
<|endoftext|>


<commit_before># Copyright 2018 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Colab-specific Google Drive integration."""

from __future__ import absolute_import as _
from __future__ import division as _
from __future__ import print_function as _

import collections as _collections
import os as _os
import signal as _signal
import socket as _socket
import subprocess as _subprocess
import sys as _sys
import tempfile as _tempfile
import uuid as _uuid

import pexpect.popen_spawn as _popen_spawn

from google.colab import output as _output

__all__ = ['flush_and_unmount', 'mount']

_Environment = _collections.namedtuple(
    '_Environment',
    ('home', 'root_dir', 'inet_family', 'dev', 'path', 'config_dir'))

def _env():
  """Create and return an _Environment to use."""
  home = _os.environ['HOME']
  root_dir = _os.path.realpath(
      _os.path.join(_os.environ['CLOUDSDK_CONFIG'], '../..'))
  inet_family = 'IPV4_ONLY'
  dev = '/dev/fuse'
  path = '/usr/local/bin:/usr/local/sbin:/usr/bin:/usr/sbin:/bin:/sbin:.'
  if len(root_dir) > 1 and not root_dir.startswith('/usr/local/google/'):
    home = _os.path.join(root_dir, home)
    inet_family = 'IPV6_ONLY'
    fum = _os.environ['HOME'].split('mount')[0] + '/mount/alloc/fusermount'
    dev = fum + '/dev/fuse'
    path = path + ':' + fum + '/bin'
  config_dir = _os.path.join(home, '.config', 'Google')
  return _Environment(
      home=home,
      root_dir=root_dir,
      inet_family=inet_family,
      dev=dev,
      path=path,
      config_dir=config_dir)

def _logs_dir():
  return _os.path.join(_env().config_dir, 'DriveFS/Logs/')

def _timeouts_path():
  return _os.path.join(_logs_dir(), 'timeouts.txt')

def flush_and_unmount(timeout_ms=24 * 60 * 60 * 1000):
  """Unmount Google Drive and flush any outstanding writes to it."""
  env = _env()
  if b'type fuse.drive' not in _subprocess.check_output(['/bin/mount']):
    print('Drive not mounted, so nothing to flush and unmount.')
    return
  drive_bin = _os.path.join(env.root_dir, 'opt/google/drive/drive')
  p = _subprocess.Popen([
      drive_bin, '--push_changes_and_quit', '--single_process',
      '--timeout_sec={}'.format(int(timeout_ms / 1000))
  ],
                        stdout=_subprocess.PIPE,
                        stderr=_subprocess.PIPE)
  out, err = p.communicate()
  if mount._DEBUG:  # pylint:disable=protected-access
    print('flush_and_unmount: out: {}\nerr: {}'.format(out, err))
  if p.returncode:
    raise ValueError('flush_and_unmount failed')

def mount(mountpoint,
          force_remount=False,
          timeout_ms=60000,
          use_metadata_server=False):
  """Mount your Google Drive at the specified mountpoint path."""

  if ' ' in mountpoint:
    raise ValueError('Mountpoint must not contain a space.')

  mountpoint = _os.path.expanduser(mountpoint)
  # If we've already mounted drive at the specified mountpoint, exit now.
  already_mounted = _os.path.isdir(_os.path.join(mountpoint, 'My Drive'))
  if not force_remount and already_mounted:
    print(
        'Drive already mounted at {mnt}; to attempt to forcibly remount, '
        'call drive.mount("{mnt}", force_remount=True).'.format(mnt=mountpoint))
    return

  env = _env()
  home = env.home
  root_dir = env.root_dir
  inet_family = env.inet_family
  dev = env.dev
  path = env.path
  config_dir = env.config_dir

  try:
    _os.makedirs(config_dir)
  except OSError:
    if not _os.path.isdir(config_dir):
      raise ValueError('{} must be a directory if present'.format(config_dir))

  # Launch an intermediate bash to manage DriveFS' I/O (b/141747058#comment6).
  prompt = u'root@{}-{}: '.format(_socket.gethostname(), _uuid.uuid4().hex)
  logfile = None
  if mount._DEBUG:  # pylint:disable=protected-access
    logfile = _sys.stdout
  d = _popen_spawn.PopenSpawn(
      '/usr/bin/setsid /bin/bash --noediting -i',  # Need -i to get prompt echo.
      timeout=120,
      maxread=int(1e6),
      encoding='utf-8',
      logfile=logfile,
      env={
          'HOME': home,
          'FUSE_DEV_NAME': dev,
          'PATH': path
      })
  d.sendline('export PS1="{}"'.format(prompt))
  d.expect(prompt)  # The new prompt.
  drive_dir = _os.path.join(root_dir, 'opt/google/drive')
  # Robustify to previously-running copies of drive. Don't only [pkill -9]
  # because that leaves enough cruft behind in the mount table that future
  # operations fail with "Transport endpoint is not connected".
  d.sendline('umount -f {mnt} || umount {mnt}; pkill -9 -x drive'.format(
      mnt=mountpoint))
  # Wait for above to be received, using the next prompt.
  d.expect(prompt)
  d.sendline('pkill -9 -f {d}/directoryprefetcher_binary'.format(d=drive_dir))
  d.expect(prompt)
  # Only check the mountpoint after potentially unmounting/pkill'ing above.
  try:
    if _os.path.islink(mountpoint):
      raise ValueError('Mountpoint must not be a symlink')
    if _os.path.isdir(mountpoint) and _os.listdir(mountpoint):
      raise ValueError('Mountpoint must not already contain files')
    if not _os.path.isdir(mountpoint) and _os.path.exists(mountpoint):
      raise ValueError('Mountpoint must either be a directory or not exist')
    normed = _os.path.normpath(mountpoint)
    if '/' in normed and not _os.path.exists(_os.path.dirname(normed)):
      raise ValueError('Mountpoint must be in a directory that exists')
  except:
    d.kill(_signal.SIGKILL)
    raise

  # Watch for success.
  success = u'google.colab.drive MOUNTED'
  success_watcher = (
      '( while `sleep 0.5`; do if [[ -d "{m}" && "$(ls -A {m})" != "" ]]; '
      'then echo "{s}"; break; fi; done ) &').format(
          m=mountpoint, s=success)
  d.sendline(success_watcher)
  d.expect(prompt)

  oauth_prompt = u'(Go to this URL in a browser: https://.*)$'
  oauth_failed = u'Authorization failed'
  domain_disabled_drivefs = u'The domain policy has disabled Drive File Stream'
  problem_and_stopped = (
      u'Drive File Stream encountered a problem and has stopped')
  drive_exited = u'drive EXITED'
  metadata_auth_arg = (
      '--metadata_server_auth_uri={metadata_server}/computeMetadata/v1 '.format(
          metadata_server=_os.environ['TBE_CREDS_ADDR'])
      if use_metadata_server else '')

  # Create a pipe for sending the oauth code to a backgrounded drive binary.
  # (popen -> no pty -> no bash job control -> can't background post-launch).
  fifo_dir = _tempfile.mkdtemp()
  fifo = _os.path.join(fifo_dir, 'drive.fifo')
  _os.mkfifo(fifo)
  # cat is needed below since the FIFO isn't opened for writing yet.
  d.sendline((
      'cat {fifo} | head -1 | ( {d}/drive '
      '--features=' + ','.join([
          'fuse_max_background:1000',
          'max_bytes_per_fetch_content_request:134217728',
          'max_read_qps:1000',
          'max_write_qps:1000',
          'max_operation_batch_size:15',
          'max_parallel_push_task_instances:10',
          'opendir_timeout_ms:{timeout_ms}',
          'shortcut_support:true',
          'virtual_folders_omit_spaces:true',
      ]) + ' '
      '--inet_family=' + inet_family + ' ' + metadata_auth_arg +
      '--preferences=trusted_root_certs_file_path:'
      '{d}/roots.pem,mount_point_path:{mnt} --console_auth 2>&1 '
      '| grep --line-buffered -E "{oauth_prompt}|{problem_and_stopped}|{oauth_failed}|{domain_disabled_drivefs}"; '
      'echo "{drive_exited}"; ) &').format(
          d=drive_dir,
          timeout_ms=timeout_ms,
          mnt=mountpoint,
          fifo=fifo,
          oauth_failed=oauth_failed,
          domain_disabled_drivefs=domain_disabled_drivefs,
          oauth_prompt=oauth_prompt,
          problem_and_stopped=problem_and_stopped,
          drive_exited=drive_exited))
  d.expect(prompt)

  # LINT.IfChange(drivetimedout)
  timeout_pattern = 'QueryManager timed out'
  # LINT.ThenChange()
  dfs_log = _os.path.join(_logs_dir(), 'drive_fs.txt')

  # TODO(b/147296819): Delete this line.
  get_code = input if _sys.version_info[0] == 3 else raw_input  # pylint: disable=undefined-variable

  wrote_to_fifo = False
  while True:
    case = d.expect([
        success,
        prompt,
        oauth_prompt,
        problem_and_stopped,
        drive_exited,
        oauth_failed,
        domain_disabled_drivefs,
    ])
    if case == 0:
      break
    elif (case == 1 or case == 3 or case == 4):
      # Prompt appearing here means something went wrong with the drive binary.
      d.kill(_signal.SIGKILL)
      extra_reason = ''
      if 0 == _subprocess.call(
          'grep -q "{}" "{}"'.format(timeout_pattern, dfs_log), shell=True):
        extra_reason = (
            ': timeout during initial read of root folder; for more info: '
            'https://research.google.com/colaboratory/faq.html#drive-timeout')
      raise ValueError('mount failed' + extra_reason)
    elif case == 2:
      # Not already authorized, so do the authorization dance.
      auth_prompt = d.match.group(1) + '\nEnter your authorization code:\n'
      with _output.use_tags('dfs-auth-dance'):
        with open(fifo, 'w') as fifo_file:
          fifo_file.write(get_code(auth_prompt) + '\n')
      wrote_to_fifo = True
    elif case == 5:
      raise ValueError('mount failed: invalid oauth code')
    elif case == 6:
      raise ValueError(
          str(domain_disabled_drivefs) +
          ': https://support.google.com/a/answer/7496409')
  if not wrote_to_fifo:
    with open(fifo, 'w') as fifo_file:
      fifo_file.write('ignored\n')
  filtered_logfile = _timeouts_path()
  d.sendline('fuser -kw "{f}" ; rm -rf "{f}"'.format(f=filtered_logfile))
  d.expect(prompt)
  filter_script = _os.path.join(drive_dir, 'drive-filter.py')
  filter_cmd = (
      """nohup bash -c 'tail -n +0 -F "{}" | """
      """python3 {} > "{}" ' < /dev/null > /dev/null 2>&1 &""").format(
          dfs_log, filter_script, filtered_logfile)
  d.sendline(filter_cmd)
  d.expect(prompt)
  if 'ENABLE_DIRECTORYPREFETCHER' in _os.environ:
    d.sendline(
        """nohup bash -c '{d}/directoryprefetcher_binary -mountpoint={mnt}' """
        """>> {log} 2>&1 &""".format(
            d=drive_dir,
            mnt=mountpoint,
            log=_os.path.join(_logs_dir(), 'dpb.txt')))
    d.expect(prompt)
  d.sendline('disown -a')
  d.expect(prompt)
  d.sendline('exit')
  assert d.wait() == 0
  _output.clear(wait=True, output_tags='dfs-auth-dance')
  print('Mounted at {}'.format(mountpoint))

mount._DEBUG = False  # pylint:disable=protected-access
<commit_msg>Internal update.<commit_after># Copyright 2018 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Colab-specific Google Drive integration."""

from __future__ import absolute_import as _
from __future__ import division as _
from __future__ import print_function as _

import collections as _collections
import os as _os
import signal as _signal
import socket as _socket
import subprocess as _subprocess
import sys as _sys
import tempfile as _tempfile
import uuid as _uuid

import pexpect.popen_spawn as _popen_spawn

from google.colab import output as _output

__all__ = ['flush_and_unmount', 'mount']

_Environment = _collections.namedtuple(
    '_Environment',
    ('home', 'root_dir', 'inet_family', 'dev', 'path', 'config_dir'))

def _env():
  """Create and return an _Environment to use."""
  home = _os.environ['HOME']
  root_dir = _os.path.realpath(
      _os.path.join(_os.environ['CLOUDSDK_CONFIG'], '../..'))
  inet_family = 'IPV4_ONLY'
  dev = '/dev/fuse'
  path = '/usr/local/bin:/usr/local/sbin:/usr/bin:/usr/sbin:/bin:/sbin:.'
  if len(root_dir) > 1 and not root_dir.startswith('/usr/local/google/'):
    home = _os.path.join(root_dir, home)
    inet_family = 'IPV6_ONLY'
    fum = _os.environ['HOME'].split('mount')[0] + '/mount/alloc/fusermount'
    if 'BORG_ALLOC_DIR' in _os.environ:
      # For support in guitar cluster invocations.
      fum = _os.environ['BORG_ALLOC_DIR'] + '/alloc/fusermount'
    dev = fum + '/dev/fuse'
    path = path + ':' + fum + '/bin'
  config_dir = _os.path.join(home, '.config', 'Google')
  return _Environment(
      home=home,
      root_dir=root_dir,
      inet_family=inet_family,
      dev=dev,
      path=path,
      config_dir=config_dir)

def _logs_dir():
  return _os.path.join(_env().config_dir, 'DriveFS/Logs/')

def _timeouts_path():
  return _os.path.join(_logs_dir(), 'timeouts.txt')

def flush_and_unmount(timeout_ms=24 * 60 * 60 * 1000):
  """Unmount Google Drive and flush any outstanding writes to it."""
  env = _env()
  if b'type fuse.drive' not in _subprocess.check_output(['/bin/mount']):
    print('Drive not mounted, so nothing to flush and unmount.')
    return
  drive_bin = _os.path.join(env.root_dir, 'opt/google/drive/drive')
  p = _subprocess.Popen([
      drive_bin, '--push_changes_and_quit', '--single_process',
      '--timeout_sec={}'.format(int(timeout_ms / 1000))
  ],
                        stdout=_subprocess.PIPE,
                        stderr=_subprocess.PIPE)
  out, err = p.communicate()
  if mount._DEBUG:  # pylint:disable=protected-access
    print('flush_and_unmount: out: {}\nerr: {}'.format(out, err))
  if p.returncode:
    raise ValueError('flush_and_unmount failed')

def mount(mountpoint,
          force_remount=False,
          timeout_ms=60000,
          use_metadata_server=False):
  """Mount your Google Drive at the specified mountpoint path."""

  if ' ' in mountpoint:
    raise ValueError('Mountpoint must not contain a space.')

  mountpoint = _os.path.expanduser(mountpoint)
  # If we've already mounted drive at the specified mountpoint, exit now.
  already_mounted = _os.path.isdir(_os.path.join(mountpoint, 'My Drive'))
  if not force_remount and already_mounted:
    print(
        'Drive already mounted at {mnt}; to attempt to forcibly remount, '
        'call drive.mount("{mnt}", force_remount=True).'.format(mnt=mountpoint))
    return

  env = _env()
  home = env.home
  root_dir = env.root_dir
  inet_family = env.inet_family
  dev = env.dev
  path = env.path
  config_dir = env.config_dir

  try:
    _os.makedirs(config_dir)
  except OSError:
    if not _os.path.isdir(config_dir):
      raise ValueError('{} must be a directory if present'.format(config_dir))

  # Launch an intermediate bash to manage DriveFS' I/O (b/141747058#comment6).
  prompt = u'root@{}-{}: '.format(_socket.gethostname(), _uuid.uuid4().hex)
  logfile = None
  if mount._DEBUG:  # pylint:disable=protected-access
    logfile = _sys.stdout
  d = _popen_spawn.PopenSpawn(
      '/usr/bin/setsid /bin/bash --noediting -i',  # Need -i to get prompt echo.
      timeout=120,
      maxread=int(1e6),
      encoding='utf-8',
      logfile=logfile,
      env={
          'HOME': home,
          'FUSE_DEV_NAME': dev,
          'PATH': path
      })
  d.sendline('export PS1="{}"'.format(prompt))
  d.expect(prompt)  # The new prompt.
  drive_dir = _os.path.join(root_dir, 'opt/google/drive')
  # Robustify to previously-running copies of drive. Don't only [pkill -9]
  # because that leaves enough cruft behind in the mount table that future
  # operations fail with "Transport endpoint is not connected".
  d.sendline('umount -f {mnt} || umount {mnt}; pkill -9 -x drive'.format(
      mnt=mountpoint))
  # Wait for above to be received, using the next prompt.
  d.expect(prompt)
  d.sendline('pkill -9 -f {d}/directoryprefetcher_binary'.format(d=drive_dir))
  d.expect(prompt)
  # Only check the mountpoint after potentially unmounting/pkill'ing above.
  try:
    if _os.path.islink(mountpoint):
      raise ValueError('Mountpoint must not be a symlink')
    if _os.path.isdir(mountpoint) and _os.listdir(mountpoint):
      raise ValueError('Mountpoint must not already contain files')
    if not _os.path.isdir(mountpoint) and _os.path.exists(mountpoint):
      raise ValueError('Mountpoint must either be a directory or not exist')
    normed = _os.path.normpath(mountpoint)
    if '/' in normed and not _os.path.exists(_os.path.dirname(normed)):
      raise ValueError('Mountpoint must be in a directory that exists')
  except:
    d.kill(_signal.SIGKILL)
    raise

  # Watch for success.
  success = u'google.colab.drive MOUNTED'
  success_watcher = (
      '( while `sleep 0.5`; do if [[ -d "{m}" && "$(ls -A {m})" != "" ]]; '
      'then echo "{s}"; break; fi; done ) &').format(
          m=mountpoint, s=success)
  d.sendline(success_watcher)
  d.expect(prompt)

  oauth_prompt = u'(Go to this URL in a browser: https://.*)$'
  oauth_failed = u'Authorization failed'
  domain_disabled_drivefs = u'The domain policy has disabled Drive File Stream'
  problem_and_stopped = (
      u'Drive File Stream encountered a problem and has stopped')
  drive_exited = u'drive EXITED'
  metadata_auth_arg = (
      '--metadata_server_auth_uri={metadata_server}/computeMetadata/v1 '.format(
          metadata_server=_os.environ['TBE_CREDS_ADDR'])
      if use_metadata_server else '')

  # Create a pipe for sending the oauth code to a backgrounded drive binary.
  # (popen -> no pty -> no bash job control -> can't background post-launch).
  fifo_dir = _tempfile.mkdtemp()
  fifo = _os.path.join(fifo_dir, 'drive.fifo')
  _os.mkfifo(fifo)
  # cat is needed below since the FIFO isn't opened for writing yet.
  d.sendline((
      'cat {fifo} | head -1 | ( {d}/drive '
      '--features=' + ','.join([
          'fuse_max_background:1000',
          'max_bytes_per_fetch_content_request:134217728',
          'max_read_qps:1000',
          'max_write_qps:1000',
          'max_operation_batch_size:15',
          'max_parallel_push_task_instances:10',
          'opendir_timeout_ms:{timeout_ms}',
          'shortcut_support:true',
          'virtual_folders_omit_spaces:true',
      ]) + ' '
      '--inet_family=' + inet_family + ' ' + metadata_auth_arg +
      '--preferences=trusted_root_certs_file_path:'
      '{d}/roots.pem,mount_point_path:{mnt} --console_auth 2>&1 '
      '| grep --line-buffered -E "{oauth_prompt}|{problem_and_stopped}|{oauth_failed}|{domain_disabled_drivefs}"; '
      'echo "{drive_exited}"; ) &').format(
          d=drive_dir,
          timeout_ms=timeout_ms,
          mnt=mountpoint,
          fifo=fifo,
          oauth_failed=oauth_failed,
          domain_disabled_drivefs=domain_disabled_drivefs,
          oauth_prompt=oauth_prompt,
          problem_and_stopped=problem_and_stopped,
          drive_exited=drive_exited))
  d.expect(prompt)

  # LINT.IfChange(drivetimedout)
  timeout_pattern = 'QueryManager timed out'
  # LINT.ThenChange()
  dfs_log = _os.path.join(_logs_dir(), 'drive_fs.txt')

  # TODO(b/147296819): Delete this line.
  get_code = input if _sys.version_info[0] == 3 else raw_input  # pylint: disable=undefined-variable

  wrote_to_fifo = False
  while True:
    case = d.expect([
        success,
        prompt,
        oauth_prompt,
        problem_and_stopped,
        drive_exited,
        oauth_failed,
        domain_disabled_drivefs,
    ])
    if case == 0:
      break
    elif (case == 1 or case == 3 or case == 4):
      # Prompt appearing here means something went wrong with the drive binary.
      d.kill(_signal.SIGKILL)
      extra_reason = ''
      if 0 == _subprocess.call(
          'grep -q "{}" "{}"'.format(timeout_pattern, dfs_log), shell=True):
        extra_reason = (
            ': timeout during initial read of root folder; for more info: '
            'https://research.google.com/colaboratory/faq.html#drive-timeout')
      raise ValueError('mount failed' + extra_reason)
    elif case == 2:
      # Not already authorized, so do the authorization dance.
      auth_prompt = d.match.group(1) + '\nEnter your authorization code:\n'
      with _output.use_tags('dfs-auth-dance'):
        with open(fifo, 'w') as fifo_file:
          fifo_file.write(get_code(auth_prompt) + '\n')
      wrote_to_fifo = True
    elif case == 5:
      raise ValueError('mount failed: invalid oauth code')
    elif case == 6:
      raise ValueError(
          str(domain_disabled_drivefs) +
          ': https://support.google.com/a/answer/7496409')
  if not wrote_to_fifo:
    with open(fifo, 'w') as fifo_file:
      fifo_file.write('ignored\n')
  filtered_logfile = _timeouts_path()
  d.sendline('fuser -kw "{f}" ; rm -rf "{f}"'.format(f=filtered_logfile))
  d.expect(prompt)
  filter_script = _os.path.join(drive_dir, 'drive-filter.py')
  filter_cmd = (
      """nohup bash -c 'tail -n +0 -F "{}" | """
      """python3 {} > "{}" ' < /dev/null > /dev/null 2>&1 &""").format(
          dfs_log, filter_script, filtered_logfile)
  d.sendline(filter_cmd)
  d.expect(prompt)
  if 'ENABLE_DIRECTORYPREFETCHER' in _os.environ:
    d.sendline(
        """nohup bash -c '{d}/directoryprefetcher_binary -mountpoint={mnt}' """
        """>> {log} 2>&1 &""".format(
            d=drive_dir,
            mnt=mountpoint,
            log=_os.path.join(_logs_dir(), 'dpb.txt')))
    d.expect(prompt)
  d.sendline('disown -a')
  d.expect(prompt)
  d.sendline('exit')
  assert d.wait() == 0
  _output.clear(wait=True, output_tags='dfs-auth-dance')
  print('Mounted at {}'.format(mountpoint))

mount._DEBUG = False  # pylint:disable=protected-access
<|endoftext|>


<commit_before>import os
import sys
import shlex
import signal
import threading
import itertools
import subprocess
from .utils import *

if sys.version_info >= (3, 0):
    import queue as Queue
else:
    import Queue

class AsyncExecutor(object):
    """
    A class to implement executing a command in subprocess, then
    read the output asynchronously.
    """
    def __init__(self):
        self._errQueue = Queue.Queue()
        self._process = None
        self._finished = False
        self._max_count = int(lfEval("g:Lf_MaxCount"))

    def _readerThread(self, fd, queue):
        try:
            for line in iter(fd.readline, b""):
                queue.put(line)
        except ValueError:
            pass
        finally:
            queue.put(None)

    def execute(self, cmd, encoding=None, cleanup=None, env=None):
        if os.name == 'nt':
            self._process = subprocess.Popen(cmd, bufsize=-1,
                                             stdin=subprocess.PIPE,
                                             stdout=subprocess.PIPE,
                                             stderr=subprocess.PIPE,
                                             shell=True,
                                             env=env,
                                             universal_newlines=False)
        else:
            self._process = subprocess.Popen(cmd, bufsize=-1,
                                             stdout=subprocess.PIPE,
                                             stderr=subprocess.PIPE,
                                             preexec_fn=os.setsid,
                                             shell=True,
                                             env=env,
                                             universal_newlines=False)

        self._finished = False

        stderr_thread = threading.Thread(target=self._readerThread,
                                         args=(self._process.stderr, self._errQueue))
        stderr_thread.daemon = True
        stderr_thread.start()

        def read(source):
            try:
                count = 0
                if encoding:
                    for line in source:
                        yield lfBytes2Str(line.rstrip(b"\r\n"), encoding)
                        if self._max_count > 0:
                            count += 1
                            if count >= self._max_count:
                                self.killProcess()
                                break
                else:
                    for line in source:
                        yield lfEncode(lfBytes2Str(line.rstrip(b"\r\n")))
                        if self._max_count > 0:
                            count += 1
                            if count >= self._max_count:
                                self.killProcess()
                                break

                err = b"".join(iter(self._errQueue.get, None))
                if err:
                    raise Exception(lfBytes2Str(err, encoding))
            except ValueError:
                pass
            finally:
                self._finished = True
                try:
                    if self._process:
                        self._process.stdout.close()
                        self._process.stderr.close()
                except IOError:
                    pass

                if cleanup:
                    cleanup()

        result = AsyncExecutor.Result(read(iter(self._process.stdout.readline, b"")))

        return result

    def killProcess(self):
        # Popen.poll always returns None, bug?
        # if self._process and not self._process.poll():
        if self._process and not self._finished:
            if os.name == 'nt':
                subprocess.Popen("TASKKILL /F /PID {pid} /T".format(pid=self._process.pid), shell=True)
            else:
                try:
                    os.killpg(os.getpgid(self._process.pid), signal.SIGTERM)
                except OSError:
                    pass

            self._process = None

    class Result(object):
        def __init__(self, iterable):
            self._g = iterable

        def __add__(self, iterable):
            self._g = itertools.chain(self._g, iterable)
            return self

        def __iadd__(self, iterable):
            self._g = itertools.chain(self._g, iterable)
            return self

        def join_left(self, iterable):
            self._g = itertools.chain(iterable, self._g)
            return self

        def __iter__(self):
            return self._g

if __name__ == "__main__":
    executor = AsyncExecutor()
    out = executor.execute("ctags -f- -R")
    print("stdout begin: ============================================")
    for i in out:
        print(repr(i))
    print("stdout end: ==============================================")
<commit_msg>optimize the asyncExecutor<commit_after>import os
import sys
import shlex
import signal
import threading
import itertools
import subprocess
from .utils import *

if sys.version_info >= (3, 0):
    import queue as Queue
else:
    import Queue

class AsyncExecutor(object):
    """
    A class to implement executing a command in subprocess, then
    read the output asynchronously.
    """
    def __init__(self):
        self._errQueue = Queue.Queue()
        self._process = None
        self._finished = False
        self._max_count = int(lfEval("g:Lf_MaxCount"))

    def _readerThread(self, fd, queue):
        try:
            for line in iter(fd.readline, b""):
                queue.put(line)
        except ValueError:
            pass
        finally:
            queue.put(None)

    def execute(self, cmd, encoding=None, cleanup=None, env=None):
        if os.name == 'nt':
            self._process = subprocess.Popen(cmd, bufsize=-1,
                                             stdin=subprocess.PIPE,
                                             stdout=subprocess.PIPE,
                                             stderr=subprocess.PIPE,
                                             shell=True,
                                             env=env,
                                             universal_newlines=False)
        else:
            self._process = subprocess.Popen(cmd, bufsize=-1,
                                             stdout=subprocess.PIPE,
                                             stderr=subprocess.PIPE,
                                             preexec_fn=os.setsid,
                                             shell=True,
                                             env=env,
                                             universal_newlines=False)

        self._finished = False

        stderr_thread = threading.Thread(target=self._readerThread,
                                         args=(self._process.stderr, self._errQueue))
        stderr_thread.daemon = True
        stderr_thread.start()

        if sys.version_info >= (3, 0):
            def read(source):
                try:
                    count = 0
                    if encoding:
                        for line in source:
                            try:
                                line.decode("ascii")
                                yield line.rstrip(b"\r\n").decode()
                            except UnicodeDecodeError:
                                yield lfBytes2Str(line.rstrip(b"\r\n"), encoding)
                            if self._max_count > 0:
                                count += 1
                                if count >= self._max_count:
                                    self.killProcess()
                                    break
                    else:
                        for line in source:
                            try:
                                line.decode("ascii")
                                yield line.rstrip(b"\r\n").decode()
                            except UnicodeDecodeError:
                                yield lfBytes2Str(line.rstrip(b"\r\n"))
                            if self._max_count > 0:
                                count += 1
                                if count >= self._max_count:
                                    self.killProcess()
                                    break

                    err = b"".join(iter(self._errQueue.get, None))
                    if err:
                        raise Exception(lfBytes2Str(err, encoding))
                except ValueError:
                    pass
                finally:
                    self._finished = True
                    try:
                        if self._process:
                            self._process.stdout.close()
                            self._process.stderr.close()
                    except IOError:
                        pass

                    if cleanup:
                        cleanup()
        else:
            def read(source):
                try:
                    count = 0
                    if encoding:
                        for line in source:
                            yield line.rstrip(b"\r\n")
                            if self._max_count > 0:
                                count += 1
                                if count >= self._max_count:
                                    self.killProcess()
                                    break
                    else:
                        for line in source:
                            try:
                                line.decode("ascii")
                                yield line.rstrip(b"\r\n")
                            except UnicodeDecodeError:
                                yield lfEncode(line.rstrip(b"\r\n"))
                            if self._max_count > 0:
                                count += 1
                                if count >= self._max_count:
                                    self.killProcess()
                                    break

                    err = b"".join(iter(self._errQueue.get, None))
                    if err:
                        raise Exception(err)
                except ValueError:
                    pass
                finally:
                    self._finished = True
                    try:
                        if self._process:
                            self._process.stdout.close()
                            self._process.stderr.close()
                    except IOError:
                        pass

                    if cleanup:
                        cleanup()

        result = AsyncExecutor.Result(read(iter(self._process.stdout.readline, b"")))

        return result

    def killProcess(self):
        # Popen.poll always returns None, bug?
        # if self._process and not self._process.poll():
        if self._process and not self._finished:
            if os.name == 'nt':
                subprocess.Popen("TASKKILL /F /PID {pid} /T".format(pid=self._process.pid), shell=True)
            else:
                try:
                    os.killpg(os.getpgid(self._process.pid), signal.SIGTERM)
                except OSError:
                    pass

            self._process = None

    class Result(object):
        def __init__(self, iterable):
            self._g = iterable

        def __add__(self, iterable):
            self._g = itertools.chain(self._g, iterable)
            return self

        def __iadd__(self, iterable):
            self._g = itertools.chain(self._g, iterable)
            return self

        def join_left(self, iterable):
            self._g = itertools.chain(iterable, self._g)
            return self

        def __iter__(self):
            return self._g

if __name__ == "__main__":
    executor = AsyncExecutor()
    out = executor.execute("ctags -f- -R")
    print("stdout begin: ============================================")
    for i in out:
        print(repr(i))
    print("stdout end: ==============================================")
<|endoftext|>


<commit_before>import os, errno, StringIO

from PIL import Image
import numpy

from django.core.exceptions import ImproperlyConfigured
from django.core.files.uploadedfile import InMemoryUploadedFile

from .settings import (
    QUAL,
    USE_PLACEHOLDIT,
    cache,
    SIZEDIMAGEFIELD_CACHE_LENGTH,
    SIZEDIMAGEFIELD_PLACEHOLDER_IMAGE
)
from .utils import (
    get_filtered_path,
    get_resized_path,
    get_image_format_from_file_extension
)

if not USE_PLACEHOLDIT:
    SIZEDIMAGEFIELD_PLACEHOLDER_FOLDER,
    SIZEDIMAGEFIELD_PLACEHOLDER_FILENAME = os.path.split(SIZEDIMAGEFIELD_PLACEHOLDER_IMAGE)

class ProcessedImage(object):

    def __init__(self, path_to_image, storage):
        self.path_to_image = path_to_image
        self.storage = storage

    def process_image(self, image, image_format, **kwargs):
        """
        Arguments:
            * `image`: a PIL Image instance
            * `width`: an int representing the intended width
            * `height`: an int representing the intended height
            * `image_format`: A valid image mime type (e.g. 'image/jpeg')

        Returns a StringIO.StringIO representation of the resized image.

        Subclasses MUST implement this method.
        """
        raise NotImplementedError('Subclasses MUST provide a `process_image` method.')

    def preprocess(self, image, image_format):
        save_kwargs = {'format':image_format}
        if hasattr(self, 'preprocess_%s' % image_format):
            image, addl_save_kwargs = getattr(self, 'preprocess_%s' % image_format)(
                image=image
            )
            save_kwargs.update(addl_save_kwargs)

        return image, save_kwargs

    def retrieve_image(self, path_to_image, prepared_path):
        if not path_to_image:
            containing_folder, filename = os.path.split(prepared_path)
            if not os.path.exists(containing_folder):
                try:
                    os.makedirs(containing_folder)
                except OSError as exc:
                    if exc.errno == errno.EEXIST and os.path.isdir(path):
                        pass
                    else:
                        raise Exception("Can't create directory: '%s' - This is probably due to a permissions issue." % (containing_folder))
            image = SIZEDIMAGEFIELD_PLACEHOLDER_IMAGE
        else:
            image = self.storage.open(path_to_image, 'r')
        file_ext = prepared_path.rsplit('.')[-1]
        image_format, mime_type = get_image_format_from_file_extension(file_ext)

        return (
            Image.open(image),
            file_ext,
            image_format,
            mime_type
        )

    def save_image(self, imagefile, save_path, file_ext, mime_type):
        image_in_memory = Image.open(
            StringIO.StringIO(
                imagefile.getvalue()
            )
        )

        file_to_save = InMemoryUploadedFile(
                            imagefile,
                            None,
                            'foo.%s' % file_ext,
                            mime_type,
                            imagefile.len,
                            None
                        )
        file_to_save.seek(0)
        self.storage.save(save_path, file_to_save)

        return True

class SizedImage(ProcessedImage, dict):

    def __init__(self, path_to_image, storage):

        if getattr(self, 'filename_key', None) is None:
            raise NotImplementedError(
                "%s instances MUST have a `filename_key` attribute" % self.__class__.__name__
            )
        super(SizedImage, self).__init__(path_to_image, storage)

    def get_filename_key(self):
        """
        Returns a string that will be used to identify the resized image.
        """
        if not getattr(self, 'filename_key'):
            raise AttributError('SizedImage subclasses must define a `filename_key` '
                                'attribute or override the `get_filename_key` method.')
        else:
            return self.filename_key

    def __setitem__(self, key, value):
        raise NotImplementedError(
            '%s instances do not allow key assignment.' % self.__class__.__name__
        )

    def __getitem__(self, key):
        try:
            width, height = [int(i) for i in key.split('x')]
        except KeyError:
            raise SizedImageDictKeyError(
                "%s keys must be in the following format: "
                "'`width`x`height`' where both `width` and `height` are "
                "integers." % self.__class__.__name__
            )

        if not self.path_to_image and USE_PLACEHOLDIT:
            resized_url = "http://placehold.it/%dx%d" % (width, height)
        else:
            resized_url = get_resized_path(
                path_to_image=self.path_to_image,
                width=width,
                height=height,
                filename_key=self.get_filename_key(),
                base_url=self.storage.base_url
            )
            if cache.get(resized_url):
                # The sized path exists in the cache so the image already exists.
                # So we `pass` to skip directly to the return statement.
                pass
            else:
                resized_image_path = get_resized_path(
                    path_to_image=self.path_to_image,
                    width=width,
                    height=height,
                    filename_key=self.get_filename_key()
                )
                if not self.storage.exists(resized_image_path):
                    image_created = self.create_resized_image(
                        path_to_image=self.path_to_image,
                        width=width,
                        height=height,
                        filename_key=self.get_filename_key()
                    )
                # Setting a super-long cache for a resized image (30 Days)
                cache.set(resized_url, 1, SIZEDIMAGEFIELD_CACHE_LENGTH)
        return resized_url

    def process_image(self, image, image_format, width, height, save_kwargs={}):
        """
        Arguments:
            * `image`: a PIL Image instance
            * `width`: value in pixels (as int) representing the intended width
            * `height`: value in pixels (as int) representing the intended height
            * `image_format`: A valid image mime type (e.g. 'image/jpeg')

        Returns a StringIO.StringIO representation of the resized image.

        Subclasses MUST implement this method.
        """
        raise NotImplementedError('Subclasses MUST provide a `process_image` method.')

    def preprocess_PNG(self, image, **kwargs):
        """
        Receives a PIL Image instance of a PNG and returns a 2-tuple:
            * [0]: Image instance with a properly processed alpha (transparency) layer that
                   is resize ready. (Found here: http://stackoverflow.com/a/9146202/1149774)
            * [1]: Empty dict ({})
        """
        premult = numpy.fromstring(image.tobytes(), dtype=numpy.uint8)
        alpha_layer = premult[3::4] / 255.0
        premult[::4] *= alpha_layer
        premult[1::4] *= alpha_layer
        premult[2::4] *= alpha_layer
        return (Image.frombytes("RGBA", image.size, premult.tostring()), {})

    def preprocess_GIF(self, image, **kwargs):
        """
        Receives a PIL Image instance of a GIF and returns 2-tuple:
            * [0]: Original Image instance (passed to `image`)
            * [1]: Dict with a transparency key (to GIF transparency layer)
        """
        return (image, {'transparency': image.info['transparency']})

    def preprocess_JPEG(self, image, **kwargs):
        """
        Receives a PIL Image instance of a JPEG and returns 2-tuple:
            * [0]: Image instance, converted to RGB
            * [1]: Dict with a quality key (mapped to the value of `QUAL` as
                   defined by the `SIZEDIMAGEFIELD_JPEG_RESIZE_QUALITY` setting)
        """
        if image.mode != 'RGB':
            image = image.convert('RGB')
        return (image, {'quality': QUAL})

    def create_resized_image(self, path_to_image, width, height, filename_key):
        """
        Creates a resized image.
        `path_to_image`: The path to the image with the media directory to resize.
            If `None`, the SIZEDIMAGEFIELD_PLACEHOLDER_IMAGE will be used.
        `width`: Width of resized image (int)
        `height`: Desired height of resized image (int)
        `filename_key`: A string that will be used in the sized image filename to signify
        what operation was done to it. Examples: 'crop' or 'scale'
        """
        resized_image_path = get_resized_path(
            path_to_image=path_to_image,
            width=width,
            height=height,
            filename_key=filename_key
        )

        image, file_ext, image_format, mime_type = self.retrieve_image(path_to_image, resized_image_path)

        image, save_kwargs = self.preprocess(image, image_format)

        imagefile = self.process_image(
            image=image,
            image_format=image_format,
            width=width,
            height=height,
            save_kwargs=save_kwargs
        )

        saved = self.save_image(imagefile, resized_image_path, file_ext, mime_type)

        return saved

class FilteredImage(ProcessedImage):
    name = None
    url = None

    def __init__(self, path_to_image, storage, filename_key):
        super(FilteredImage, self).__init__(path_to_image, storage)
        self.name = get_filtered_path(
            path_to_image=self.path_to_image,
            filename_key=filename_key
        )
        self.url = self.storage.url(self.name)

    def add_sizedimage_attrs(self, sizedimage_registry):
        for attr_name, sizedimage_cls in sizedimage_registry.iteritems():
            setattr(
                self,
                attr_name,
                sizedimage_cls(
                    path_to_image=self.name,
                    storage=self.storage
                )
            )

    def process_filter(self, image, image_format, save_kwargs={}):
        raise NotImplementedError('Subclasses MUST provide a `process_filter` method.')

    def create_filtered_image(self, path_to_image, prepared_path):
        """
        Creates a resized image.
        `path_to_image`: The path to the image with the media directory to resize.
        `prepared_path`:
        """

        image, file_ext, image_format, mime_type = self.retrieve_image(path_to_image, prepared_path)
        #print image.__class__
        image, save_kwargs = self.preprocess(image, image_format)
        imagefile = self.process_filter(image, image_format, save_kwargs)
        saved = self.save_image(imagefile, prepared_path, file_ext, mime_type)

        return saved

class Filters(dict):

    def __init__(self, sizedimagefield, registry):
        self.sizedimagefield = sizedimagefield
        self.registry = registry

    def __getitem__(self, key):
        try:
            prepped_filter = dict.__getitem__(self, key)
        except KeyError:
            if key not in self.registry._filter_registry:
                return None
            else:
                filtered_path = get_filtered_path(
                    path_to_image=self.sizedimagefield.name,
                    filename_key=key
                )
                filter_cls = self.registry._filter_registry[key]
                prepped_filter = filter_cls(
                    path_to_image=self.sizedimagefield.name,
                    storage=self.sizedimagefield.storage,
                    filename_key=key
                )
                if cache.get(filtered_path):
                    # The sized path exists in the cache so the image already exists.
                    # So we `pass` to skip directly to the return statement.
                    pass
                else:
                    if not self.sizedimagefield.storage.exists(filtered_path):
                        image_created = prepped_filter.create_filtered_image(
                            path_to_image=self.sizedimagefield.name,
                            prepared_path=filtered_path
                        )

                        # Setting a super-long cache for a resized image (30 Days)
                        cache.set(self.sizedimagefield.name, 1, SIZEDIMAGEFIELD_CACHE_LENGTH)
                for attr_name, sizedimage_cls in self.registry._sizedimage_registry.iteritems():
                    setattr(
                        prepped_filter,
                        attr_name,
                        sizedimage_cls(
                            path_to_image=filtered_path,
                            storage=self.sizedimagefield.storage
                        )
                    )
                self[key] = prepped_filter

        return dict.__getitem__(self, key)
<commit_msg>Ensuring cache entry is properly set<commit_after>import os, errno, StringIO

from PIL import Image
import numpy

from django.core.exceptions import ImproperlyConfigured
from django.core.files.uploadedfile import InMemoryUploadedFile

from .settings import (
    QUAL,
    USE_PLACEHOLDIT,
    cache,
    SIZEDIMAGEFIELD_CACHE_LENGTH,
    SIZEDIMAGEFIELD_PLACEHOLDER_IMAGE
)
from .utils import (
    get_filtered_path,
    get_resized_path,
    get_image_format_from_file_extension
)

if not USE_PLACEHOLDIT:
    SIZEDIMAGEFIELD_PLACEHOLDER_FOLDER,
    SIZEDIMAGEFIELD_PLACEHOLDER_FILENAME = os.path.split(SIZEDIMAGEFIELD_PLACEHOLDER_IMAGE)

class ProcessedImage(object):

    def __init__(self, path_to_image, storage):
        self.path_to_image = path_to_image
        self.storage = storage

    def process_image(self, image, image_format, **kwargs):
        """
        Arguments:
            * `image`: a PIL Image instance
            * `width`: an int representing the intended width
            * `height`: an int representing the intended height
            * `image_format`: A valid image mime type (e.g. 'image/jpeg')

        Returns a StringIO.StringIO representation of the resized image.

        Subclasses MUST implement this method.
        """
        raise NotImplementedError('Subclasses MUST provide a `process_image` method.')

    def preprocess(self, image, image_format):
        save_kwargs = {'format':image_format}
        if hasattr(self, 'preprocess_%s' % image_format):
            image, addl_save_kwargs = getattr(self, 'preprocess_%s' % image_format)(
                image=image
            )
            save_kwargs.update(addl_save_kwargs)

        return image, save_kwargs

    def retrieve_image(self, path_to_image, prepared_path):
        if not path_to_image:
            containing_folder, filename = os.path.split(prepared_path)
            if not os.path.exists(containing_folder):
                try:
                    os.makedirs(containing_folder)
                except OSError as exc:
                    if exc.errno == errno.EEXIST and os.path.isdir(path):
                        pass
                    else:
                        raise Exception("Can't create directory: '%s' - This is probably due to a permissions issue." % (containing_folder))
            image = SIZEDIMAGEFIELD_PLACEHOLDER_IMAGE
        else:
            image = self.storage.open(path_to_image, 'r')
        file_ext = prepared_path.rsplit('.')[-1]
        image_format, mime_type = get_image_format_from_file_extension(file_ext)

        return (
            Image.open(image),
            file_ext,
            image_format,
            mime_type
        )

    def save_image(self, imagefile, save_path, file_ext, mime_type):
        image_in_memory = Image.open(
            StringIO.StringIO(
                imagefile.getvalue()
            )
        )

        file_to_save = InMemoryUploadedFile(
                            imagefile,
                            None,
                            'foo.%s' % file_ext,
                            mime_type,
                            imagefile.len,
                            None
                        )
        file_to_save.seek(0)
        self.storage.save(save_path, file_to_save)

        return True

class SizedImage(ProcessedImage, dict):

    def __init__(self, path_to_image, storage):

        if getattr(self, 'filename_key', None) is None:
            raise NotImplementedError(
                "%s instances MUST have a `filename_key` attribute" % self.__class__.__name__
            )
        super(SizedImage, self).__init__(path_to_image, storage)

    def get_filename_key(self):
        """
        Returns a string that will be used to identify the resized image.
        """
        if not getattr(self, 'filename_key'):
            raise AttributError('SizedImage subclasses must define a `filename_key` '
                                'attribute or override the `get_filename_key` method.')
        else:
            return self.filename_key

    def __setitem__(self, key, value):
        raise NotImplementedError(
            '%s instances do not allow key assignment.' % self.__class__.__name__
        )

    def __getitem__(self, key):
        try:
            width, height = [int(i) for i in key.split('x')]
        except KeyError:
            raise SizedImageDictKeyError(
                "%s keys must be in the following format: "
                "'`width`x`height`' where both `width` and `height` are "
                "integers." % self.__class__.__name__
            )

        if not self.path_to_image and USE_PLACEHOLDIT:
            resized_url = "http://placehold.it/%dx%d" % (width, height)
        else:
            resized_url = get_resized_path(
                path_to_image=self.path_to_image,
                width=width,
                height=height,
                filename_key=self.get_filename_key(),
                base_url=self.storage.base_url
            )
            if cache.get(resized_url):
                # The sized path exists in the cache so the image already exists.
                # So we `pass` to skip directly to the return statement.
                pass
            else:
                resized_image_path = get_resized_path(
                    path_to_image=self.path_to_image,
                    width=width,
                    height=height,
                    filename_key=self.get_filename_key()
                )
                if not self.storage.exists(resized_image_path):
                    image_created = self.create_resized_image(
                        path_to_image=self.path_to_image,
                        width=width,
                        height=height,
                        filename_key=self.get_filename_key()
                    )
                # Setting a super-long cache for a resized image (30 Days)
                cache.set(resized_url, 1, SIZEDIMAGEFIELD_CACHE_LENGTH)
        return resized_url

    def process_image(self, image, image_format, width, height, save_kwargs={}):
        """
        Arguments:
            * `image`: a PIL Image instance
            * `width`: value in pixels (as int) representing the intended width
            * `height`: value in pixels (as int) representing the intended height
            * `image_format`: A valid image mime type (e.g. 'image/jpeg')

        Returns a StringIO.StringIO representation of the resized image.

        Subclasses MUST implement this method.
        """
        raise NotImplementedError('Subclasses MUST provide a `process_image` method.')

    def preprocess_PNG(self, image, **kwargs):
        """
        Receives a PIL Image instance of a PNG and returns a 2-tuple:
            * [0]: Image instance with a properly processed alpha (transparency) layer that
                   is resize ready. (Found here: http://stackoverflow.com/a/9146202/1149774)
            * [1]: Empty dict ({})
        """
        premult = numpy.fromstring(image.tobytes(), dtype=numpy.uint8)
        alpha_layer = premult[3::4] / 255.0
        premult[::4] *= alpha_layer
        premult[1::4] *= alpha_layer
        premult[2::4] *= alpha_layer
        return (Image.frombytes("RGBA", image.size, premult.tostring()), {})

    def preprocess_GIF(self, image, **kwargs):
        """
        Receives a PIL Image instance of a GIF and returns 2-tuple:
            * [0]: Original Image instance (passed to `image`)
            * [1]: Dict with a transparency key (to GIF transparency layer)
        """
        return (image, {'transparency': image.info['transparency']})

    def preprocess_JPEG(self, image, **kwargs):
        """
        Receives a PIL Image instance of a JPEG and returns 2-tuple:
            * [0]: Image instance, converted to RGB
            * [1]: Dict with a quality key (mapped to the value of `QUAL` as
                   defined by the `SIZEDIMAGEFIELD_JPEG_RESIZE_QUALITY` setting)
        """
        if image.mode != 'RGB':
            image = image.convert('RGB')
        return (image, {'quality': QUAL})

    def create_resized_image(self, path_to_image, width, height, filename_key):
        """
        Creates a resized image.
        `path_to_image`: The path to the image with the media directory to resize.
            If `None`, the SIZEDIMAGEFIELD_PLACEHOLDER_IMAGE will be used.
        `width`: Width of resized image (int)
        `height`: Desired height of resized image (int)
        `filename_key`: A string that will be used in the sized image filename to signify
        what operation was done to it. Examples: 'crop' or 'scale'
        """
        resized_image_path = get_resized_path(
            path_to_image=path_to_image,
            width=width,
            height=height,
            filename_key=filename_key
        )

        image, file_ext, image_format, mime_type = self.retrieve_image(path_to_image, resized_image_path)

        image, save_kwargs = self.preprocess(image, image_format)

        imagefile = self.process_image(
            image=image,
            image_format=image_format,
            width=width,
            height=height,
            save_kwargs=save_kwargs
        )

        saved = self.save_image(imagefile, resized_image_path, file_ext, mime_type)

        return saved

class FilteredImage(ProcessedImage):
    name = None
    url = None

    def __init__(self, path_to_image, storage, filename_key):
        super(FilteredImage, self).__init__(path_to_image, storage)
        self.name = get_filtered_path(
            path_to_image=self.path_to_image,
            filename_key=filename_key
        )
        self.url = self.storage.url(self.name)

    def add_sizedimage_attrs(self, sizedimage_registry):
        for attr_name, sizedimage_cls in sizedimage_registry.iteritems():
            setattr(
                self,
                attr_name,
                sizedimage_cls(
                    path_to_image=self.name,
                    storage=self.storage
                )
            )

    def process_filter(self, image, image_format, save_kwargs={}):
        raise NotImplementedError('Subclasses MUST provide a `process_filter` method.')

    def create_filtered_image(self, path_to_image, prepared_path):
        """
        Creates a resized image.
        `path_to_image`: The path to the image with the media directory to resize.
        `prepared_path`:
        """

        image, file_ext, image_format, mime_type = self.retrieve_image(path_to_image, prepared_path)
        #print image.__class__
        image, save_kwargs = self.preprocess(image, image_format)
        imagefile = self.process_filter(image, image_format, save_kwargs)
        saved = self.save_image(imagefile, prepared_path, file_ext, mime_type)

        return saved

class Filters(dict):

    def __init__(self, sizedimagefield, registry):
        self.sizedimagefield = sizedimagefield
        self.registry = registry

    def __getitem__(self, key):
        try:
            prepped_filter = dict.__getitem__(self, key)
        except KeyError:
            if key not in self.registry._filter_registry:
                return None
            else:
                filtered_path = get_filtered_path(
                    path_to_image=self.sizedimagefield.name,
                    filename_key=key
                )
                filter_cls = self.registry._filter_registry[key]
                prepped_filter = filter_cls(
                    path_to_image=self.sizedimagefield.name,
                    storage=self.sizedimagefield.storage,
                    filename_key=key
                )
                if cache.get(filtered_path):
                    # The sized path exists in the cache so the image already exists.
                    # So we `pass` to skip directly to the return statement.
                    pass
                else:
                    if not self.sizedimagefield.storage.exists(filtered_path):
                        image_created = prepped_filter.create_filtered_image(
                            path_to_image=self.sizedimagefield.name,
                            prepared_path=filtered_path
                        )

                    # Setting a super-long cache for a resized image (30 Days)
                    cache.set(filtered_path, 1, SIZEDIMAGEFIELD_CACHE_LENGTH)
                for attr_name, sizedimage_cls in self.registry._sizedimage_registry.iteritems():
                    setattr(
                        prepped_filter,
                        attr_name,
                        sizedimage_cls(
                            path_to_image=filtered_path,
                            storage=self.sizedimagefield.storage
                        )
                    )
                self[key] = prepped_filter

        return dict.__getitem__(self, key)
<|endoftext|>