1mod native_kernel;
2use std::{fmt::Debug, future::Future, path::PathBuf, sync::Arc};
3
4use futures::{
5 channel::mpsc::{self, Receiver},
6 future::Shared,
7 stream,
8};
9use gpui::{App, Entity, Task, Window};
10use language::LanguageName;
11pub use native_kernel::*;
12
13mod remote_kernels;
14use project::{Project, ProjectPath, Toolchains, WorktreeId};
15pub use remote_kernels::*;
16
17use anyhow::Result;
18use jupyter_protocol::JupyterKernelspec;
19use runtimelib::{ExecutionState, JupyterMessage, KernelInfoReply};
20use ui::{Icon, IconName, SharedString};
21
22pub type JupyterMessageChannel = stream::SelectAll<Receiver<JupyterMessage>>;
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub enum KernelSpecification {
26 Remote(RemoteKernelSpecification),
27 Jupyter(LocalKernelSpecification),
28 PythonEnv(LocalKernelSpecification),
29}
30
31impl KernelSpecification {
32 pub fn name(&self) -> SharedString {
33 match self {
34 Self::Jupyter(spec) => spec.name.clone().into(),
35 Self::PythonEnv(spec) => spec.name.clone().into(),
36 Self::Remote(spec) => spec.name.clone().into(),
37 }
38 }
39
40 pub fn type_name(&self) -> SharedString {
41 match self {
42 Self::Jupyter(_) => "Jupyter".into(),
43 Self::PythonEnv(_) => "Python Environment".into(),
44 Self::Remote(_) => "Remote".into(),
45 }
46 }
47
48 pub fn path(&self) -> SharedString {
49 SharedString::from(match self {
50 Self::Jupyter(spec) => spec.path.to_string_lossy().to_string(),
51 Self::PythonEnv(spec) => spec.path.to_string_lossy().to_string(),
52 Self::Remote(spec) => spec.url.to_string(),
53 })
54 }
55
56 pub fn language(&self) -> SharedString {
57 SharedString::from(match self {
58 Self::Jupyter(spec) => spec.kernelspec.language.clone(),
59 Self::PythonEnv(spec) => spec.kernelspec.language.clone(),
60 Self::Remote(spec) => spec.kernelspec.language.clone(),
61 })
62 }
63
64 pub fn icon(&self, cx: &App) -> Icon {
65 let lang_name = match self {
66 Self::Jupyter(spec) => spec.kernelspec.language.clone(),
67 Self::PythonEnv(spec) => spec.kernelspec.language.clone(),
68 Self::Remote(spec) => spec.kernelspec.language.clone(),
69 };
70
71 file_icons::FileIcons::get(cx)
72 .get_icon_for_type(&lang_name.to_lowercase(), cx)
73 .map(Icon::from_path)
74 .unwrap_or(Icon::new(IconName::ReplNeutral))
75 }
76}
77
78pub fn python_env_kernel_specifications(
79 project: &Entity<Project>,
80 worktree_id: WorktreeId,
81 cx: &mut App,
82) -> impl Future<Output = Result<Vec<KernelSpecification>>> + use<> {
83 let python_language = LanguageName::new("Python");
84 let toolchains = project.read(cx).available_toolchains(
85 ProjectPath {
86 worktree_id,
87 path: Arc::from("".as_ref()),
88 },
89 python_language,
90 cx,
91 );
92 let background_executor = cx.background_executor().clone();
93
94 async move {
95 let (toolchains, user_toolchains) = if let Some(Toolchains {
96 toolchains,
97 root_path: _,
98 user_toolchains,
99 }) = toolchains.await
100 {
101 (toolchains, user_toolchains)
102 } else {
103 return Ok(Vec::new());
104 };
105
106 let kernelspecs = user_toolchains
107 .into_values()
108 .flatten()
109 .chain(toolchains.toolchains)
110 .map(|toolchain| {
111 background_executor.spawn(async move {
112 let python_path = toolchain.path.to_string();
113
114 // Check if ipykernel is installed
115 let ipykernel_check = util::command::new_smol_command(&python_path)
116 .args(&["-c", "import ipykernel"])
117 .output()
118 .await;
119
120 if ipykernel_check.is_ok() && ipykernel_check.unwrap().status.success() {
121 // Create a default kernelspec for this environment
122 let default_kernelspec = JupyterKernelspec {
123 argv: vec![
124 python_path.clone(),
125 "-m".to_string(),
126 "ipykernel_launcher".to_string(),
127 "-f".to_string(),
128 "{connection_file}".to_string(),
129 ],
130 display_name: toolchain.name.to_string(),
131 language: "python".to_string(),
132 interrupt_mode: None,
133 metadata: None,
134 env: None,
135 };
136
137 Some(KernelSpecification::PythonEnv(LocalKernelSpecification {
138 name: toolchain.name.to_string(),
139 path: PathBuf::from(&python_path),
140 kernelspec: default_kernelspec,
141 }))
142 } else {
143 None
144 }
145 })
146 });
147
148 let kernel_specs = futures::future::join_all(kernelspecs)
149 .await
150 .into_iter()
151 .flatten()
152 .collect();
153
154 anyhow::Ok(kernel_specs)
155 }
156}
157
158pub trait RunningKernel: Send + Debug {
159 fn request_tx(&self) -> mpsc::Sender<JupyterMessage>;
160 fn working_directory(&self) -> &PathBuf;
161 fn execution_state(&self) -> &ExecutionState;
162 fn set_execution_state(&mut self, state: ExecutionState);
163 fn kernel_info(&self) -> Option<&KernelInfoReply>;
164 fn set_kernel_info(&mut self, info: KernelInfoReply);
165 fn force_shutdown(&mut self, window: &mut Window, cx: &mut App) -> Task<anyhow::Result<()>>;
166}
167
168#[derive(Debug, Clone)]
169pub enum KernelStatus {
170 Idle,
171 Busy,
172 Starting,
173 Error,
174 ShuttingDown,
175 Shutdown,
176 Restarting,
177}
178
179impl KernelStatus {
180 pub fn is_connected(&self) -> bool {
181 matches!(self, KernelStatus::Idle | KernelStatus::Busy)
182 }
183}
184
185impl ToString for KernelStatus {
186 fn to_string(&self) -> String {
187 match self {
188 KernelStatus::Idle => "Idle".to_string(),
189 KernelStatus::Busy => "Busy".to_string(),
190 KernelStatus::Starting => "Starting".to_string(),
191 KernelStatus::Error => "Error".to_string(),
192 KernelStatus::ShuttingDown => "Shutting Down".to_string(),
193 KernelStatus::Shutdown => "Shutdown".to_string(),
194 KernelStatus::Restarting => "Restarting".to_string(),
195 }
196 }
197}
198
199#[derive(Debug)]
200pub enum Kernel {
201 RunningKernel(Box<dyn RunningKernel>),
202 StartingKernel(Shared<Task<()>>),
203 ErroredLaunch(String),
204 ShuttingDown,
205 Shutdown,
206 Restarting,
207}
208
209impl From<&Kernel> for KernelStatus {
210 fn from(kernel: &Kernel) -> Self {
211 match kernel {
212 Kernel::RunningKernel(kernel) => match kernel.execution_state() {
213 ExecutionState::Idle => KernelStatus::Idle,
214 ExecutionState::Busy => KernelStatus::Busy,
215 },
216 Kernel::StartingKernel(_) => KernelStatus::Starting,
217 Kernel::ErroredLaunch(_) => KernelStatus::Error,
218 Kernel::ShuttingDown => KernelStatus::ShuttingDown,
219 Kernel::Shutdown => KernelStatus::Shutdown,
220 Kernel::Restarting => KernelStatus::Restarting,
221 }
222 }
223}
224
225impl Kernel {
226 pub fn status(&self) -> KernelStatus {
227 self.into()
228 }
229
230 pub fn set_execution_state(&mut self, status: &ExecutionState) {
231 if let Kernel::RunningKernel(running_kernel) = self {
232 running_kernel.set_execution_state(status.clone());
233 }
234 }
235
236 pub fn set_kernel_info(&mut self, kernel_info: &KernelInfoReply) {
237 if let Kernel::RunningKernel(running_kernel) = self {
238 running_kernel.set_kernel_info(kernel_info.clone());
239 }
240 }
241
242 pub fn is_shutting_down(&self) -> bool {
243 match self {
244 Kernel::Restarting | Kernel::ShuttingDown => true,
245 Kernel::RunningKernel(_)
246 | Kernel::StartingKernel(_)
247 | Kernel::ErroredLaunch(_)
248 | Kernel::Shutdown => false,
249 }
250 }
251}