Refactor shell switching related code (#6258)

* Refactor shell switching related code

Signed-off-by: nibon7 <nibon7@163.com>

* add tests

Signed-off-by: nibon7 <nibon7@163.com>

* fix tests

Signed-off-by: nibon7 <nibon7@163.com>

crates/nu-command/src/shells/g.rs

@@ -1,4 +1,4 @@
-use super::{get_current_shell, get_last_shell, get_shells};
+use super::{get_current_shell, get_shells, switch_shell, SwitchTo};
 use nu_engine::{current_dir, CallExt};
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
@@ -50,47 +50,16 @@ impl Command for GotoShell {
 
         match new_shell {
             Some(shell_span) => {
-                let index = if shell_span.item == "-" {
+                if shell_span.item == "-" {
-                    get_last_shell(engine_state, stack)
+                    switch_shell(engine_state, stack, call, shell_span.span, SwitchTo::Last)
                 } else {
-                    shell_span
+                    let n = shell_span
                         .item
                         .parse::<usize>()
-                        .map_err(|_| ShellError::NotFound(shell_span.span))?
+                        .map_err(|_| ShellError::NotFound(shell_span.span))?;
-                };
 
-                let new_path = shells
+                    switch_shell(engine_state, stack, call, shell_span.span, SwitchTo::Nth(n))
-                    .get(index)
+                }
-                    .ok_or(ShellError::NotFound(shell_span.span))?
-                    .to_owned();
-
-                let current_shell = get_current_shell(engine_state, stack);
-
-                stack.add_env_var(
-                    "NUSHELL_SHELLS".into(),
-                    Value::List {
-                        vals: shells,
-                        span: call.head,
-                    },
-                );
-                stack.add_env_var(
-                    "NUSHELL_CURRENT_SHELL".into(),
-                    Value::Int {
-                        val: index as i64,
-                        span: call.head,
-                    },
-                );
-                stack.add_env_var(
-                    "NUSHELL_LAST_SHELL".into(),
-                    Value::Int {
-                        val: current_shell as i64,
-                        span: call.head,
-                    },
-                );
-
-                stack.add_env_var("PWD".into(), new_path);
-
-                Ok(PipelineData::new(call.head))
             }
             None => {
                 let current_shell = get_current_shell(engine_state, stack);

crates/nu-command/src/shells/mod.rs

@@ -9,11 +9,20 @@ pub use enter::Enter;
 pub use exit::Exit;
 pub use g::GotoShell;
 pub use n::NextShell;
+use nu_engine::current_dir;
+use nu_protocol::ast::Call;
 use nu_protocol::engine::{EngineState, Stack};
-use nu_protocol::Value;
+use nu_protocol::{PipelineData, ShellError, Span, Value};
 pub use p::PrevShell;
 pub use shells_::Shells;
 
+enum SwitchTo {
+    Next,
+    Prev,
+    Last,
+    Nth(usize),
+}
+
 pub fn get_shells(engine_state: &EngineState, stack: &mut Stack, cwd: Value) -> Vec<Value> {
     let shells = stack.get_env_var(engine_state, "NUSHELL_SHELLS");
     let shells = if let Some(v) = shells {
@@ -43,3 +52,74 @@ fn get_last_shell(engine_state: &EngineState, stack: &mut Stack) -> usize {
         0
     }
 }
+
+fn switch_shell(
+    engine_state: &EngineState,
+    stack: &mut Stack,
+    call: &Call,
+    span: Span,
+    switch_to: SwitchTo,
+) -> Result<PipelineData, ShellError> {
+    let cwd = current_dir(engine_state, stack)?;
+    let cwd = Value::String {
+        val: cwd.to_string_lossy().to_string(),
+        span: call.head,
+    };
+
+    let shells = get_shells(engine_state, stack, cwd);
+    let current_shell = get_current_shell(engine_state, stack);
+
+    let new_shell = match switch_to {
+        SwitchTo::Next => {
+            let mut new_shell = current_shell + 1;
+
+            if new_shell == shells.len() {
+                new_shell = 0;
+            }
+
+            new_shell
+        }
+        SwitchTo::Prev => {
+            if current_shell == 0 {
+                shells.len() - 1
+            } else {
+                current_shell - 1
+            }
+        }
+        SwitchTo::Last => get_last_shell(engine_state, stack),
+        SwitchTo::Nth(n) => n,
+    };
+
+    let new_path = shells
+        .get(new_shell)
+        .ok_or(ShellError::NotFound(span))?
+        .to_owned();
+
+    stack.add_env_var(
+        "NUSHELL_SHELLS".into(),
+        Value::List {
+            vals: shells,
+            span: call.head,
+        },
+    );
+
+    stack.add_env_var(
+        "NUSHELL_CURRENT_SHELL".into(),
+        Value::Int {
+            val: new_shell as i64,
+            span: call.head,
+        },
+    );
+
+    stack.add_env_var(
+        "NUSHELL_LAST_SHELL".into(),
+        Value::Int {
+            val: current_shell as i64,
+            span: call.head,
+        },
+    );
+
+    stack.add_env_var("PWD".into(), new_path);
+
+    Ok(PipelineData::new(call.head))
+}

crates/nu-command/src/shells/n.rs

@@ -1,8 +1,7 @@
-use super::{get_current_shell, get_shells};
+use super::{switch_shell, SwitchTo};
-use nu_engine::current_dir;
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
-use nu_protocol::{Category, Example, PipelineData, ShellError, Signature, Value};
+use nu_protocol::{Category, Example, PipelineData, ShellError, Signature};
 
 /// Source a file for environment variables.
 #[derive(Clone)]
@@ -28,48 +27,7 @@ impl Command for NextShell {
         call: &Call,
         _input: PipelineData,
     ) -> Result<PipelineData, ShellError> {
-        let cwd = current_dir(engine_state, stack)?;
+        switch_shell(engine_state, stack, call, call.head, SwitchTo::Next)
-        let cwd = Value::String {
-            val: cwd.to_string_lossy().to_string(),
-            span: call.head,
-        };
-
-        let shells = get_shells(engine_state, stack, cwd);
-        let mut current_shell = get_current_shell(engine_state, stack);
-        stack.add_env_var(
-            "NUSHELL_LAST_SHELL".into(),
-            Value::Int {
-                val: current_shell as i64,
-                span: call.head,
-            },
-        );
-
-        current_shell += 1;
-
-        if current_shell == shells.len() {
-            current_shell = 0;
-        }
-
-        let new_path = shells[current_shell].clone();
-
-        stack.add_env_var(
-            "NUSHELL_SHELLS".into(),
-            Value::List {
-                vals: shells,
-                span: call.head,
-            },
-        );
-        stack.add_env_var(
-            "NUSHELL_CURRENT_SHELL".into(),
-            Value::Int {
-                val: current_shell as i64,
-                span: call.head,
-            },
-        );
-
-        stack.add_env_var("PWD".into(), new_path);
-
-        Ok(PipelineData::new(call.head))
     }
 
     fn examples(&self) -> Vec<Example> {

crates/nu-command/src/shells/p.rs

@@ -1,8 +1,7 @@
-use super::{get_current_shell, get_shells};
+use super::{switch_shell, SwitchTo};
-use nu_engine::current_dir;
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
-use nu_protocol::{Category, Example, PipelineData, ShellError, Signature, Value};
+use nu_protocol::{Category, Example, PipelineData, ShellError, Signature};
 
 /// Source a file for environment variables.
 #[derive(Clone)]
@@ -28,48 +27,7 @@ impl Command for PrevShell {
         call: &Call,
         _input: PipelineData,
     ) -> Result<PipelineData, ShellError> {
-        let cwd = current_dir(engine_state, stack)?;
+        switch_shell(engine_state, stack, call, call.head, SwitchTo::Prev)
-        let cwd = Value::String {
-            val: cwd.to_string_lossy().to_string(),
-            span: call.head,
-        };
-
-        let shells = get_shells(engine_state, stack, cwd);
-        let mut current_shell = get_current_shell(engine_state, stack);
-        stack.add_env_var(
-            "NUSHELL_LAST_SHELL".into(),
-            Value::Int {
-                val: current_shell as i64,
-                span: call.head,
-            },
-        );
-
-        if current_shell == 0 {
-            current_shell = shells.len() - 1;
-        } else {
-            current_shell -= 1;
-        }
-
-        let new_path = shells[current_shell].clone();
-
-        stack.add_env_var(
-            "NUSHELL_SHELLS".into(),
-            Value::List {
-                vals: shells,
-                span: call.head,
-            },
-        );
-        stack.add_env_var(
-            "NUSHELL_CURRENT_SHELL".into(),
-            Value::Int {
-                val: current_shell as i64,
-                span: call.head,
-            },
-        );
-
-        stack.add_env_var("PWD".into(), new_path);
-
-        Ok(PipelineData::new(call.head))
     }
 
     fn examples(&self) -> Vec<Example> {

crates/nu-command/tests/commands/mod.rs

@@ -41,9 +41,11 @@ mod math;
 mod merge;
 mod mkdir;
 mod move_;
+mod n;
 mod network;
 mod nu_check;
 mod open;
+mod p;
 mod parse;
 mod path;
 mod prepend;

crates/nu-command/tests/commands/n.rs

@@ -0,0 +1,31 @@
+use nu_test_support::{nu, pipeline, playground::Playground};
+
+#[test]
+fn switch_to_next_shell_1() {
+    Playground::setup("switch_to_next_shell_1", |dirs, sandbox| {
+        sandbox.mkdir("foo").mkdir("bar");
+
+        let actual = nu!(
+            cwd: dirs.test(),
+            pipeline(
+            r#"enter foo; enter ../bar; n; g | get active.0"#
+        ));
+
+        assert_eq!(actual.out, "true");
+    })
+}
+
+#[test]
+fn switch_to_next_shell_2() {
+    Playground::setup("switch_to_next_shell_2", |dirs, sandbox| {
+        sandbox.mkdir("foo").mkdir("bar");
+
+        let actual = nu!(
+            cwd: dirs.test(),
+            pipeline(
+            r#"enter foo; enter ../bar; n; n; g | get active.1"#
+        ));
+
+        assert_eq!(actual.out, "true");
+    })
+}

crates/nu-command/tests/commands/p.rs

@@ -0,0 +1,31 @@
+use nu_test_support::{nu, pipeline, playground::Playground};
+
+#[test]
+fn switch_to_prev_shell_1() {
+    Playground::setup("switch_to_next_shell_1", |dirs, sandbox| {
+        sandbox.mkdir("foo").mkdir("bar");
+
+        let actual = nu!(
+            cwd: dirs.test(),
+            pipeline(
+            r#"enter foo; enter ../bar; p; g | get active.1"#
+        ));
+
+        assert_eq!(actual.out, "true");
+    })
+}
+
+#[test]
+fn switch_to_prev_shell_2() {
+    Playground::setup("switch_to_next_shell_2", |dirs, sandbox| {
+        sandbox.mkdir("foo").mkdir("bar");
+
+        let actual = nu!(
+            cwd: dirs.test(),
+            pipeline(
+            r#"enter foo; enter ../bar; p; p; p; g | get active.2"#
+        ));
+
+        assert_eq!(actual.out, "true");
+    })
+}