diff --git a/cli/openbb_cli/controllers/base_controller.py b/cli/openbb_cli/controllers/base_controller.py index ab5f70be386..a6b858d1686 100644 --- a/cli/openbb_cli/controllers/base_controller.py +++ b/cli/openbb_cli/controllers/base_controller.py @@ -201,7 +201,7 @@ class BaseController(metaclass=ABCMeta): # Single command fed, process else: try: - (known_args, other_args) = self.parser.parse_known_args( + known_args, other_args = self.parser.parse_known_args( shlex.split(an_input) ) except Exception as exc: @@ -622,7 +622,7 @@ class BaseController(metaclass=ABCMeta): system_clear() try: - (ns_parser, l_unknown_args) = parser.parse_known_args(other_args) + ns_parser, l_unknown_args = parser.parse_known_args(other_args) except SystemExit: # In case the command has required argument that isn't specified session.console.print("\n") @@ -768,11 +768,45 @@ class BaseController(metaclass=ABCMeta): ), -1, ) - # Split comma-separated arguments, except for the argument at routine_args_index + # Collect indices whose values should NOT be comma-split because + # the provider may accept a comma-separated string (e.g. --symbol + # AAPL,MSFT). Handles --flag value, --flag=value, -f value, and + # multi-value flags (nargs="+" / nargs="*" / nargs=N). + no_split_indices: set[int] = set() + if 0 <= routine_args_index < len(other_args): + no_split_indices.add(routine_args_index) + + for i, arg in enumerate(other_args): + if not arg.startswith("-"): + continue + # Handle --flag=value by extracting the flag portion. + flag_part = arg.split("=", 1)[0] if "=" in arg else arg + for action in parser._actions: # pylint: disable=protected-access + if flag_part in action.option_strings and action.nargs != 0: + if "=" in arg: + # Value is embedded in the same token. + no_split_indices.add(i) + elif action.nargs in ("+", "*") or ( + isinstance(action.nargs, int) and action.nargs > 1 + ): + # Multi-value flag: protect all consecutive + # non-flag tokens after the flag. + j = i + 1 + while j < len(other_args) and not other_args[j].startswith( + "-" + ): + no_split_indices.add(j) + j += 1 + # Single-value flag: protect the next token. + elif i + 1 < len(other_args): + no_split_indices.add(i + 1) + break + + # Split comma-separated arguments only for positional / unflagged values. other_args = [ part for index, arg in enumerate(other_args) - for part in (arg.split(",") if index != routine_args_index else [arg]) + for part in ([arg] if index in no_split_indices else arg.split(",")) ] # Check if the action has optional choices, if yes, remove them @@ -780,7 +814,7 @@ class BaseController(metaclass=ABCMeta): if getattr(action, "optional_choices", None): action.choices = None - (ns_parser, l_unknown_args) = parser.parse_known_args(other_args) + ns_parser, l_unknown_args = parser.parse_known_args(other_args) if export_allowed in [ "raw_data_only", diff --git a/cli/tests/test_controllers_base_controller.py b/cli/tests/test_controllers_base_controller.py index ebf0fb3e9b1..6cec53bbfd7 100644 --- a/cli/tests/test_controllers_base_controller.py +++ b/cli/tests/test_controllers_base_controller.py @@ -1,5 +1,6 @@ """Test the base controller.""" +import argparse from unittest.mock import MagicMock, patch import pytest @@ -77,3 +78,114 @@ def test_call_exit(): with patch.object(controller, "save_class", MagicMock()): controller.queue = ["quit"] controller.call_exit(None) + + +@pytest.fixture +def mock_base_session(): + """Mock the session for parse_known_args_and_warn tests.""" + with patch("openbb_cli.controllers.base_controller.session") as mock_session: + mock_session.settings.USE_CLEAR_AFTER_CMD = False + yield mock_session + + +def _make_parser(*args_spec): + """Create an argparse parser from a list of add_argument kwargs.""" + parser = argparse.ArgumentParser(add_help=False) + for spec in args_spec: + flags = spec.pop("flags") + parser.add_argument(*flags, **spec) + return parser + + +def test_comma_split_flagged_value_not_split(mock_base_session): + """Simple test: --symbol AAPL,MSFT must stay as one value.""" + parser = _make_parser({"flags": ["--symbol", "-s"], "dest": "symbol", "type": str}) + result = BaseController.parse_known_args_and_warn(parser, ["--symbol", "AAPL,MSFT"]) + assert result is not None + assert result.symbol == "AAPL,MSFT" + + +def test_comma_split_short_flag_not_split(mock_base_session): + """Short flag -s AAPL,MSFT must also stay as one value.""" + parser = _make_parser({"flags": ["--symbol", "-s"], "dest": "symbol", "type": str}) + result = BaseController.parse_known_args_and_warn(parser, ["-s", "AAPL,MSFT"]) + assert result is not None + assert result.symbol == "AAPL,MSFT" + + +def test_comma_split_equals_syntax_not_split(mock_base_session): + """--symbol=AAPL,MSFT must not be split.""" + parser = _make_parser({"flags": ["--symbol", "-s"], "dest": "symbol", "type": str}) + result = BaseController.parse_known_args_and_warn(parser, ["--symbol=AAPL,MSFT"]) + assert result is not None + assert result.symbol == "AAPL,MSFT" + + +def test_comma_split_nargs_plus_all_values_protected(mock_base_session): + """nargs='+': all consecutive values after --symbols are protected.""" + parser = _make_parser( + {"flags": ["--symbols"], "dest": "symbols", "nargs": "+", "type": str} + ) + result = BaseController.parse_known_args_and_warn( + parser, ["--symbols", "AAPL,MSFT", "GOOG,AMZN"] + ) + assert result is not None + assert result.symbols == ["AAPL,MSFT", "GOOG,AMZN"] + + +def test_comma_split_nargs_star_values_protected(mock_base_session): + """nargs='*': consecutive values after --tags are protected.""" + parser = _make_parser( + {"flags": ["--tags"], "dest": "tags", "nargs": "*", "type": str} + ) + result = BaseController.parse_known_args_and_warn(parser, ["--tags", "a,b", "c,d"]) + assert result is not None + assert result.tags == ["a,b", "c,d"] + + +def test_comma_split_nargs_int_values_protected(mock_base_session): + """nargs=2: both values after --pair are protected.""" + parser = _make_parser( + {"flags": ["--pair"], "dest": "pair", "nargs": 2, "type": str} + ) + result = BaseController.parse_known_args_and_warn(parser, ["--pair", "a,b", "c,d"]) + assert result is not None + assert result.pair == ["a,b", "c,d"] + + +def test_comma_split_store_true_not_confused(mock_base_session): + """store_true flags (nargs=0) should not protect the next token.""" + parser = _make_parser( + {"flags": ["--symbol", "-s"], "dest": "symbol", "type": str}, + {"flags": ["--raw"], "dest": "raw", "action": "store_true", "default": False}, + ) + result = BaseController.parse_known_args_and_warn( + parser, ["--raw", "--symbol", "AAPL,MSFT"] + ) + assert result is not None + assert result.raw is True + assert result.symbol == "AAPL,MSFT" + + +def test_comma_split_no_comma_values_unchanged(mock_base_session): + """Values without commas pass through unaffected.""" + parser = _make_parser({"flags": ["--symbol", "-s"], "dest": "symbol", "type": str}) + result = BaseController.parse_known_args_and_warn(parser, ["--symbol", "AAPL"]) + assert result is not None + assert result.symbol == "AAPL" + + +def test_comma_split_multiple_flags_each_protected(mock_base_session): + """Multiple flags each protect their own values independently.""" + parser = _make_parser( + {"flags": ["--symbol", "-s"], "dest": "symbol", "type": str}, + {"flags": ["--raw"], "dest": "raw", "action": "store_true", "default": False}, + {"flags": ["--provider"], "dest": "provider", "type": str}, + ) + result = BaseController.parse_known_args_and_warn( + parser, + ["--symbol", "AAPL,MSFT", "--provider", "yfinance,polygon"], + ) + assert result is not None + assert result.symbol == "AAPL,MSFT" + assert result.provider == "yfinance,polygon"