jupyter display support (#5)

Author: marcrasi

.gitignore

@@ -1 +1,3 @@
 venv
+
+**/*.pyc

EnableIPythonDisplay.swift

@@ -0,0 +1,82 @@
+// Copyright 2018 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// Hooks IPython to the KernelCommunicator, so that it can send display
+/// messages to Jupyter.
+
+import Python
+
+// Workaround SR-7757.
+#if canImport(Darwin)
+import func Darwin.C.dlopen
+#elseif canImport(Glibc)
+import func Glibc.dlopen
+#else
+#error("Cannot import Darwin or Glibc!")
+#endif
+dlopen("libpython2.7.so", RTLD_NOW | RTLD_GLOBAL)
+
+enum IPythonDisplay {
+  static var socket: PythonObject = Python.None
+  static var shell: PythonObject = Python.None
+
+}
+
+extension IPythonDisplay {
+  private static func bytes(_ py: PythonObject) -> [CChar] {
+    // faster not-yet-introduced method
+    // return py.swiftBytes!
+
+    // slow placeholder implementation
+    return py.map { el in
+      return CChar(bitPattern: UInt8(Python.ord(el))!)
+    }
+  }
+
+  private static func updateParentMessage(
+      to parentMessage: KernelCommunicator.ParentMessage) {
+    let json = Python.import("json")
+    IPythonDisplay.shell.set_parent(json.loads(parentMessage.json))
+  }
+
+  private static func consumeDisplayMessages()
+      -> [KernelCommunicator.JupyterDisplayMessage] {
+    let displayMessages = IPythonDisplay.socket.messages.map {
+      KernelCommunicator.JupyterDisplayMessage(parts: $0.map { bytes($0) })
+    }
+    IPythonDisplay.socket.messages = []
+    return displayMessages
+  }
+
+  static func enable() {
+    if IPythonDisplay.shell != Python.None {
+      print("Warning: IPython display already enabled.")
+      return
+    }
+
+    let swift_shell = Python.import("swift_shell")
+    let socketAndShell = swift_shell.create_shell(
+      username: JupyterKernel.communicator.jupyterSession.username,
+      session_id: JupyterKernel.communicator.jupyterSession.id,
+      key: JupyterKernel.communicator.jupyterSession.key)
+    IPythonDisplay.socket = socketAndShell[0]
+    IPythonDisplay.shell = socketAndShell[1]
+
+    JupyterKernel.communicator.handleParentMessage(updateParentMessage)
+    JupyterKernel.communicator.afterSuccessfulExecution(
+      run: consumeDisplayMessages)
+  }
+}
+
+IPythonDisplay.enable()

KernelCommunicator.swift

@@ -0,0 +1,79 @@
+// Copyright 2018 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// A struct with functions that the kernel and the code running inside the
+/// kernel use to talk to each other.
+///
+/// Note that it would be more Jupyter-y for the communication to happen over
+/// ZeroMQ. This is not currently possible, because ZeroMQ sends messages
+/// asynchronously using IO threads, and LLDB pauses those IO threads, which
+/// prevents them from sending the messages.
+public struct KernelCommunicator {
+  private var afterSuccessfulExecutionHandlers: [() -> [JupyterDisplayMessage]]
+  private var parentMessageHandlers: [(ParentMessage) -> ()]
+
+  public let jupyterSession: JupyterSession
+
+  init(jupyterSession: JupyterSession) {
+    self.afterSuccessfulExecutionHandlers = []
+    self.parentMessageHandlers = []
+    self.jupyterSession = jupyterSession
+  }
+
+  /// Register a handler to run after the kernel successfully executes a cell
+  /// of user code. The handler may return messages. These messages will be
+  /// send to the Jupyter client.
+  public mutating func afterSuccessfulExecution(
+      run handler: @escaping () -> [JupyterDisplayMessage]) {
+    afterSuccessfulExecutionHandlers.append(handler)
+  }
+
+  /// Register a handler to run when the parent message changes.
+  public mutating func handleParentMessage(
+      _ handler: @escaping (ParentMessage) -> ()) {
+    parentMessageHandlers.append(handler)
+  }
+
+  /// The kernel calls this after successfully executing a cell of user code.
+  public func triggerAfterSuccessfulExecution() -> [JupyterDisplayMessage] {
+    return afterSuccessfulExecutionHandlers.flatMap { $0() }
+  }
+
+  /// The kernel calls this when the parent message changes.
+  public mutating func updateParentMessage(
+      to parentMessage: ParentMessage) {
+    for parentMessageHandler in parentMessageHandlers {
+      parentMessageHandler(parentMessage)
+    }
+  }
+
+  /// A single serialized display message for the Jupyter client.
+  public struct JupyterDisplayMessage {
+    public let parts: [[CChar]]
+  }
+
+  /// ParentMessage identifies the request that causes things to happen.
+  /// This lets Jupyter, for example, know which cell to display graphics
+  /// messages in.
+  public struct ParentMessage {
+    let json: String
+  }
+
+  /// The data necessary to identify and sign outgoing jupyter messages.
+  public struct JupyterSession {
+    let id: String
+    let key: String
+    let username: String
+  }
+}

README.md

@@ -1,6 +1,6 @@
 # Swift-Jupyter
 
-This is a Jupyter Kernel for Swift, intended to make it possible to use Juptyer
+This is a Jupyter Kernel for Swift, intended to make it possible to use Jupyter
 with the [Swift for Tensorflow](https://github.com/tensorflow/swift) project.
 
 This kernel is currently very barebones and experimental.
@@ -16,7 +16,9 @@ virtualenv venv
 pip2 install jupyter # Must use python2, because LLDB doesn't support python3.
 ```
 
-Install a Swift toolchain ([see instructions here](https://github.com/tensorflow/swift/blob/master/Installation.md)).
+Install a Swift toolchain. For example, install the [Swift for Tensorflow] toolchain.
+
+[Swift for Tensorflow]: https://github.com/tensorflow/swift/blob/master/Installation.md
 
 Optionally [install SourceKitten](https://github.com/jpsim/SourceKitten) (this enables code completion).
 
@@ -30,6 +32,66 @@ Now run `jupyter notebook`, and it should have a Swift kernel.
 
 # Usage Instructions
 
+## Rich output
+
+You can call Python libaries using [Swift's Python interop] to display rich
+output in your Swift notebooks. (Eventually, we'd like to support Swift
+libraries the produce rich output too!)
+
+Prerequisites:
+
+* You must use a Swift toolchain that has Python interop. As of July 2018,
+  only the [Swift for Tensorflow] toolchain has Python interop.
+
+* Install the `ipykernel` Python library, and any other Python libraries
+  that you want output from (such as `matplotlib` or `pandas`) on your
+  system Python. (Do not install them on the virtualenv from the Swift-Jupyter
+  installation instructions. Swift's Python interop talks to your system
+  Python.)
+
+After taking care of the prerequisites, run
+`%include "EnableIPythonDisplay.swift"` in your Swift notebook. Now you should
+be able to display rich output! For example:
+
+```swift
+let np = Python.import("numpy")
+let plt = Python.import("matplotlib.pyplot")
+IPythonDisplay.shell.enable_matplotlib("inline")
+```
+
+```swift
+let time = np.arange(0, 10, 0.01)
+let amplitude = np.exp(-0.1 * time)
+let position = amplitude * np.sin(3 * time)
+
+plt.figure(figsize: [15, 10])
+
+plt.plot(time, position)
+plt.plot(time, amplitude)
+plt.plot(time, -amplitude)
+
+plt.xlabel("time (s)")
+plt.ylabel("position (m)")
+plt.title("Oscillations")
+
+plt.show()
+```
+
+![Screenshot of running the above two snippets of code in Jupyter](./screenshots/display_matplotlib.png)
+
+```swift
+let display = Python.import("IPython.display")
+let pd = Python.import("pandas")
+```
+
+```swift
+display.display(pd.DataFrame.from_records([["col 1": 3, "col 2": 5], ["col 1": 8, "col 2": 2]]))
+```
+
+![Screenshot of running the above two snippets of code in Jupyter](./screenshots/display_pandas.png)
+
+[Swift's Python interop]: https://github.com/tensorflow/swift/blob/master/docs/PythonInteroperability.md
+
 ## %include directives
 
 `%include` directives let you include code from files. To use them, put a line