Merge branch 'master' into traff/zip_helpers
[idea/community.git] / python / ipnb / src / org / jetbrains / plugins / ipnb / configuration / IpnbConnectionManager.java
1 package org.jetbrains.plugins.ipnb.configuration;
2
3 import com.google.common.collect.ImmutableMap;
4 import com.google.common.collect.Lists;
5 import com.intellij.execution.ExecutionException;
6 import com.intellij.execution.RunContentExecutor;
7 import com.intellij.execution.configurations.GeneralCommandLine;
8 import com.intellij.execution.process.KillableColoredProcessHandler;
9 import com.intellij.execution.process.UnixProcessManager;
10 import com.intellij.openapi.application.ApplicationManager;
11 import com.intellij.openapi.components.ProjectComponent;
12 import com.intellij.openapi.diagnostic.Logger;
13 import com.intellij.openapi.module.Module;
14 import com.intellij.openapi.options.ShowSettingsUtil;
15 import com.intellij.openapi.project.Project;
16 import com.intellij.openapi.projectRoots.Sdk;
17 import com.intellij.openapi.roots.ProjectFileIndex;
18 import com.intellij.openapi.ui.InputValidator;
19 import com.intellij.openapi.ui.MessageType;
20 import com.intellij.openapi.ui.Messages;
21 import com.intellij.openapi.ui.popup.Balloon;
22 import com.intellij.openapi.ui.popup.BalloonBuilder;
23 import com.intellij.openapi.ui.popup.JBPopupFactory;
24 import com.intellij.openapi.util.Computable;
25 import com.intellij.openapi.util.Pair;
26 import com.intellij.openapi.util.text.StringUtil;
27 import com.intellij.openapi.vfs.VirtualFile;
28 import com.intellij.ui.HyperlinkAdapter;
29 import com.intellij.util.Alarm;
30 import com.intellij.util.io.HttpRequests;
31 import com.jetbrains.python.PythonHelper;
32 import com.jetbrains.python.packaging.PyPackage;
33 import com.jetbrains.python.packaging.PyPackageManager;
34 import com.jetbrains.python.sdk.PythonSdkType;
35 import org.jetbrains.annotations.NotNull;
36 import org.jetbrains.annotations.Nullable;
37 import org.jetbrains.plugins.ipnb.editor.IpnbFileEditor;
38 import org.jetbrains.plugins.ipnb.editor.panels.code.IpnbCodePanel;
39 import org.jetbrains.plugins.ipnb.format.cells.output.IpnbOutputCell;
40 import org.jetbrains.plugins.ipnb.protocol.IpnbConnection;
41 import org.jetbrains.plugins.ipnb.protocol.IpnbConnectionListenerBase;
42 import org.jetbrains.plugins.ipnb.protocol.IpnbConnectionV3;
43
44 import javax.swing.event.HyperlinkEvent;
45 import java.io.IOException;
46 import java.net.URI;
47 import java.net.URISyntaxException;
48 import java.util.ArrayList;
49 import java.util.HashMap;
50 import java.util.List;
51 import java.util.Map;
52
53 public final class IpnbConnectionManager implements ProjectComponent {
54   private static final Logger LOG = Logger.getInstance(IpnbConnectionManager.class);
55   private final Project myProject;
56   private Map<String, IpnbConnection> myKernels = new HashMap<String, IpnbConnection>();
57   private Map<String, IpnbCodePanel> myUpdateMap = new HashMap<String, IpnbCodePanel>();
58
59   public IpnbConnectionManager(final Project project) {
60     myProject = project;
61   }
62
63   public static IpnbConnectionManager getInstance(Project project) {
64     return project.getComponent(IpnbConnectionManager.class);
65   }
66
67   public void executeCell(@NotNull final IpnbCodePanel codePanel) {
68     final IpnbFileEditor fileEditor = codePanel.getFileEditor();
69     final VirtualFile virtualFile = fileEditor.getVirtualFile();
70     final String path = virtualFile.getPath();
71     if (!myKernels.containsKey(path)) {
72       startConnection(codePanel, fileEditor, path);
73     }
74     else {
75       IpnbConnection connection = myKernels.get(path);
76       if (!connection.isAlive()) {
77         myKernels.remove(path);
78         startConnection(codePanel, fileEditor, path);
79       }
80       else {
81         final String messageId = connection.execute(codePanel.getCell().getSourceAsString());
82         myUpdateMap.put(messageId, codePanel);
83       }
84     }
85   }
86
87   private void startConnection(@NotNull final IpnbCodePanel codePanel, final IpnbFileEditor fileEditor, final String path) {
88     String url = IpnbSettings.getInstance(myProject).getURL();
89     if (StringUtil.isEmptyOrSpaces(url)) {
90       url = IpnbSettings.DEFAULT_URL;
91     }
92
93     boolean connectionStarted = startConnection(codePanel, path, url, false);
94     if (!connectionStarted) {
95       final String finalUrl = url;
96       url = showDialogUrl(url);
97       if (url == null) return;
98       IpnbSettings.getInstance(myProject).setURL(url);
99
100       ApplicationManager.getApplication().executeOnPooledThread(new Runnable() {
101         @Override
102         public void run() {
103           final boolean serverStarted = startIpythonServer(finalUrl, fileEditor);
104           if (!serverStarted) {
105             return;
106           }
107           ApplicationManager.getApplication().invokeLater(new Runnable() {
108             @Override
109             public void run() {
110               new Alarm(Alarm.ThreadToUse.SWING_THREAD).addRequest(new Runnable() {
111                 @Override
112                 public void run() {
113                   startConnection(codePanel, path, finalUrl, true);
114                 }
115               }, 3000);
116             }
117           });
118         }
119       });
120     }
121   }
122
123   @Nullable
124   private static String showDialogUrl(@NotNull final String initialUrl) {
125     final String url = Messages.showInputDialog("IPython Notebook URL:", "Start IPython Notebook", null, initialUrl,
126                                                 new InputValidator() {
127                                                   @Override
128                                                   public boolean checkInput(String inputString) {
129                                                     try {
130                                                       final URI uri = new URI(inputString);
131                                                       if (uri.getPort() == -1 || StringUtil.isEmptyOrSpaces(uri.getHost())) {
132                                                         return false;
133                                                       }
134                                                     }
135                                                     catch (URISyntaxException e) {
136                                                       return false;
137                                                     }
138                                                     return !inputString.isEmpty();
139                                                   }
140
141                                                   @Override
142                                                   public boolean canClose(String inputString) {
143                                                     return true;
144                                                   }
145                                                 });
146     return url == null ? null : StringUtil.trimEnd(url, "/");
147   }
148
149   private boolean startConnection(@NotNull final IpnbCodePanel codePanel, @NotNull final String path, @NotNull final String urlString,
150                                   final boolean showNotification) {
151     try {
152       final IpnbConnectionListenerBase listener = new IpnbConnectionListenerBase() {
153         @Override
154         public void onOpen(@NotNull IpnbConnection connection) {
155           final String messageId = connection.execute(codePanel.getCell().getSourceAsString());
156           myUpdateMap.put(messageId, codePanel);
157         }
158
159         @Override
160         public void onOutput(@NotNull IpnbConnection connection,
161                              @NotNull String parentMessageId) {
162           if (!myUpdateMap.containsKey(parentMessageId)) return;
163           final IpnbCodePanel cell = myUpdateMap.get(parentMessageId);
164           cell.getCell().setPromptNumber(connection.getExecCount());
165           //noinspection unchecked
166           cell.updatePanel(null, (List<IpnbOutputCell>)connection.getOutput().clone());
167         }
168
169         @Override
170         public void onPayload(@Nullable String payload, @NotNull String parentMessageId) {
171           if (!myUpdateMap.containsKey(parentMessageId)) return;
172           final IpnbCodePanel cell = myUpdateMap.remove(parentMessageId);
173           if (payload != null) {
174             //noinspection unchecked
175             cell.updatePanel(payload, null);
176           }
177         }
178       };
179
180       HttpRequests.request(urlString + "/api").connect(new HttpRequests.RequestProcessor<Object>() {
181         @Override
182         public Object process(@NotNull HttpRequests.Request request) throws IOException {
183           final IpnbConnection connection;
184           try {
185             if (request.isSuccessful()) {
186               connection = new IpnbConnectionV3(urlString, listener);
187             }
188             else {
189               connection = new IpnbConnection(urlString, listener);
190             }
191             myKernels.put(path, connection);
192           }
193           catch (URISyntaxException e) {
194             if (showNotification) {
195               showWarning(codePanel.getFileEditor(),
196                           "Please, check IPython Notebook URL in <a href=\"\">Settings->Tools->IPython Notebook</a>",
197                           new IpnbSettingsAdapter());
198               LOG.warn("IPython Notebook connection refused: " + e.getMessage());
199             }
200           }
201           return null;
202         }
203       });
204     }
205     catch (IOException e) {
206       if (showNotification) {
207         LOG.warn("IPython Notebook connection refused: " + e.getMessage());
208       }
209       return false;
210     }
211     return true;
212   }
213
214   public void interruptKernel(@NotNull final String filePath) {
215     if (!myKernels.containsKey(filePath)) return;
216     final IpnbConnection connection = myKernels.get(filePath);
217     try {
218       connection.interrupt();
219     }
220     catch (IOException e) {
221       LOG.warn("Failed to interrupt kernel " + filePath);
222       LOG.warn(e.getMessage());
223     }
224   }
225
226   public void reloadKernel(@NotNull final String filePath) {
227     if (!myKernels.containsKey(filePath)) return;
228     final IpnbConnection connection = myKernels.get(filePath);
229     try {
230       connection.reload();
231     }
232     catch (IOException e) {
233       LOG.warn("Failed to reload kernel " + filePath);
234       LOG.warn(e.getMessage());
235     }
236   }
237
238   private static void showWarning(@NotNull final IpnbFileEditor fileEditor, @NotNull final String message,
239                                   @Nullable final HyperlinkAdapter listener) {
240     BalloonBuilder balloonBuilder = JBPopupFactory.getInstance().createHtmlTextBalloonBuilder(
241       message, null, MessageType.WARNING.getPopupBackground(), listener);
242     final Balloon balloon = balloonBuilder.createBalloon();
243     ApplicationManager.getApplication().invokeLater(new Runnable() {
244       @Override
245       public void run() {
246         balloon.showInCenterOf(fileEditor.getRunCellButton());
247       }
248     });
249   }
250
251   private boolean startIpythonServer(@NotNull final String url, @NotNull final IpnbFileEditor fileEditor) {
252     final Module module = ProjectFileIndex.SERVICE.getInstance(myProject).getModuleForFile(fileEditor.getVirtualFile());
253     if (module == null) return false;
254     final Sdk sdk = PythonSdkType.findPythonSdk(module);
255     if (sdk == null) {
256       showWarning(fileEditor, "Please check Python Interpreter in Settings->Python Interpreter", null);
257       return false;
258     }
259     try {
260       final PyPackage ipythonPackage = PyPackageManager.getInstance(sdk).findPackage("ipython", false);
261       if (ipythonPackage == null) {
262         showWarning(fileEditor, "Add IPython to the interpreter of the current project.", null);
263         return false;
264       }
265     }
266     catch (ExecutionException ignored) {
267     }
268     final Map<String, String> env = ImmutableMap.of("PYCHARM_EP_DIST", "ipython", "PYCHARM_EP_NAME", "ipython");
269
270     final Pair<String, String> hostPort = getHostPortFromUrl(url);
271     if (hostPort == null) {
272       showWarning(fileEditor, "Please, check IPython Notebook URL in <a href=\"\">Settings->Tools->IPython Notebook</a>",
273                   new IpnbSettingsAdapter());
274       return false;
275     }
276     final String homePath = sdk.getHomePath();
277     if (homePath == null) {
278       showWarning(fileEditor, "Python Sdk is invalid, please check Python Interpreter in Settings->Python Interpreter", null);
279       return false;
280     }
281     String ipython = findIPythonRunner(homePath);
282     Map<String, String> env = null;
283     if (ipython == null) {
284       ipython = PythonHelpersLocator.getHelperPath("pycharm/pycharm_load_entry_point.py");
285       env = ImmutableMap.of("PYCHARM_EP_DIST", "ipython", "PYCHARM_EP_NAME", "ipython");
286     }
287
288     final ArrayList<String> parameters = Lists.newArrayList(homePath, ipython, "notebook", "--no-browser");
289     if (hostPort.getFirst() != null) {
290       parameters.add("--ip");
291       parameters.add(hostPort.getFirst());
292     }
293     if (hostPort.getSecond() != null) {
294       parameters.add("--port");
295       parameters.add(hostPort.getSecond());
296     }
297     final GeneralCommandLine commandLine = new GeneralCommandLine(parameters).withWorkDirectory(myProject.getBasePath());
298     if (env != null) {
299       commandLine.withEnvironment(env);
300     }
301
302     try {
303       final KillableColoredProcessHandler processHandler = new KillableColoredProcessHandler(commandLine) {
304         @Override
305         protected void doDestroyProcess() {
306           super.doDestroyProcess();
307           UnixProcessManager.sendSigIntToProcessTree(getProcess());
308         }
309
310         @Override
311         public boolean isSilentlyDestroyOnClose() {
312           return true;
313         }
314       };
315       processHandler.setShouldDestroyProcessRecursively(true);
316       ApplicationManager.getApplication().invokeLater(new Runnable() {
317         @Override
318         public void run() {
319           new RunContentExecutor(myProject, processHandler)
320             .withTitle("IPython Notebook")
321             .withStop(new Runnable() {
322               @Override
323               public void run() {
324                 processHandler.destroyProcess();
325                 UnixProcessManager.sendSigIntToProcessTree(processHandler.getProcess());
326               }
327             }, new Computable<Boolean>() {
328               @Override
329               public Boolean compute() {
330                 return !processHandler.isProcessTerminated();
331               }
332             })
333             .withRerun(new Runnable() {
334               @Override
335               public void run() {
336                 startIpythonServer(url, fileEditor);
337               }
338             })
339             .run();
340         }
341       });
342       return true;
343     }
344     catch (ExecutionException e) {
345       return false;
346     }
347   }
348
349   @Nullable
350   private static String findIPythonRunner(String homePath) {
351     for (String name : Lists.newArrayList("ipython", "ipython-script.py")) {
352       String runnerPath = PythonSdkType.getExecutablePath(homePath, name);
353       if (runnerPath != null) {
354         return runnerPath;
355       }
356     }
357
358     return null;
359   }
360
361   @Nullable
362   public static Pair<String, String> getHostPortFromUrl(@NotNull String url) {
363     try {
364       final URI uri = new URI(url);
365       final int port = uri.getPort();
366       return Pair.create(uri.getHost(), port == -1 ? null : String.valueOf(port));
367     }
368     catch (URISyntaxException e) {
369       return null;
370     }
371   }
372
373   public void projectOpened() {
374   }
375
376
377   public void projectClosed() {
378     shutdownKernels();
379   }
380
381   private void shutdownKernels() {
382     for (IpnbConnection connection : myKernels.values()) {
383       if (!connection.isAlive()) continue;
384       connection.shutdown();
385       try {
386         connection.close();
387       }
388       catch (IOException e) {
389         LOG.error(e);
390       }
391       catch (InterruptedException e) {
392         LOG.error(e);
393       }
394     }
395     myKernels.clear();
396   }
397
398   @NotNull
399   public String getComponentName() {
400     return "IpnbConnectionManager";
401   }
402
403   public void initComponent() {
404   }
405
406   public void disposeComponent() {
407     shutdownKernels();
408   }
409
410   class IpnbSettingsAdapter extends HyperlinkAdapter {
411     @Override
412     protected void hyperlinkActivated(HyperlinkEvent e) {
413       ShowSettingsUtil.getInstance().showSettingsDialog(myProject, "IPython Notebook");
414     }
415   }
416 }