import fnmatch


def filter_files(files, suffix):
    """
    Filter the list to contain only files matching the given `suffix`.

    Parameters
    ----------
    files : list
        List of files to be filtered.
    suffix : str
        Return only files matching this suffix.

    Returns
    -------
    list
        List of files.

    """
    # make sure files is a list
    if not isinstance(files, list):
        files = [files]
    # no suffix given, return the list unaltered
    if suffix is None:
        return files
    # filter the files with the given suffix
    file_list = []
    if isinstance(suffix, list):
        # a list of suffices is given
        for s in suffix:
            file_list.extend(fnmatch.filter(files, "*%s" % s))
    else:
        # a single suffix is given
        file_list.extend(fnmatch.filter(files, "*%s" % suffix))
    # return the filtered list
    return file_list