query.py

#!/usr/bin/env python
#
# query.py - The FileTreeQuery class
#
# Author: Paul McCarthy <pauldmccarthy@gmail.com>
# Author: Michiel Cottaar <michiel.cottaar@.ndcn.ox.ac.uk>
#
"""This module contains the :class:`FileTreeQuery` class, which can be used to
search for files in a directory described by a :class:`.FileTree`. A
``FileTreeQuery`` object returns :class:`Match` objects which each represent a
file that is described by the ``FileTree``, and which is present in the
directory.

The following utility functions, used by the ``FileTreeQuery`` class, are also
defined in this module:

.. autosummary::
   :nosignatures:

   scan
   allVariables
"""


import              logging
import              collections
import functools as ft

import os.path as op
from typing import Dict, List, Tuple

import numpy as np

from . import FileTree


log = logging.getLogger(__name__)


class FileTreeQuery(object):
    """The ``FileTreeQuery`` class uses a :class:`.FileTree` to search
    a directory for files which match a specific query.

    A ``FileTreeQuery`` scans the contents of a directory which is described
    by a :class:`.FileTree`, and identifies all file types (a.k.a. *templates*
    or *short names*) that are present, and the values of variables within each
    short name that are present. The :meth:`query` method can be used to
    retrieve files which match a specific template, and variable values.

    The :meth:`query` method returns a collection of :class:`Match` objects,
    each of which represents one file which matches the query.

    Example usage::

        >>> from fsl.utils.filetree import FileTree, FileTreeQuery

        >>> tree  = FileTree.read('bids_raw', './my_bids_data')
        >>> query = FileTreeQuery(tree)

        >>> query.axes('anat_image')
        ['acq', 'ext', 'modality', 'participant', 'rec', 'run_index',
         'session']

        >>> query.variables('anat_image')
        {'acq': [None],
         'ext': ['.nii.gz'],
         'modality': ['T1w', 'T2w'],
         'participant': ['01', '02', '03'],
         'rec': [None],
         'run_index': [None, '01', '02', '03'],
         'session': [None]}

        >>> query.query('anat_image', participant='01')
        [Match(./my_bids_data/sub-01/anat/sub-01_T1w.nii.gz),
         Match(./my_bids_data/sub-01/anat/sub-01_T2w.nii.gz)]


    Matches for templates contained within sub-trees are referred to by
    constructing a hierarchical path from the sub-tree template name(s),
    and the template name - see the :meth:`Match.full_name` method.
    """


    def __init__(self, tree):
        """Create a ``FileTreeQuery``. The contents of the tree directory are
        scanned via the :func:`scan` function, which may take some time for
        large data sets.

        :arg tree: The :class:`.FileTree` object
        """
        # Hard-code into the templates any pre-defined variables
        tree = tree.partial_fill()

        # Find all files present in the directory
        # (as Match objects), and find all variables,
        # plus their values, and all templates,
        # that are present in the directory.
        matches               = scan(tree)
        allvars, templatevars = allVariables(tree, matches)

        # Now we are going to build a series of ND
        # arrays to store Match objects. We create
        # one array for each template. Each axis
        # in an array corresponds to a variable
        # present in files of that template type,
        # and each position along an axis corresponds
        # to one value of that variable.
        #
        # These arrays will be used to store and
        # retrieve Match objects - given a template
        # and a set of variable values, we can
        # quickly find the corresponding Match
        # object (or objects).

        # matcharrays contains {template : ndarray}
        # mappings, and varidxs contains
        # {template : {varvalue : index}} mappings
        matcharrays = {}
        varidxs     = {}

        for template, tvars in templatevars.items():

            tvarlens = [len(allvars[v]) for v in tvars]

            # "Scalar" match objects - templates
            # which have no variables, and for
            # which zero or one file is present
            if len(tvarlens) == 0:
                tvarlens = 1

            # An ND array for this short
            # name. Each element is a
            # Match object, or nan.
            matcharray    = np.zeros(tvarlens, dtype=np.object)
            matcharray[:] = np.nan

            # indices into the match array
            # for each variable value
            tvaridxs = {}
            for v in tvars:
                tvaridxs[v] = {n : i for i, n in enumerate(allvars[v])}

            matcharrays[template] = matcharray
            varidxs[    template] = tvaridxs

        # Populate the match arrays
        for match in matches:
            tvars    = templatevars[match.full_name]
            tvaridxs = varidxs[     match.full_name]
            tarr     = matcharrays[ match.full_name]
            idx      = []

            if len(match.variables) == 0:
                idx = [0]
            else:
                for var in tvars:
                    val = match.variables[var]
                    idx.append(tvaridxs[var][val])

            tarr[tuple(idx)] = match

        self.__tree          = tree
        self.__allvars       = allvars
        self.__templatevars  = templatevars
        self.__matches       = matches
        self.__matcharrays   = matcharrays
        self.__varidxs       = varidxs


    def axes(self, template) -> List[str]:
        """Returns a list containing the names of variables present in files
        of the given ``template`` type, in the same order of the axes of
        :class:`Match` arrays that are returned by the :meth:`query` method.
        """
        return self.__templatevars[template]


    def variables(self, template=None) -> Dict[str, List]:
        """Return a dict of ``{variable : [values]}`` mappings.
        This dict describes all variables and their possible values in
        the tree.

        If a ``template`` is specified, only variables which are present in
        files of that ``template`` type are returned.
        """
        if template is None:
            return {var : list(vals) for var, vals in self.__allvars.items()}
        else:
            varnames = self.__templatevars[template]
            return {var : list(self.__allvars[var]) for var in varnames}


    @property
    def tree(self):
        """Returns the :class:`.FileTree` associated with this
        ``FileTreeQuery``.
        """
        return self.__tree


    @property
    def templates(self) -> List[str]:
        """Returns a list containing all templates of the ``FileTree`` that
        are present in the directory.
        """
        return list(self.__templatevars.keys())


    def query(self, template, asarray=False, **variables):
        """Search for files of the given ``template``, which match
        the specified ``variables``. All hits are returned for variables
        that are unspecified.

        :arg template: Template of files to search for.

        :arg asarray:  If ``True``, the relevant :class:`Match` objects are
                       returned in a in a ND ``numpy.array`` where each
                       dimension corresponds to a variable for the
                       ``templates`` in question (as returned by
                       :meth:`axes`). Otherwise (the default), they are
                       returned in a list.

        All other arguments are assumed to be ``variable=value`` pairs,
        used to restrict which matches are returned. All values are returned
        for variables that are not specified, or variables which are given a
        value of ``'*'``.

        :returns: A list  of ``Match`` objects, (or a ``numpy.array`` if
                  ``asarray=True``).
        """

        varnames    = list(variables.keys())
        allvarnames = self.__templatevars[template]
        varidxs     = self.__varidxs[     template]
        matcharray  = self.__matcharrays[ template]
        slc         = []

        for var in allvarnames:

            if var in varnames: val = variables[var]
            else:               val = '*'

            # We're using np.newaxis to retain
            # the full dimensionality of the
            # array, so that the axis labels
            # returned by the axes() method
            # are valid.
            if val == '*': slc.append(slice(None))
            else:          slc.extend([np.newaxis, varidxs[var][val]])

        result = matcharray[tuple(slc)]

        if asarray: return result
        else:       return [m for m in result.flat if isinstance(m, Match)]


@ft.total_ordering
class Match(object):
    """A ``Match`` object represents a file with a name matching a template in
    a ``FileTree``.  The :func:`scan` function and :meth:`FileTree.query`
    method both return ``Match`` objects.
    """


    def __init__(self, filename, template, tree, variables):
        """Create a ``Match`` object. All arguments are added as attributes.

        :arg filename:   name of existing file
        :arg template:   template identifier
        :arg tree:       :class:`.FileTree` which contains this ``Match``
        :arg variables:  Dictionary of ``{variable : value}`` mappings
                         containing all variables present in the file name.
        """
        self.__filename   = filename
        self.__template   = template
        self.__tree       = tree
        self.__variables  = dict(variables)


    @property
    def filename(self):
        return self.__filename


    @property
    def template(self):
        return self.__template


    @property
    def full_name(self):
        """The ``full_name`` of a ``Match`` is a combination of the
        ``template`` (i.e. the matched template), and the name(s) of
        the relevant ``FileTree`` objects.

        It allows one to unamiguously identify the location of a ``Match``
        in a ``FileTree`` hierarchy, where the same ``short_name`` may be
        used in different sub-trees.
        """

        def parents(tree):
            if tree.parent is None:
                return []
            else:
                return [tree.parent] + parents(tree.parent)

        trees = [self.tree] + parents(self.tree)

        # Drop the root tree
        trees = list(reversed(trees))[1:]

        return '/'.join([t.name for t in trees] + [self.template])


    @property
    def tree(self):
        return self.__tree


    @property
    def variables(self):
        return dict(self.__variables)


    def __eq__(self, other):
        return (isinstance(other, Match)            and
                self.filename   == other.filename   and
                self.template   == other.template   and
                self.tree       is other.tree       and
                self.variables  == other.variables)


    def __lt__(self, other):
        return isinstance(other, Match) and self.filename < other.filename


    def __repr__(self):
        """Returns a string representation of this ``Match``. """
        return 'Match({}: {})'.format(self.full_name, self.filename)


    def __str__(self):
        """Returns a string representation of this ``Match``. """
        return repr(self)


def scan(tree : FileTree) -> List[Match]:
    """Scans the directory of the given ``FileTree`` to find all files which
    match a tree template.

    :arg tree: :class:`.FileTree` to scan
    :returns:  list of :class:`Match` objects
    """

    matches = []
    for template in tree.templates:

        for variables in tree.get_all_vars(template, glob_vars='all'):

            filename = tree.update(**variables).get(template)

            if not op.isfile(filename):
                continue

            matches.append(Match(filename, template, tree, variables))

    for tree_name, sub_tree in tree.sub_trees.items():
        matches.extend(scan(sub_tree))

    return matches


def allVariables(
        tree    : FileTree,
        matches : List[Match]) -> Tuple[Dict[str, List], Dict[str, List]]:
    """Identifies the ``FileTree`` variables which are actually represented
    in files in the directory.

    :arg filetree: The ``FileTree`` object
    :arg matches:  list of ``Match`` objects (e.g. as returned by :func:`scan`)

    :returns: a tuple containing two dicts:

               - A dict of ``{ variable : [values] }`` mappings containing all
                 variables and their possible values present in the given list
                 of ``Match`` objects.

               - A dict of ``{ full_name : [variables] }`` mappings,
                 containing the variables which are relevant to each template.
    """
    allvars      = collections.defaultdict(set)
    alltemplates = {}

    for m in matches:

        if m.full_name not in alltemplates:
            alltemplates[m.full_name] = set()

        for var, val in m.variables.items():
            allvars[     var]        .add(val)
            alltemplates[m.full_name].add(var)

    # allow us to compare None with strings
    def key(v):
        if v is None: return ''
        else:         return v

    allvars      = {var : list(sorted(vals, key=key))
                    for var, vals in allvars.items()}
    alltemplates = {tn  : list(sorted(vars))
                    for tn, vars in alltemplates.items()}

    return allvars, alltemplates