Ensured shared environment for 'initial-env' command executions

Author: fritzrehde

Cargo.lock

@@ -29,6 +29,7 @@ ranges = "0.3.3"
 tokio = { version = "1.33.0", features = ["full"] }
 futures = "0.3.29"
 ansi-to-tui = "3.1.0"
+once_cell = "1.18.0"
 
 # Config for 'cargo dist'
 [workspace.metadata.dist]

Cargo.toml

@@ -185,6 +185,12 @@ The environment variable `lines` set to all selected lines, or if none are selec
 All set environment variables `ENV` will be made available in all future spawned commands/processes, including the watched command, any executed subcommands, as well as commands executed in `set-env` operations.
 If multiple lines are selected, they will be separated by newlines in `lines`.
 
+### State management
+
+The `set-env` and `unset-env` operations allow you to manage state through environment variables.
+Additionally, you can use the `initial-env` option to specify a list of `set-env` commands that will be executed **before** the first execution of the watched command.
+This powerful combination allows you to set some initial state with `initial-env`, reference that state directly in the watched command, and update the state with keybindings at runtime with `set-env`.
+
 ### Formatting with Field Separators and Field Selections
 
 `watchbind` supports some extra formatting features reminiscent of the Unix `cut` command:

README.md

@@ -33,7 +33,7 @@ impl Operations {
     }
 }
 
-#[derive(IntoIterator, Eq, Ord, PartialEq, PartialOrd, From)]
+#[derive(IntoIterator, Eq, Ord, PartialEq, PartialOrd, From, Clone)]
 pub struct OperationsParsed(#[into_iterator(ref)] Vec<OperationParsed>);
 
 impl TryFrom<Vec<String>> for OperationsParsed {

src/config/keybindings/operations/mod.rs

@@ -14,7 +14,7 @@ use tokio::sync::Mutex;
 /// The version of Operation used for parsing and displaying. The reason we
 /// can't parse directly into Operation is because any operations that execute
 /// something need to receive access to the globally set environment variables.
-#[derive(FromStr, Display, PartialEq, PartialOrd, Eq, Ord)]
+#[derive(FromStr, Display, PartialEq, PartialOrd, Eq, Ord, Clone)]
 #[display(style = "kebab-case")]
 pub enum OperationParsed {
     Exit,
@@ -114,13 +114,13 @@ impl Operation {
             Self::Reload => return Ok(RequestedAction::ReloadWatchedCommand),
             Self::Exit => return Ok(RequestedAction::Exit),
             Self::ExecuteNonBlocking(non_blocking_cmd) => {
-                state.add_lines_to_env().await?;
+                state.add_cursor_and_selected_lines_to_env().await;
                 non_blocking_cmd.execute().await?;
+                state.remove_cursor_and_selected_lines_from_env().await;
             }
             Self::ExecuteBlocking(blocking_cmd) => {
-                state.add_lines_to_env().await?;
+                state.add_cursor_and_selected_lines_to_env().await;
 
-                // TODO: these clones are preventable by using Arc<> (I think Arc<Mutex> isn't required because executing them doesn't mutate them)
                 let blocking_cmd = Arc::clone(blocking_cmd);
                 let event_tx = event_tx.clone();
                 tokio::spawn(async move {
@@ -133,7 +133,7 @@ impl Operation {
                 return Ok(RequestedAction::ExecutingBlockingSubcommand);
             }
             Self::ExecuteTUI(tui_cmd) => {
-                state.add_lines_to_env().await?;
+                state.add_cursor_and_selected_lines_to_env().await;
 
                 // Create channels for waiting until TUI has actually been hidden.
                 let (tui_hidden_tx, mut tui_hidden_rx) = mpsc::channel(1);
@@ -153,7 +153,7 @@ impl Operation {
                 return Ok(RequestedAction::ExecutingTUISubcommand(tui_hidden_tx));
             }
             Self::SetEnv(env_variable, blocking_cmd) => {
-                state.add_lines_to_env().await?;
+                state.add_cursor_and_selected_lines_to_env().await;
 
                 let blocking_cmd = blocking_cmd.clone();
                 let env_variable = env_variable.clone();
@@ -179,7 +179,8 @@ impl Operation {
         Ok(RequestedAction::Continue)
     }
 
-    /// Convert the parsed form into the normal, runtime Operation form.
+    /// Convert the parsed form into the normal, runtime Operation form. The
+    /// `env_variables` is required so it can be passed to the `SetEnv` command.
     pub fn from_parsed(parsed: OperationParsed, env_variables: &Arc<Mutex<EnvVariables>>) -> Self {
         match parsed {
             OperationParsed::Exit => Self::Exit,

src/config/keybindings/operations/operation.rs

@@ -28,7 +28,7 @@ pub struct Config {
     pub keybindings_parsed: KeybindingsParsed,
     pub header_lines: usize,
     pub fields: Fields,
-    pub initial_env_variables: OperationsParsed,
+    pub initial_env_ops: OperationsParsed,
 }
 
 impl Config {
@@ -76,7 +76,7 @@ impl TryFrom<TomlConfig> for Config {
 
         Ok(Self {
             log_file: toml.log_file,
-            initial_env_variables: toml.initial_env_variables.unwrap_or_default().try_into()?,
+            initial_env_ops: toml.initial_env_vars.unwrap_or_default().try_into()?,
             watched_command: match toml.watched_command {
                 Some(command) => command,
                 None => bail!("A command must be provided via command line or config file"),
@@ -100,7 +100,7 @@ pub struct TomlConfig {
     log_file: Option<PathBuf>,
 
     #[serde(rename = "initial-env")]
-    initial_env_variables: Option<Vec<String>>,
+    initial_env_vars: Option<Vec<String>>,
 
     watched_command: Option<String>,
     interval: Option<f64>,
@@ -151,7 +151,7 @@ impl TomlConfig {
     fn merge(self, other: Self) -> Self {
         Self {
             log_file: self.log_file.or(other.log_file),
-            initial_env_variables: self.initial_env_variables.or(other.initial_env_variables),
+            initial_env_vars: self.initial_env_vars.or(other.initial_env_vars),
             watched_command: self.watched_command.or(other.watched_command),
             interval: self.interval.or(other.interval),
             non_cursor_non_header_fg: self
@@ -182,7 +182,7 @@ impl From<ClapConfig> for TomlConfig {
     fn from(clap: ClapConfig) -> Self {
         Self {
             log_file: clap.log_file,
-            initial_env_variables: clap.initial_env_variables,
+            initial_env_vars: clap.initial_env_vars,
             watched_command: clap.watched_command.map(|s| s.join(" ")),
             interval: clap.interval,
             non_cursor_non_header_fg: clap.non_cursor_non_header_fg,
@@ -248,9 +248,9 @@ pub struct ClapConfig {
     #[arg(short, long, value_name = "FILE")]
     log_file: Option<PathBuf>,
 
-    /// Command to watch by executing periodically
+    /// Comman-separated `set-env` operations to execute before first watched command execution
     #[arg(long = "initial-env", value_name = "LIST", value_delimiter = ',')]
-    initial_env_variables: Option<Vec<String>>,
+    initial_env_vars: Option<Vec<String>>,
 
     /// Command to watch by executing periodically
     #[arg(trailing_var_arg(true))]

src/config/mod.rs

@@ -157,15 +157,18 @@ impl UI {
     async fn new(config: Config) -> Result<(Self, PollingState)> {
         let terminal_manager = Tui::new()?;
 
-        let env_variables = EnvVariables::generate_initial(config.initial_env_variables).await?;
+        // Create `State`.
         let keybindings_str = config.keybindings_parsed.to_string();
-        let state = State::new(
+        let mut state = State::new(
             config.header_lines,
             config.fields,
             config.styles,
             keybindings_str,
-            env_variables,
+            EnvVariables::new(),
         );
+        state
+            .generate_initial_env_vars(config.initial_env_ops)
+            .await?;
 
         // TODO: room for optimization: we can probably get away with much smaller buffer sizes for some of our channels
 
@@ -269,6 +272,9 @@ impl UI {
                     Event::TUISubcommandCompleted(potential_error) => {
                         potential_error?;
 
+                        // Remove temporary env vars that were added just for execution.
+                        self.state.remove_cursor_and_selected_lines_from_env().await;
+
                         self.tui.restore()?;
                         log::info!("Watchbind's TUI is shown.");
 
@@ -311,13 +317,18 @@ impl UI {
                 },
                 BlockingState::BlockedExecutingSubcommand => match event {
                     Event::CommandOutput(lines) => {
+                        // TODO: it's up for discussion if we really want this behaviour, need to find use-cases against first
+
                         // We handle new output lines, but don't exit the
                         // blocking state.
                         self.state.update_lines(lines?)?;
                     }
                     Event::SubcommandCompleted(potential_error) => {
                         potential_error?;
 
+                        // Remove temporary env vars that were added just for execution.
+                        self.state.remove_cursor_and_selected_lines_from_env().await;
+
                         if let ControlFlow::Exit = self.conclude_blocking().await? {
                             break 'event_loop;
                         }
@@ -336,7 +347,10 @@ impl UI {
                         self.state.update_lines(lines?)?;
                     }
                     Event::SubcommandForEnvCompleted(new_env_variables) => {
-                        self.state.set_env(new_env_variables?).await;
+                        // Remove temporary env vars that were added just for execution.
+                        self.state.remove_cursor_and_selected_lines_from_env().await;
+
+                        self.state.set_envs(new_env_variables?).await;
 
                         if let ControlFlow::Exit = self.conclude_blocking().await? {
                             break 'event_loop;

src/ui/mod.rs

@@ -2,7 +2,7 @@ mod env_variable;
 
 use crate::{
     command::CommandBuilder,
-    config::{OperationParsed, OperationsParsed},
+    config::{Operation, OperationParsed, Operations, OperationsParsed},
 };
 use anyhow::{bail, Result};
 use itertools::Itertools;
@@ -15,7 +15,11 @@ pub use env_variable::EnvVariable;
 pub struct EnvVariables(HashMap<EnvVariable, String>);
 
 impl EnvVariables {
-    // TODO: maybe extend to also allow the general execution of normal commands before the initial watched command is executed (if there are use-cases for that)
+    /// Create a new empty structure of environment variables.
+    pub fn new() -> Self {
+        Self(HashMap::new())
+    }
+
     /// Receive parsed operations, but only execute the "set-env" operations.
     pub async fn generate_initial(value: OperationsParsed) -> Result<Self> {
         // TODO: consider trying to use async iterators to do this in one iterator pass (instead of the mut hashmap) once stable
@@ -36,13 +40,41 @@ impl EnvVariables {
         Ok(Self(map))
     }
 
+    /// Receive parsed operations, but only execute the "set-env" operations.
+    pub async fn generate_initial2(&mut self, operations: Operations) -> Result<()> {
+        // TODO: consider trying to use async iterators to do this in one iterator pass (instead of the mut hashmap) once stable
+        for op in operations.into_iter() {
+            match op {
+                Operation::SetEnv(env_variable, blocking_cmd) => {
+                    log::info!("executing cmd");
+
+                    let cmd_output = blocking_cmd.execute().await?;
+
+                    log::info!("finished executing cmd");
+
+                    self.0.insert(env_variable, cmd_output);
+                }
+                // other_op => bail!("Invalid operation for env variable setup (only \"set-env\" operations allowed here): {}", other_op),
+                _ => bail!("Only \"set-env\" operations allowed here."),
+            }
+        }
+        Ok(())
+    }
+
     pub fn merge_new_envs(&mut self, env_variables: Self) {
         self.0.extend(env_variables.0);
     }
 
+    // TODO: expose EnvVariableValue type instead of String
+
+    /// Add an environment variable mapping.
+    pub fn set_env(&mut self, env_var: EnvVariable, value: String) {
+        self.0.insert(env_var, value);
+    }
+
     /// Unset/remove the specified environment variable.
-    pub fn unset_env(&mut self, env: &EnvVariable) {
-        self.0.remove(env);
+    pub fn unset_env(&mut self, env_var: &EnvVariable) {
+        self.0.remove(env_var);
     }
 
     /// Write formatted version (insert elastic tabstops) to a buffer.