1  //! Extensions to the athenacl's dialog module.
  2  
  3  use rfd::FileDialog;
  4  use rustpython_vm::{pymodule, VirtualMachine};
  5  
  6  pub(crate) fn make_module(vm: &mut VirtualMachine) {
  7      vm.add_native_module("dialogExt", Box::new(_inner::make_module));
  8  }
  9  
 10  #[pymodule]
 11  pub(super) mod _inner {
 12      use std::env;
 13      use std::str;
 14  
 15      use super::*;
 16      use rustpython_vm::{convert::ToPyObject, PyResult};
 17  
 18      #[pyfunction(name = "promptChooseDir")]
 19      pub(crate) fn prompt_choose_dir(
 20          title: String,
 21          initial_dir: String,
 22          vm: &VirtualMachine,
 23      ) -> PyResult {
 24          prompt_dialog(title, initial_dir, PromptType::ChooseDir, vm)
 25      }
 26  
 27      #[pyfunction(name = "promptChooseFile")]
 28      pub(crate) fn prompt_choose_file(
 29          title: String,
 30          initial_dir: String,
 31          vm: &VirtualMachine,
 32      ) -> PyResult {
 33          prompt_dialog(title, initial_dir, PromptType::ChooseFile, vm)
 34      }
 35  
 36      #[pyfunction(name = "promptSaveFile")]
 37      pub(crate) fn prompt_save_file(
 38          title: String,
 39          initial_dir: String,
 40          vm: &VirtualMachine,
 41      ) -> PyResult {
 42          prompt_dialog(title, initial_dir, PromptType::SaveFile, vm)
 43      }
 44  
 45      fn prompt_dialog(
 46          title: String,
 47          initial_dir: String,
 48          prompt_type: PromptType,
 49          vm: &VirtualMachine,
 50      ) -> PyResult {
 51          let initial_dir = if initial_dir.is_empty() {
 52              match env::current_dir() {
 53                  Ok(path) => path.to_string_lossy().to_string(),
 54                  Err(e) => {
 55                      eprint!("{}", e);
 56                      let res = vec![
 57                          vm.ctx.new_str("").to_pyobject(vm),
 58                          vm.ctx.new_int(0).to_pyobject(vm),
 59                      ];
 60                      return Ok(vm.ctx.new_tuple(res).into());
 61                  }
 62              }
 63          } else {
 64              initial_dir
 65          };
 66  
 67          let title = if title.is_empty() {
 68              "Select directory".to_string()
 69          } else {
 70              title
 71          };
 72  
 73          let fd = FileDialog::new()
 74              .set_title(title)
 75              .set_directory(initial_dir)
 76              .set_can_create_directories(true);
 77  
 78          let res = match prompt_type {
 79              PromptType::ChooseDir => fd.pick_folder(),
 80              PromptType::ChooseFile => fd.pick_file(),
 81              PromptType::SaveFile => fd.save_file(),
 82          };
 83  
 84          let res = match res {
 85              Some(path) => vec![
 86                  vm.ctx
 87                      .new_str(path.to_string_lossy().to_string())
 88                      .to_pyobject(vm),
 89                  vm.ctx.new_int(1).to_pyobject(vm),
 90              ],
 91              None => vec![
 92                  vm.ctx.new_str("").to_pyobject(vm),
 93                  vm.ctx.new_int(0).to_pyobject(vm),
 94              ],
 95          };
 96  
 97          Ok(vm.ctx.new_tuple(res).into())
 98      }
 99  
100      #[derive(Clone, Copy)]
101      enum PromptType {
102          ChooseDir,
103          ChooseFile,
104          SaveFile,
105      }
106  }