Add support for def-env and export def-env (#887)

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

@@ -0,0 +1,38 @@
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{Category, PipelineData, Signature, SyntaxShape};
+
+#[derive(Clone)]
+pub struct DefEnv;
+
+impl Command for DefEnv {
+    fn name(&self) -> &str {
+        "def-env"
+    }
+
+    fn usage(&self) -> &str {
+        "Define a custom command, which participates in the caller environment"
+    }
+
+    fn signature(&self) -> nu_protocol::Signature {
+        Signature::build("def-env")
+            .required("def_name", SyntaxShape::String, "definition name")
+            .required("params", SyntaxShape::Signature, "parameters")
+            .required(
+                "block",
+                SyntaxShape::Block(Some(vec![])),
+                "body of the definition",
+            )
+            .category(Category::Core)
+    }
+
+    fn run(
+        &self,
+        _engine_state: &EngineState,
+        _stack: &mut Stack,
+        call: &Call,
+        _input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        Ok(PipelineData::new(call.head))
+    }
+}

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

@@ -0,0 +1,38 @@
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{Category, PipelineData, Signature, SyntaxShape};
+
+#[derive(Clone)]
+pub struct ExportDefEnv;
+
+impl Command for ExportDefEnv {
+    fn name(&self) -> &str {
+        "export def-env"
+    }
+
+    fn usage(&self) -> &str {
+        "Define a custom command that participates in the environment and export it from a module"
+    }
+
+    fn signature(&self) -> nu_protocol::Signature {
+        Signature::build("export def-env")
+            .required("name", SyntaxShape::String, "definition name")
+            .required("params", SyntaxShape::Signature, "parameters")
+            .required(
+                "block",
+                SyntaxShape::Block(Some(vec![])),
+                "body of the definition",
+            )
+            .category(Category::Core)
+    }
+
+    fn run(
+        &self,
+        _engine_state: &EngineState,
+        _stack: &mut Stack,
+        call: &Call,
+        _input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        Ok(PipelineData::new(call.head))
+    }
+}

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

@@ -1,11 +1,13 @@
 mod alias;
 mod debug;
 mod def;
+mod def_env;
 mod describe;
 mod do_;
 mod echo;
 mod export;
 mod export_def;
+mod export_def_env;
 mod export_env;
 mod for_;
 mod help;
@@ -24,11 +26,13 @@ mod version;
 pub use alias::Alias;
 pub use debug::Debug;
 pub use def::Def;
+pub use def_env::DefEnv;
 pub use describe::Describe;
 pub use do_::Do;
 pub use echo::Echo;
 pub use export::ExportCommand;
 pub use export_def::ExportDef;
+pub use export_def_env::ExportDefEnv;
 pub use export_env::ExportEnv;
 pub use for_::For;
 pub use help::Help;

crates/nu-command/src/default_context.rs

@@ -28,11 +28,13 @@ pub fn create_default_context(cwd: impl AsRef<Path>) -> EngineState {
             Alias,
             Debug,
             Def,
+            DefEnv,
             Describe,
             Do,
             Echo,
             ExportCommand,
             ExportDef,
+            ExportDefEnv,
             ExportEnv,
             For,
             Help,

crates/nu-engine/src/eval.rs

@@ -133,7 +133,15 @@ fn eval_call(
                 }
             }
         }
-        eval_block(engine_state, &mut callee_stack, block, input)
+        let result = eval_block(engine_state, &mut callee_stack, block, input);
+        if block.redirect_env {
+            for env_vars in callee_stack.env_vars {
+                for (var, value) in env_vars {
+                    caller_stack.add_env_var(var, value)
+                }
+            }
+        }
+        result
     } else {
         // We pass caller_stack here with the knowledge that internal commands
         // are going to be specifically looking for global state in the stack

crates/nu-parser/src/parse_keywords.rs

@@ -30,7 +30,7 @@ pub fn parse_def_predecl(working_set: &mut StateWorkingSet, spans: &[Span]) -> O
         (name, spans)
     };
 
-    if name == b"def" && spans.len() >= 4 {
+    if (name == b"def" || name == b"def-env") && spans.len() >= 4 {
         let (name_expr, ..) = parse_string(working_set, spans[1]);
         let name = name_expr.as_string();
 
@@ -235,7 +235,9 @@ pub fn parse_def(
 
     // Checking that the function is used with the correct name
     // Maybe this is not necessary but it is a sanity check
-    if working_set.get_span_contents(spans[0]) != b"def" {
+
+    let def_call = working_set.get_span_contents(spans[0]).to_vec();
+    if def_call != b"def" && def_call != b"def-env" {
         return (
             garbage_statement(spans),
             Some(ParseError::UnknownState(
@@ -248,7 +250,7 @@ pub fn parse_def(
     // Parsing the spans and checking that they match the register signature
     // Using a parsed call makes more sense than checking for how many spans are in the call
     // Also, by creating a call, it can be checked if it matches the declaration signature
-    let (call, call_span) = match working_set.find_decl(b"def") {
+    let (call, call_span) = match working_set.find_decl(&def_call) {
         None => {
             return (
                 garbage_statement(spans),
@@ -333,6 +335,7 @@ pub fn parse_def(
             let captures = find_captures_in_block(working_set, block, &mut seen, &mut seen_decls);
 
             let mut block = working_set.get_block_mut(block_id);
+            block.redirect_env = def_call == b"def-env";
             block.captures = captures;
         } else {
             error = error.or_else(|| {
@@ -544,6 +547,74 @@ pub fn parse_export(
                     None
                 }
             }
+            b"def-env" => {
+                let lite_command = LiteCommand {
+                    comments: lite_command.comments.clone(),
+                    parts: spans[1..].to_vec(),
+                };
+                let (stmt, err) = parse_def(working_set, &lite_command);
+                error = error.or(err);
+
+                let export_def_decl_id = if let Some(id) = working_set.find_decl(b"export def-env")
+                {
+                    id
+                } else {
+                    return (
+                        garbage_statement(spans),
+                        None,
+                        Some(ParseError::InternalError(
+                            "missing 'export def-env' command".into(),
+                            export_span,
+                        )),
+                    );
+                };
+
+                // Trying to warp the 'def' call into the 'export def' in a very clumsy way
+                if let Statement::Pipeline(ref pipe) = stmt {
+                    if let Some(Expression {
+                        expr: Expr::Call(ref def_call),
+                        ..
+                    }) = pipe.expressions.get(0)
+                    {
+                        call = def_call.clone();
+
+                        call.head = span(&spans[0..=1]);
+                        call.decl_id = export_def_decl_id;
+                    } else {
+                        error = error.or_else(|| {
+                            Some(ParseError::InternalError(
+                                "unexpected output from parsing a definition".into(),
+                                span(&spans[1..]),
+                            ))
+                        });
+                    }
+                } else {
+                    error = error.or_else(|| {
+                        Some(ParseError::InternalError(
+                            "unexpected output from parsing a definition".into(),
+                            span(&spans[1..]),
+                        ))
+                    });
+                };
+
+                if error.is_none() {
+                    let decl_name = working_set.get_span_contents(spans[2]);
+                    let decl_name = trim_quotes(decl_name);
+                    if let Some(decl_id) = working_set.find_decl(decl_name) {
+                        Some(Exportable::Decl(decl_id))
+                    } else {
+                        error = error.or_else(|| {
+                            Some(ParseError::InternalError(
+                                "failed to find added declaration".into(),
+                                span(&spans[1..]),
+                            ))
+                        });
+                        None
+                    }
+                } else {
+                    None
+                }
+            }
             b"env" => {
                 if let Some(id) = working_set.find_decl(b"export env") {
                     call.decl_id = id;
@@ -700,7 +771,7 @@ pub fn parse_module_block(
                 let name = working_set.get_span_contents(pipeline.commands[0].parts[0]);
 
                 let (stmt, err) = match name {
-                    b"def" => {
+                    b"def" | b"def-env" => {
                         let (stmt, err) = parse_def(working_set, &pipeline.commands[0]);
 
                         (stmt, err)

crates/nu-parser/src/parser.rs

@@ -3517,7 +3517,7 @@ pub fn parse_statement(
     let name = working_set.get_span_contents(lite_command.parts[0]);
 
     match name {
-        b"def" => parse_def(working_set, lite_command),
+        b"def" | b"def-env" => parse_def(working_set, lite_command),
         b"let" => parse_let(working_set, &lite_command.parts),
         b"for" => {
             let (expr, err) = parse_for(working_set, &lite_command.parts);

crates/nu-protocol/src/ast/block.rs

@@ -9,6 +9,7 @@ pub struct Block {
     pub signature: Box<Signature>,
     pub stmts: Vec<Statement>,
     pub captures: Vec<VarId>,
+    pub redirect_env: bool,
 }
 
 impl Block {
@@ -47,6 +48,7 @@ impl Block {
             signature: Box::new(Signature::new("")),
             stmts: vec![],
             captures: vec![],
+            redirect_env: false,
         }
     }
 }
@@ -60,6 +62,7 @@ where
             signature: Box::new(Signature::new("")),
             stmts: stmts.collect(),
             captures: vec![],
+            redirect_env: false,
         }
     }
 }

src/tests/test_engine.rs

@@ -190,3 +190,27 @@ fn let_sees_in_variable2() -> TestResult {
         "3",
     )
 }
+
+#[test]
+fn def_env() -> TestResult {
+    run_test(
+        r#"def-env bob [] { let-env BAR = BAZ }; bob; $env.BAR"#,
+        "BAZ",
+    )
+}
+
+#[test]
+fn not_def_env() -> TestResult {
+    fail_test(
+        r#"def bob [] { let-env BAR = BAZ }; bob; $env.BAR"#,
+        "did you mean",
+    )
+}
+
+#[test]
+fn export_def_env() -> TestResult {
+    run_test(
+        r#"module foo { export def-env bob [] { let-env BAR = BAZ } }; use foo bob; bob; $env.BAR"#,
+        "BAZ",
+    )
+}