vantage6-algorithm-tools 5.0.3rc5__tar.gz → 5.0.4rc1__tar.gz

{vantage6_algorithm_tools-5.0.3rc5 → vantage6_algorithm_tools-5.0.4rc1}/PKG-INFO

@@ -1,24 +1,20 @@
 Metadata-Version: 2.5
 Name: vantage6-algorithm-tools
-Version: 5.0.3rc5
+Version: 5.0.4rc1
 Summary: Vantage6 algorithm tools
 Author: Vantage6 Team
 Maintainer-email: Frank Martin <f.martin@iknl.nl>, Bart van Beusekom <b.vanbeusekom@iknl.nl>
-License: MIT
+License: Apache-2.0
 Requires-Python: >=3.13
 Requires-Dist: openpyxl>=3.0.0
 Requires-Dist: pandas>=2.2.3
 Requires-Dist: psycopg2-binary==2.9.10
 Requires-Dist: pyarrow>=16.1.0
 Requires-Dist: pyfiglet==1.0.4
-Requires-Dist: pyjwt==2.13.0
+Requires-Dist: pyjwt==2.15.1
 Requires-Dist: sparqlwrapper>=2.0.0
 Requires-Dist: sqlalchemy>=1.4.46
-Requires-Dist: vantage6-common==5.0.3rc5
-Provides-Extra: dev
-Requires-Dist: black; extra == 'dev'
-Requires-Dist: coverage==7.10.2; extra == 'dev'
-Requires-Dist: pre-commit; extra == 'dev'
+Requires-Dist: vantage6-common==5.0.4rc1
 Description-Content-Type: text/markdown
 
 <h1 align="center">

{vantage6_algorithm_tools-5.0.3rc5 → vantage6_algorithm_tools-5.0.4rc1}/pyproject.toml

@@ -4,10 +4,10 @@ build-backend = "hatchling.build"
 
 [project]
 name = "vantage6-algorithm-tools"
-version = "5.0.3rc5"
+version = "5.0.4rc1"
 description = "Vantage6 algorithm tools"
 readme = "README.md"
-license = {text = "MIT"}
+license = {text = "Apache-2.0"}
 authors = [
     {name = "Vantage6 Team"}
 ]
@@ -23,18 +23,11 @@ dependencies = [
     # running SQL queries on a postgres database.
     "psycopg2-binary==2.9.10",
     "pyarrow>=16.1.0",
-    "pyjwt==2.13.0",
+    "pyjwt==2.15.1",
     "pyfiglet==1.0.4",
     "SPARQLWrapper>=2.0.0",
     "sqlalchemy>=1.4.46",
-    "vantage6-common==5.0.3rc5",
-]
-
-[project.optional-dependencies]
-dev = [
-    "coverage==7.10.2",
-    "black",
-    "pre-commit",
+    "vantage6-common==5.0.4rc1",
 ]
 
 [tool.hatch.build.targets.wheel]

vantage6_algorithm_tools-5.0.4rc1/tests/preprocessing/test_expression.py

@@ -0,0 +1,207 @@
+import inspect
+import unittest
+
+import numpy as np
+import pandas as pd
+
+from vantage6.algorithm.tools.exceptions import UserInputError
+
+from vantage6.algorithm.preprocessing.column import assign_column, redefine_column
+from vantage6.algorithm.preprocessing.expression import (
+    MAX_EXPRESSION_LENGTH,
+    evaluate_expression,
+)
+from vantage6.algorithm.preprocessing.filtering import select_rows
+
+# The decorators require a running algorithm container, so test the wrapped functions
+_assign_column = inspect.unwrap(assign_column)
+_redefine_column = inspect.unwrap(redefine_column)
+_select_rows = inspect.unwrap(select_rows)
+
+
+def _df() -> pd.DataFrame:
+    return pd.DataFrame(
+        {
+            "a": [1, 2, 3, 4, 5],
+            "b": [6, 7, 8, 9, 10],
+            "c": [11, 12, 13, 14, 15],
+            "s": ["x", "y", "x", "z", "y"],
+        }
+    )
+
+
+class TestEvaluateExpression(unittest.TestCase):
+    def assert_series_equal(self, expression: str, expected: pd.Series) -> None:
+        pd.testing.assert_series_equal(
+            evaluate_expression(_df(), expression), expected, check_names=False
+        )
+
+    def test_arithmetic(self):
+        df = _df()
+        self.assert_series_equal("a + b * 2", df.a + df.b * 2)
+        self.assert_series_equal("(a - b) / c", (df.a - df.b) / df.c)
+        self.assert_series_equal("c // a % 3", df.c // df.a % 3)
+        self.assert_series_equal("-a", -df.a)
+        self.assert_series_equal("+a", df.a)
+
+    def test_comparisons(self):
+        df = _df()
+        self.assert_series_equal("a > 2", df.a > 2)
+        self.assert_series_equal("a == 3", df.a == 3)
+        self.assert_series_equal("a != 3", df.a != 3)
+        self.assert_series_equal("s == 'x'", df.s == "x")
+        self.assert_series_equal("2 <= a < 5", (df.a >= 2) & (df.a < 5))
+
+    def test_membership(self):
+        df = _df()
+        self.assert_series_equal("a in [1, 5]", df.a.isin([1, 5]))
+        self.assert_series_equal("s not in ('x', 'y')", ~df.s.isin(["x", "y"]))
+
+    def test_logical_operators(self):
+        df = _df()
+        self.assert_series_equal("a > 2 and b < 10", (df.a > 2) & (df.b < 10))
+        self.assert_series_equal("a > 4 or c < 12", (df.a > 4) | (df.c < 12))
+        self.assert_series_equal("not a > 2", ~(df.a > 2))
+        self.assert_series_equal("(a > 2) & (b < 10)", (df.a > 2) & (df.b < 10))
+        self.assert_series_equal("(a > 4) | (c < 12)", (df.a > 4) | (df.c < 12))
+        self.assert_series_equal("~(a > 2)", ~(df.a > 2))
+
+    def test_functions(self):
+        df = _df()
+        for name, function in [
+            ("abs", np.abs),
+            ("sqrt", np.sqrt),
+            ("log", np.log),
+            ("log2", np.log2),
+            ("log10", np.log10),
+            ("exp", np.exp),
+            ("floor", np.floor),
+            ("ceil", np.ceil),
+        ]:
+            with self.subTest(function=name):
+                self.assert_series_equal(f"{name}(a / 3)", function(df.a / 3))
+        self.assert_series_equal("round(a / 3)", np.round(df.a / 3))
+        self.assert_series_equal("round(a / 3, 2)", np.round(df.a / 3, 2))
+        self.assert_series_equal("sqrt(log(b) + 1)", np.sqrt(np.log(df.b) + 1))
+
+    def test_text(self):
+        df = _df()
+        self.assert_series_equal("s + '_1'", df.s + "_1")
+        for expression in [
+            "s * 999999999",
+            "'x' * 999999999",
+            "999999999 * s",
+            "s - 'x'",
+        ]:
+            with (
+                self.subTest(expression=expression),
+                self.assertRaises(UserInputError),
+            ):
+                evaluate_expression(df, expression)
+
+    def test_constants(self):
+        self.assertEqual(evaluate_expression(_df(), "1 + 2"), 3)
+        self.assertEqual(evaluate_expression(_df(), "not True"), False)
+        self.assertEqual(evaluate_expression(_df(), "  1  "), 1)
+
+    def test_rejects_code_execution(self):
+        rejected = [
+            # attribute chain to reach the os module
+            "a.__class__.__init__.__globals__.get('sys').modules.get('os').getpid()",
+            "a.sum()",
+            "a.values",
+            "log(a).real",
+            "__import__('os')",
+            "eval('1')",
+            "exec('1')",
+            "open('/etc/passwd')",
+            "getattr(a, 'sum')()",
+            "(lambda: 1)()",
+            "log(a)(1)",
+        ]
+        for expression in rejected:
+            with (
+                self.subTest(expression=expression),
+                self.assertRaises(UserInputError),
+            ):
+                evaluate_expression(_df(), expression)
+
+    def test_rejects_unsupported_syntax(self):
+        rejected = [
+            "a ** 2",
+            "a @ b",
+            "a << 1",
+            "a[0]",
+            "a is None",
+            "[1, 2]",
+            "{'x': 1}",
+            "f'{a}'",
+            "a if b else c",
+            "b'bytes'",
+            "1j",
+            "...",
+            "a in b",
+            "a in [b]",
+            "x := 1",
+            "@a",
+            "`a`",
+            "log(a, 2)",
+            "log()",
+            "log(a=1)",
+            "round(a, 1.5)",
+            "round(a, True)",
+            "round(a, b)",
+            "round(a, 1, 2)",
+            "unknown_column + 1",
+            "",
+            "a +",
+        ]
+        for expression in rejected:
+            with (
+                self.subTest(expression=expression),
+                self.assertRaises(UserInputError),
+            ):
+                evaluate_expression(_df(), expression)
+
+    def test_rejects_non_string(self):
+        with self.assertRaises(UserInputError):
+            evaluate_expression(_df(), 1)
+
+    def test_maximum_length(self):
+        expression = "a" + " + a" * ((MAX_EXPRESSION_LENGTH - 1) // 4)
+        self.assertLessEqual(len(expression), MAX_EXPRESSION_LENGTH)
+        evaluate_expression(_df(), expression)
+        with self.assertRaises(UserInputError):
+            evaluate_expression(_df(), expression + " + a")
+
+
+class TestPreprocessingFunctions(unittest.TestCase):
+    def test_assign_column(self):
+        df = _assign_column(_df(), "d", "a + b")
+        self.assertEqual(df["d"].tolist(), [7, 9, 11, 13, 15])
+
+    def test_assign_column_rejects_code(self):
+        with self.assertRaises(UserInputError):
+            _assign_column(_df(), "d", "a.sum()")
+
+    def test_redefine_column(self):
+        df = _redefine_column(_df(), "c", "a * b")
+        self.assertEqual(df["c"].tolist(), [6, 14, 24, 36, 50])
+
+    def test_redefine_column_rejects_code(self):
+        with self.assertRaises(UserInputError):
+            _redefine_column(_df(), "c", "a.sum()")
+
+    def test_select_rows(self):
+        df = _select_rows(_df(), "a > 2 and s != 'z'")
+        self.assertEqual(df["a"].tolist(), [3, 5])
+        self.assertEqual(df.index.tolist(), [2, 4])
+
+    def test_select_rows_rejects_code(self):
+        with self.assertRaises(UserInputError):
+            _select_rows(_df(), "a.sum() > 0")
+
+    def test_select_rows_requires_boolean(self):
+        for query in ["a + 1", "True", "s"]:
+            with self.subTest(query=query), self.assertRaises(UserInputError):
+                _select_rows(_df(), query)

{vantage6_algorithm_tools-5.0.3rc5 → vantage6_algorithm_tools-5.0.4rc1}/vantage6/algorithm/mock/client.py

@@ -40,8 +40,11 @@ class MockBaseClient:
             "_download_run_data_from_hq",
             "_fetch_and_decrypt_run_data",
             "_multi_page_request",
+            "_obtain_client_credentials_tokens",
             "_refresh_token",
+            "_store_access_token",
             "_upload_run_data_to_hq",
+            "AlgorithmStoreSubClientBase",
             "auth_url",
             "authenticate",
             "check_if_blob_store_enabled",
@@ -59,6 +62,7 @@ class MockBaseClient:
             "hq_url",
             "session_id",
             "setup_encryption",
+            "store_headers",
             "store_id",
             "study_id",
             "token",
@@ -510,6 +514,7 @@ class MockUserClient(MockBaseClient):
         self.set_missing_attributes(
             [
                 "_get_logger",
+                "_obtain_tokens_per_audience",
                 "auth_client",
                 "auth_realm",
                 "authenticate_service_account",

{vantage6_algorithm_tools-5.0.3rc5 → vantage6_algorithm_tools-5.0.4rc1}/vantage6/algorithm/preprocessing/column.py

@@ -10,6 +10,7 @@ from vantage6.algorithm.tools.error_handling import handle_data_errors
 from vantage6.algorithm.tools.exceptions import UserInputError
 
 from vantage6.algorithm.decorator.action import preprocessing
+from vantage6.algorithm.preprocessing.expression import evaluate_expression
 
 
 @handle_data_errors
@@ -68,7 +69,14 @@ def assign_column(df: pd.DataFrame, column_name: str, expression: str) -> pd.Dat
     column_name : str
         The name of the new column.
     expression : str
-        The expression used to create the new column.
+        The expression used to create the new column, e.g. ``weight / (height *
+        height)``. Supported are column names, numbers, strings, ``True``, ``False``,
+        ``None``, the operators ``+ - * / // %`` (only ``+`` for text), comparisons
+        (``== != < <= > >=``), ``in``/``not in`` a list of constants,
+        ``and or not & | ~`` and the functions
+        ``abs, sqrt, log, log2, log10, exp, floor, ceil, round``. Other functions,
+        attribute access, indexing and ``**`` are not supported. The expression may
+        be at most 250 characters long.
 
     Returns
     -------
@@ -92,7 +100,7 @@ def assign_column(df: pd.DataFrame, column_name: str, expression: str) -> pd.Dat
     if column_name in df.columns:
         raise UserInputError(f"Column '{column_name}' already exists.")
 
-    df[column_name] = df.eval(expression)
+    df[column_name] = evaluate_expression(df, expression)
     return df
 
 
@@ -111,7 +119,14 @@ def redefine_column(
     column_name : str
         The name of the existing column to modify.
     expression : str
-        The expression used to modify the existing column.
+        The expression used to modify the existing column, e.g. ``round(age / 10)``.
+        Supported are column names, numbers, strings, ``True``, ``False``, ``None``,
+        the operators ``+ - * / // %`` (only ``+`` for text), comparisons
+        (``== != < <= > >=``), ``in``/``not in`` a list of constants,
+        ``and or not & | ~`` and the functions
+        ``abs, sqrt, log, log2, log10, exp, floor, ceil, round``. Other functions,
+        attribute access, indexing and ``**`` are not supported. The expression may
+        be at most 250 characters long.
 
     Returns
     -------
@@ -136,7 +151,7 @@ def redefine_column(
     if column_name not in df.columns:
         raise UserInputError(f"Column '{column_name}' does not exist.")
 
-    df[column_name] = df.eval(expression)
+    df[column_name] = evaluate_expression(df, expression)
     return df
 
 

vantage6_algorithm_tools-5.0.4rc1/vantage6/algorithm/preprocessing/expression.py

@@ -0,0 +1,259 @@
+"""
+Safe evaluation of expressions on pandas DataFrames.
+
+Expressions are provided by users as task arguments, so they must not be evaluated with
+``pandas.DataFrame.eval`` or ``pandas.DataFrame.query``: those allow attribute access and
+method calls, and can be used to execute arbitrary code. Instead, the expression is
+parsed with :mod:`ast` and evaluated by walking the syntax tree, supporting only:
+
+- column names, e.g. ``age`` (columns must be valid Python identifiers)
+- numbers, strings, ``True``, ``False`` and ``None``
+- arithmetic: ``+``, ``-``, ``*``, ``/``, ``//``, ``%`` (only ``+`` for text)
+- comparisons: ``==``, ``!=``, ``<``, ``<=``, ``>``, ``>=``, also chained (``1 < a < 5``)
+- membership in a list of constants: ``a in [1, 2]``, ``a not in ['x', 'y']``
+- logical operators: ``and``, ``or``, ``not``, ``&``, ``|``, ``~``
+- the functions ``abs``, ``sqrt``, ``log``, ``log2``, ``log10``, ``exp``, ``floor``,
+  ``ceil`` and ``round`` (``round(a)`` or ``round(a, 2)``)
+
+Anything else, such as other functions, attribute access (``a.b``), indexing
+(``a[0]``) and exponentiation (``**``), is rejected. Expressions may be at most
+``MAX_EXPRESSION_LENGTH`` characters long.
+"""
+
+import ast
+import operator
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+import pandas as pd
+
+from vantage6.algorithm.tools.exceptions import UserInputError
+
+MAX_EXPRESSION_LENGTH = 250
+
+_BINARY_OPERATORS: dict[type[ast.operator], Callable] = {
+    ast.Add: operator.add,
+    ast.Sub: operator.sub,
+    ast.Mult: operator.mul,
+    ast.Div: operator.truediv,
+    ast.FloorDiv: operator.floordiv,
+    ast.Mod: operator.mod,
+    ast.BitAnd: operator.and_,
+    ast.BitOr: operator.or_,
+}
+
+_COMPARISON_OPERATORS: dict[type[ast.cmpop], Callable] = {
+    ast.Eq: operator.eq,
+    ast.NotEq: operator.ne,
+    ast.Lt: operator.lt,
+    ast.LtE: operator.le,
+    ast.Gt: operator.gt,
+    ast.GtE: operator.ge,
+}
+
+# Functions that may be called in expressions, with the number of arguments they take
+_FUNCTIONS: dict[str, tuple[Callable, int]] = {
+    "abs": (np.abs, 1),
+    "sqrt": (np.sqrt, 1),
+    "log": (np.log, 1),
+    "log2": (np.log2, 1),
+    "log10": (np.log10, 1),
+    "exp": (np.exp, 1),
+    "floor": (np.floor, 1),
+    "ceil": (np.ceil, 1),
+}
+# `round` is handled separately, as its second argument must be an integer constant
+_SUPPORTED_FUNCTIONS = sorted([*_FUNCTIONS, "round"])
+
+_CONSTANT_TYPES = (int, float, str, bool, type(None))
+
+
+def evaluate_expression(df: pd.DataFrame, expression: str) -> Any:
+    """
+    Evaluate an expression on a DataFrame.
+
+    See the module docstring for the supported syntax.
+
+    Parameters
+    ----------
+    df : pandas.DataFrame
+        The DataFrame whose columns may be used in the expression.
+    expression : str
+        The expression to evaluate.
+
+    Returns
+    -------
+    Any
+        The result of the expression: usually a pandas Series, or a scalar if the
+        expression does not refer to any column.
+
+    Raises
+    ------
+    UserInputError
+        If the expression is too long, cannot be parsed or uses unsupported syntax.
+    """
+    if not isinstance(expression, str):
+        raise UserInputError("Expression must be a string.")
+    if len(expression) > MAX_EXPRESSION_LENGTH:
+        raise UserInputError(
+            f"Expression is too long (maximum is {MAX_EXPRESSION_LENGTH} characters)."
+        )
+    try:
+        tree = ast.parse(expression.strip(), mode="eval")
+    except (SyntaxError, ValueError, RecursionError, MemoryError):
+        raise UserInputError("Expression could not be parsed.") from None
+    return _Evaluator(df).visit(tree.body)
+
+
+class _Evaluator:
+    """Evaluate the nodes of a parsed expression that are explicitly supported."""
+
+    def __init__(self, df: pd.DataFrame) -> None:
+        self.df = df
+
+    def visit(self, node: ast.AST) -> Any:
+        if isinstance(node, ast.Name):
+            return self._visit_name(node)
+        if isinstance(node, ast.Constant):
+            return self._visit_constant(node)
+        if isinstance(node, ast.BinOp):
+            return self._visit_binop(node)
+        if isinstance(node, ast.UnaryOp):
+            return self._visit_unaryop(node)
+        if isinstance(node, ast.BoolOp):
+            return self._visit_boolop(node)
+        if isinstance(node, ast.Compare):
+            return self._visit_compare(node)
+        if isinstance(node, ast.Call):
+            return self._visit_call(node)
+        raise UserInputError(_unsupported(node))
+
+    def _visit_name(self, node: ast.Name) -> pd.Series:
+        if node.id not in self.df.columns:
+            raise UserInputError(f"Column '{node.id}' does not exist.")
+        return self.df[node.id]
+
+    @staticmethod
+    def _visit_constant(node: ast.Constant) -> Any:
+        if not isinstance(node.value, _CONSTANT_TYPES):
+            raise UserInputError(_unsupported(node))
+        return node.value
+
+    def _visit_binop(self, node: ast.BinOp) -> Any:
+        op = _BINARY_OPERATORS.get(type(node.op))
+        if op is None:
+            raise UserInputError(_unsupported(node.op))
+        left, right = self.visit(node.left), self.visit(node.right)
+        # only allow joining text, as e.g. repeating text (`'x' * 999999999`) could
+        # use excessive memory
+        if not isinstance(node.op, ast.Add) and (_is_text(left) or _is_text(right)):
+            raise UserInputError("Only '+' is supported for text.")
+        return op(left, right)
+
+    def _visit_unaryop(self, node: ast.UnaryOp) -> Any:
+        operand = self.visit(node.operand)
+        if isinstance(node.op, ast.USub):
+            return -operand
+        if isinstance(node.op, ast.UAdd):
+            return +operand
+        if isinstance(node.op, (ast.Not, ast.Invert)):
+            # like pandas, `not` and `~` are element-wise negations of booleans
+            if isinstance(operand, (bool, np.bool_)):
+                return not operand
+            return ~operand
+        raise UserInputError(_unsupported(node.op))
+
+    def _visit_boolop(self, node: ast.BoolOp) -> Any:
+        # like pandas, `and` and `or` are evaluated element-wise
+        op = operator.and_ if isinstance(node.op, ast.And) else operator.or_
+        result = self.visit(node.values[0])
+        for value in node.values[1:]:
+            result = op(result, self.visit(value))
+        return result
+
+    def _visit_compare(self, node: ast.Compare) -> Any:
+        result = None
+        left = self.visit(node.left)
+        for op, comparator in zip(node.ops, node.comparators):
+            if isinstance(op, (ast.In, ast.NotIn)):
+                right = self._visit_constant_list(comparator)
+                comparison = _isin(left, right)
+                if isinstance(op, ast.NotIn):
+                    comparison = (
+                        not comparison if isinstance(comparison, bool) else ~comparison
+                    )
+            else:
+                compare = _COMPARISON_OPERATORS.get(type(op))
+                if compare is None:
+                    raise UserInputError(_unsupported(op))
+                right = self.visit(comparator)
+                comparison = compare(left, right)
+            # chained comparisons (e.g. `1 < a < 5`) must all be true
+            result = comparison if result is None else operator.and_(result, comparison)
+            left = right
+        return result
+
+    def _visit_constant_list(self, node: ast.AST) -> list:
+        if not isinstance(node, (ast.List, ast.Tuple)) or not all(
+            isinstance(element, ast.Constant) for element in node.elts
+        ):
+            raise UserInputError(
+                "Only a list of constants is supported after 'in' or 'not in'."
+            )
+        return [self._visit_constant(element) for element in node.elts]
+
+    def _visit_call(self, node: ast.Call) -> Any:
+        if (
+            not isinstance(node.func, ast.Name)
+            or node.func.id not in _SUPPORTED_FUNCTIONS
+        ):
+            raise UserInputError(
+                f"Only the functions {', '.join(_SUPPORTED_FUNCTIONS)} are supported."
+            )
+        name = node.func.id
+        if node.keywords:
+            raise UserInputError(
+                f"Function '{name}' does not accept keyword arguments."
+            )
+        if name == "round":
+            return self._visit_round(node)
+        function, n_args = _FUNCTIONS[name]
+        if len(node.args) != n_args:
+            raise UserInputError(f"Function '{name}' takes {n_args} argument(s).")
+        return function(*(self.visit(arg) for arg in node.args))
+
+    def _visit_round(self, node: ast.Call) -> Any:
+        if len(node.args) not in (1, 2):
+            raise UserInputError("Function 'round' takes 1 or 2 arguments.")
+        decimals = 0
+        if len(node.args) == 2:
+            arg = node.args[1]
+            if not (
+                isinstance(arg, ast.Constant)
+                and isinstance(arg.value, int)
+                and not isinstance(arg.value, bool)
+            ):
+                raise UserInputError(
+                    "The number of decimals of 'round' must be an integer."
+                )
+            decimals = arg.value
+        return np.round(self.visit(node.args[0]), decimals)
+
+
+def _is_text(value: Any) -> bool:
+    if isinstance(value, pd.Series):
+        return pd.api.types.is_string_dtype(value) or pd.api.types.is_object_dtype(
+            value
+        )
+    return isinstance(value, str)
+
+
+def _isin(value: Any, options: list) -> Any:
+    if isinstance(value, pd.Series):
+        return value.isin(options)
+    return value in options
+
+
+def _unsupported(node: ast.AST) -> str:
+    return f"'{type(node).__name__}' is not supported in expressions."

{vantage6_algorithm_tools-5.0.3rc5 → vantage6_algorithm_tools-5.0.4rc1}/vantage6/algorithm/preprocessing/filtering.py

@@ -6,25 +6,31 @@ pandas DataFrames, such as selecting and filtering rows and columns.
 import pandas as pd
 
 from vantage6.algorithm.tools.error_handling import handle_data_errors
+from vantage6.algorithm.tools.exceptions import UserInputError
 
 from vantage6.algorithm.decorator.action import preprocessing
+from vantage6.algorithm.preprocessing.expression import evaluate_expression
 
 
 @handle_data_errors
 @preprocessing
 def select_rows(df: pd.DataFrame, query: str) -> pd.DataFrame:
     """
-    Select rows from the data based on a query. It uses the
-    pandas.DataFrame.query function to filter the data.
-
-    See the documentation of that function for more information on the query syntax.
+    Select the rows of the data for which the query is true.
 
     Parameters
     ----------
     df : pd.DataFrame
         The data to filter.
     query : str
-        The query string to filter on.
+        The condition that rows must meet, e.g. ``age >= 18 and sex == 'F'``.
+        Supported are column names, numbers, strings, ``True``, ``False``, ``None``,
+        the operators ``+ - * / // %`` (only ``+`` for text), comparisons
+        (``== != < <= > >=``), ``in``/``not in`` a list of constants,
+        ``and or not & | ~`` and the functions
+        ``abs, sqrt, log, log2, log10, exp, floor, ceil, round``. Other functions,
+        attribute access, indexing and ``**`` are not supported. The query may be at
+        most 250 characters long.
 
     Returns
     -------
@@ -73,9 +79,15 @@ def select_rows(df: pd.DataFrame, query: str) -> pd.DataFrame:
     3  4   9  14
     4  5  10  15
 
-    for more examples, see the documentation of pandas.DataFrame.query
+    >>> select_rows(df, "a in [1, 5]")
+       a   b   c
+    0  1   6  11
+    4  5  10  15
     """
-    return df.query(query)
+    mask = evaluate_expression(df, query)
+    if not isinstance(mask, pd.Series) or not pd.api.types.is_bool_dtype(mask):
+        raise UserInputError("The query must evaluate to true or false for each row.")
+    return df[mask]
 
 
 @handle_data_errors