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