diff --git a/nemo_run/cli/cli_parser.py b/nemo_run/cli/cli_parser.py index 76d8cca1..c4d45056 100644 --- a/nemo_run/cli/cli_parser.py +++ b/nemo_run/cli/cli_parser.py @@ -518,8 +518,8 @@ def _contains_unsafe_operations(self, node: ast.AST) -> bool: return self._contains_unsafe_operations(node.operand) elif isinstance(node, (ast.List, ast.Tuple, ast.Set, ast.Dict)): return any(self._contains_unsafe_operations(elt) for elt in ast.iter_child_nodes(node)) - elif isinstance(node, (ast.Num, ast.Str, ast.Bytes, ast.NameConstant, ast.Ellipsis)): - # Allow basic literals + elif isinstance(node, ast.Constant): + # Allow basic literals (numbers, strings, bytes, True/False/None, Ellipsis) return False return True diff --git a/nemo_run/cli/lazy.py b/nemo_run/cli/lazy.py index f6785dbb..deaf2ce3 100644 --- a/nemo_run/cli/lazy.py +++ b/nemo_run/cli/lazy.py @@ -833,8 +833,8 @@ def _load_entrypoint_from_script(script_content: str, module_name: str = "__dyna if ( isinstance(node.test.left, ast.Name) and node.test.left.id == "__name__" - and isinstance(node.test.comparators[0], ast.Str) - and node.test.comparators[0].s == "__main__" + and isinstance(node.test.comparators[0], ast.Constant) + and node.test.comparators[0].value == "__main__" ): main_block = node.body break diff --git a/test/cli/test_ast_constants.py b/test/cli/test_ast_constants.py new file mode 100644 index 00000000..f0fb36c8 --- /dev/null +++ b/test/cli/test_ast_constants.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import ast + +import pytest + +from nemo_run.cli.cli_parser import PythonicParser + + +class TestPythonicParserLiteralSafety: + @pytest.fixture + def parser(self): + return PythonicParser() + + @pytest.mark.parametrize( + "expression", + [ + "1", + "'hello'", + "b'hello'", + "True", + "None", + "...", + ], + ) + def test_literal_expressions_are_safe(self, parser, expression): + """Literal expressions previously checked via deprecated ast.* types.""" + assert parser._contains_unsafe_operations(ast.parse(expression, mode="eval").body) is False