<commit_before>import ast
from os import path
from copy import deepcopy
from operator import lt, gt
from itertools import takewhile
import fnmatch

from .code_util import DocStringInheritor

__all__ = [
    'keywords', 'EvaluateQueries', 'eval_queries',
    'queries_syntax_check', 'queries_preprocess',
    'TractQuerierSyntaxError', 'TractQuerierLabelNotFound'
]

keywords = [
    'and',
    'or',
    'not in',
    'not',
    'only',
    'endpoints_in',
    'both_endpoints_in',
    'anterior_of',
    'posterior_of',
    'medial_of',
    'lateral_of',
    'inferior_of',
    'superior_of',
]


class FiberQueryInfo(object):

    r"""
    Information about a processed query

    Attribute
    ---------
        tracts : set
            set of tract indices resulting from the query
        labels : set
            set of labels resulting by the query
        tracts_endpoints : (set, set)
            sets of labels of where the tract endpoints are
    """

    def __init__(self, tracts=None, labels=None, tracts_endpoints=None):
        if tracts is None:
            tracts = set()
        if labels is None:
            labels = set()
        if tracts_endpoints is None:
            tracts_endpoints = (set(), set())
        self.tracts = tracts
        self.labels = labels
        self.tracts_endpoints = tracts_endpoints

    def __getattribute__(self, name):
        if name in (
            'update', 'intersection_update', 'union', 'intersection',
            'difference', 'difference_update'
        ):
            return self.set_operation(name)
        else:
            return object.__getattribute__(self, name)

    def copy(self):
        return FiberQueryInfo(
            self.tracts.copy(), self.labels.copy(),
            (self.tracts_endpoints[0].copy(), self.tracts_endpoints[1].copy()),
        )

    def set_operation(self, name):
        def operation(tract_query_info):
            tracts_op = getattr(self.tracts, name)
            if name == 'intersection':
                name_labels = 'union'
            elif name == 'intersection_update':
                name_labels = 'update'
            else:
                name_labels = name
            labels_op = getattr(self.labels, name_labels)

            new_tracts = tracts_op(tract_query_info.tracts)
            new_labels = labels_op(tract_query_info.labels)

            new_tracts_endpoints = (
                getattr(self.tracts_endpoints[0], name)(
                    tract_query_info.tracts_endpoints[0]
                ),
                getattr(self.tracts_endpoints[1], name)(
                    tract_query_info.tracts_endpoints[1]
                )
            )

            if name.endswith('update'):
                return self
            else:
                return FiberQueryInfo(
                    new_tracts, new_labels,
                    new_tracts_endpoints,
                )

        return operation


class EvaluateQueries(ast.NodeVisitor):

    r"""
    This class implements the parser to process
    White Matter Query Language modules. By inheriting from
    :py:mod:`ast.NodeVisitor` it uses a syntax close to the
    python language.

    Every node expression visitor has the following signature

    Parameters
    ----------
    node : ast.Node

    Returns
    -------
    tracts : set
        numbers of the tracts that result of this
        query

    labels : set
        numbers of the labels that are traversed by
        the tracts resulting from this query

    """
    __metaclass__ = DocStringInheritor

    relative_terms = [
        'anterior_of',
        'posterior_of',
        'medial_of',
        'lateral_of',
        'inferior_of',
        'superior_of'
    ]

    def __init__(
        self,
        tractography_spatial_indexing,
    ):
        self.tractography_spatial_indexing = tractography_spatial_indexing

        self.evaluated_queries_info = {}
        self.queries_to_save = set()

        self.evaluating_endpoints = False

    def visit_Module(self, node):
        for line in node.body:
            self.visit(line)

    def visit_Compare(self, node):
        if any(not isinstance(op, ast.NotIn) for op in node.ops):
            raise TractQuerierSyntaxError(
                "Invalid syntax in query line %d" % node.lineno
            )

        query_info = self.visit(node.left).copy()
        for value in node.comparators:
            query_info_ = self.visit(value)
            query_info.difference_update(query_info_)

        return query_info

    def visit_BoolOp(self, node):
        query_info = self.visit(node.values[0])
        query_info = query_info.copy()

        if isinstance(node.op, ast.Or):
            for value in node.values[1:]:
                query_info_ = self.visit(value)
                query_info.update(query_info_)

        elif isinstance(node.op, ast.And):
            for value in node.values[1:]:
                query_info_ = self.visit(value)
                query_info.intersection_update(query_info_)

        else:
            return self.generic_visit(node)

        return query_info

    def visit_BinOp(self, node):
        info_left = self.visit(node.left)
        info_right = self.visit(node.right)
        if isinstance(node.op, ast.Add):
            return info_left.union(info_right)
        if isinstance(node.op, ast.Mult):
            return info_left.intersection(info_right)
        if isinstance(node.op, ast.Sub):
            return (
                info_left.difference(info_right)
            )
        else:
            return self.generic_visit(node)

    def visit_UnaryOp(self, node):
        query_info = self.visit(node.operand)
        if isinstance(node.op, ast.Invert):
            return FiberQueryInfo(
                set(
                    tract for tract in query_info.tracts
                    if (
                        self.tractography_spatial_indexing.
                        crossing_tracts_labels[tract].
                        issubset(query_info.labels)
                    )
                ),
                query_info.labels
            )
        elif isinstance(node.op, ast.UAdd):
            return query_info
        elif isinstance(node.op, ast.USub) or isinstance(node.op, ast.Not):
            all_labels = set(
                self.tractography_spatial_indexing.
                crossing_labels_tracts.keys()
            )
            all_labels.difference_update(query_info.labels)
            all_tracts = set().union(*tuple(
                (
                    self.tractography_spatial_indexing.
                    crossing_labels_tracts[label]
                    for label in all_labels
                )
            ))

            new_info = FiberQueryInfo(all_tracts, all_labels)
            return new_info
        else:
            raise TractQuerierSyntaxError(
                "Syntax error in query line %d" % node.lineno)

    def visit_Str(self, node):
        query_info = FiberQueryInfo()
        for name in fnmatch.filter(self.evaluated_queries_info.keys(), node.s):
            query_info.update(self.evaluated_queries_info[name])
        return query_info

    def visit_Call(self, node):
        # Single string argument function
        if (
            isinstance(node.func, ast.Name) and
            len(node.args) == 1 and
            len(node.args) == 1 and
            node.starargs is None and
            node.keywords == [] and
            node.kwargs is None
        ):
            if (node.func.id.lower() == 'only'):
                query_info = self.visit(node.args[0])

                only_tracts = set(
                    tract for tract in query_info.tracts
                    if (
                        self.tractography_spatial_indexing.
                        crossing_tracts_labels[tract].
                        issubset(query_info.labels)
                    )
                )
                only_endpoints = tuple((
                    set(
                        tract for tract in query_info.tracts_endpoints[i]
                        if (
                            self.tractography_spatial_indexing.
                            ending_tracts_labels[i][tract] in query_info.labels
                        )
                    )
                    for i in (0, 1)
                ))
                return FiberQueryInfo(
                    only_tracts,
                    query_info.labels,
                    only_endpoints
                )
            elif (node.func.id.lower() == 'endpoints_in'):
                query_info = self.visit(node.args[0])
                new_tracts = (
                    query_info.tracts_endpoints[0].
                    union(query_info.tracts_endpoints[1])
                )
                return FiberQueryInfo(
                    new_tracts, query_info.labels,
                    query_info.tracts_endpoints
                )
            elif (node.func.id.lower() == 'both_endpoints_in'):
                query_info = self.visit(node.args[0])
                new_tracts = (
                    query_info.tracts_endpoints[0].
                    intersection(query_info.tracts_endpoints[1])
                )
                return FiberQueryInfo(
                    new_tracts, query_info.labels,
                    query_info.tracts_endpoints
                )
            elif (
                node.func.id.lower() == 'save' and
                isinstance(node.args, ast.Str)
            ):
                self.queries_to_save.add(node.args[0].s)
                return
            elif node.func.id.lower() in self.relative_terms:
                return self.process_relative_term(node)

        raise TractQuerierSyntaxError("Invalid query in line %d" % node.lineno)

    def process_relative_term(self, node):
        r"""
        Processes the relative terms

        * anterior_of
        * posterior_of
        * superior_of
        * inferior_of
        * medial_of<commit_msg>Cleaned code and made metaclass 2to3 compatible<commit_after>import ast
from os import path
from copy import deepcopy
from operator import lt, gt
from itertools import takewhile
import fnmatch

from .code_util import DocStringInheritor
from six import add_metaclass


__all__ = [
    'keywords', 'EvaluateQueries', 'eval_queries',
    'queries_syntax_check', 'queries_preprocess',
    'TractQuerierSyntaxError', 'TractQuerierLabelNotFound'
]

keywords = [
    'and',
    'or',
    'not in',
    'not',
    'only',
    'endpoints_in',
    'both_endpoints_in',
    'anterior_of',
    'posterior_of',
    'medial_of',
    'lateral_of',
    'inferior_of',
    'superior_of',
]


class FiberQueryInfo(object):

    r"""
    Information about a processed query

    Attribute
    ---------
        tracts : set
            set of tract indices resulting from the query
        labels : set
            set of labels resulting by the query
        tracts_endpoints : (set, set)
            sets of labels of where the tract endpoints are
    """

    def __init__(self, tracts=None, labels=None, tracts_endpoints=None):
        if tracts is None:
            tracts = set()
        if labels is None:
            labels = set()
        if tracts_endpoints is None:
            tracts_endpoints = (set(), set())
        self.tracts = tracts
        self.labels = labels
        self.tracts_endpoints = tracts_endpoints

    def __getattribute__(self, name):
        if name in (
            'update', 'intersection_update', 'union', 'intersection',
            'difference', 'difference_update'
        ):
            return self.set_operation(name)
        else:
            return object.__getattribute__(self, name)

    def copy(self):
        return FiberQueryInfo(
            self.tracts.copy(), self.labels.copy(),
            (self.tracts_endpoints[0].copy(), self.tracts_endpoints[1].copy()),
        )

    def set_operation(self, name):
        def operation(tract_query_info):
            tracts_op = getattr(self.tracts, name)
            if name == 'intersection':
                name_labels = 'union'
            elif name == 'intersection_update':
                name_labels = 'update'
            else:
                name_labels = name
            labels_op = getattr(self.labels, name_labels)

            new_tracts = tracts_op(tract_query_info.tracts)
            new_labels = labels_op(tract_query_info.labels)

            new_tracts_endpoints = (
                getattr(self.tracts_endpoints[0], name)(
                    tract_query_info.tracts_endpoints[0]
                ),
                getattr(self.tracts_endpoints[1], name)(
                    tract_query_info.tracts_endpoints[1]
                )
            )
            
            if name.endswith('update'):
                return self
            else:
                return FiberQueryInfo(
                    new_tracts, new_labels,
                    new_tracts_endpoints,
                )

        return operation


@add_metaclass(DocStringInheritor)
class EvaluateQueries(ast.NodeVisitor):

    r"""
    This class implements the parser to process
    White Matter Query Language modules. By inheriting from
    :py:mod:`ast.NodeVisitor` it uses a syntax close to the
    python language.

    Every node expression visitor has the following signature

    Parameters
    ----------
    node : ast.Node

    Returns
    -------
    tracts : set
        numbers of the tracts that result of this
        query

    labels : set
        numbers of the labels that are traversed by
        the tracts resulting from this query

    """

    relative_terms = [
        'anterior_of',
        'posterior_of',
        'medial_of',
        'lateral_of',
        'inferior_of',
        'superior_of'
    ]

    def __init__(
        self,
        tractography_spatial_indexing,
    ):
        self.tractography_spatial_indexing = tractography_spatial_indexing

        self.evaluated_queries_info = {}
        self.queries_to_save = set()

        self.evaluating_endpoints = False

    def visit_Module(self, node):
        for line in node.body:
            self.visit(line)

    def visit_Compare(self, node):
        if any(not isinstance(op, ast.NotIn) for op in node.ops):
            raise TractQuerierSyntaxError(
                "Invalid syntax in query line %d" % node.lineno
            )

        query_info = self.visit(node.left).copy()
        for value in node.comparators:
            query_info_ = self.visit(value)
            query_info.difference_update(query_info_)

        return query_info

    def visit_BoolOp(self, node):
        query_info = self.visit(node.values[0])
        query_info = query_info.copy()

        if isinstance(node.op, ast.Or):
            for value in node.values[1:]:
                query_info_ = self.visit(value)
                query_info.update(query_info_)

        elif isinstance(node.op, ast.And):
            for value in node.values[1:]:
                query_info_ = self.visit(value)
                query_info.intersection_update(query_info_)

        else:
            return self.generic_visit(node)

        return query_info

    def visit_BinOp(self, node):
        info_left = self.visit(node.left)
        info_right = self.visit(node.right)
        if isinstance(node.op, ast.Add):
            return info_left.union(info_right)
        if isinstance(node.op, ast.Mult):
            return info_left.intersection(info_right)
        if isinstance(node.op, ast.Sub):
            return (
                info_left.difference(info_right)
            )
        else:
            return self.generic_visit(node)

    def visit_UnaryOp(self, node):
        query_info = self.visit(node.operand)
        if isinstance(node.op, ast.Invert):
            return FiberQueryInfo(
                set(
                    tract for tract in query_info.tracts
                    if (
                        self.tractography_spatial_indexing.
                        crossing_tracts_labels[tract].
                        issubset(query_info.labels)
                    )
                ),
                query_info.labels
            )
        elif isinstance(node.op, ast.UAdd):
            return query_info
        elif isinstance(node.op, ast.USub) or isinstance(node.op, ast.Not):
            all_labels = set(
                self.tractography_spatial_indexing.
                crossing_labels_tracts.keys()
            )
            all_labels.difference_update(query_info.labels)
            all_tracts = set().union(*tuple(
                (
                    self.tractography_spatial_indexing.
                    crossing_labels_tracts[label]
                    for label in all_labels
                )
            ))

            new_info = FiberQueryInfo(all_tracts, all_labels)
            return new_info
        else:
            raise TractQuerierSyntaxError(
                "Syntax error in query line %d" % node.lineno)

    def visit_Str(self, node):
        query_info = FiberQueryInfo()
        for name in fnmatch.filter(self.evaluated_queries_info.keys(), node.s):
            query_info.update(self.evaluated_queries_info[name])
        return query_info

    def visit_Call(self, node):
        # Single string argument function
        if (
            isinstance(node.func, ast.Name) and
            len(node.args) == 1 and
            len(node.args) == 1 and
            node.starargs is None and
            node.keywords == [] and
            node.kwargs is None
        ):
            if (node.func.id.lower() == 'only'):
                query_info = self.visit(node.args[0])

                only_tracts = set(
                    tract for tract in query_info.tracts
                    if (
                        self.tractography_spatial_indexing.
                        crossing_tracts_labels[tract].
                        issubset(query_info.labels)
                    )
                )
                only_endpoints = tuple((
                    set(
                        tract for tract in query_info.tracts_endpoints[i]
                        if (
                            self.tractography_spatial_indexing.
                            ending_tracts_labels[i][tract] in query_info.labels
                        )
                    )
                    for i in (0, 1)
                ))
                return FiberQueryInfo(
                    only_tracts,
                    query_info.labels,
                    only_endpoints
                )
            elif (node.func.id.lower() == 'endpoints_in'):
                query_info = self.visit(node.args[0])
                new_tracts = query_info.tracts_endpoints[0].union(query_info.tracts_endpoints[1])
                return FiberQueryInfo(new_tracts, query_info.labels, query_info.tracts_endpoints)
            elif (node.func.id.lower() == 'both_endpoints_in'):
                query_info = self.visit(node.args[0])
                new_tracts = (
                    query_info.tracts_endpoints[0].
                    intersection(query_info.tracts_endpoints[1])
                )
                return FiberQueryInfo(
                    new_tracts, query_info.labels,
                    query_info.tracts_endpoints
                )
            elif (
                node.func.id.lower() == 'save' and
                isinstance(node.args, ast.Str)
            ):
                self.queries_to_save.add(node.args[0].s)
                return
            elif node.func.id.lower() in self.relative_terms:
                return self.process_relative_term(node)

        raise TractQuerierSyntaxError("Invalid query in line %d" % node.lineno)

    def process_relative_term(self, node):
        r"""
        Processes the relative terms

        * anterior_of
        * posterior_of
        * superior_of
        * inferior_of
        * medial_of
{'fraction_non_alphanumeric': 0.08259735046166199, 'fraction_numerical': 0.0022079486150140505, 'mean_word_length': 1.5625562556255626, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1343232', 'n_tokens_mistral': 5647, 'n_tokens_neox': 5321, 'n_words': 1150}