def convert_str_to_ascii(input_str):
    """Convert a list of strings into an ascii encoded numpy array.

    Converts a string or list of strings to a numpy array. The array size is
    defined by the size of string plus one. This is needed for ascii to str
    convertion in lua using ffi.string() which expects a 0 at the end of an
    array.

    If a list of strings is used, the size of the array is defined by the size
    of the longest string (plus one), and zero padded to maitain the array
    shape.

    Parameters
    ----------
    input_str : str/list/tuple
        String or list of strings to convert to an ascii array.

    Returns
    -------
    np.ndarray
        Single/multi-dimensional array of ASCII encoded strings.

    Examples
    --------
    Example1: Convert a string to a numpy array encoded into ASCII values.

    >>> from dbcollection.utils.string_ascii import convertstr_to_ascii
    >>> convertstr_to_ascii('string1')
    array([115, 116, 114, 105, 110, 103,  49,   0], dtype=uint8)

    Example2: Convert a list of lists into an ASCII array.

    >>> from dbcollection.utils.string_ascii import convertstr_to_ascii
    >>> convertstr_to_ascii(['string1', 'string2', 'string3'])
    array([[115, 116, 114, 105, 110, 103,  49,   0],
        [115, 116, 114, 105, 110, 103,  50,   0],
        [115, 116, 114, 105, 110, 103,  51,   0]], dtype=uint8)

    """
    assert isinstance(input_str, (list, tuple, str)), 'Must input a valid input string.'
    assert any(input_str), 'Must input a non-empty list'
    if isinstance(input_str, str):
        input_str = [input_str]

    # get max size of the list strings
    max_size = max([len(a) for a in input_str])

    # allocate array
    ascii_array = np.zeros([len(input_str), max_size + 1], dtype=np.uint8)

    # iteratively copy data to the array
    for i, val in enumerate(input_str):
        ascii_array[i, :len(val)] = str_to_ascii(val)

    if len(input_str) > 1:
        return ascii_array
    else:
        return ascii_array[0] def pad_list(listA, val=-1, length=None):
    """Pad list of lists with 'val' such that all lists have the same length.

    Parameters
    ----------
    listA : list
        List of lists of different sizes.
    val : number, optional
        Value to pad the lists.
    length : number, optional
        Total length of the list.

    Returns
    -------
    list
        A list of lists with the same same.

    Examples
    --------
    Pad an uneven list of lists with a value.

    >>> from dbcollection.utils.pad import pad_list
    >>> pad_list([[0,1,2,3],[45,6],[7,8],[9]])  # pad with -1 (default)
    [[0, 1, 2, 3], [4, 5, 6, -1], [7, 8, -1, -1], [9-1, -1, -1]]
    >>> pad_list([[1,2],[3,4]])  # does nothing
    [[1, 2], [3, 4]]
    >>> pad_list([[],[1],[3,4,5]], 0)  # pad lists with 0
    [[0, 0, 0], [1, 0, 0], [3, 4, 5]]
    >>> pad_list([[],[1],[3,4,5]], 0, 6)  # pad lists with 0 of size 6
    [[0, 0, 0, 0, 0, 0], [1, 0, 0, 0, 0, 0], [3, 4, 5, 0, 0, 0]]

    """
    # pad list with zeros in order to have all lists of the same size
    assert isinstance(listA, list), 'Input must be a list. Got {}, expected {}' \
                                    .format(type(listA), type(list))

    # get size of the biggest list
    if length:
        max_size = length
    else:
        max_size = len(max(listA, key=len))

    # pad all lists with the a padding value
    return [l + [val] * int(max_size - len(l)) for l in listA] class Classification(BaseTask):
    """Cifar10 Classification preprocessing functions."""

    # metadata filename
    filename_h5 = 'classification'

    # extracted file names
    data_files = [
        "batches.meta",
        "data_batch_1",
        "data_batch_2",
        "data_batch_3",
        "data_batch_4",
        "data_batch_5",
        "test_batch"
    ]

    def load_data(self):
        """
        Loads data from annotation files.
        """
        loader = DatasetAnnotationLoader(
            data_files=self.data_files,
            data_path=self.data_path,
            cache_path=self.cache_path,
            verbose=self.verbose
        )
        yield {"train": loader.load_train_data()}
        yield {"test": loader.load_test_data()}

    def process_set_metadata(self, data, set_name):
        """
        Saves the metadata of a set.
        """
        args = {
            "data": data,
            "set_name": set_name,
            "hdf5_manager": self.hdf5_manager,
            "verbose": self.verbose
        }

        # Fields
        if self.verbose:
            print('\n==> Setting up the data fields:')
        ClassLabelField(**args).process()
        image_ids = ImageField(**args).process()
        label_ids = LabelIdField(**args).process()
        ObjectFieldNamesField(**args).process()
        ObjectIdsField(**args).process(image_ids, label_ids)

        # Lists
        if self.verbose:
            print('\n==> Setting up ordered lists:')
        ImagesPerClassList(**args).process() class DatasetAnnotationLoader:
    """Annotation's data loader for the cifar10 dataset (train/test)."""

    def __init__(self, data_files, data_path, cache_path, verbose):
        self.data_files = data_files
        self.data_path = data_path
        self.cache_path = cache_path
        self.verbose = verbose

    def load_train_data(self):
        """Loads the train set annotation data from disk
        and returns it as a dictionary."""
        return self.load_data_set(is_test=False)

    def load_test_data(self):
        """Loads the test set annotation data from disk
        and returns it as a dictionary."""
        return self.load_data_set(is_test=True)

    def load_data_set(self, is_test):
        """Fetches the train/test data."""
        assert isinstance(is_test, bool), "Must input a valid boolean input."
        images, labels, class_names = self.load_data_annotations(is_test)
        return {
            "classes": class_names,
            "images": images,
            "labels": labels
        }

    def load_data_annotations(self, is_test):
        """Loads the data from the annotations' files."""
        assert isinstance(is_test, bool), "Must input a valid boolean input."
        data_path = os.path.join(self.data_path, 'cifar-10-batches-py')
        class_names = self.get_class_names(data_path)
        if is_test:
            data, labels = self.get_data_test(data_path)
        else:
            data, labels = self.get_data_train(data_path)
        data = np.transpose(data, (0, 2, 3, 1))  # NxHxWxC
        return data, labels, class_names

    def get_class_names(self, path):
        """Returns the class names/labels."""
        assert path, "Must input a valid path."
        filename = os.path.join(path, self.data_files[0])
        return self.load_annotation_file(filename)

    def load_annotation_file(self, path):
        """Reads the data from annotation file from disk."""
        return load_pickle(path)

    def get_data_train(self, path):
        assert path, "Must input a valid path."
        batch1 = self.load_annotation_file(os.path.join(path, self.data_files[1]))
        batch2 = self.load_annotation_file(os.path.join(path, self.data_files[2]))
        batch3 = self.load_annotation_file(os.path.join(path, self.data_files[3]))
        batch4 = self.load_annotation_file(os.path.join(path, self.data_files[4]))
        batch5 = self.load_annotation_file(os.path.join(path, self.data_files[5]))

        data = np.concatenate(
            (batch1['data'], batch2['data'], batch3['data'], batch4['data'], batch5['data']),
            axis=0
        )
        data = self.reshape_array(data, 50000)
        labels = np.concatenate(
            (batch1['labels'], batch2['labels'], batch3['labels'], batch4['labels'], batch5['labels']),
            axis=0
        )

        return data, labels

    def reshape_array(self, data, size_array):
        """Reshapes the numpy array to a fixed format."""
        return data.reshape(size_array, 3, 32, 32)

    def get_data_test(self, path):
        assert path, "Must input a valid path."
        batch = self.load_annotation_file(os.path.join(path, self.data_files[6]))
        data = self.reshape_array(batch['data'], 10000)
        labels = np.array(batch['labels'], dtype=np.uint8)
        return data, labels class ClassLabelField(BaseField):
    """Class label names' field metadata process/save class."""

    @display_message_processing('class labels')
    def process(self):
        """Processes and saves the classes metadata to hdf5."""
        class_names = self.get_class_names()
        self.save_field_to_hdf5(
            set_name=self.set_name,
            field='classes',
            data=str2ascii(class_names),
            dtype=np.uint8,
            fillvalue=0
        )

    def get_class_names(self):
        """Returns a list of class names."""
        return self.data['classes'] class ImageField(BaseField):
    """Images' data field process/save class."""

    @display_message_processing('images')
    def process(self):
        """Processes and saves the images metadata to hdf5."""
        images, image_ids = self.get_images()
        self.save_field_to_hdf5(
            set_name=self.set_name,
            field='images',
            data=images,
            dtype=np.uint8,
            fillvalue=-1
        )
        return image_ids

    def get_images(self):
        """Returns a np.ndarray of images and a list
        of image ids for each row of 'object_ids' field."""
        images = self.data['images']
        image_ids = list(range(len(images)))
        return images, image_ids class LabelIdField(BaseField):
    """Label id field metadata process/save class."""

    @display_message_processing('labels')
    def process(self):
        """Processes and saves the labels metadata to hdf5."""
        labels, label_ids = self.get_labels()
        self.save_field_to_hdf5(
            set_name=self.set_name,
            field='labels',
            data=labels,
            dtype=np.uint8,
            fillvalue=0
        )
        return label_ids

    def get_labels(self):
        """Returns a np.ndarray of labels and a list
        of label ids for each row of 'object_ids' field."""
        labels = self.data['labels']
        label_ids = list(range(len(labels)))
        return labels, label_ids class ObjectFieldNamesField(BaseField):
    """Object field names metadata process/save class."""

    def process(self):
        """Processes and saves the labels metadata to hdf5."""
        self.save_field_to_hdf5(
            set_name=self.set_name,
            field='object_fields',
            data=str2ascii(['images', 'labels']),
            dtype=np.uint8,
            fillvalue=0
        ) class ObjectIdsField(BaseField):
    """Object ids' field metadata process/save class."""

    def process(self, image_ids, label_ids):
        """Processes and saves the object ids metadata to hdf5."""
        # images, labels
        object_ids = [[image_ids[i], label_ids[i]] for i, _ in enumerate(label_ids)]
        self.save_field_to_hdf5(
            set_name=self.set_name,
            field='object_ids',
            data=np.array(object_ids, dtype=np.int32),
            dtype=np.int32,
            fillvalue=-1
        ) class ImagesPerClassList(BaseField):
    """Images per class list metadata process/save class."""

    @display_message_processing('images per class list')
    def process(self):
        """Processes and saves the list ids metadata to hdf5."""
        images_per_class = self.get_image_ids_per_class()
        images_per_class_array = self.convert_list_to_array(images_per_class)
        self.save_field_to_hdf5(
            set_name=self.set_name,
            field='list_images_per_class',
            data=images_per_class_array,
            dtype=np.int32,
            fillvalue=-1
        )

    def get_image_ids_per_class(self):
        """Returns a list of lists of image filename ids per class id."""
        images_per_class = []
        labels = self.data['labels']
        unique_labels = np.unique(labels)
        for label in unique_labels:
            images_idx = np.where(labels == label)[0].tolist()
            images_per_class.append(images_idx)
        return images_per_class

    def convert_list_to_array(self, list_ids):
        """Pads a list of listsand converts it into a numpy.ndarray."""
        padded_list = pad_list(list_ids, val=-1)
        return np.array(padded_list, dtype=np.int32)
import os
import sys
import pytest
import numpy as np
from numpy.testing import assert_array_equal
from dbcollection.utils.string_ascii import convert_str_to_ascii as str2ascii
from dbcollection.utils.pad import pad_list
from dbcollection.datasets.cifar.cifar10.classification import (
    Classification,
    DatasetAnnotationLoader,
    ClassLabelField,
    ImageField,
    LabelIdField,
    ObjectFieldNamesField,
    ObjectIdsField,
    ImagesPerClassList
)

        mock_classlabel_class.process()

        mock_get_class.assert_called_once_with()
        assert mock_save_hdf5.called
        # **disabled until I find a way to do assert calls with numpy arrays**
        # mock_save_hdf5.assert_called_once_with(
        #     set_name='train',
        #     field='classes',
        #     data=str2ascii(dummy_names),
        #     dtype=np.uint8,
        #     fillvalue=-1
        # )

    def test_get_class_names(self, mocker, mock_classlabel_class, test_data_loaded):
        class_names = mock_classlabel_class.get_class_names()

        assert class_names == test_data_loaded['classes']


class TestImageField:
    """Unit tests for the ImageField class."""

    @staticmethod
    @pytest.fixture()
    def mock_image_class(field_kwargs):
        return ImageField(**field_kwargs)

    def test_process(self, mocker, mock_image_class):
        dummy_images = np.random.rand(5,3, 32, 32)
        dummy_ids = list(range(5))
        mock_get_images = mocker.patch.object(ImageField, "get_images", return_value=(dummy_images, dummy_ids))
        mock_save_hdf5 = mocker.patch.object(ImageField, "save_field_to_hdf5")

        image_ids = mock_image_class.process()

        assert image_ids == dummy_ids
        mock_get_images.assert_called_once_with()
        assert mock_save_hdf5.called
        # **disabled until I find a way to do assert calls with numpy arrays**
        # mock_save_hdf5.assert_called_once_with(
        #     set_name='train',
        #     field='images',
        #     data=dummy_images,
        #     dtype=np.uint8,
        #     fillvalue=-1
        # )

    def test_get_images(self, mocker, mock_image_class, test_data_loaded):
        images, image_ids = mock_image_class.get_images()

        assert_array_equal(images, test_data_loaded['images'])
        assert image_ids == list(range(len(images)))


class TestLabelIdField:
    """Unit tests for the LabelIdField class."""

    @staticmethod
    @pytest.fixture()
    def mock_label_class(field_kwargs):
        return LabelIdField(**field_kwargs)

    def test_process(self, mocker, mock_label_class):
        dummy_labels = np.array(range(10))
        dummy_ids = list(range(10))
        mock_get_labels = mocker.patch.object(LabelIdField, "get_labels", return_value=(dummy_labels, dummy_ids))
        mock_save_hdf5 = mocker.patch.object(LabelIdField, "save_field_to_hdf5")

        label_ids = mock_label_class.process()

        assert label_ids == dummy_ids
        mock_get_labels.assert_called_once_with()
        assert mock_save_hdf5.called
        # **disabled until I find a way to do assert calls with numpy arrays**
        # mock_save_hdf5.assert_called_once_with(
        #     set_name='train',
        #     field='labels',
        #     data=dummy_labels,
        #     dtype=np.uint8,
        #     fillvalue=0
        # )

    def test_get_images(self, mocker, mock_label_class, test_data_loaded):
        labels, label_ids = mock_label_class.get_labels()

        assert_array_equal(labels, test_data_loaded['labels'])
        assert label_ids == list(range(len(labels)))


class TestObjectFieldNamesField:
    """Unit tests for the ObjectFieldNamesField class."""

    @staticmethod
    @pytest.fixture()
    def mock_objfields_class(field_kwargs):
        return ObjectFieldNamesField(**field_kwargs)

    def test_process(self, mocker, mock_objfields_class):
        mock_save_hdf5 = mocker.patch.object(ObjectFieldNamesField, "save_field_to_hdf5")

        mock_objfields_class.process()

        assert mock_save_hdf5.called
        # **disabled until I find a way to do assert calls with numpy arrays**
        # mock_save_hdf5.assert_called_once_with(
        #     set_name='train',
        #     field='object_fields',
        #     data=str2ascii(['images', 'labels']),
        #     dtype=np.uint8,
        #     fillvalue=0
        # )


class TestObjectIdsField:
    """Unit tests for the ObjectIdsField class."""

    @staticmethod
    @pytest.fixture()
    def mock_objfids_class(field_kwargs):

        return ObjectIdsField(**field_kwargs)