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