Repository: joblib/joblib
Pull request #1327 (2022-09-12 07:53:00)

FIX parse pre-dispatch with AST instead of calling eval
I realized https://github.com/joblib/joblib/pull/1321 doesn't actually fix the issue, since `eval("exec('import os')", {}, {})` would still work.

This PR uses AST and limits the operations to the very basic ones.

This should be an actual fix for https://github.com/joblib/joblib/issues/1128

cc @ogrisel 

codecov[bot]: # [Codecov](https://codecov.io/gh/joblib/joblib/pull/1327?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) Report
> Merging [#1327](https://codecov.io/gh/joblib/joblib/pull/1327?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) (74eedad) into [master](https://codecov.io/gh/joblib/joblib/commit/1f00a1c30a107342bfb77ef39c0d65f5991eb1bc?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) (1f00a1c) will **decrease** coverage by `6.85%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #1327      +/-   ##
==========================================
- Coverage   93.91%   87.05%   -6.86%     
==========================================
  Files          50       52       +2     
  Lines        7275     7298      +23     
==========================================
- Hits         6832     6353     -479     
- Misses        443      945     +502     
```


| [Impacted Files](https://codecov.io/gh/joblib/joblib/pull/1327?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) | Coverage Δ | |
|---|---|---|
| [joblib/\_utils.py](https://codecov.io/gh/joblib/joblib/pull/1327/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL191dGlscy5weQ==) | `100.00% <100.00%> (ø)` | |
| [joblib/parallel.py](https://codecov.io/gh/joblib/joblib/pull/1327/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3BhcmFsbGVsLnB5) | `94.97% <100.00%> (-1.58%)` | :arrow_down: |
| [joblib/test/test\_utils.py](https://codecov.io/gh/joblib/joblib/pull/1327/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Rlc3QvdGVzdF91dGlscy5weQ==) | `100.00% <100.00%> (ø)` | |
| [joblib/test/test\_dask.py](https://codecov.io/gh/joblib/joblib/pull/1327/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Rlc3QvdGVzdF9kYXNrLnB5) | `3.53% <0.00%> (-87.64%)` | :arrow_down: |
| [joblib/\_dask.py](https://codecov.io/gh/joblib/joblib/pull/1327/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL19kYXNrLnB5) | `24.39% <0.00%> (-68.30%)` | :arrow_down: |
| [joblib/backports.py](https://codecov.io/gh/joblib/joblib/pull/1327/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL2JhY2twb3J0cy5weQ==) | `52.52% <0.00%> (-17.18%)` | :arrow_down: |
| [joblib/test/common.py](https://codecov.io/gh/joblib/joblib/pull/1327/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Rlc3QvY29tbW9uLnB5) | `67.24% <0.00%> (-10.35%)` | :arrow_down: |
| [joblib/pool.py](https://codecov.io/gh/joblib/joblib/pull/1327/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL3Bvb2wucHk=) | `83.73% <0.00%> (-4.88%)` | :arrow_down: |
| [joblib/\_parallel\_backends.py](https://codecov.io/gh/joblib/joblib/pull/1327/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL19wYXJhbGxlbF9iYWNrZW5kcy5weQ==) | `88.23% <0.00%> (-4.78%)` | :arrow_down: |
| [joblib/\_memmapping\_reducer.py](https://codecov.io/gh/joblib/joblib/pull/1327/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib#diff-am9ibGliL19tZW1tYXBwaW5nX3JlZHVjZXIucHk=) | `92.50% <0.00%> (-3.38%)` | :arrow_down: |
| ... and [12 more](https://codecov.io/gh/joblib/joblib/pull/1327/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=joblib)

adrinjalali: > Could you please add a new test for a pre_dispatch expression which includes the n_jobs placeholder in it?

Doesn't this already cover it?

```python
@parametrize('n_tasks, n_jobs, pre_dispatch, batch_size',
             [(2, 2, 'all', 'auto'),
              (2, 2, 'n_jobs', 'auto'),
              (10, 2, 'n_jobs', 'auto'),
              (517, 2, 'n_jobs', 'auto'),
              (10, 2, 'n_jobs', 'auto'),
              (10, 4, 'n_jobs', 'auto'),
              (200, 12, 'n_jobs', 'auto'),
              (25, 12, '2 * n_jobs', 1),
              (250, 12, 'all', 1),
              (250, 12, '2 * n_jobs', 7),
              (200, 12, '2 * n_jobs', 'auto')])
def test_dispatch_race_condition(n_tasks, n_jobs, pre_dispatch, batch_size):
    # Check that using (async-)dispatch does not yield a race condition on the
    # iterable generator that is not thread-safe natively.
    # This is a non-regression test for the "Pool seems closed" class of error
    params = {'n_jobs': n_jobs, 'pre_dispatch': pre_dispatch,
              'batch_size': batch_size}
    expected = [square(i) for i in range(n_tasks)]
    results = Parallel(**params)(delayed(square)(i) for i in range(n_tasks))
    assert results == expected
```
ogrisel: Indeed, my `git grep` was too restrictive.

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 62bf7ed..26742a4 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -30,7 +30,7 @@ Development version
 
 - Fix a security issue where ``eval(pre_dispatch)`` could potentially run
   arbitrary code. Now only basic numerics are supported.
-  https://github.com/joblib/joblib/pull/1321
+  https://github.com/joblib/joblib/pull/1327
 
 - Vendor cloudpickle 2.2.0 which adds support for PyPy 3.8+.
 
diff --git a/joblib/_utils.py b/joblib/_utils.py
new file mode 100644
index 0000000..2dbd4f6
--- /dev/null
+++ b/joblib/_utils.py
@@ -0,0 +1,44 @@
+# Adapted from https://stackoverflow.com/a/9558001/2536294
+
+import ast
+import operator as op
+
+# supported operators
+operators = {
+    ast.Add: op.add,
+    ast.Sub: op.sub,
+    ast.Mult: op.mul,
+    ast.Div: op.truediv,
+    ast.FloorDiv: op.floordiv,
+    ast.Mod: op.mod,
+    ast.Pow: op.pow,
+    ast.USub: op.neg,
+}
+
+
+def eval_expr(expr):
+    """
+    >>> eval_expr('2*6')
+    12
+    >>> eval_expr('2**6')
+    64
+    >>> eval_expr('1 + 2*3**(4) / (6 + -7)')
+    -161.0
+    """
+    try:
+        return eval_(ast.parse(expr, mode="eval").body)
+    except (TypeError, SyntaxError, KeyError) as e:
+        raise ValueError(
+            f"{expr!r} is not a valid or supported arithmetic expression."
+        ) from e
+
+
+def eval_(node):
+    if isinstance(node, ast.Num):  # <number>
+        return node.n
+    elif isinstance(node, ast.BinOp):  # <left> <operator> <right>
+        return operators[type(node.op)](eval_(node.left), eval_(node.right))
+    elif isinstance(node, ast.UnaryOp):  # <operator> <operand> e.g., -1
+        return operators[type(node.op)](eval_(node.operand))
+    else:
+        raise TypeError(node)
diff --git a/joblib/parallel.py b/joblib/parallel.py
index 1c2fe18..6e7b1b1 100644
--- a/joblib/parallel.py
+++ b/joblib/parallel.py
@@ -27,6 +27,7 @@ from ._parallel_backends import (FallbackToBackend, MultiprocessingBackend,
                                  ThreadingBackend, SequentialBackend,
                                  LokyBackend)
 from .externals.cloudpickle import dumps, loads
+from ._utils import eval_expr
 
 # Make sure that those two classes are part of the public joblib.parallel API
 # so that 3rd party backend implementers can import them from here.
@@ -1051,10 +1052,8 @@ class Parallel(Logger):
         else:
             self._original_iterator = iterator
             if hasattr(pre_dispatch, 'endswith'):
-                pre_dispatch = eval(
-                    pre_dispatch,
-                    {"n_jobs": n_jobs, "__builtins__": {}},  # globals
-                    {}  # locals
+                pre_dispatch = eval_expr(
+                    pre_dispatch.replace("n_jobs", str(n_jobs))
                 )
             self._pre_dispatch_amount = pre_dispatch = int(pre_dispatch)
 


Test patch:
diff --git a/joblib/test/test_utils.py b/joblib/test/test_utils.py
new file mode 100644
index 0000000..4999a21
--- /dev/null
+++ b/joblib/test/test_utils.py
@@ -0,0 +1,27 @@
+import pytest
+
+from joblib._utils import eval_expr
+
+
+@pytest.mark.parametrize(
+    "expr",
+    ["exec('import os')", "print(1)", "import os", "1+1; import os", "1^1"],
+)
+def test_eval_expr_invalid(expr):
+    with pytest.raises(
+        ValueError, match="is not a valid or supported arithmetic"
+    ):
+        eval_expr(expr)
+
+
+@pytest.mark.parametrize(
+    "expr, result",
+    [
+        ("2*6", 12),
+        ("2**6", 64),
+        ("1 + 2*3**(4) / (6 + -7)", -161.0),
+        ("(20 // 3) % 5", 1),
+    ],
+)
+def test_eval_expr_valid(expr, result):
+    assert eval_expr(expr) == result