import pickle


def _load_pkl_dict(path, key, idx=None, dtype='float32'):
    """Helper function to load pickled data.

    Parameters
    ----------
    path : :obj:`str`
        full file name including `.pkl` extention
    key : :obj:`str`
        data is returned from this key of the pickled dictionary
    idx : :obj:`int` or :obj:`NoneType`
        if :obj:`NoneType` return all data, else return data from this index
    dtype : :obj:`str`
        numpy data type of data

    Returns
    -------
    :obj:`list` of :obj:`numpy.ndarray` if :obj:`idx=None`
    :obj:`numpy.ndarray` is :obj:`idx=int`

    """
    with open(path, 'rb') as f:
        data_dict = pickle.load(f)

    if idx is None:
        samp = [data.astype(dtype) for data in data_dict[key]]
    else:
        samp = [data_dict[key][idx].astype(dtype)]

    return samp