diff --git a/repl/src/main/resources/fake_shell.py b/repl/src/main/resources/fake_shell.py index 2a99185b1..cd7129264 100644 --- a/repl/src/main/resources/fake_shell.py +++ b/repl/src/main/resources/fake_shell.py @@ -46,9 +46,27 @@ global_dict = {} job_context = None local_tmp_dir_path = None +ip_completer = None +p_completer = None TOP_FRAME_REGEX = re.compile(r'\s*File "".*in ') + +def initialize_completer(): + global ip_completer,p_completer + try: + __IPYTHON__ + from IPython.core.completer import IPCompleter + ip = get_ipython() + ip_completer = IPCompleter(ip, global_namespace=global_dict) + except NameError: + try: + from rlcompleter import Completer + p_completer = Completer(namespace=global_dict) + except: + pass + + def execute_reply(status, content): return { 'msg_type': 'execute_reply', @@ -195,6 +213,20 @@ def processBypassJob(self, serialized_job): def addFile(self, uri_path): job_context.sc.addFile(uri_path) + def complete(self, code, cursor_pos=None): + results = [] + if ip_completer is not None: + results = ip_completer.complete(code, cursor_pos=cursor_pos)[1] + elif p_completer is not None: + state = 0 + result = None + while state == 0 or result is not None: + result = p_completer.complete(code, state) + if result is not None: + results.append(result) + state += 1 + return ",".join(results) + def addPyFile(self, uri_path): job_context.sc.addPyFile(uri_path) @@ -555,6 +587,7 @@ def main(): sys.stderr = UnicodeDecodingStringIO() spark_major_version = os.getenv("LIVY_SPARK_MAJOR_VERSION") + initialize_completer() try: listening_port = 0 if os.environ.get("LIVY_TEST") != "true": diff --git a/repl/src/main/scala/org/apache/livy/repl/PySparkJobProcessor.scala b/repl/src/main/scala/org/apache/livy/repl/PySparkJobProcessor.scala index 95b91cffb..1154cc108 100644 --- a/repl/src/main/scala/org/apache/livy/repl/PySparkJobProcessor.scala +++ b/repl/src/main/scala/org/apache/livy/repl/PySparkJobProcessor.scala @@ -24,4 +24,6 @@ trait PySparkJobProcessor { def addPyFile(path: String) def getLocalTmpDirPath: String + + def complete(code: String, cursor_pos: Int): String } diff --git a/repl/src/main/scala/org/apache/livy/repl/PythonInterpreter.scala b/repl/src/main/scala/org/apache/livy/repl/PythonInterpreter.scala index bab2bfa7b..75d8593b7 100644 --- a/repl/src/main/scala/org/apache/livy/repl/PythonInterpreter.scala +++ b/repl/src/main/scala/org/apache/livy/repl/PythonInterpreter.scala @@ -299,6 +299,11 @@ private class PythonInterpreter( } } + override def complete(code: String, cursor: Int): Array[String] = { + val r = pysparkJobProcessor.complete(code, cursor) + r.split(",") + } + private def updatePythonGatewayPort(port: Int): Unit = { // The python gateway port can be 0 only when LivyConf.TEST_MODE is true // Py4j 0.10 has different API signature for "getCallbackClient", use reflection to handle it.