jupyter_websocket_client/
client.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
use anyhow::{Context, Result};
use async_tungstenite::{async_std::connect_async, tungstenite::http::Response};
use serde::{Deserialize, Serialize};
use url::Url;

use crate::websocket::JupyterWebSocket;

pub struct RemoteServer {
    pub base_url: String,
    pub token: String,
}

// Only `id` is used right now, but other fields will be useful when pulling up a listing later
#[derive(Debug, Serialize, Deserialize)]
pub struct Kernel {
    pub id: String,
    pub name: String,
    pub last_activity: String,
    pub execution_state: String,
    pub connections: u64,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Session {
    pub id: String,
    pub path: String,
    pub name: String,
    #[serde(rename = "type")]
    pub session_type: String,
    pub kernel: Kernel,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct NewSession {
    pub path: String,
    pub name: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct KernelSpec {
    pub name: String,
    pub spec: jupyter_serde::JupyterKernelspec,
    pub resources: std::collections::HashMap<String, String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct KernelLaunchRequest {
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub path: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct HelpLink {
    pub text: String,
    pub url: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct KernelSpecsResponse {
    pub default: String,
    pub kernelspecs: std::collections::HashMap<String, KernelSpec>,
}

fn api_url(base_url: &str, path: &str) -> String {
    format!(
        "{}/api/{}",
        base_url.trim_end_matches('/'),
        path.trim_start_matches('/')
    )
}

impl RemoteServer {
    pub fn from_url(url: &str) -> Result<Self> {
        let parsed_url = Url::parse(url).context("Failed to parse Jupyter URL")?;
        let base_url = format!(
            "{}://{}{}{}",
            parsed_url.scheme(),
            parsed_url.host_str().unwrap_or("localhost"),
            parsed_url
                .port()
                .map(|p| format!(":{}", p))
                .unwrap_or_default(),
            parsed_url.path().trim_end_matches("/tree")
        );

        let token = parsed_url
            .query_pairs()
            .find(|(key, _)| key == "token")
            .map(|(_, value)| value.into_owned())
            .ok_or_else(|| anyhow::anyhow!("Token not found in URL"))?;

        Ok(Self { base_url, token })
    }

    pub fn api_url(&self, path: &str) -> String {
        api_url(&self.base_url, path)
    }

    /// Connect to a kernel by ID
    ///
    /// ```rust
    /// use jupyter_websocket_client::RemoteServer;
    ///
    /// use jupyter_protocol::{KernelInfoRequest, JupyterMessageContent};
    ///
    /// // Import the sink and stream extensions to allow splitting the socket into a writer and reader pair
    /// use futures::{SinkExt as _, StreamExt as _};
    ///
    /// pub async fn connect_kernel() -> anyhow::Result<()> {
    ///     let server = RemoteServer::from_url(
    ///         "http://127.0.0.1:8888/lab?token=f487535a46268da4a0752c0e162c873b721e33a9e6ec8390"
    ///     )?;
    ///
    ///     // You'll need to launch a kernel and get a kernel ID using your own HTTP
    ///     // request library
    ///     let kernel_id = "1057-1057-1057-1057";
    ///
    ///     let (kernel_socket, response) = server.connect_to_kernel(kernel_id).await?;
    ///
    ///     let (mut w, mut r) = kernel_socket.split();
    ///
    ///     w.send(KernelInfoRequest {}.into()).await?;
    ///
    ///     while let Some(response) = r.next().await.transpose()? {
    ///         match response.content {
    ///             JupyterMessageContent::KernelInfoReply(kernel_info_reply) => {
    ///                 println!("Received kernel_info_reply");
    ///                 println!("{:?}", kernel_info_reply);
    ///                 break;
    ///             }
    ///             other => {
    ///                 println!("Received");
    ///                 println!("{:?}", other);
    ///             }
    ///         }
    ///     }
    ///
    ///     Ok(())
    /// }
    /// ```
    pub async fn connect_to_kernel(
        &self,
        kernel_id: &str,
    ) -> Result<(JupyterWebSocket, Response<Option<Vec<u8>>>)> {
        let ws_url = format!(
            "{}?token={}",
            api_url(&self.base_url, &format!("kernels/{}/channels", kernel_id))
                .replace("http", "ws"),
            self.token
        );

        let response = connect_async(ws_url).await;

        let (ws_stream, response) = response?;

        Ok((JupyterWebSocket { inner: ws_stream }, response))
    }
}