jupyter_protocol/
connection_info.rs

1//! Defines structures and functions for Jupyter kernel connection information.
2//!
3//! This module provides types and utilities for working with Jupyter kernel
4//! connection information, including the `ConnectionInfo` struct and related
5//! functionality for formatting URLs and serializing/deserializing connection data.
6//!
7//! The main struct, `ConnectionInfo`, encapsulates all necessary information for
8//! establishing a connection with a Jupyter kernel, including IP address, ports,
9//! transport protocol, and authentication details.
10//! Defines structures and functions for Jupyter kernel connection information.
11//!
12//! This module provides types and utilities for working with Jupyter kernel
13//! connection information, including the `ConnectionInfo` struct and related
14//! functionality for formatting URLs and serializing/deserializing connection data.
15//!
16//! The main struct, `ConnectionInfo`, encapsulates all necessary information for
17//! establishing a connection with a Jupyter kernel, including IP address, ports,
18//! transport protocol, and authentication details.
19//!
20//! # Examples
21//!
22//! ```rust
23//! use jupyter_protocol::connection_info::{ConnectionInfo, Transport};
24//!
25//! let info = ConnectionInfo {
26//!     ip: "127.0.0.1".to_string(),
27//!     transport: Transport::TCP,
28//!     shell_port: 6767,
29//!     iopub_port: 6768,
30//!     stdin_port: 6769,
31//!     control_port: 6770,
32//!     hb_port: 6771,
33//!     key: "secret_key".to_string(),
34//!     signature_scheme: "hmac-sha256".to_string(),
35//!     kernel_name: Some("python3".to_string()),
36//! };
37//!
38//! assert_eq!(info.shell_url(), "tcp://127.0.0.1:6767");
39//! ```
40use serde::{Deserialize, Serialize};
41
42/// Represents the transport protocol used for Jupyter kernel communication.
43///
44/// This enum is used to specify whether the kernel should use IPC (Inter-Process Communication)
45/// or TCP (Transmission Control Protocol) for its network communications.
46#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
47#[serde(rename_all = "lowercase")]
48pub enum Transport {
49    IPC,
50    TCP,
51}
52
53impl std::fmt::Display for Transport {
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        match self {
56            Transport::IPC => write!(f, "ipc"),
57            Transport::TCP => write!(f, "tcp"),
58        }
59    }
60}
61
62/// Represents the runtime connection information for a Jupyter kernel.
63///
64/// This struct contains all the necessary information for a Jupyter client
65/// to connect to a kernel, including ports, transport protocol, and authentication details.
66///
67/// # Fields
68///
69/// * `ip` - The IP address of the kernel.
70/// * `transport` - The transport protocol (TCP or IPC).
71/// * `shell_port` - The port number for the shell channel.
72/// * `iopub_port` - The port number for the IOPub channel.
73/// * `stdin_port` - The port number for the stdin channel.
74/// * `control_port` - The port number for the control channel.
75/// * `hb_port` - The port number for the heartbeat channel.
76/// * `key` - The authentication key.
77/// * `signature_scheme` - The signature scheme used for message authentication.
78/// * `kernel_name` - An optional name for the kernel.
79///
80/// # Example
81///
82/// ```
83/// use jupyter_protocol::connection_info::{ConnectionInfo, Transport};
84///
85/// let info = ConnectionInfo {
86///     ip: "127.0.0.1".to_string(),
87///     transport: Transport::TCP,
88///     shell_port: 6767,
89///     iopub_port: 6768,
90///     stdin_port: 6790,
91///     control_port: 6791,
92///     hb_port: 6792,
93///     key: "secret_key".to_string(),
94///     signature_scheme: "hmac-sha256".to_string(),
95///     kernel_name: Some("python3".to_string()),
96/// };
97///
98/// assert_eq!(info.shell_url(), "tcp://127.0.0.1:6767");
99/// ```
100#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
101pub struct ConnectionInfo {
102    pub ip: String,
103    pub transport: Transport,
104    pub shell_port: u16,
105    pub iopub_port: u16,
106    pub stdin_port: u16,
107    pub control_port: u16,
108    pub hb_port: u16,
109    pub key: String,
110    pub signature_scheme: String,
111    // Ignore if not present
112    #[serde(skip_serializing_if = "Option::is_none")]
113    pub kernel_name: Option<String>,
114}
115
116/// Constructs a URL string from the given transport, IP address, and port.
117///
118/// This is a helper function used internally to create formatted URL strings
119/// for various Jupyter communication channels.
120///
121/// # Arguments
122///
123/// * `transport` - The transport protocol (`Transport::TCP` or `Transport::IPC`).
124/// * `ip` - The IP address as a string.
125/// * `port` - The port number.
126///
127/// # Returns
128///
129/// A `String` containing the formatted URL.
130fn form_url(transport: &Transport, ip: &str, port: u16) -> String {
131    format!("{}://{}:{}", transport, ip, port)
132}
133
134/// Provides methods to generate formatted URLs for various Jupyter communication channels.
135impl ConnectionInfo {
136    /// Formats the URL for the IOPub channel.
137    ///
138    /// # Returns
139    ///
140    /// A `String` containing the formatted URL for the IOPub channel.
141    /// format the iopub url for a ZeroMQ connection
142    pub fn iopub_url(&self) -> String {
143        form_url(&self.transport, &self.ip, self.iopub_port)
144    }
145
146    /// format the shell url for a ZeroMQ connection
147    /// Formats the URL for the shell channel.
148    ///
149    /// # Returns
150    ///
151    /// A `String` containing the formatted URL for the shell channel.
152    pub fn shell_url(&self) -> String {
153        form_url(&self.transport, &self.ip, self.shell_port)
154    }
155
156    /// format the stdin url for a ZeroMQ connection
157    /// Formats the URL for the stdin channel.
158    ///
159    /// # Returns
160    ///
161    /// A `String` containing the formatted URL for the stdin channel.
162    pub fn stdin_url(&self) -> String {
163        form_url(&self.transport, &self.ip, self.stdin_port)
164    }
165
166    /// format the control url for a ZeroMQ connection
167    /// Formats the URL for the control channel.
168    ///
169    /// # Returns
170    ///
171    /// A `String` containing the formatted URL for the control channel.
172    pub fn control_url(&self) -> String {
173        form_url(&self.transport, &self.ip, self.control_port)
174    }
175
176    /// format the heartbeat url for a ZeroMQ connection
177    /// Formats the URL for the heartbeat channel.
178    ///
179    /// # Returns
180    ///
181    /// A `String` containing the formatted URL for the heartbeat channel.
182    pub fn hb_url(&self) -> String {
183        form_url(&self.transport, &self.ip, self.hb_port)
184    }
185}
186
187#[cfg(test)]
188mod test {
189    use super::*;
190
191    #[test]
192    fn test_connection_info_urls() {
193        let connection_info = ConnectionInfo {
194            ip: "127.0.0.1".to_string(),
195            transport: Transport::TCP,
196            shell_port: 6767,
197            iopub_port: 6768,
198            stdin_port: 6769,
199            control_port: 6770,
200            hb_port: 6771,
201            key: "test_key".to_string(),
202            signature_scheme: "hmac-sha256".to_string(),
203            kernel_name: Some("test_kernel".to_string()),
204        };
205
206        assert_eq!(connection_info.shell_url(), "tcp://127.0.0.1:6767");
207        assert_eq!(connection_info.iopub_url(), "tcp://127.0.0.1:6768");
208        assert_eq!(connection_info.stdin_url(), "tcp://127.0.0.1:6769");
209        assert_eq!(connection_info.control_url(), "tcp://127.0.0.1:6770");
210        assert_eq!(connection_info.hb_url(), "tcp://127.0.0.1:6771");
211
212        let ipc_connection_info = ConnectionInfo {
213            transport: Transport::IPC,
214            ..connection_info
215        };
216
217        assert_eq!(ipc_connection_info.shell_url(), "ipc://127.0.0.1:6767");
218        assert_eq!(ipc_connection_info.iopub_url(), "ipc://127.0.0.1:6768");
219        assert_eq!(ipc_connection_info.stdin_url(), "ipc://127.0.0.1:6769");
220        assert_eq!(ipc_connection_info.control_url(), "ipc://127.0.0.1:6770");
221        assert_eq!(ipc_connection_info.hb_url(), "ipc://127.0.0.1:6771");
222    }
223
224    #[test]
225    fn test_parse_connection_info() {
226        let json_str = r#"
227        {
228            "shell_port": 53380,
229            "iopub_port": 53381,
230            "stdin_port": 53382,
231            "control_port": 53383,
232            "hb_port": 53384,
233            "ip": "127.0.0.1",
234            "key": "e733b584-1d43845bc7d8d11a60df6363",
235            "transport": "tcp",
236            "signature_scheme": "hmac-sha256",
237            "kernel_name": "anaconda",
238            "jupyter_session": "/Users/kylekelley/Untitled3.ipynb"
239        }"#;
240
241        let connection_info: ConnectionInfo = serde_json::from_str(json_str).unwrap();
242
243        assert_eq!(connection_info.shell_port, 53380);
244        assert_eq!(connection_info.iopub_port, 53381);
245        assert_eq!(connection_info.stdin_port, 53382);
246        assert_eq!(connection_info.control_port, 53383);
247        assert_eq!(connection_info.hb_port, 53384);
248        assert_eq!(connection_info.ip, "127.0.0.1");
249        assert_eq!(connection_info.key, "e733b584-1d43845bc7d8d11a60df6363");
250        assert_eq!(connection_info.transport, Transport::TCP);
251        assert_eq!(connection_info.signature_scheme, "hmac-sha256");
252        assert_eq!(connection_info.kernel_name, Some("anaconda".to_string()));
253    }
254}