[lldb] Support CommandInterpreter print callbacks

lldb/bindings/python/python-swigsafecast.swig

@@ -9,6 +9,10 @@ PythonObject SWIGBridge::ToSWIGWrapper(std::unique_ptr<lldb::SBValue> value_sb)
   return ToSWIGHelper(value_sb.release(), SWIGTYPE_p_lldb__SBValue);
 }
 
+PythonObject SWIGBridge::ToSWIGWrapper(std::unique_ptr<lldb::SBCommandReturnObject> result_up) {
+  return ToSWIGHelper(result_up.release(), SWIGTYPE_p_lldb__SBCommandReturnObject);
+}
+
 PythonObject SWIGBridge::ToSWIGWrapper(lldb::ValueObjectSP value_sp) {
   return ToSWIGWrapper(std::unique_ptr<lldb::SBValue>(new lldb::SBValue(value_sp)));
 }

lldb/bindings/python/python-typemaps.swig

@@ -476,6 +476,25 @@ template <> bool SetNumberFromPyObject<double>(double &number, PyObject *obj) {
   $1 = $1 || PyCallable_Check(reinterpret_cast<PyObject *>($input));
 }
 
+// For lldb::SBCommandPrintCallback
+%typemap(in) (lldb::SBCommandPrintCallback callback, void *baton) {
+  if (!($input == Py_None ||
+        PyCallable_Check(reinterpret_cast<PyObject *>($input)))) {
+    PyErr_SetString(PyExc_TypeError, "Need a callable object or None!");
+    SWIG_fail;
+  }
+
+  // Don't lose the callback reference.
+  Py_INCREF($input);
+  $1 = LLDBSwigPythonCallPythonCommandPrintCallback;
+  $2 = $input;
+}
+
+%typemap(typecheck) (lldb::SBCommandPrintCallback callback, void *baton) {
+  $1 = $input == Py_None;
+  $1 = $1 || PyCallable_Check(reinterpret_cast<PyObject *>($input));
+}
+
 %typemap(in) (lldb::CommandOverrideCallback callback, void *baton) {
   if (!($input == Py_None ||
         PyCallable_Check(reinterpret_cast<PyObject *>($input)))) {

lldb/bindings/python/python-wrapper.swig

@@ -727,7 +727,7 @@ lldb_private::python::SWIGBridge::LLDBSwigPythonHandleOptionArgumentCompletionFo
     dict_sp->AddBooleanItem("no-completion", true);
     return dict_sp;
   }
-    
+
 
   // Convert the return dictionary to a DictionarySP.
   StructuredData::ObjectSP result_obj_sp = result.CreateStructuredObject();
@@ -753,7 +753,7 @@ bool lldb_private::python::SWIGBridge::LLDBSwigPythonCallParsedCommandObject(
   auto pfunc = self.ResolveName<PythonCallable>("__call__");
 
   if (!pfunc.IsAllocated()) {
-    cmd_retobj.AppendError("Could not find '__call__' method in implementation class"); 
+    cmd_retobj.AppendError("Could not find '__call__' method in implementation class");
     return false;
   }
 
@@ -1012,6 +1012,26 @@ static void LLDBSwigPythonCallPythonLogOutputCallback(const char *str,
   }
 }
 
+// For CommandPrintCallback functions
+static CommandReturnObjectCallbackResult LLDBSwigPythonCallPythonCommandPrintCallback(SBCommandReturnObject& result, void *callback_baton) {
+  SWIG_Python_Thread_Block swig_thread_block;
+
+  PyErr_Cleaner py_err_cleaner(true);
+
+  PythonObject result_arg = SWIGBridge::ToSWIGWrapper(
+      std::make_unique<SBCommandReturnObject>(result));
+  PythonCallable callable =
+      Retain<PythonCallable>(reinterpret_cast<PyObject *>(callback_baton));
+
+  if (!callable.IsValid())
+    return eCommandReturnObjectPrintCallbackSkipped;
+
+  PythonObject callback_result = callable(result_arg);
+
+  long long ret_val = unwrapOrSetPythonException(As<long long>(callback_result));
+  return (CommandReturnObjectCallbackResult)ret_val;
+}
+
 // For DebuggerTerminateCallback functions
 static void LLDBSwigPythonCallPythonSBDebuggerTerminateCallback(lldb::user_id_t debugger_id,
                                                       void *baton) {

lldb/include/lldb/API/SBCommandInterpreter.h

@@ -247,13 +247,13 @@ class SBCommandInterpreter {
                                        lldb::SBStringList &matches,
                                        lldb::SBStringList &descriptions);
 
-  /// Returns whether an interrupt flag was raised either by the SBDebugger - 
+  /// Returns whether an interrupt flag was raised either by the SBDebugger -
   /// when the function is not running on the RunCommandInterpreter thread, or
   /// by SBCommandInterpreter::InterruptCommand if it is.  If your code is doing
-  /// interruptible work, check this API periodically, and interrupt if it 
+  /// interruptible work, check this API periodically, and interrupt if it
   /// returns true.
   bool WasInterrupted() const;
-  
+
   /// Interrupts the command currently executing in the RunCommandInterpreter
   /// thread.
   ///
@@ -331,6 +331,8 @@ class SBCommandInterpreter {
   /// this list. Otherwise this list is empty.
   SBStructuredData GetTranscript();
 
+  void SetPrintCallback(lldb::SBCommandPrintCallback callback, void *baton);
+
 protected:
   friend class lldb_private::CommandPluginInterfaceImplementation;
 

lldb/include/lldb/API/SBDefines.h

@@ -144,6 +144,9 @@ typedef bool (*SBBreakpointHitCallback)(void *baton, lldb::SBProcess &process,
 typedef void (*SBDebuggerDestroyCallback)(lldb::user_id_t debugger_id,
                                           void *baton);
 
+typedef CommandReturnObjectCallbackResult (*SBCommandPrintCallback)(
+    lldb::SBCommandReturnObject &result, void *baton);
+
 typedef lldb::SBError (*SBPlatformLocateModuleCallback)(
     void *baton, const lldb::SBModuleSpec &module_spec,
     lldb::SBFileSpec &module_file_spec, lldb::SBFileSpec &symbol_file_spec);

lldb/include/lldb/Interpreter/CommandInterpreter.h

@@ -16,6 +16,7 @@
 #include "lldb/Interpreter/CommandObject.h"
 #include "lldb/Interpreter/ScriptInterpreter.h"
 #include "lldb/Utility/Args.h"
+#include "lldb/Utility/Baton.h"
 #include "lldb/Utility/Broadcaster.h"
 #include "lldb/Utility/CompletionRequest.h"
 #include "lldb/Utility/Event.h"
@@ -253,6 +254,10 @@ class CommandInterpreter : public Broadcaster,
     eCommandTypesAllThem = 0xFFFF  //< all commands
   };
 
+  using CommandReturnObjectCallback =
+      std::function<lldb::CommandReturnObjectCallbackResult(
+          CommandReturnObject &)>;
+
   // The CommandAlias and CommandInterpreter both have a hand in
   // substituting for alias commands.  They work by writing special tokens
   // in the template form of the Alias command, and then detecting them when the
@@ -664,6 +669,8 @@ class CommandInterpreter : public Broadcaster,
     ++m_command_usages[cmd_obj.GetCommandName()];
   }
 
+  void SetPrintCallback(CommandReturnObjectCallback callback);
+
   llvm::json::Value GetStatistics();
   const StructuredData::Array &GetTranscript() const;
 
@@ -774,6 +781,9 @@ class CommandInterpreter : public Broadcaster,
   std::vector<uint32_t> m_command_source_flags;
   CommandInterpreterRunResult m_result;
 
+  /// An optional callback to handle printing the CommandReturnObject.
+  CommandReturnObjectCallback m_print_callback;
+
   // The exit code the user has requested when calling the 'quit' command.
   // No value means the user hasn't set a custom exit code so far.
   std::optional<int> m_quit_exit_code;

lldb/include/lldb/lldb-enumerations.h

@@ -1368,6 +1368,15 @@ enum Severity {
   eSeverityInfo, // Equivalent to Remark used in clang.
 };
 
+/// Callback return value, indicating whether it handled printing the
+/// CommandReturnObject or deferred doing so to the CommandInterpreter.
+enum CommandReturnObjectCallbackResult {
+  /// The callback deferred printing the command return object.
+  eCommandReturnObjectPrintCallbackSkipped = 0,
+  /// The callback handled printing the command return object.
+  eCommandReturnObjectPrintCallbackHandled = 1,
+};
+
 } // namespace lldb
 
 #endif // LLDB_LLDB_ENUMERATIONS_H

lldb/source/API/SBCommandInterpreter.cpp

@@ -98,8 +98,8 @@ SBCommandInterpreter::SBCommandInterpreter(const SBCommandInterpreter &rhs)
 
 SBCommandInterpreter::~SBCommandInterpreter() = default;
 
-const SBCommandInterpreter &SBCommandInterpreter::
-operator=(const SBCommandInterpreter &rhs) {
+const SBCommandInterpreter &
+SBCommandInterpreter::operator=(const SBCommandInterpreter &rhs) {
   LLDB_INSTRUMENT_VA(this, rhs);
 
   m_opaque_ptr = rhs.m_opaque_ptr;
@@ -151,7 +151,7 @@ bool SBCommandInterpreter::WasInterrupted() const {
 
 bool SBCommandInterpreter::InterruptCommand() {
   LLDB_INSTRUMENT_VA(this);
-  
+
   return (IsValid() ? m_opaque_ptr->InterruptCommand() : false);
 }
 
@@ -222,8 +222,7 @@ void SBCommandInterpreter::HandleCommandsFromFile(
   if (override_context.get())
     m_opaque_ptr->HandleCommandsFromFile(tmp_spec,
                                          override_context.get()->Lock(true),
-                                         options.ref(),
-                                         result.ref());
+                                         options.ref(), result.ref());
 
   else
     m_opaque_ptr->HandleCommandsFromFile(tmp_spec, options.ref(), result.ref());
@@ -649,7 +648,8 @@ SBCommand::operator bool() const {
 const char *SBCommand::GetName() {
   LLDB_INSTRUMENT_VA(this);
 
-  return (IsValid() ? ConstString(m_opaque_sp->GetCommandName()).AsCString() : nullptr);
+  return (IsValid() ? ConstString(m_opaque_sp->GetCommandName()).AsCString()
+                    : nullptr);
 }
 
 const char *SBCommand::GetHelp() {
@@ -743,3 +743,15 @@ void SBCommand::SetFlags(uint32_t flags) {
   if (IsValid())
     m_opaque_sp->GetFlags().Set(flags);
 }
+
+void SBCommandInterpreter::SetPrintCallback(
+    lldb::SBCommandPrintCallback callback, void *baton) {
+  LLDB_INSTRUMENT_VA(this, callback, baton);
+
+  if (m_opaque_ptr)
+    m_opaque_ptr->SetPrintCallback(
+        [callback, baton](lldb_private::CommandReturnObject &result) {
+          SBCommandReturnObject sb_result(result);
+          return callback(sb_result, baton);
+        });
+}

lldb/source/Interpreter/CommandInterpreter.cpp

@@ -3186,30 +3186,40 @@ void CommandInterpreter::IOHandlerInputComplete(IOHandler &io_handler,
   if ((result.Succeeded() &&
        io_handler.GetFlags().Test(eHandleCommandFlagPrintResult)) ||
       io_handler.GetFlags().Test(eHandleCommandFlagPrintErrors)) {
-    // Display any inline diagnostics first.
-    const bool inline_diagnostics = !result.GetImmediateErrorStream() &&
-                                    GetDebugger().GetShowInlineDiagnostics();
-    if (inline_diagnostics) {
-      unsigned prompt_len = m_debugger.GetPrompt().size();
-      if (auto indent = result.GetDiagnosticIndent()) {
-        std::string diags =
-            result.GetInlineDiagnosticString(prompt_len + *indent);
-        PrintCommandOutput(io_handler, diags, true);
+    auto DefaultPrintCallback = [&](const CommandReturnObject &result) {
+      // Display any inline diagnostics first.
+      const bool inline_diagnostics = !result.GetImmediateErrorStream() &&
+                                      GetDebugger().GetShowInlineDiagnostics();
+      if (inline_diagnostics) {
+        unsigned prompt_len = m_debugger.GetPrompt().size();
+        if (auto indent = result.GetDiagnosticIndent()) {
+          std::string diags =
+              result.GetInlineDiagnosticString(prompt_len + *indent);
+          PrintCommandOutput(io_handler, diags, true);
+        }
       }
-    }
 
-    // Display any STDOUT/STDERR _prior_ to emitting the command result text.
-    GetProcessOutput();
+      // Display any STDOUT/STDERR _prior_ to emitting the command result text.
+      GetProcessOutput();
 
-    if (!result.GetImmediateOutputStream()) {
-      llvm::StringRef output = result.GetOutputString();
-      PrintCommandOutput(io_handler, output, true);
-    }
+      if (!result.GetImmediateOutputStream()) {
+        llvm::StringRef output = result.GetOutputString();
+        PrintCommandOutput(io_handler, output, true);
+      }
 
-    // Now emit the command error text from the command we just executed.
-    if (!result.GetImmediateErrorStream()) {
-      std::string error = result.GetErrorString(!inline_diagnostics);
-      PrintCommandOutput(io_handler, error, false);
+      // Now emit the command error text from the command we just executed.
+      if (!result.GetImmediateErrorStream()) {
+        std::string error = result.GetErrorString(!inline_diagnostics);
+        PrintCommandOutput(io_handler, error, false);
+      }
+    };
+
+    if (m_print_callback) {
+      const auto callback_result = m_print_callback(result);
+      if (callback_result == eCommandReturnObjectPrintCallbackSkipped)
+        DefaultPrintCallback(result);
+    } else {
+      DefaultPrintCallback(result);
     }
   }
 
@@ -3660,3 +3670,8 @@ llvm::json::Value CommandInterpreter::GetStatistics() {
 const StructuredData::Array &CommandInterpreter::GetTranscript() const {
   return m_transcript;
 }
+
+void CommandInterpreter::SetPrintCallback(
+    CommandReturnObjectCallback callback) {
+  m_print_callback = callback;
+}

lldb/source/Plugins/ScriptInterpreter/Python/SWIGPythonBridge.h

@@ -81,6 +81,8 @@ template <typename T> class ScopedPythonObject : PythonObject {
 class SWIGBridge {
 public:
   static PythonObject ToSWIGWrapper(std::unique_ptr<lldb::SBValue> value_sb);
+  static PythonObject
+  ToSWIGWrapper(std::unique_ptr<lldb::SBCommandReturnObject> result_up);
   static PythonObject ToSWIGWrapper(lldb::ValueObjectSP value_sp);
   static PythonObject ToSWIGWrapper(lldb::TargetSP target_sp);
   static PythonObject ToSWIGWrapper(lldb::ProcessSP process_sp);
@@ -190,12 +192,11 @@ class SWIGBridge {
                                   lldb::DebuggerSP debugger, const char *args,
                                   lldb_private::CommandReturnObject &cmd_retobj,
                                   lldb::ExecutionContextRefSP exe_ctx_ref_sp);
-  static bool
-  LLDBSwigPythonCallParsedCommandObject(PyObject *implementor,
-                                  lldb::DebuggerSP debugger,  
-                                  StructuredDataImpl &args_impl,
-                                  lldb_private::CommandReturnObject &cmd_retobj,
-                                  lldb::ExecutionContextRefSP exe_ctx_ref_sp);
+  static bool LLDBSwigPythonCallParsedCommandObject(
+      PyObject *implementor, lldb::DebuggerSP debugger,
+      StructuredDataImpl &args_impl,
+      lldb_private::CommandReturnObject &cmd_retobj,
+      lldb::ExecutionContextRefSP exe_ctx_ref_sp);
 
   static std::optional<std::string>
   LLDBSwigPythonGetRepeatCommandForScriptedCommand(PyObject *implementor,

lldb/test/API/python_api/interpreter_callback/Makefile

@@ -0,0 +1,3 @@
+C_SOURCES := main.c
+
+include Makefile.rules

lldb/test/API/python_api/interpreter_callback/TestCommandInterepterPrintCallback.py

@@ -0,0 +1,66 @@
+import lldb
+from lldbsuite.test.decorators import *
+from lldbsuite.test.lldbtest import *
+from lldbsuite.test import lldbutil
+
+
+class CommandInterepterPrintCallbackTest(TestBase):
+    NO_DEBUG_INFO_TESTCASE = True
+
+    def run_command_interpreter_with_output_file(self, out_filename, input_str):
+        with open(out_filename, "w") as f:
+            self.dbg.SetOutputFileHandle(f, False)
+            self.dbg.SetInputString(input_str)
+            opts = lldb.SBCommandInterpreterRunOptions()
+            self.dbg.RunCommandInterpreter(True, False, opts, 0, False, False)
+
+    def test_command_interpreter_print_callback(self):
+        """Test the command interpreter print callback."""
+        self.build()
+        exe = self.getBuildArtifact("a.out")
+
+        target = self.dbg.CreateTarget(exe)
+        self.assertTrue(target, VALID_TARGET)
+
+        lldbutil.run_to_source_breakpoint(
+            self, "// Break here", lldb.SBFileSpec("main.c")
+        )
+
+        out_filename = self.getBuildArtifact("output")
+        ci = self.dbg.GetCommandInterpreter()
+        called = False
+
+        # The string we'll be looking for in the command output.
+        needle = "Show a list of all debugger commands"
+
+        # Test registering a callback that handles the printing. Make sure the
+        # result is passed to the callback and that we don't print the result.
+        def handling_callback(return_object):
+            nonlocal called
+            called = True
+            self.assertIn(needle, return_object.GetOutput())
+            return lldb.eCommandReturnObjectPrintCallbackHandled
+
+        ci.SetPrintCallback(handling_callback)
+        self.assertFalse(called)
+        self.run_command_interpreter_with_output_file(out_filename, "help help\n")
+        with open(out_filename, "r") as f:
+            self.assertNotIn(needle, f.read())
+
+        # Test registering a callback that defers the printing to lldb. Make
+        # sure the result is passed to the callback and that the result is
+        # printed by lldb.
+        def non_handling_callback(return_object):
+            nonlocal called
+            called = True
+            self.assertIn(needle, return_object.GetOutput())
+            return lldb.eCommandReturnObjectPrintCallbackSkipped
+
+        called = False
+        ci.SetPrintCallback(non_handling_callback)
+        self.assertFalse(called)
+        self.run_command_interpreter_with_output_file(out_filename, "help help\n")
+        self.assertTrue(called)
+
+        with open(out_filename, "r") as f:
+            self.assertIn(needle, f.read())