Add JediSettings proxy

Integrations/src/main/java/io/deephaven/integrations/python/PythonDeephavenSession.java

@@ -327,6 +327,10 @@ public Object unwrapObject(Object object) {
         return object;
     }
 
+    public PyObject scope() {
+        return scope.mainGlobals().unwrap();
+    }
+
     interface PythonScriptSessionModule extends Closeable {
         PyObject create_change_list(PyObject from, PyObject to);
 

engine/table/src/main/java/io/deephaven/engine/util/DelegatingScriptSession.java

@@ -33,6 +33,10 @@ public DelegatingScriptSession(final ScriptSession delegate) {
         this.delegate = Objects.requireNonNull(delegate);
     }
 
+    public ScriptSession delegate() {
+        return delegate;
+    }
+
     private Changes contextualizeChanges(final Changes diff) {
         knownVariables.removeAll(diff.removed.keySet());
 

open-api/lang-tools/src/test/groovy/io/deephaven/lang/completion/ChunkerCompletionHandlerTest.groovy

@@ -160,8 +160,8 @@ b = 2
 c = 3
 """
         String src2 = "t = "
-        p.update(uri, "0", [ makeChange(0, 0, src1) ])
-        p.update(uri, "1", [ makeChange(3, 0, src2) ])
+        p.update(uri, 0, [ makeChange(0, 0, src1) ])
+        p.update(uri, 1, [ makeChange(3, 0, src2) ])
         doc = p.finish(uri)
 
         VariableProvider variables = Mock(VariableProvider) {

py/embedded-server/setup.py

@@ -56,6 +56,6 @@ def normalize_version(version):
     install_requires=[
         'jpy>=0.13.0',
         "java-utilities",
-        f"deephaven-core=={__normalized_version__}",
+        f"deephaven-core[autocomplete]=={__normalized_version__}",
     ]
 )

py/server/deephaven/completer/_completer.py

@@ -65,7 +65,7 @@ def is_enabled(self) -> bool:
     def can_jedi(self) -> bool:
         return self.__can_jedi
 
-    def do_completion(self, uri: str, version: int, line: int, col: int) -> list[list[Any]]:
+    def do_completion(self, scope: dict, uri: str, version: int, line: int, col: int) -> list[list[Any]]:
         if not self._versions[uri] == version:
             # if you aren't the newest completion, you get nothing, quickly
             return []
@@ -75,7 +75,7 @@ def do_completion(self, uri: str, version: int, line: int, col: int) -> list[lis
         # The Script completer is static analysis only, so we should actually be feeding it a whole document at once.
 
         completer = Script if self.__mode == CompleterMode.safe else Interpreter
-        completions = completer(txt, [globals()]).complete(line, col)
+        completions = completer(txt, [scope]).complete(line, col)
         # for now, a simple sorting based on number of preceding _
         # we may want to apply additional sorting to each list before combining
         results: list = []

py/server/setup.py

@@ -62,6 +62,9 @@ def normalize_version(version):
         # TODO(deephaven-core#3082): Remove numba dependency workarounds
         'numba; python_version < "3.11"',
     ],
+    extras_require = {
+        "autocomplete": ["jedi==0.18.2"],
+    },
     entry_points={
         'deephaven.plugin': ['registration_cls = deephaven.pandasplugin:PandasPluginRegistration']
     }

server/src/main/java/io/deephaven/server/console/ConsoleServiceGrpcImpl.java

@@ -24,13 +24,15 @@
 import io.deephaven.proto.backplane.grpc.TypedTicket;
 import io.deephaven.proto.backplane.script.grpc.*;
 import io.deephaven.server.console.completer.JavaAutoCompleteObserver;
+import io.deephaven.server.console.completer.JediSettings;
 import io.deephaven.server.console.completer.PythonAutoCompleteObserver;
 import io.deephaven.server.session.SessionCloseableObserver;
 import io.deephaven.server.session.SessionService;
 import io.deephaven.server.session.SessionState;
 import io.deephaven.server.session.SessionState.ExportBuilder;
 import io.deephaven.server.session.TicketRouter;
 import io.grpc.stub.StreamObserver;
+import org.jpy.PyModule;
 import org.jpy.PyObject;
 
 import javax.inject.Inject;
@@ -48,6 +50,9 @@ public class ConsoleServiceGrpcImpl extends ConsoleServiceGrpc.ConsoleServiceImp
     public static final boolean REMOTE_CONSOLE_DISABLED =
             Configuration.getInstance().getBooleanWithDefault("deephaven.console.disable", false);
 
+    public static final boolean AUTOCOMPLETE_DISABLED =
+            Configuration.getInstance().getBooleanWithDefault("deephaven.console.autocomplete.disable", false);
+
     public static final boolean QUIET_AUTOCOMPLETE_ERRORS =
             Configuration.getInstance().getBooleanWithDefault("deephaven.console.autocomplete.quiet", true);
 
@@ -257,21 +262,28 @@ public void bindTableToVariable(BindTableToVariableRequest request,
     public StreamObserver<AutoCompleteRequest> autoCompleteStream(
             StreamObserver<AutoCompleteResponse> responseObserver) {
         return GrpcUtil.rpcWrapper(log, responseObserver, () -> {
+            if (AUTOCOMPLETE_DISABLED) {
+                return new NoopAutoCompleteObserver(responseObserver);
+            }
             final SessionState session = sessionService.getCurrentSession();
             if (PythonDeephavenSession.SCRIPT_TYPE.equals(scriptSessionProvider.get().scriptType())) {
-                PyObject[] settings = new PyObject[1];
+                JediSettings[] settings = new JediSettings[1];
                 safelyExecute(() -> {
-                    final ScriptSession scriptSession = scriptSessionProvider.get();
-                    scriptSession.evaluateScript("from deephaven.completer import jedi_settings");
-                    settings[0] = (PyObject) scriptSession.getVariable("jedi_settings");
+                    try (final PyModule pyModule = PyModule.importModule("deephaven.completer")) {
+                        settings[0] = pyModule.getAttribute("jedi_settings").createProxy(JediSettings.class);
+                    }
                 });
-                boolean canJedi = settings[0] != null && settings[0].call("can_jedi").getBooleanValue();
+                boolean canJedi = settings[0] != null && settings[0].can_jedi();
                 log.info().append(canJedi ? "Using jedi for python autocomplete"
                         : "No jedi dependency available in python environment; disabling autocomplete.").endl();
-                return canJedi ? new PythonAutoCompleteObserver(responseObserver, scriptSessionProvider, session)
-                        : new NoopAutoCompleteObserver(responseObserver);
+                if (!canJedi) {
+                    if (settings[0] != null) {
+                        settings[0].close();
+                    }
+                    return new NoopAutoCompleteObserver(responseObserver);
+                }
+                return new PythonAutoCompleteObserver(responseObserver, session, settings[0]);
             }
-
             return new JavaAutoCompleteObserver(session, responseObserver);
         });
     }

server/src/main/java/io/deephaven/server/console/completer/JediSettings.java

@@ -0,0 +1,25 @@
+package io.deephaven.server.console.completer;
+
+import org.jpy.PyObject;
+
+import java.io.Closeable;
+
+public interface JediSettings extends Closeable {
+
+    void open_doc(String text, String uri, int version);
+
+    String get_doc(String uri);
+
+    void update_doc(String document, String uri, int version);
+
+    void close_doc(String uri);
+
+    boolean is_enabled();
+
+    PyObject do_completion(PyObject scope, String uri, int version, int line, int character);
+
+    boolean can_jedi();
+
+    @Override
+    void close();
+}

server/src/main/java/io/deephaven/server/console/completer/PythonAutoCompleteObserver.java

@@ -1,23 +1,39 @@
 package io.deephaven.server.console.completer;
 
 import com.google.rpc.Code;
+import io.deephaven.engine.util.DelegatingScriptSession;
 import io.deephaven.engine.util.ScriptSession;
 import io.deephaven.extensions.barrage.util.GrpcUtil;
+import io.deephaven.integrations.python.PythonDeephavenSession;
 import io.deephaven.internal.log.LoggerFactory;
 import io.deephaven.io.logger.Logger;
 import io.deephaven.lang.completion.ChunkerCompleter;
 import io.deephaven.lang.parse.CompletionParser;
-import io.deephaven.proto.backplane.script.grpc.*;
+import io.deephaven.proto.backplane.script.grpc.AutoCompleteRequest;
+import io.deephaven.proto.backplane.script.grpc.AutoCompleteResponse;
+import io.deephaven.proto.backplane.script.grpc.ChangeDocumentRequest;
+import io.deephaven.proto.backplane.script.grpc.CloseDocumentRequest;
+import io.deephaven.proto.backplane.script.grpc.CompletionItem;
+import io.deephaven.proto.backplane.script.grpc.DocumentRange;
+import io.deephaven.proto.backplane.script.grpc.GetCompletionItemsRequest;
+import io.deephaven.proto.backplane.script.grpc.GetCompletionItemsResponse;
+import io.deephaven.proto.backplane.script.grpc.OpenDocumentRequest;
+import io.deephaven.proto.backplane.script.grpc.Position;
+import io.deephaven.proto.backplane.script.grpc.TextDocumentItem;
+import io.deephaven.proto.backplane.script.grpc.TextEdit;
+import io.deephaven.proto.backplane.script.grpc.VersionedTextDocumentIdentifier;
 import io.deephaven.server.console.ConsoleServiceGrpcImpl;
 import io.deephaven.server.session.SessionState;
 import io.deephaven.util.SafeCloseable;
 import io.grpc.stub.StreamObserver;
 import org.jpy.PyObject;
 
 import javax.inject.Provider;
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 
+import static io.deephaven.extensions.barrage.util.GrpcUtil.safelyExecute;
 import static io.deephaven.extensions.barrage.util.GrpcUtil.safelyExecuteLocked;
 
 /**
@@ -31,45 +47,44 @@ public class PythonAutoCompleteObserver implements StreamObserver<AutoCompleteRe
      * We only log timing for completions that take longer than, currently, 100ms
      */
     private static final long HUNDRED_MS_IN_NS = 100_000_000;
-    private final Provider<ScriptSession> scriptSession;
     private final SessionState session;
     private final StreamObserver<AutoCompleteResponse> responseObserver;
 
-    public PythonAutoCompleteObserver(StreamObserver<AutoCompleteResponse> responseObserver,
-            Provider<ScriptSession> scriptSession, final SessionState session) {
-        this.scriptSession = scriptSession;
+    private JediSettings jediSettings;
+
+    public PythonAutoCompleteObserver(
+            StreamObserver<AutoCompleteResponse> responseObserver, SessionState session, JediSettings jediSettings) {
         this.session = session;
         this.responseObserver = responseObserver;
+        this.jediSettings = jediSettings;
     }
 
     @Override
     @SuppressWarnings("DuplicatedCode")
     public void onNext(AutoCompleteRequest value) {
+        if (jediSettings == null) {
+            throw GrpcUtil.statusRuntimeException(Code.INTERNAL, "jediSettings already closed");
+        }
         switch (value.getRequestCase()) {
             case OPEN_DOCUMENT: {
                 final OpenDocumentRequest openDoc = value.getOpenDocument();
                 final TextDocumentItem doc = openDoc.getTextDocument();
-                PyObject completer = (PyObject) scriptSession.get().getVariable("jedi_settings");
-                completer.callMethod("open_doc", doc.getText(), doc.getUri(), doc.getVersion());
+                jediSettings.open_doc(doc.getText(), doc.getUri(), doc.getVersion());
                 break;
             }
             case CHANGE_DOCUMENT: {
                 ChangeDocumentRequest request = value.getChangeDocument();
                 final VersionedTextDocumentIdentifier text = request.getTextDocument();
-
-                PyObject completer = (PyObject) scriptSession.get().getVariable("jedi_settings");
                 String uri = text.getUri();
                 int version = text.getVersion();
-                String document = completer.callMethod("get_doc", text.getUri()).getStringValue();
-
+                String document = jediSettings.get_doc(text.getUri());
                 final List<ChangeDocumentRequest.TextDocumentContentChangeEvent> changes =
                         request.getContentChangesList();
                 document = CompletionParser.updateDocumentChanges(uri, version, document, changes);
                 if (document == null) {
                     return;
                 }
-
-                completer.callMethod("update_doc", document, uri, version);
+                jediSettings.update_doc(document, uri, version);
                 break;
             }
             case GET_COMPLETION_ITEMS: {
@@ -86,13 +101,11 @@ public void onNext(AutoCompleteRequest value) {
             }
             case CLOSE_DOCUMENT: {
                 CloseDocumentRequest request = value.getCloseDocument();
-                PyObject completer = (PyObject) scriptSession.get().getVariable("jedi_settings");
-                completer.callMethod("close_doc", request.getTextDocument().getUri());
+                jediSettings.close_doc(request.getTextDocument().getUri());
                 break;
             }
             case REQUEST_NOT_SET: {
-                throw GrpcUtil.statusRuntimeException(Code.INVALID_ARGUMENT,
-                        "Autocomplete command missing request");
+                throw GrpcUtil.statusRuntimeException(Code.INVALID_ARGUMENT, "Autocomplete command missing request");
             }
         }
     }
@@ -102,9 +115,7 @@ private void getCompletionItems(GetCompletionItemsRequest request,
             StreamObserver<AutoCompleteResponse> responseObserver) {
         final ScriptSession scriptSession = exportedConsole.get();
         try (final SafeCloseable ignored = scriptSession.getExecutionContext().open()) {
-
-            PyObject completer = (PyObject) scriptSession.getVariable("jedi_settings");
-            boolean canJedi = completer.callMethod("is_enabled").getBooleanValue();
+            boolean canJedi = jediSettings.is_enabled();
             if (!canJedi) {
                 log.trace().append("Ignoring completion request because jedi is disabled").endl();
                 // send back an empty, failed response...
@@ -121,13 +132,16 @@ private void getCompletionItems(GetCompletionItemsRequest request,
             final long startNano = System.nanoTime();
 
             if (log.isTraceEnabled()) {
-                String text = completer.call("get_doc", doc.getUri()).getStringValue();
+                String text = jediSettings.get_doc(doc.getUri());
                 log.trace().append("Completion version ").append(doc.getVersion())
                         .append(" has source code:").append(text).endl();
             }
-            final PyObject results = completer.callMethod("do_completion", doc.getUri(), doc.getVersion(),
-                    // our java is 0-indexed lines, 1-indexed chars. jedi is 1-indexed-both.
-                    // we'll keep that translation ugliness to the in-java result-processing.
+
+            final PyObject scope =
+                    ((PythonDeephavenSession) ((DelegatingScriptSession) scriptSession).delegate()).scope();
+            // our java is 0-indexed lines, 1-indexed chars. jedi is 1-indexed-both.
+            // we'll keep that translation ugliness to the in-java result-processing.
+            final PyObject results = jediSettings.do_completion(scope, doc.getUri(), doc.getVersion(),
                     pos.getLine() + 1, pos.getCharacter());
             if (!results.isList()) {
                 throw new UnsupportedOperationException(
@@ -228,13 +242,25 @@ private String toMillis(final long totalNanos) {
     @Override
     public void onError(Throwable t) {
         // ignore, client doesn't need us, will be cleaned up later
+        safelyExecute(this::closeJedi);
     }
 
     @Override
     public void onCompleted() {
         // just hang up too, browser will reconnect if interested
-        synchronized (responseObserver) {
+        safelyExecuteLocked(responseObserver, () -> {
             responseObserver.onCompleted();
+            closeJedi();
+        });
+    }
+
+    private void closeJedi() {
+        if (jediSettings != null) {
+            try {
+                jediSettings.close();
+            } finally {
+                jediSettings = null;
+            }
         }
     }
 }