Use argparse to process arguments to /source

changelog.md

@@ -25,6 +25,7 @@ Internal
 * Upgrade `coverage` dev dependency to v7.15.4.
 * Upgrade `pytest-cov` dev dependency to v7.1.0.
 * Upgrade `wcwidth` dependency to v0.8.3.
+* Use `argparse` to process arguments to `/source`.
 
 
 2.19.0 (2026/09/02)

mycli/client_commands.py

@@ -20,7 +20,6 @@
 from mycli.packages.special.source import (
     SOURCE_HELP_ROWS,
     parse_source_arguments,
-    parse_source_filename,
     source_special_command_is_safe,
 )
 from mycli.packages.sqlresult import SQLResult
@@ -291,20 +290,16 @@ def change_db(self, arg: str, **_) -> Generator[SQLResult, None, None]:
 
     def execute_from_file(self, arg: str, **_) -> Generator[SQLResult, None, None]:
         try:
-            filename, allow_special, show_queries, page_output, throttle, show_help = parse_source_arguments(arg)
+            source_arguments = parse_source_arguments(arg)
         except ValueError as error:
             yield SQLResult(status=str(error), is_error=True)
             return
-        if show_help:
+        if source_arguments.show_help:
             yield SQLResult(header=['Argument', 'Description'], rows=SOURCE_HELP_ROWS)
             return
-        if page_output:
+        if source_arguments.page_output:
             yield SQLResult(command={'name': 'source_page'})
-        try:
-            filename = parse_source_filename(filename)
-        except ValueError as error:
-            yield SQLResult(status=str(error), is_error=True)
-            return
+        filename = source_arguments.filename
         if not filename:
             yield SQLResult(status="Missing required argument: filename.  See /source --help.", is_error=True)
             return
@@ -330,7 +325,7 @@ def execute_from_file(self, arg: str, **_) -> Generator[SQLResult, None, None]:
 
                 special_query = query.rstrip(';')
                 if special.is_special_command(special_query):
-                    if not allow_special:
+                    if not source_arguments.allow_special:
                         yield SQLResult(
                             status='Special commands are not supported without /source --special.',
                             is_error=True,
@@ -343,10 +338,10 @@ def execute_from_file(self, arg: str, **_) -> Generator[SQLResult, None, None]:
                             is_error=True,
                         )
                         return
-                    if executed_statement and throttle > 0:
-                        time.sleep(throttle)
-                    if show_queries:
-                        if page_output:
+                    if executed_statement and source_arguments.throttle > 0:
+                        time.sleep(source_arguments.throttle)
+                    if source_arguments.show_queries:
+                        if source_arguments.page_output:
                             yield SQLResult(command={'name': 'source_show', 'text': special_query})
                         else:
                             click.secho(f'> {special_query}')
@@ -356,10 +351,10 @@ def execute_from_file(self, arg: str, **_) -> Generator[SQLResult, None, None]:
 
                 if self.destructive_warning and confirm_destructive_query(self.destructive_keywords, query) is False:
                     continue
-                if executed_statement and throttle > 0:
-                    time.sleep(throttle)
-                if show_queries:
-                    if page_output:
+                if executed_statement and source_arguments.throttle > 0:
+                    time.sleep(source_arguments.throttle)
+                if source_arguments.show_queries:
+                    if source_arguments.page_output:
                         yield SQLResult(command={'name': 'source_show', 'text': query})
                     else:
                         click.secho(f'> {query}')

mycli/packages/completion_engine.py

@@ -818,7 +818,23 @@ def suggest_special(text: str) -> list[dict[str, Any]]:
         '/source',
     ]:
         source_options = list(SOURCE_OPTIONS)
-        source_arguments = _arg.split()
+        try:
+            source_arguments = shlex.split(_arg, posix=False)
+        except ValueError:
+            source_filename = _arg
+            while leading_arguments := source_filename.split(maxsplit=1):
+                if leading_arguments[0] in SOURCE_BOOLEAN_OPTIONS:
+                    source_filename = leading_arguments[1] if len(leading_arguments) == 2 else ''
+                    continue
+                if leading_arguments[0] == '--throttle' and len(leading_arguments) == 2:
+                    throttle_arguments = leading_arguments[1].split(maxsplit=1)
+                    source_filename = throttle_arguments[1] if len(throttle_arguments) == 2 else ''
+                    continue
+                if leading_arguments[0].startswith('--throttle='):
+                    source_filename = leading_arguments[1] if len(leading_arguments) == 2 else ''
+                    continue
+                break
+            return [{'type': 'file_name', 'quote_spaces': True, 'source_filename': source_filename}]
         if not source_arguments:
             return [
                 {'type': 'special_subcommand', 'subcommands': source_options},
@@ -827,15 +843,22 @@ def suggest_special(text: str) -> list[dict[str, Any]]:
 
         used_options: set[str] = set()
         argument_index = 0
+        filename_index: int | None = None
+        options_ended = False
         while argument_index < len(source_arguments):
             argument = source_arguments[argument_index]
-            if argument in SOURCE_BOOLEAN_OPTIONS:
+            if not options_ended and argument == '--':
+                options_ended = True
+                used_options.update(source_options)
+                argument_index += 1
+                continue
+            if not options_ended and argument in SOURCE_BOOLEAN_OPTIONS:
                 used_options.add(argument)
                 argument_index += 1
                 continue
-            if argument == '--help':
+            if not options_ended and argument == '--help':
                 return []
-            if argument == '--throttle':
+            if not options_ended and argument == '--throttle':
                 used_options.add(argument)
                 argument_index += 1
                 if argument_index >= len(source_arguments):
@@ -844,30 +867,31 @@ def suggest_special(text: str) -> list[dict[str, Any]]:
                     return []
                 argument_index += 1
                 continue
-            if argument.startswith('--throttle='):
+            if not options_ended and argument.startswith('--throttle='):
                 used_options.add('--throttle')
                 if argument == '--throttle=' or not text[-1].isspace():
                     return []
                 argument_index += 1
                 continue
-            break
+            if not options_ended and argument.startswith('-'):
+                break
+            if filename_index is not None:
+                return []
+            filename_index = argument_index
+            argument_index += 1
         remaining_options = [option for option in source_options if option not in used_options]
-        source_filename = _arg
-        for _index in range(argument_index):
-            parsed_argument = source_filename.split(maxsplit=1)
-            source_filename = parsed_argument[1] if len(parsed_argument) == 2 else ''
+        source_filename = source_arguments[filename_index] if filename_index is not None else ''
         file_suggestion = {'type': 'file_name', 'quote_spaces': True, 'source_filename': source_filename}
 
         if argument_index < len(source_arguments):
-            if source_arguments[argument_index].startswith('-'):
-                return [{'type': 'special_subcommand', 'subcommands': remaining_options}]
-            return [file_suggestion]
+            return [{'type': 'special_subcommand', 'subcommands': remaining_options}]
         if not text[-1].isspace():
-            return []
+            return [file_suggestion] if filename_index == len(source_arguments) - 1 else []
         suggestions: list[dict[str, Any]] = []
         if remaining_options:
             suggestions.append({'type': 'special_subcommand', 'subcommands': remaining_options})
-        suggestions.append(file_suggestion)
+        if filename_index is None:
+            suggestions.append(file_suggestion)
         return suggestions
 
     if cmd.lower() in [

mycli/packages/hybrid_redirection.py

@@ -9,26 +9,26 @@
 from mycli.packages.special.delimitercommand import DelimiterCommand
 from mycli.packages.special.source import (
     parse_source_arguments,
-    parse_source_filename,
 )
 
 logger = logging.getLogger(__name__)
 delimiter_command = DelimiterCommand()
 SOURCE_COMMAND_PATTERN = re.compile(r'^([/]?source|[/\\]\.)\s+', re.IGNORECASE)
-SOURCE_OPTIONS_PATTERN = re.compile(
-    r'^([/]?source|[/\\]\.)\s+'
-    r'(?P<options>(?:(?:--special|--show|--page)\s+|--throttle(?:=[^\s]+|\s+[^\s]+)\s+)*)',
-    re.IGNORECASE,
-)
+SOURCE_OPTION_PATTERN = re.compile(r'(?<!\S)--(?:special|show|page|help|throttle)(?==|\s|$)', re.IGNORECASE)
+SOURCE_OPTION_TERMINATOR_PATTERN = re.compile(r'(?<!\S)--(?=\s|$)')
+HYBRID_OPERATOR_PATTERN = re.compile(r'\$(?:>>?|\|)')
 
 
 def tokenize_command(command: str) -> list[sqlglot.Token]:
     """Tokenize a command without treating source options as SQL comments."""
-    source_match = SOURCE_OPTIONS_PATTERN.match(command)
-    if source_match:
-        options_start, options_end = source_match.span('options')
-        options = command[options_start:options_end].replace('-', '_')
-        command = command[:options_start] + options + command[options_end:]
+    if SOURCE_COMMAND_PATTERN.match(command):
+        operator_match = HYBRID_OPERATOR_PATTERN.search(command)
+        source_end = operator_match.start() if operator_match is not None else len(command)
+        source_part = command[:source_end]
+        terminator_match = SOURCE_OPTION_TERMINATOR_PATTERN.search(source_part)
+        options_end = terminator_match.start() if terminator_match is not None else len(source_part)
+        masked_options = SOURCE_OPTION_PATTERN.sub(lambda match: match.group().replace('-', '_'), source_part[:options_end])
+        command = masked_options + source_part[options_end:] + command[source_end:]
     return sqlglot.tokenize(command)
 
 
@@ -68,8 +68,7 @@ def find_sql_part(
     if SOURCE_COMMAND_PATTERN.match(sql_part):
         source_arg_str = SOURCE_COMMAND_PATTERN.sub('', sql_part)
         try:
-            filename, _allow_special, _show_queries, _page_output, _throttle, _show_help = parse_source_arguments(source_arg_str)
-            filename = parse_source_filename(filename)
+            filename = parse_source_arguments(source_arg_str).filename
         except ValueError:
             return ''
         if not filename:

mycli/packages/special/__init__.py

@@ -55,10 +55,7 @@
     write_pipe_once,
     write_tee,
 )
-from mycli.packages.special.source import (
-    parse_source_arguments,
-    parse_source_filename,
-)
+from mycli.packages.special.source import parse_source_arguments
 
 if not os.environ.get('MYCLI_LLM_OFF'):
     from mycli.packages.special.llm import (
@@ -125,7 +122,6 @@ def sql_using_llm(*args, **kwargs):  # type: ignore[no-redef, misc]
     'list_tables',
     'open_external_editor',
     'parse_source_arguments',
-    'parse_source_filename',
     'parse_special_command',
     'ping',
     'register_special_command',