Source code for zcollection.expression

# Copyright (c) 2023 CNES
#
# All rights reserved. Use of this source code is governed by a
# BSD-style license that can be found in the LICENSE file.
"""
Handles the partition selection expressions
===========================================
"""
from __future__ import annotations

from typing import Any, ClassVar
import ast
import dataclasses
import types




@dataclasses.dataclass
class Expression:
    """Partitioning expressions.

    Args:
        expression: The expression to be evaluated
    Raises:
        NameError: If a variable is not defined.
    Example:
        >>> expr = Expression("year==2000 and month==1 and day in range(1, 12)")
    """
    #: Compiled expression to be evaluated
    code: types.CodeType

    #: Known data members
    __slots__: tuple[str, ...] = ('code', )

    #: The builtins that are allowed in the expression.
    BUILTINS: ClassVar[dict[str, Any]] = {'range': range}

    def __init__(self, expression: str) -> None:
        self.code = compile(ast.parse(expression, mode='eval'), ' ', 'eval')



    def __call__(self, variables: dict[str, Any]) -> Any:
        try:
            __locals: dict[str, Any] = {
                name: variables[name]
                for name in self.code.co_names if name not in self.BUILTINS
            }
            # pylint: disable=eval-used
            # The eval function is used here to evaluate a simple expression.
            # The only builtin functions allowed is the range function.
            return eval(self.code, {'__builtins__': self.BUILTINS}, __locals)
            # pylint: enable=eval-used
        except KeyError as err:
            raise NameError(f'name {err!s} is not defined') from err