Add the ability for Script based commands to specify their "repeat command" (#94823)

jimingham · web-flow · commit 77d131eddb6c · 2024-07-03T10:39:34.000-07:00
Among other things, returning an empty string as the repeat command
disables auto-repeat, which can be useful for state-changing commands.

There's one remaining refinement to this setup, which is that for parsed
script commands, it should be possible to change an option value, or add
a new option value that wasn't originally specified, then ask lldb "make
this back into a command string". That would make doing fancy things
with repeat commands easier.

That capability isn't present in the lldb_private side either, however.
So that's for a next iteration.

I haven't added this to the docs on adding commands yet. I wanted to
make sure this was an acceptable approach before I spend the time to do
that.
diff --git a/lldb/bindings/python/python-wrapper.swig b/lldb/bindings/python/python-wrapper.swig
@@ -728,6 +728,28 @@ bool lldb_private::python::SWIGBridge::LLDBSwigPythonCallCommandObject(
   return true;
 }
 
+std::optional<std::string>
+lldb_private::python::SWIGBridge::LLDBSwigPythonGetRepeatCommandForScriptedCommand(PyObject *implementor,
+                                               std::string &command) {
+  PyErr_Cleaner py_err_cleaner(true);
+
+  PythonObject self(PyRefType::Borrowed, implementor);
+  auto pfunc = self.ResolveName<PythonCallable>("get_repeat_command");
+  // If not implemented, repeat the exact command.
+  if (!pfunc.IsAllocated())
+    return std::nullopt;
+
+  PythonString command_str(command);
+  PythonObject result = pfunc(command_str);
+
+  // A return of None is the equivalent of nullopt - means repeat
+  // the command as is:
+  if (result.IsNone())
+    return std::nullopt;
+
+  return result.Str().GetString().str();
+}
+
 #include "lldb/Interpreter/CommandReturnObject.h"
 
 bool lldb_private::python::SWIGBridge::LLDBSwigPythonCallParsedCommandObject(
diff --git a/lldb/docs/use/python-reference.rst b/lldb/docs/use/python-reference.rst
@@ -562,6 +562,18 @@ which should implement the following interface:
           this call should return the short help text for this command[1]
       def get_long_help(self):
           this call should return the long help text for this command[1]
+      def get_repeat_command(self, command):
+          The auto-repeat command is what will get executed when the user types just
+          a return at the next prompt after this command is run.  Even if your command
+          was run because it was specified as a repeat command, that invocation will still
+          get asked for IT'S repeat command, so you can chain a series of repeats, for instance
+          to implement a pager.
+
+          The command argument is the command that is about to be executed.
+
+          If this call returns None, then the ordinary repeat mechanism will be used
+          If this call returns an empty string, then auto-repeat is disabled
+          If this call returns any other string, that will be the repeat command [1]
 
 [1] This method is optional.
 
diff --git a/lldb/examples/python/cmdtemplate.py b/lldb/examples/python/cmdtemplate.py
@@ -19,7 +19,7 @@ class FrameStatCommand(ParsedCommand):
 
     @classmethod
     def register_lldb_command(cls, debugger, module_name):
-        ParsedCommandBase.do_register_cmd(cls, debugger, module_name)
+        ParsedCommand.do_register_cmd(cls, debugger, module_name)
         print(
             'The "{0}" command has been installed, type "help {0}" or "{0} '
             '--help" for detailed help.'.format(cls.program)
@@ -72,6 +72,10 @@ def setup_command_definition(self):
             default = True,
         )
 
+    def get_repeat_command(self, args):
+        """As an example, make the command not auto-repeat:"""
+        return ""
+
     def get_short_help(self):
         return "Example command for use in debugging"
 
diff --git a/lldb/include/lldb/Interpreter/CommandObject.h b/lldb/include/lldb/Interpreter/CommandObject.h
@@ -297,6 +297,10 @@ class CommandObject : public std::enable_shared_from_this<CommandObject> {
   /// \param[in] current_command_args
   ///    The command arguments.
   ///
+  /// \param[in] index
+  ///    This is for internal use - it is how the completion request is tracked
+  ///    in CommandObjectMultiword, and should otherwise be ignored.
+  ///
   /// \return
   ///     std::nullopt if there is no special repeat command - it will use the
   ///     current command line.
diff --git a/lldb/include/lldb/Interpreter/ScriptInterpreter.h b/lldb/include/lldb/Interpreter/ScriptInterpreter.h
@@ -439,6 +439,12 @@ class ScriptInterpreter : public PluginInterface {
     return false;
   }
 
+  virtual std::optional<std::string>
+  GetRepeatCommandForScriptedCommand(StructuredData::GenericSP impl_obj_sp,
+                                     Args &args) {
+    return std::nullopt;
+  }
+
   virtual bool RunScriptFormatKeyword(const char *impl_function,
                                       Process *process, std::string &output,
                                       Status &error) {
diff --git a/lldb/source/Commands/CommandObjectCommands.cpp b/lldb/source/Commands/CommandObjectCommands.cpp
@@ -1142,6 +1142,15 @@ class CommandObjectScriptingObjectRaw : public CommandObjectRaw {
 
   ScriptedCommandSynchronicity GetSynchronicity() { return m_synchro; }
 
+  std::optional<std::string> GetRepeatCommand(Args &args,
+                                              uint32_t index) override {
+    ScriptInterpreter *scripter = GetDebugger().GetScriptInterpreter();
+    if (!scripter)
+      return std::nullopt;
+
+    return scripter->GetRepeatCommandForScriptedCommand(m_cmd_obj_sp, args);
+  }
+
   llvm::StringRef GetHelp() override {
     if (m_fetched_help_short)
       return CommandObjectRaw::GetHelp();
@@ -1588,7 +1597,9 @@ class CommandObjectScriptingObjectParsed : public CommandObjectParsed {
       options.ForEach(add_element);
       return error;
     }
-    
+
+    size_t GetNumOptions() { return m_num_options; }
+
   private:
     struct EnumValueStorage {
       EnumValueStorage() {
@@ -1827,6 +1838,15 @@ class CommandObjectScriptingObjectParsed : public CommandObjectParsed {
 
   ScriptedCommandSynchronicity GetSynchronicity() { return m_synchro; }
 
+  std::optional<std::string> GetRepeatCommand(Args &args,
+                                              uint32_t index) override {
+    ScriptInterpreter *scripter = GetDebugger().GetScriptInterpreter();
+    if (!scripter)
+      return std::nullopt;
+
+    return scripter->GetRepeatCommandForScriptedCommand(m_cmd_obj_sp, args);
+  }
+
   llvm::StringRef GetHelp() override {
     if (m_fetched_help_short)
       return CommandObjectParsed::GetHelp();
@@ -1857,9 +1877,14 @@ class CommandObjectScriptingObjectParsed : public CommandObjectParsed {
       SetHelpLong(docstring);
     return CommandObjectParsed::GetHelpLong();
   }
-  
-  Options *GetOptions() override { return &m_options; }
 
+  Options *GetOptions() override {
+    // CommandObjectParsed requires that a command with no options return
+    // nullptr.
+    if (m_options.GetNumOptions() == 0)
+      return nullptr;
+    return &m_options;
+  }
 
 protected:
   void DoExecute(Args &args,
diff --git a/lldb/source/Commands/CommandObjectThread.cpp b/lldb/source/Commands/CommandObjectThread.cpp
@@ -132,7 +132,7 @@ class CommandObjectThreadBacktrace : public CommandObjectIterateOverThreads {
   Options *GetOptions() override { return &m_options; }
 
   std::optional<std::string> GetRepeatCommand(Args &current_args,
-                                              uint32_t idx) override {
+                                              uint32_t index) override {
     llvm::StringRef count_opt("--count");
     llvm::StringRef start_opt("--start");
 
diff --git a/lldb/source/Plugins/ScriptInterpreter/Python/SWIGPythonBridge.h b/lldb/source/Plugins/ScriptInterpreter/Python/SWIGPythonBridge.h
@@ -206,6 +206,10 @@ class SWIGBridge {
                                   lldb_private::CommandReturnObject &cmd_retobj,
                                   lldb::ExecutionContextRefSP exe_ctx_ref_sp);
 
+  static std::optional<std::string>
+  LLDBSwigPythonGetRepeatCommandForScriptedCommand(PyObject *implementor,
+                                                   std::string &command);
+
   static bool LLDBSwigPythonCallModuleInit(const char *python_module_name,
                                            const char *session_dictionary_name,
                                            lldb::DebuggerSP debugger);
diff --git a/lldb/source/Plugins/ScriptInterpreter/Python/ScriptInterpreterPython.cpp b/lldb/source/Plugins/ScriptInterpreter/Python/ScriptInterpreterPython.cpp
@@ -2708,6 +2708,33 @@ bool ScriptInterpreterPythonImpl::RunScriptBasedParsedCommand(
   return ret_val;
 }
 
+std::optional<std::string>
+ScriptInterpreterPythonImpl::GetRepeatCommandForScriptedCommand(
+    StructuredData::GenericSP impl_obj_sp, Args &args) {
+  if (!impl_obj_sp || !impl_obj_sp->IsValid())
+    return std::nullopt;
+
+  lldb::DebuggerSP debugger_sp = m_debugger.shared_from_this();
+
+  if (!debugger_sp.get())
+    return std::nullopt;
+
+  std::optional<std::string> ret_val;
+
+  {
+    Locker py_lock(this, Locker::AcquireLock | Locker::NoSTDIN,
+                   Locker::FreeLock);
+
+    StructuredData::ArraySP args_arr_sp(new StructuredData::Array());
+
+    // For scripting commands, we send the command string:
+    std::string command;
+    args.GetQuotedCommandString(command);
+    ret_val = SWIGBridge::LLDBSwigPythonGetRepeatCommandForScriptedCommand(
+        static_cast<PyObject *>(impl_obj_sp->GetValue()), command);
+  }
+  return ret_val;
+}
 
 /// In Python, a special attribute __doc__ contains the docstring for an object
 /// (function, method, class, ...) if any is defined Otherwise, the attribute's
diff --git a/lldb/source/Plugins/ScriptInterpreter/Python/ScriptInterpreterPythonImpl.h b/lldb/source/Plugins/ScriptInterpreter/Python/ScriptInterpreterPythonImpl.h
@@ -160,13 +160,16 @@ class ScriptInterpreterPythonImpl : public ScriptInterpreterPython {
       lldb_private::CommandReturnObject &cmd_retobj, Status &error,
       const lldb_private::ExecutionContext &exe_ctx) override;
 
-    virtual bool RunScriptBasedParsedCommand(
-      StructuredData::GenericSP impl_obj_sp, Args& args,
+  bool RunScriptBasedParsedCommand(
+      StructuredData::GenericSP impl_obj_sp, Args &args,
       ScriptedCommandSynchronicity synchronicity,
       lldb_private::CommandReturnObject &cmd_retobj, Status &error,
       const lldb_private::ExecutionContext &exe_ctx) override;
 
-  
+  std::optional<std::string>
+  GetRepeatCommandForScriptedCommand(StructuredData::GenericSP impl_obj_sp,
+                                     Args &args) override;
+
   Status GenerateFunction(const char *signature, const StringList &input,
                           bool is_callback) override;
 
diff --git a/lldb/test/API/commands/command/script/add/TestAddParsedCommand.py b/lldb/test/API/commands/command/script/add/TestAddParsedCommand.py
@@ -16,6 +16,11 @@ class ParsedCommandTestCase(TestBase):
     def test(self):
         self.pycmd_tests()
 
+    def setUp(self):
+        TestBase.setUp(self)
+        self.stdin_path = self.getBuildArtifact("stdin.txt")
+        self.stdout_path = self.getBuildArtifact("stdout.txt")
+
     def check_help_options(self, cmd_name, opt_list, substrs=[]):
         """
         Pass the command name in cmd_name and a vector of the short option, type & long option.
@@ -29,9 +34,40 @@ def check_help_options(self, cmd_name, opt_list, substrs=[]):
             else:
                 (short_opt, type, long_opt) = elem
                 substrs.append(f"-{short_opt} <{type}> ( --{long_opt} <{type}> )")
-        print(f"Opt Vec\n{substrs}")
         self.expect("help " + cmd_name, substrs=substrs)
 
+    def run_one_repeat(self, commands, expected_num_errors):
+        with open(self.stdin_path, "w") as input_handle:
+            input_handle.write(commands)
+
+        in_fileH = open(self.stdin_path, "r")
+        self.dbg.SetInputFileHandle(in_fileH, False)
+
+        out_fileH = open(self.stdout_path, "w")
+        self.dbg.SetOutputFileHandle(out_fileH, False)
+        self.dbg.SetErrorFileHandle(out_fileH, False)
+
+        options = lldb.SBCommandInterpreterRunOptions()
+        options.SetEchoCommands(False)
+        options.SetPrintResults(True)
+        options.SetPrintErrors(True)
+        options.SetAllowRepeats(True)
+
+        n_errors, quit_requested, has_crashed = self.dbg.RunCommandInterpreter(
+            True, False, options, 0, False, False
+        )
+
+        in_fileH.close()
+        out_fileH.close()
+
+        results = None
+        with open(self.stdout_path, "r") as out_fileH:
+            results = out_fileH.read()
+
+        self.assertEqual(n_errors, expected_num_errors)
+
+        return results
+
     def pycmd_tests(self):
         source_dir = self.getSourceDir()
         test_file_path = os.path.join(source_dir, "test_commands.py")
@@ -168,9 +204,6 @@ def cleanup():
         num_completions = interp.HandleCompletionWithDescriptions(
             cmd_str, len(cmd_str) - 1, 0, 1000, matches, descriptions
         )
-        print(
-            f"First: {matches.GetStringAtIndex(0)}\nSecond: {matches.GetStringAtIndex(1)}\nThird: {matches.GetStringAtIndex(2)}"
-        )
         self.assertEqual(num_completions, 1, "Only one completion for source file")
         self.assertEqual(matches.GetSize(), 2, "The first element is the complete line")
         self.assertEqual(
@@ -197,3 +230,23 @@ def cleanup():
             "two-args 'First Argument' 'Second Argument'",
             substrs=["0: First Argument", "1: Second Argument"],
         )
+
+        # Now make sure get_repeat_command works properly:
+
+        # no-args turns off auto-repeat
+        results = self.run_one_repeat("no-args\n\n", 1)
+        self.assertIn("No auto repeat", results, "Got auto-repeat error")
+
+        # one-args does the normal repeat
+        results = self.run_one_repeat("one-arg-no-opt ONE_ARG\n\n", 0)
+        self.assertEqual(results.count("ONE_ARG"), 2, "We did a normal repeat")
+
+        # two-args adds an argument:
+        results = self.run_one_repeat("two-args FIRST_ARG SECOND_ARG\n\n", 0)
+        self.assertEqual(
+            results.count("FIRST_ARG"), 2, "Passed first arg to both commands"
+        )
+        self.assertEqual(
+            results.count("SECOND_ARG"), 2, "Passed second arg to both commands"
+        )
+        self.assertEqual(results.count("THIRD_ARG"), 1, "Passed third arg in repeat")
diff --git a/lldb/test/API/commands/command/script/add/test_commands.py b/lldb/test/API/commands/command/script/add/test_commands.py
@@ -32,6 +32,12 @@ def __call__(self, debugger, args_array, exe_ctx, result):
             )
 
 
+# Use these to make sure that get_repeat_command sends the right
+# command.
+no_args_repeat = None
+one_arg_repeat = None
+two_arg_repeat = None
+
 class NoArgsCommand(ReportingCmd):
     program = "no-args"
 
@@ -96,6 +102,12 @@ def setup_command_definition(self):
             default="foo",
         )
 
+    def get_repeat_command(self, command):
+        # No auto-repeat
+        global no_args_repeat
+        no_args_repeat = command
+        return ""
+
     def get_short_help(self):
         return "Example command for use in debugging"
 
@@ -118,6 +130,12 @@ def setup_command_definition(self):
             [self.ov_parser.make_argument_element(lldb.eArgTypeSourceFile, "plain")]
         )
 
+    def get_repeat_command(self, command):
+        # Repeat the current command
+        global one_arg_repeat
+        one_arg_repeat = command
+        return None
+
     def get_short_help(self):
         return "Example command for use in debugging"
 
@@ -187,8 +205,13 @@ def setup_command_definition(self):
             ]
         )
 
+    def get_repeat_command(self, command):
+        global two_arg_repeat
+        two_arg_repeat = command
+        return command + " THIRD_ARG"
+
     def get_short_help(self):
-        return "Example command for use in debugging"
+        return "This is my short help string"
 
     def get_long_help(self):
         return self.help_string
diff --git a/lldb/unittests/ScriptInterpreter/Python/PythonTestSuite.cpp b/lldb/unittests/ScriptInterpreter/Python/PythonTestSuite.cpp
@@ -200,6 +200,12 @@ bool lldb_private::python::SWIGBridge::LLDBSwigPythonCallParsedCommandObject(
   return false;
 }
 
+std::optional<std::string>
+LLDBSwigPythonGetRepeatCommandForScriptedCommand(PyObject *implementor,
+                                                 std::string &command) {
+  return std::nullopt;
+}
+
 bool lldb_private::python::SWIGBridge::LLDBSwigPythonCallModuleInit(
     const char *python_module_name, const char *session_dictionary_name,
     lldb::DebuggerSP debugger) {
