pub mod call; pub mod class; pub mod config; use std::ffi::{CString, OsStr}; use std::fmt::Display; use std::path::Path; use std::sync::OnceLock; use std::time::SystemTime; use std::{collections::HashMap, path::PathBuf}; use colored::Colorize; use pyo3::prelude::*; use pyo3::types::PyTuple; use tracing::{event, span, warn, Level}; use crate::MainStatus; #[derive(Debug, Clone)] pub struct PyStatus { pub files: PyPlugins, pub config: config::PluginConfigFile, } pub type PyPlugins = HashMap; pub type RawPyPlugin = (PathBuf, Option, String); #[allow(non_upper_case_globals)] static mut PyPluginStatus: OnceLock = OnceLock::new(); impl PyStatus { pub fn init() { let config = config::PluginConfigFile::default_init().expect("初始化 Python 插件配置文件失败"); let status = PyStatus { files: HashMap::new(), config, }; let _ = unsafe { PyPluginStatus.get_or_init(|| status) }; } pub fn get() -> &'static PyStatus { unsafe { PyPluginStatus.get().unwrap() } } pub fn get_mut() -> &'static mut PyStatus { unsafe { PyPluginStatus.get_mut().unwrap() } } /// 添加一个插件 pub fn add_file(&mut self, path: PathBuf, plugin: PyPlugin) { self.files.insert(path, plugin); } /// 重新加载一个插件 pub fn reload_plugin(&mut self, plugin_name: &str) -> bool { let plugin = self.files.iter_mut().find_map(|(_, plugin)| { if plugin.get_id() == plugin_name { Some(plugin) } else { None } }); if let Some(plugin) = plugin { plugin.reload_from_file() } else { event!(Level::WARN, "没有找到插件: {}", plugin_name); false } } /// 删除一个插件 pub fn delete_file(&mut self, path: &PathBuf) -> Option { self.files.remove(path) } pub fn get_status(&self, pluging_id: &str) -> Option { self.files.iter().find_map(|(_, plugin)| { if plugin.get_id() == pluging_id { return Some(plugin.enabled); } None }) } pub fn set_status(&mut self, pluging_id: &str, status: bool) { self.files.iter_mut().for_each(|(_, plugin)| { if plugin.get_id() == pluging_id { plugin.enabled = status; } }); } pub fn verify_file(&self, path: &PathBuf) -> bool { self.files.get(path).is_some_and(|plugin| plugin.verifiy()) } pub fn display() -> String { format!( "Python 插件 {{ {} }}", Self::get() .files .values() .map(|v| v.to_string()) .collect::>() .join("\n") ) } } pub fn get_py_err_traceback(py_err: &PyErr) -> String { Python::with_gil(|py| match py_err.traceback(py) { Some(traceback) => match traceback.format() { Ok(trace) => trace, Err(e) => format!("{:?}", e), }, None => "".to_string(), }) .red() .to_string() } #[derive(Debug, Clone)] pub struct PyPlugin { pub file_path: PathBuf, pub changed_time: Option, pub py_module: Py, pub enabled: bool, } impl PyPlugin { /// 从文件创建一个新的 pub fn new_from_path(path: &PathBuf) -> Option { let raw_file = load_py_file(path); match raw_file { Ok(raw_file) => match Self::try_from(raw_file) { Ok(plugin) => Some(plugin), Err(e) => { warn!( "加载 Python 插件文件{:?}: {:?} 失败\n{}", path, e, get_py_err_traceback(&e) ); None } }, Err(e) => { warn!("加载插件 {:?}: {:?} 失败", path, e); None } } } /// 从文件更新 pub fn reload_from_file(&mut self) -> bool { let raw_file = load_py_file(&self.file_path); match raw_file { Ok(raw_file) => match Self::try_from(raw_file) { Ok(plugin) => { self.py_module = plugin.py_module; self.changed_time = plugin.changed_time; self.enabled = PyStatus::get().config.get_status(&self.get_id()); event!(Level::INFO, "更新 Python 插件文件 {:?} 完成", self.file_path); true } Err(e) => { warn!( "更新 Python 插件文件{:?}: {:?} 失败\n{}", self.file_path, e, get_py_err_traceback(&e) ); false } }, Err(e) => { warn!("更新插件 {:?}: {:?} 失败", self.file_path, e); false } } } /// 检查文件是否被修改 pub fn verifiy(&self) -> bool { match get_change_time(&self.file_path) { None => false, Some(time) => { if let Some(changed_time) = self.changed_time { time.eq(&changed_time) } else { true } } } } pub fn get_id(&self) -> String { plugin_path_as_id(&self.file_path) } } impl Display for PyPlugin { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}({:?})-{}", self.get_id(), self.file_path, self.enabled) } } pub const CONFIG_DATA_NAME: &str = "CONFIG_DATA"; impl TryFrom for PyPlugin { type Error = PyErr; fn try_from(value: RawPyPlugin) -> Result { let (path, changed_time, content) = value; let py_module = match py_module_from_code(&content, &path) { Ok(module) => module, Err(e) => { warn!("加载 Python 插件: {:?} 失败", e); return Err(e); } }; Python::with_gil(|py| { let module = py_module.bind(py); if let Ok(config_func) = call::get_func(module, "on_config") { match config_func.call0() { Ok(config) => { if config.is_instance_of::() { let (config, default) = config.extract::<(String, String)>().unwrap(); let base_path = MainStatus::global_config().py().config_path; let mut base_path: PathBuf = PathBuf::from(base_path); if !base_path.exists() { event!(Level::WARN, "python 插件路径不存在, 创建: {:?}", base_path); std::fs::create_dir_all(&base_path)?; } base_path.push(&config); let config_value = if base_path.exists() { event!( Level::INFO, "加载 {:?} 的配置文件 {:?} 中", path, base_path ); let content = std::fs::read_to_string(&base_path)?; toml::from_str(&content) } else { event!( Level::WARN, "配置文件 {:?} 不存在, 创建默认配置", base_path ); // 写入默认配置 std::fs::write(base_path, &default)?; toml::from_str(&default) }; match config_value { Ok(config) => { let py_config = Bound::new(py, class::ConfigDataPy::new(config)); if let Err(e) = py_config { event!(Level::WARN, "添加配置文件信息失败: {:?}", e); return Err(e); } let py_config = py_config.unwrap(); // 先判定一下原来有没有 if let Ok(true) = module.hasattr(CONFIG_DATA_NAME) { // get 过来, 后面直接覆盖, 这里用于发个警告 match module.getattr(CONFIG_DATA_NAME) { Ok(old_config) => { // 先判断是不是 None, 直接忽略掉 None // 毕竟有可能有占位 if !old_config.is_none() { warn!( "Python 插件 {:?} 的配置文件信息已经存在\n原始内容: {}", path, old_config ); } } Err(e) => { warn!( "Python 插件 {:?} 的配置文件信息已经存在, 但获取失败:{:?}", path, e ); } } } match module.setattr(CONFIG_DATA_NAME, py_config) { Ok(()) => Ok(PyPlugin { file_path: path, changed_time, py_module: module.into_py(py), enabled: true, }), Err(e) => { warn!( "Python 插件 {:?} 的配置文件信息设置失败:{:?}", path, e ); Err(PyErr::new::( format!( "Python 插件 {:?} 的配置文件信息设置失败:{:?}", path, e ), )) } } } Err(e) => { warn!( "加载 Python 插件 {:?} 的配置文件信息时失败:{:?}", path, e ); Err(PyErr::new::(format!( "加载 Python 插件 {:?} 的配置文件信息时失败:{:?}", path, e ))) } } } else if config.is_none() { // 没有配置文件 Ok(PyPlugin { file_path: path, changed_time, py_module: module.into_py(py), enabled: true, }) } else { warn!( "加载 Python 插件 {:?} 的配置文件信息时失败:返回的不是 [str, str]", path ); Err(PyErr::new::( "返回的不是 [str, str]".to_string(), )) } } Err(e) => { warn!("加载 Python 插件 {:?} 的配置文件信息时失败:{:?}", path, e); Err(e) } } } else { Ok(PyPlugin { file_path: path, changed_time, py_module: module.into_py(py), enabled: true, }) } }) } } /// 插件路径转换为 id pub fn plugin_path_as_id(path: &Path) -> String { path.file_name() .unwrap_or_default() .to_str() .unwrap_or("decode-failed") .to_string() } pub fn load_py_plugins(path: &PathBuf) { let plugins = PyStatus::get_mut(); if path.exists() { event!(Level::INFO, "找到位于 {:?} 的插件", path); // 搜索所有的 py 文件 和 文件夹单层下面的 py 文件 match path.read_dir() { Err(e) => { event!(Level::WARN, "读取插件路径失败 {:?}", e); } Ok(dir) => { for entry in dir { let entry = entry.unwrap(); let path = entry.path(); if let Some(ext) = path.extension() { if ext == "py" { if let Some(plugin) = PyPlugin::new_from_path(&path) { plugins.add_file(path, plugin); } } } } } } } else { event!(Level::WARN, "插件加载目录不存在: {:?}", path); } plugins.config.read_status_from_default(); plugins.config.sync_status_to_config(); event!( Level::INFO, "python 插件目录: {:?} 加载完成, 加载到 {} 个插件", path, plugins.files.len() ); } pub fn get_change_time(path: &Path) -> Option { path.metadata().ok()?.modified().ok() } pub fn py_module_from_code(content: &str, path: &Path) -> PyResult> { Python::with_gil(|py| -> PyResult> { let module: PyResult> = PyModule::from_code( py, CString::new(content).unwrap().as_c_str(), CString::new(path.to_string_lossy().as_bytes()).unwrap().as_c_str(), CString::new(path.file_name().unwrap().to_string_lossy().as_bytes()) .unwrap() .as_c_str(), // !!!! 请注意, 一定要给他一个名字, cpython 会自动把后面的重名模块覆盖掉前面的 ) .map(|module| module.into()); module }) } /// 传入文件路径 /// 返回 hash 和 文件内容 pub fn load_py_file(path: &PathBuf) -> std::io::Result { let changed_time = get_change_time(path); let content = std::fs::read_to_string(path)?; Ok((path.clone(), changed_time, content)) } fn init_py_with_env_path(path: &str) { unsafe { #[cfg(target_os = "linux")] use std::os::unix::ffi::OsStrExt; #[cfg(target_os = "windows")] use std::os::windows::ffi::OsStrExt; let mut config = std::mem::zeroed::(); let config_ptr = &mut config as *mut pyo3::ffi::PyConfig; // 初始化配置 // pyo3::ffi::PyConfig_InitIsolatedConfig(config_ptr); pyo3::ffi::PyConfig_InitPythonConfig(config_ptr); let wide_path = OsStr::new(path).encode_wide().chain(Some(0)).collect::>(); // 设置 prefix 和 exec_prefix pyo3::ffi::PyConfig_SetString(config_ptr, &mut config.prefix as *mut _, wide_path.as_ptr()); pyo3::ffi::PyConfig_SetString( config_ptr, &mut config.exec_prefix as *mut _, wide_path.as_ptr(), ); // 使用 Py_InitializeFromConfig 初始化 python let status = pyo3::ffi::Py_InitializeFromConfig(&config as *const _); pyo3::ffi::PyEval_SaveThread(); // 清理配置 pyo3::ffi::PyConfig_Clear(config_ptr); match status._type { pyo3::ffi::_PyStatus_TYPE::_PyStatus_TYPE_OK => { event!(Level::INFO, "根据配置初始化 python 完成"); } pyo3::ffi::_PyStatus_TYPE::_PyStatus_TYPE_EXIT => { event!(Level::ERROR, "初始化 python 时发生错误: EXIT"); } pyo3::ffi::_PyStatus_TYPE::_PyStatus_TYPE_ERROR => { event!(Level::ERROR, "初始化 python 时发生错误: ERROR"); pyo3::ffi::Py_ExitStatusException(status); } } } } /// Python 侧初始化 pub fn init_py() { // 从 全局配置中获取 python 插件路径 let span = span!(Level::INFO, "py init"); let _enter = span.enter(); let plugin_path = MainStatus::global_config().py().plugin_path; let cli_args = std::env::args().collect::>(); if cli_args.contains(&"-env".to_string()) { let env_path = cli_args.iter().find(|&arg| arg != "-env").expect("未找到 -env 参数的值"); event!(Level::INFO, "找到 -env 参数: {} 正在初始化", env_path); // 判断一下是否有 VIRTUAL_ENV 环境变量 if let Ok(virtual_env) = std::env::var("VIRTUAL_ENV") { event!(Level::WARN, "找到 VIRTUAL_ENV 环境变量: {} 将会被 -env 参数覆盖", virtual_env); } init_py_with_env_path(env_path); } else { // 根据 VIRTUAL_ENV 环境变量 进行一些处理 match std::env::var("VIRTUAL_ENV") { Ok(virtual_env) => { event!(Level::INFO, "找到 VIRTUAL_ENV 环境变量: {} 正在初始化", virtual_env); init_py_with_env_path(&virtual_env); } Err(_) => { event!(Level::INFO, "未找到 VIRTUAL_ENV 环境变量, 正常初始化"); pyo3::prepare_freethreaded_python(); event!(Level::INFO, "prepare_freethreaded_python 完成"); } } } PyStatus::init(); let plugin_path = PathBuf::from(plugin_path); load_py_plugins(&plugin_path); event!(Level::DEBUG, "python 插件列表: {}", PyStatus::display()); event!(Level::INFO, "python 初始化完成") } pub async fn post_py() -> anyhow::Result<()> { let status = PyStatus::get_mut(); status.config.sync_status_to_config(); status.config.write_to_default()?; stop_tasks().await; unsafe { if !pyo3::ffi::Py_FinalizeEx() == 0 { event!(Level::ERROR, "Python 退出失败 (不过应该无所谓)"); } } Ok(()) } async fn stop_tasks() { if call::PY_TASKS.lock().await.is_empty() { return; } let waiter = tokio::spawn(async { call::PY_TASKS.lock().await.join_all().await; }); tokio::select! { _ = waiter => { event!(Level::INFO, "Python 任务完成"); } _ = tokio::signal::ctrl_c() => { call::PY_TASKS.lock().await.cancel_all(); event!(Level::INFO, "Python 任务被中断"); } } }