Extract code from line magics for attribute completion (ipython#14982)

Fixes ipython#14981

Author: krassowski

.github/workflows/downstream.yml

@@ -88,4 +88,4 @@ jobs:
     - name: Test pyflyby (IPython integration only)
       run: |
         cd ../pyflyby
-        pytest tests/test_interactive.py -k 'not test_timeit_complete_autoimport_member_1'
+        pytest tests/test_interactive.py

IPython/core/completer.py

@@ -218,7 +218,7 @@
     EvaluationContext,
     _validate_policy_overrides,
 )
-from IPython.core.error import TryNext
+from IPython.core.error import TryNext, UsageError
 from IPython.core.inputtransformer2 import ESC_MAGIC
 from IPython.core.latex_symbols import latex_symbols, reverse_latex_symbol
 from IPython.testing.skipdoctest import skip_doctest
@@ -1219,6 +1219,10 @@ def _strip_code_before_operator(self, code: str) -> str:
         else:
             return code
 
+    def _extract_code(self, line: str):
+        """No-op in Completer, but can be used in subclasses to customise behaviour"""
+        return line
+
     def _attr_matches(
         self, text: str, include_prefix: bool = True
     ) -> tuple[Sequence[str], str]:
@@ -2244,6 +2248,57 @@ def file_matcher(self, context: CompletionContext) -> SimpleMatcherResult:
             "suppress": False,
         }
 
+    def _extract_code(self, line: str) -> str:
+        """Extract code from magics if any."""
+
+        if not line:
+            return line
+        maybe_magic, *rest = line.split(maxsplit=1)
+        if not rest:
+            return line
+        args = rest[0]
+        known_magics = self.shell.magics_manager.lsmagic()
+        line_magics = known_magics["line"]
+        magic_name = maybe_magic.lstrip(self.magic_escape)
+        if magic_name not in line_magics:
+            return line
+
+        if not maybe_magic.startswith(self.magic_escape):
+            all_variables = [*self.namespace.keys(), *self.global_namespace.keys()]
+            if magic_name in all_variables:
+                # short circuit if we see a line starting with say `time`
+                # but time is defined as a variable (in addition to being
+                # a magic). In these cases users need to use explicit `%time`.
+                return line
+
+        magic_method = line_magics[magic_name]
+
+        try:
+            if magic_name == "timeit":
+                opts, stmt = magic_method.__self__.parse_options(
+                    args,
+                    "n:r:tcp:qov:",
+                    posix=False,
+                    strict=False,
+                    preserve_non_opts=True,
+                )
+                return stmt
+            elif magic_name == "prun":
+                opts, stmt = magic_method.__self__.parse_options(
+                    args, "D:l:rs:T:q", list_all=True, posix=False
+                )
+                return stmt
+            elif hasattr(magic_method, "parser") and getattr(
+                magic_method, "has_arguments", False
+            ):
+                # e.g. %debug, %time
+                args, extra = magic_method.parser.parse_argstring(args, partial=True)
+                return " ".join(extra)
+        except UsageError:
+            return line
+
+        return line
+
     @context_matcher()
     def magic_matcher(self, context: CompletionContext) -> SimpleMatcherResult:
         """Match magics."""
@@ -2255,7 +2310,7 @@ def magic_matcher(self, context: CompletionContext) -> SimpleMatcherResult:
         line_magics = lsm['line']
         cell_magics = lsm['cell']
         pre = self.magic_escape
-        pre2 = pre+pre
+        pre2 = pre + pre
 
         explicit_magic = text.startswith(pre)
 
@@ -2619,6 +2674,7 @@ def _is_in_string_or_comment(self, text):
     def python_matcher(self, context: CompletionContext) -> SimpleMatcherResult:
         """Match attributes or global python names"""
         text = context.text_until_cursor
+        text = self._extract_code(text)
         completion_type = self._determine_completion_context(text)
         if completion_type == self._CompletionContextType.ATTRIBUTE:
             try:
@@ -3476,7 +3532,7 @@ def _complete(self, *, cursor_line, cursor_pos, line_buffer=None, text=None,
             full_text=full_text,
             cursor_position=cursor_pos,
             cursor_line=cursor_line,
-            token=text,
+            token=self._extract_code(text),
             limit=MATCHES_LIMIT,
         )
 

IPython/core/magic_arguments.py

@@ -161,10 +161,12 @@ def error(self, message):
         """
         raise UsageError(message)
 
-    def parse_argstring(self, argstring):
+    def parse_argstring(self, argstring, *, partial=False):
         """ Split a string into an argument list and parse that argument list.
         """
         argv = arg_split(argstring)
+        if partial:
+            return self.parse_known_args(argv)
         return self.parse_args(argv)
 
 
@@ -190,10 +192,10 @@ def construct_parser(magic_func):
     return parser
 
 
-def parse_argstring(magic_func, argstring):
+def parse_argstring(magic_func, argstring, *, partial=False):
     """ Parse the string of arguments for the given magic function.
     """
-    return magic_func.parser.parse_argstring(argstring)
+    return magic_func.parser.parse_argstring(argstring, partial=partial)
 
 
 def real_name(magic_func):

IPython/core/magics/execution.py

@@ -317,6 +317,7 @@ def prun(self, parameter_s='', cell=None):
             the magic line is always left unmodified.
 
         """
+        # TODO: port to magic_arguments as currently this is duplicated in IPCompleter._extract_code
         opts, arg_str = self.parse_options(parameter_s, 'D:l:rs:T:q',
                                            list_all=True, posix=False)
         if cell is not None:
@@ -441,11 +442,10 @@ def pdb(self, parameter_s=''):
         Set break point at LINE in FILE.
         """
     )
-    @magic_arguments.argument('statement', nargs='*',
-        help="""
-        Code to run in debugger.
-        You can omit this in cell magic mode.
-        """
+    @magic_arguments.kwds(
+        epilog="""
+      Any remaining arguments will be treated as code to run in the debugger.
+    """
     )
     @no_var_expand
     @line_cell_magic
@@ -454,12 +454,12 @@ def debug(self, line="", cell=None, local_ns=None):
         """Activate the interactive debugger.
 
         This magic command support two ways of activating debugger.
-        One is to activate debugger before executing code.  This way, you
+        One is to activate debugger before executing code. This way, you
         can set a break point, to step through the code from the point.
         You can use this mode by giving statements to execute and optionally
         a breakpoint.
 
-        The other one is to activate debugger in post-mortem mode.  You can
+        The other one is to activate debugger in post-mortem mode. You can
         activate this mode simply running %debug without any argument.
         If an exception has just occurred, this lets you inspect its stack
         frames interactively.  Note that this will always work only on the last
@@ -475,9 +475,9 @@ def debug(self, line="", cell=None, local_ns=None):
             the magic line is always left unmodified.
 
         """
-        args = magic_arguments.parse_argstring(self.debug, line)
+        args, extra = magic_arguments.parse_argstring(self.debug, line, partial=True)
 
-        if not (args.breakpoint or args.statement or cell):
+        if not (args.breakpoint or extra or cell):
             self._debug_post_mortem()
         elif not (args.breakpoint or cell):
             # If there is no breakpoints, the line is just code to execute
@@ -486,7 +486,7 @@ def debug(self, line="", cell=None, local_ns=None):
             # Here we try to reconstruct the code from the output of
             # parse_argstring. This might not work if the code has spaces
             # For example this fails for `print("a b")`
-            code = "\n".join(args.statement)
+            code = " ".join(extra)
             if cell:
                 code += "\n" + cell
             self._debug_exec(code, args.breakpoint, local_ns)
@@ -1138,6 +1138,7 @@ def timeit(self, line='', cell=None, local_ns=None):
         does not matter as long as results from timeit.py are not mixed with
         those from ``%timeit``."""
 
+        # TODO: port to magic_arguments as currently this is duplicated in IPCompleter._extract_code
         opts, stmt = self.parse_options(
             line, "n:r:tcp:qov:", posix=False, strict=False, preserve_non_opts=True
         )
@@ -1266,6 +1267,11 @@ def timeit(self, line='', cell=None, local_ns=None):
         dest="no_raise_error",
         help="If given, don't re-raise exceptions",
     )
+    @magic_arguments.kwds(
+        epilog="""
+      Any remaining arguments will be treated as code to run.
+    """
+    )
     @skip_doctest
     @needs_local_scope
     @line_cell_magic
@@ -1337,34 +1343,10 @@ def time(self, line="", cell=None, local_ns=None):
                 Wall time: 0.00 s
                 Compiler : 0.78 s
         """
-        line_present = False
-        # Try to parse --no-raise-error if present, else ignore unrecognized args
-        try:
-            args = magic_arguments.parse_argstring(self.time, line)
-        except UsageError as e:
-            # Only ignore UsageError if caused by unrecognized arguments
-            # We'll manually check for --no-raise-error and remove it from line
-            line_present = True
-
-            # Check if --no-raise-error is present
-            no_raise_error = "--no-raise-error" in line
-
-            if no_raise_error:
-                # Remove --no-raise-error while preserving the rest of the line structure
-                line = re.sub(r"\s*--no-raise-error\s*", " ", line).strip()
-                # Clean up any double spaces
-                line = re.sub(r"\s+", " ", line)
-
-            class Args:
-                def __init__(self, no_raise_error):
-                    self.no_raise_error = no_raise_error
-
-            args = Args(no_raise_error)
-        else:
-            if not hasattr(args, "no_raise_error"):
-                args.no_raise_error = False
+        args, extra = magic_arguments.parse_argstring(self.time, line, partial=True)
+        line = " ".join(extra)
 
-        if line_present and cell:
+        if line and cell:
             raise UsageError("Can't use statement directly after '%%time'!")
 
         if cell:

tests/test_completer.py

@@ -905,6 +905,66 @@ def _bar_cellm(line, cell):
         self.assertNotIn("%_bar_cellm", matches)
         self.assertIn("%%_bar_cellm", matches)
 
+    def test_line_magics_with_code_argument(self):
+        ip = get_ipython()
+        c = ip.Completer
+        c.use_jedi = False
+
+        # attribute completion
+        text, matches = c.complete("%timeit -n 2 -r 1 float.as_integer")
+        self.assertEqual(matches, [".as_integer_ratio"])
+
+        text, matches = c.complete("%debug --breakpoint test float.as_integer")
+        self.assertEqual(matches, [".as_integer_ratio"])
+
+        text, matches = c.complete("%time --no-raise-error float.as_integer")
+        self.assertEqual(matches, [".as_integer_ratio"])
+
+        text, matches = c.complete("%prun -l 0.5 -r float.as_integer")
+        self.assertEqual(matches, [".as_integer_ratio"])
+
+        # implicit magics
+        text, matches = c.complete("timeit -n 2 -r 1 float.as_integer")
+        self.assertEqual(matches, [".as_integer_ratio"])
+
+        # built-ins completion
+        text, matches = c.complete("%timeit -n 2 -r 1 flo")
+        self.assertEqual(matches, ["float"])
+
+        # dict completion
+        text, matches = c.complete("%timeit -n 2 -r 1 {'my_key': 1}['my")
+        self.assertEqual(matches, ["my_key"])
+
+        # invalid arguments - should not throw
+        text, matches = c.complete("%timeit -n 2 -r 1 -invalid float.as_integer")
+        self.assertEqual(matches, [])
+
+        text, matches = c.complete("%debug --invalid float.as_integer")
+        self.assertEqual(matches, [])
+
+    def test_line_magics_with_code_argument_shadowing(self):
+        ip = get_ipython()
+        c = ip.Completer
+        c.use_jedi = False
+
+        # shadow
+        ip.run_cell("timeit = 1")
+
+        # should not suggest on implict magic when shadowed
+        text, matches = c.complete("timeit -n 2 -r 1 flo")
+        self.assertEqual(matches, [])
+
+        # should suggest on explicit magic
+        text, matches = c.complete("%timeit -n 2 -r 1 flo")
+        self.assertEqual(matches, ["float"])
+
+        # remove shadow
+        del ip.user_ns["timeit"]
+
+        # should suggest on implicit magic after shadow removal
+        text, matches = c.complete("timeit -n 2 -r 1 flo")
+        self.assertEqual(matches, ["float"])
+
     def test_magic_completion_order(self):
         ip = get_ipython()
         c = ip.Completer