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};
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 smol::process::Command;
20use ui::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
65pub fn python_env_kernel_specifications(
66 project: &Model<Project>,
67 worktree_id: WorktreeId,
68 cx: &mut AppContext,
69) -> impl Future<Output = Result<Vec<KernelSpecification>>> {
70 let python_language = LanguageName::new("Python");
71 let toolchains = project
72 .read(cx)
73 .available_toolchains(worktree_id, python_language, cx);
74 let background_executor = cx.background_executor().clone();
75
76 async move {
77 let toolchains = if let Some(toolchains) = toolchains.await {
78 toolchains
79 } else {
80 return Ok(Vec::new());
81 };
82
83 let kernelspecs = toolchains.toolchains.into_iter().map(|toolchain| {
84 background_executor.spawn(async move {
85 let python_path = toolchain.path.to_string();
86
87 // Check if ipykernel is installed
88 let ipykernel_check = Command::new(&python_path)
89 .args(&["-c", "import ipykernel"])
90 .output()
91 .await;
92
93 if ipykernel_check.is_ok() && ipykernel_check.unwrap().status.success() {
94 // Create a default kernelspec for this environment
95 let default_kernelspec = JupyterKernelspec {
96 argv: vec![
97 python_path.clone(),
98 "-m".to_string(),
99 "ipykernel_launcher".to_string(),
100 "-f".to_string(),
101 "{connection_file}".to_string(),
102 ],
103 display_name: toolchain.name.to_string(),
104 language: "python".to_string(),
105 interrupt_mode: None,
106 metadata: None,
107 env: None,
108 };
109
110 Some(KernelSpecification::PythonEnv(LocalKernelSpecification {
111 name: toolchain.name.to_string(),
112 path: PathBuf::from(&python_path),
113 kernelspec: default_kernelspec,
114 }))
115 } else {
116 None
117 }
118 })
119 });
120
121 let kernel_specs = futures::future::join_all(kernelspecs)
122 .await
123 .into_iter()
124 .flatten()
125 .collect();
126
127 anyhow::Ok(kernel_specs)
128 }
129}
130
131pub trait RunningKernel: Send + Debug {
132 fn request_tx(&self) -> mpsc::Sender<JupyterMessage>;
133 fn working_directory(&self) -> &PathBuf;
134 fn execution_state(&self) -> &ExecutionState;
135 fn set_execution_state(&mut self, state: ExecutionState);
136 fn kernel_info(&self) -> Option<&KernelInfoReply>;
137 fn set_kernel_info(&mut self, info: KernelInfoReply);
138 fn force_shutdown(&mut self) -> anyhow::Result<()>;
139}
140
141#[derive(Debug, Clone)]
142pub enum KernelStatus {
143 Idle,
144 Busy,
145 Starting,
146 Error,
147 ShuttingDown,
148 Shutdown,
149 Restarting,
150}
151
152impl KernelStatus {
153 pub fn is_connected(&self) -> bool {
154 match self {
155 KernelStatus::Idle | KernelStatus::Busy => true,
156 _ => false,
157 }
158 }
159}
160
161impl ToString for KernelStatus {
162 fn to_string(&self) -> String {
163 match self {
164 KernelStatus::Idle => "Idle".to_string(),
165 KernelStatus::Busy => "Busy".to_string(),
166 KernelStatus::Starting => "Starting".to_string(),
167 KernelStatus::Error => "Error".to_string(),
168 KernelStatus::ShuttingDown => "Shutting Down".to_string(),
169 KernelStatus::Shutdown => "Shutdown".to_string(),
170 KernelStatus::Restarting => "Restarting".to_string(),
171 }
172 }
173}
174
175#[derive(Debug)]
176pub enum Kernel {
177 RunningKernel(Box<dyn RunningKernel>),
178 StartingKernel(Shared<Task<()>>),
179 ErroredLaunch(String),
180 ShuttingDown,
181 Shutdown,
182 Restarting,
183}
184
185impl From<&Kernel> for KernelStatus {
186 fn from(kernel: &Kernel) -> Self {
187 match kernel {
188 Kernel::RunningKernel(kernel) => match kernel.execution_state() {
189 ExecutionState::Idle => KernelStatus::Idle,
190 ExecutionState::Busy => KernelStatus::Busy,
191 },
192 Kernel::StartingKernel(_) => KernelStatus::Starting,
193 Kernel::ErroredLaunch(_) => KernelStatus::Error,
194 Kernel::ShuttingDown => KernelStatus::ShuttingDown,
195 Kernel::Shutdown => KernelStatus::Shutdown,
196 Kernel::Restarting => KernelStatus::Restarting,
197 }
198 }
199}
200
201impl Kernel {
202 pub fn status(&self) -> KernelStatus {
203 self.into()
204 }
205
206 pub fn set_execution_state(&mut self, status: &ExecutionState) {
207 if let Kernel::RunningKernel(running_kernel) = self {
208 running_kernel.set_execution_state(status.clone());
209 }
210 }
211
212 pub fn set_kernel_info(&mut self, kernel_info: &KernelInfoReply) {
213 if let Kernel::RunningKernel(running_kernel) = self {
214 running_kernel.set_kernel_info(kernel_info.clone());
215 }
216 }
217
218 pub fn is_shutting_down(&self) -> bool {
219 match self {
220 Kernel::Restarting | Kernel::ShuttingDown => true,
221 Kernel::RunningKernel(_)
222 | Kernel::StartingKernel(_)
223 | Kernel::ErroredLaunch(_)
224 | Kernel::Shutdown => false,
225 }
226 }
227}