Allow let/let-env to see custom command input (nushell#854)

crates/nu-command/src/core_commands/let_.rs

@@ -1,4 +1,4 @@
-use nu_engine::eval_expression;
+use nu_engine::eval_expression_with_input;
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
 use nu_protocol::{Category, Example, PipelineData, Signature, SyntaxShape};
@@ -31,7 +31,7 @@ impl Command for Let {
         engine_state: &EngineState,
         stack: &mut Stack,
         call: &Call,
-        _input: PipelineData,
+        input: PipelineData,
     ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
         let var_id = call.positional[0]
             .as_var()
@@ -41,11 +41,11 @@ impl Command for Let {
             .as_keyword()
             .expect("internal error: missing keyword");
 
-        let rhs = eval_expression(engine_state, stack, keyword_expr)?;
+        let rhs = eval_expression_with_input(engine_state, stack, keyword_expr, input, false)?;
 
         //println!("Adding: {:?} to {}", rhs, var_id);
 
-        stack.add_var(var_id, rhs);
+        stack.add_var(var_id, rhs.into_value(call.head));
         Ok(PipelineData::new(call.head))
     }
 

crates/nu-command/src/env/let_env.rs

@@ -1,4 +1,4 @@
-use nu_engine::{current_dir, eval_expression};
+use nu_engine::{current_dir, eval_expression_with_input};
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
 use nu_protocol::{Category, PipelineData, Signature, SyntaxShape, Value};
@@ -31,7 +31,7 @@ impl Command for LetEnv {
         engine_state: &EngineState,
         stack: &mut Stack,
         call: &Call,
-        _input: PipelineData,
+        input: PipelineData,
     ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
         let env_var = call.positional[0]
             .as_string()
@@ -41,7 +41,8 @@ impl Command for LetEnv {
             .as_keyword()
             .expect("internal error: missing keyword");
 
-        let rhs = eval_expression(engine_state, stack, keyword_expr)?;
+        let rhs = eval_expression_with_input(engine_state, stack, keyword_expr, input, false)?
+            .into_value(call.head);
 
         if env_var == "PWD" {
             let cwd = current_dir(engine_state, stack)?;

crates/nu-engine/src/eval.rs

@@ -405,6 +405,48 @@ pub fn eval_expression(
     }
 }
 
+/// Checks the expression to see if it's a internal or external call. If so, passes the input
+/// into the call and gets out the result
+/// Otherwise, invokes the expression
+pub fn eval_expression_with_input(
+    engine_state: &EngineState,
+    stack: &mut Stack,
+    expr: &Expression,
+    mut input: PipelineData,
+    last_expression: bool,
+) -> Result<PipelineData, ShellError> {
+    match expr {
+        Expression {
+            expr: Expr::Call(call),
+            ..
+        } => {
+            input = eval_call(engine_state, stack, call, input)?;
+        }
+        Expression {
+            expr: Expr::ExternalCall(head, args),
+            ..
+        } => {
+            input = eval_external(engine_state, stack, head, args, input, last_expression)?;
+        }
+
+        Expression {
+            expr: Expr::Subexpression(block_id),
+            ..
+        } => {
+            let block = engine_state.get_block(*block_id);
+
+            // FIXME: protect this collect with ctrl-c
+            input = eval_subexpression(engine_state, stack, block, input)?;
+        }
+
+        elem => {
+            input = eval_expression(engine_state, stack, elem)?.into_pipeline_data();
+        }
+    }
+
+    Ok(input)
+}
+
 pub fn eval_block(
     engine_state: &EngineState,
     stack: &mut Stack,
@@ -415,31 +457,13 @@ pub fn eval_block(
     for (stmt_idx, stmt) in block.stmts.iter().enumerate() {
         if let Statement::Pipeline(pipeline) = stmt {
             for (i, elem) in pipeline.expressions.iter().enumerate() {
-                match elem {
-                    Expression {
-                        expr: Expr::Call(call),
-                        ..
-                    } => {
-                        input = eval_call(engine_state, stack, call, input)?;
-                    }
-                    Expression {
-                        expr: Expr::ExternalCall(head, args),
-                        ..
-                    } => {
-                        input = eval_external(
-                            engine_state,
-                            stack,
-                            head,
-                            args,
-                            input,
-                            i == pipeline.expressions.len() - 1,
-                        )?;
-                    }
-
-                    elem => {
-                        input = eval_expression(engine_state, stack, elem)?.into_pipeline_data();
-                    }
-                }
+                input = eval_expression_with_input(
+                    engine_state,
+                    stack,
+                    elem,
+                    input,
+                    i == pipeline.expressions.len() - 1,
+                )?
             }
         }
 
@@ -511,25 +535,8 @@ pub fn eval_subexpression(
 ) -> Result<PipelineData, ShellError> {
     for stmt in block.stmts.iter() {
         if let Statement::Pipeline(pipeline) = stmt {
-            for elem in pipeline.expressions.iter() {
-                match elem {
-                    Expression {
-                        expr: Expr::Call(call),
-                        ..
-                    } => {
-                        input = eval_call(engine_state, stack, call, input)?;
-                    }
-                    Expression {
-                        expr: Expr::ExternalCall(head, args),
-                        ..
-                    } => {
-                        input = eval_external(engine_state, stack, head, args, input, false)?;
-                    }
-
-                    elem => {
-                        input = eval_expression(engine_state, stack, elem)?.into_pipeline_data();
-                    }
-                }
+            for expr in pipeline.expressions.iter() {
+                input = eval_expression_with_input(engine_state, stack, expr, input, false)?
             }
         }
     }

crates/nu-engine/src/lib.rs

@@ -9,5 +9,5 @@ pub use call_ext::CallExt;
 pub use column::get_columns;
 pub use documentation::{generate_docs, get_brief_help, get_documentation, get_full_help};
 pub use env::*;
-pub use eval::{eval_block, eval_expression, eval_operator};
+pub use eval::{eval_block, eval_expression, eval_expression_with_input, eval_operator};
 pub use glob_from::glob_from;

src/tests/test_engine.rs

@@ -166,3 +166,11 @@ fn divide_filesize() -> TestResult {
 fn date_comparison() -> TestResult {
     run_test(r#"(date now) < ((date now) + 2min)"#, "true")
 }
+
+#[test]
+fn let_sees_input() -> TestResult {
+    run_test(
+        r#"def c [] { let x = str length; $x }; "hello world" | c"#,
+        "11",
+    )
+}