[lldb] Support Python imports relative the to the current file being sourced

Make it possible to use a relative path in command script import to the
location of the file being sourced. This allows the user to put Python
scripts next to LLDB command files and importing them without having to
specify an absolute path.

To enable this behavior pass `-c` to `command script import`. The
argument can only be used when sourcing the command from a file.

rdar://68310384

Differential revision: https://reviews.llvm.org/D89334

3 files changed, 61 insertions, 0 deletions

diff --git a/lldb/test/Shell/ScriptInterpreter/Python/Inputs/hello.split b/lldb/test/Shell/ScriptInterpreter/Python/Inputs/hello.split
new file mode 100644
index 0000000..510d5f7
--- /dev/null
+++ b/lldb/test/Shell/ScriptInterpreter/Python/Inputs/hello.split
@@ -0,0 +1,10 @@
+#--- hello.in
+command script import -c baz.hello
+#--- hello.py
+import lldb
+
+def hello(debugger, command, result, internal_dict):
+    print("Hello, World!")
+
+def __lldb_init_module(debugger, internal_dict):
+    debugger.HandleCommand('command script add -f baz.hello.hello hello')
diff --git a/lldb/test/Shell/ScriptInterpreter/Python/Inputs/relative.split b/lldb/test/Shell/ScriptInterpreter/Python/Inputs/relative.split
new file mode 100644
index 0000000..7a5c89c
--- /dev/null
+++ b/lldb/test/Shell/ScriptInterpreter/Python/Inputs/relative.split
@@ -0,0 +1,20 @@
+#--- magritte.in
+command script import magritte
+#--- magritte.py
+import lldb
+
+def magritte(debugger, command, result, internal_dict):
+    print("Ceci n'est pas une pipe")
+
+def __lldb_init_module(debugger, internal_dict):
+    debugger.HandleCommand('command script add -f magritte.magritte magritte')
+#--- zip.in
+command script import -c zip
+#--- zip.py
+import lldb
+
+def zip(debugger, command, result, internal_dict):
+    print("95126")
+
+def __lldb_init_module(debugger, internal_dict):
+    debugger.HandleCommand('command script add -f zip.zip zip')
diff --git a/lldb/test/Shell/ScriptInterpreter/Python/command_relative_import.test b/lldb/test/Shell/ScriptInterpreter/Python/command_relative_import.test
new file mode 100644
index 0000000..cd401cb
--- /dev/null
+++ b/lldb/test/Shell/ScriptInterpreter/Python/command_relative_import.test
@@ -0,0 +1,31 @@
+# REQUIRES: python
+
+# RUN: rm -rf %t && mkdir -p %t/foo/bar/baz
+# RUN: split-file %S/Inputs/relative.split %t/foo
+# RUN: split-file %S/Inputs/hello.split %t/foo/bar
+# RUN: mv %t/foo/bar/hello.py %t/foo/bar/baz
+# RUN: echo 'command source %t/foo/bar/hello.in' >> %t/foo/zip.in
+
+# RUN: %lldb --script-language python \
+# RUN:    -o 'command source %t/foo/magritte.in' \
+# RUN:    -o 'command source %t/foo/zip.in' \
+# RUN:    -o 'command source %t/foo/magritte.in' \
+# RUN;    -o 'zip' \
+# RUN:    -o 'hello'
+# RUN     -o 'magritte' 2>&1 | FileCheck %s
+
+# The first time importing 'magritte' fails because we didn't pass -c.
+# CHECK: ModuleNotFoundError: No module named 'magritte'
+# CHECK-NOT: Ceci n'est pas une pipe
+# CHECK: 95126
+# CHECK: Hello, World!
+# The second time importing 'magritte' works, even without passing -c because
+# we added '%t/foo' to the Python path when importing 'zip'.
+# CHECK: Ceci n'est pas une pipe
+
+# Cannot use `-o` here because the driver puts the commands in a file and
+# sources them.
+command script import -c %t/foo/magritte.py
+quit
+# RUN: cat %s | %lldb --script-language python 2>&1 | FileCheck %s --check-prefix ERROR
+# ERROR: error: command script import -c can only be specified from a command file