[fud2] Experimental support for managing python environment (#2265)

docs/running-calyx/fud2/index.md

@@ -30,6 +30,18 @@ Add these lines:
 base = "<path to calyx checkout>"
 ```
 
+### Environment Setup
+
+Some parts of Calyx and `fud2` require setting up and installing various python packages. With Python removing support for installing packages system wide, it's recommended to install relevant packages into a python virtual environment. `fud2` can set up this environment for you and instruct `fud2` to automatically run relevant tools in the correct virtual environment.
+
+To do this, simply run:
+
+    $ fud2 env init
+
+There may be some cases where you want to manually interact with the python virtual environment. The virtual environment is installed to `$XDG_DATA_HOME/fud2/venv` (usually `~/.local/share/fud2/venv`). You can activate the virtual environment in your current shell with:
+
+    $ fud2 env activate
+
 Now you're ready to use fud2.
 
 [ninja]: https://ninja-build.org

fud2/Cargo.toml

@@ -22,6 +22,8 @@ fud-core = { path = "fud-core", version = "0.0.2", features = ["egg_planner"] }
 anyhow.workspace = true
 manifest-dir-macros = "0.1"
 include_dir = "0.7"
+argh.workspace = true
+toml_edit = "0.22.20"
 
 [lib]
 name = "fud2"

fud2/fud-core/Cargo.toml

@@ -32,4 +32,4 @@ toml_edit = { version = "0.22.20", features = ["serde"] }
 rand_chacha = "0.3.1"
 
 [features]
-egg_planner = ["dep:egg"]
+egg_planner = ["dep:egg"]

fud2/fud-core/src/cli.rs

@@ -1,3 +1,4 @@
+pub use crate::cli_ext::{CliExt, FakeCli, FromArgFn, RedactArgFn};
 use crate::config;
 use crate::exec::{plan, Driver, Request, StateRef};
 use crate::run::Run;
@@ -6,7 +7,7 @@ use argh::FromArgs;
 use camino::Utf8PathBuf;
 use figment::providers::Serialized;
 use itertools::Itertools;
-use std::fmt::Display;
+use std::fmt::{Debug, Display};
 use std::fs;
 use std::str::FromStr;
 
@@ -120,9 +121,9 @@ pub struct RegisterCommand {
 }
 
 /// supported subcommands
-#[derive(FromArgs, PartialEq, Debug)]
+#[derive(FromArgs)]
 #[argh(subcommand)]
-pub enum Subcommand {
+pub enum Subcommand<T: CliExt> {
     /// edit the configuration file
     EditConfig(EditConfig),
 
@@ -134,13 +135,16 @@ pub enum Subcommand {
 
     /// register a plugin
     Register(RegisterCommand),
+
+    #[argh(dynamic)]
+    Extended(FakeCli<T>),
 }
 
 #[derive(FromArgs)]
 /// A generic compiler driver.
-struct FakeArgs {
+pub struct FakeArgs<T: CliExt> {
     #[argh(subcommand)]
-    pub sub: Option<Subcommand>,
+    pub sub: Option<Subcommand<T>>,
 
     /// the input file
     #[argh(positional)]
@@ -223,9 +227,9 @@ fn get_states_with_errors(
     Ok(states)
 }
 
-fn from_states(
+fn from_states<T: CliExt>(
     driver: &Driver,
-    args: &FakeArgs,
+    args: &FakeArgs<T>,
 ) -> anyhow::Result<Vec<StateRef>> {
     get_states_with_errors(
         driver,
@@ -237,7 +241,10 @@ fn from_states(
     )
 }
 
-fn to_state(driver: &Driver, args: &FakeArgs) -> anyhow::Result<Vec<StateRef>> {
+fn to_state<T: CliExt>(
+    driver: &Driver,
+    args: &FakeArgs<T>,
+) -> anyhow::Result<Vec<StateRef>> {
     get_states_with_errors(
         driver,
         &args.to,
@@ -248,7 +255,10 @@ fn to_state(driver: &Driver, args: &FakeArgs) -> anyhow::Result<Vec<StateRef>> {
     )
 }
 
-fn get_request(driver: &Driver, args: &FakeArgs) -> anyhow::Result<Request> {
+fn get_request<T: CliExt>(
+    driver: &Driver,
+    args: &FakeArgs<T>,
+) -> anyhow::Result<Request> {
     // The default working directory (if not specified) depends on the mode.
     let workdir = args.dir.clone().unwrap_or_else(|| match args.mode {
         Mode::Generate | Mode::Run => {
@@ -363,9 +373,41 @@ fn register_plugin(
     Ok(())
 }
 
-/// Given the name of a Driver, returns a config based on that name and CLI arguments.
-pub fn config_from_cli(name: &str) -> anyhow::Result<figment::Figment> {
-    let args: FakeArgs = argh::from_env();
+pub trait CliStart<T: CliExt> {
+    /// Given the name of a Driver, returns a config based on that name and CLI arguments.
+    fn config_from_cli(name: &str) -> anyhow::Result<figment::Figment>;
+
+    /// Given a driver and config, start the CLI.
+    fn cli(driver: &Driver, config: &figment::Figment) -> anyhow::Result<()>;
+}
+
+/// Default CLI that provides an interface to core actions.
+pub struct DefaultCli;
+
+impl CliStart<()> for DefaultCli {
+    fn config_from_cli(name: &str) -> anyhow::Result<figment::Figment> {
+        config_from_cli_ext::<()>(name)
+    }
+
+    fn cli(driver: &Driver, config: &figment::Figment) -> anyhow::Result<()> {
+        cli_ext::<()>(driver, config)
+    }
+}
+
+impl<T: CliExt> CliStart<T> for T {
+    fn config_from_cli(name: &str) -> anyhow::Result<figment::Figment> {
+        config_from_cli_ext::<T>(name)
+    }
+
+    fn cli(driver: &Driver, config: &figment::Figment) -> anyhow::Result<()> {
+        cli_ext::<T>(driver, config)
+    }
+}
+
+fn config_from_cli_ext<T: CliExt>(
+    name: &str,
+) -> anyhow::Result<figment::Figment> {
+    let args: FakeArgs<T> = argh::from_env();
     let mut config = config::load_config(name);
 
     // Use `--set` arguments to override configuration values.
@@ -382,9 +424,11 @@ pub fn config_from_cli(name: &str) -> anyhow::Result<figment::Figment> {
     Ok(config)
 }
 
-pub fn cli(driver: &Driver, config: &figment::Figment) -> anyhow::Result<()> {
-    let args: FakeArgs = argh::from_env();
-
+fn cli_ext<T: CliExt>(
+    driver: &Driver,
+    config: &figment::Figment,
+) -> anyhow::Result<()> {
+    let args: FakeArgs<T> = argh::from_env();
     // Configure logging.
     env_logger::Builder::new()
         .format_timestamp(None)
@@ -407,6 +451,9 @@ pub fn cli(driver: &Driver, config: &figment::Figment) -> anyhow::Result<()> {
         Some(Subcommand::Register(cmd)) => {
             return register_plugin(driver, cmd);
         }
+        Some(Subcommand::Extended(cmd)) => {
+            return cmd.0.run(driver);
+        }
         None => {}
     }
 

fud2/fud-core/src/cli_ext.rs

@@ -0,0 +1,137 @@
+use std::sync::OnceLock;
+
+use crate::Driver;
+
+/// Fn type representing the redact_arg function required for implementing `argh::FromArgs`
+pub type RedactArgFn =
+    fn(&[&str], &[&str]) -> Result<Vec<String>, argh::EarlyExit>;
+
+/// Fn type representing the from_arg function required for implementing `argh::FromArgs`
+pub type FromArgFn<T> = fn(&[&str], &[&str]) -> Result<T, argh::EarlyExit>;
+
+/// Trait for extending the cli provided by `fud_core::cli`.
+///
+/// Below is an example of how to use this trait to add a subcommand named `test` to
+/// the `fud_core` cli.
+///
+/// ```rust
+/// /// some test command
+/// #[derive(FromArgs)]
+/// #[argh(subcommand, name = "test")]
+/// pub struct TestCommand {
+///     /// some arg
+///     #[argh(positional)]
+///     arg: String,
+/// }
+///
+/// pub enum TestExt {
+///     Test(TestCommand)
+/// }
+///
+/// impl CliExt for Fud2CliExt {
+///     fn run(&self, driver: &fud_core::Driver) -> anyhow::Result<()> {
+///         match &self {
+///             Fud2CliExt::Test(cmd) => {
+///                 println!("hi there: {}", cmd.arg);
+///                 Ok(())
+///             }
+///         }
+///     }
+///
+///     fn inner_command_info() -> Vec<CommandInfo> {
+///         vec![CommandInfo {
+///             name: "test",
+///             description: "test command",
+///         }]
+///     }
+///
+///     fn inner_redact_arg_values() -> Vec<(&'static str, RedactArgFn)> {
+///         vec![("test", TestCommand::redact_arg_values)]
+///     }
+///
+///     fn inner_from_args() -> Vec<(&'static str, FromArgFn<Self>)> {
+///         vec![("test", |cmd_name, args| {
+///             TestCommand::from_args(cmd_name, args).map(Self::Test)
+///         })]
+///     }
+/// }
+/// ```
+pub trait CliExt: Sized {
+    /// Action to execute when this subcommand is provided to the cli
+    fn run(&self, driver: &Driver) -> anyhow::Result<()>;
+
+    /// Provides the command names and descriptions for all subcommands in the cli
+    /// extension.
+    fn inner_command_info() -> Vec<argh::CommandInfo>;
+
+    /// Forward `redact_arg_values` parsing of subcommands to `fud_core::cli` parsing.
+    fn inner_redact_arg_values() -> Vec<(&'static str, RedactArgFn)>;
+
+    /// Forward `from_args` parsing of subcommands to `fud_core::cli` parsing.
+    fn inner_from_args() -> Vec<(&'static str, FromArgFn<Self>)>;
+}
+
+/// Wrapper type over types that implement `CliExt`. This is needed so that we can
+/// implement the foreign trait `argh::DynamicSubCommand` on a user provided `CliExt`.
+pub struct FakeCli<T: CliExt>(pub T);
+
+impl<T: CliExt> argh::DynamicSubCommand for FakeCli<T> {
+    fn commands() -> &'static [&'static argh::CommandInfo] {
+        static RET: OnceLock<Vec<&'static argh::CommandInfo>> = OnceLock::new();
+        RET.get_or_init(|| {
+            T::inner_command_info()
+                .into_iter()
+                .map(|cmd_info| &*Box::leak(Box::new(cmd_info)))
+                .collect()
+        })
+    }
+
+    fn try_redact_arg_values(
+        command_name: &[&str],
+        args: &[&str],
+    ) -> Option<Result<Vec<String>, argh::EarlyExit>> {
+        for (reg_name, f) in T::inner_redact_arg_values() {
+            if let Some(&name) = command_name.last() {
+                if name == reg_name {
+                    return Some(f(command_name, args));
+                }
+            }
+        }
+        None
+    }
+
+    fn try_from_args(
+        command_name: &[&str],
+        args: &[&str],
+    ) -> Option<Result<Self, argh::EarlyExit>> {
+        for (reg_name, f) in T::inner_from_args() {
+            if let Some(&name) = command_name.last() {
+                if name == reg_name {
+                    return Some(f(command_name, args).map(FakeCli));
+                }
+            }
+        }
+        None
+    }
+}
+
+/// The default CliExt used if none is provided. This doesn't define any new commands.
+impl CliExt for () {
+    fn inner_command_info() -> Vec<argh::CommandInfo> {
+        vec![]
+    }
+
+    fn inner_redact_arg_values(
+    ) -> Vec<(&'static str, crate::cli_ext::RedactArgFn)> {
+        vec![]
+    }
+
+    fn inner_from_args() -> Vec<(&'static str, crate::cli_ext::FromArgFn<Self>)>
+    {
+        vec![]
+    }
+
+    fn run(&self, _driver: &Driver) -> anyhow::Result<()> {
+        Ok(())
+    }
+}

fud2/fud-core/src/config.rs

@@ -47,6 +47,16 @@ pub fn config_path(name: &str) -> std::path::PathBuf {
     config_path
 }
 
+/// Location of the data directory
+pub fn data_dir(name: &str) -> std::path::PathBuf {
+    // The configuration is usually at `~/.config/driver_name.toml`.
+    let config_base = env::var("XDG_DATA_HOME").unwrap_or_else(|_| {
+        let home = env::var("HOME").expect("$HOME not set");
+        home + "/.local/share"
+    });
+    Path::new(&config_base).join(name)
+}
+
 /// Get raw configuration data with some default options.
 pub fn default_config() -> Figment {
     Figment::from(Serialized::defaults(GlobalConfig::default()))

fud2/fud-core/src/lib.rs

@@ -1,4 +1,5 @@
 pub mod cli;
+mod cli_ext;
 pub mod config;
 pub mod exec;
 pub mod run;