From 5f14faf4b42434da69880abb3251dc29dc302a31 Mon Sep 17 00:00:00 2001
From: JT <jonathan.d.turner@gmail.com>
Date: Tue, 12 Oct 2021 17:49:17 +1300
Subject: [PATCH] Custom switch support

---
 crates/nu-engine/src/eval.rs           | 19 ++++++++++--
 crates/nu-parser/src/parser.rs         | 42 ++++++++++++++++++++++----
 crates/nu-protocol/src/syntax_shape.rs |  4 +++
 src/tests.rs                           | 32 ++++++++++++++++++++
 4 files changed, 88 insertions(+), 9 deletions(-)

diff --git a/crates/nu-engine/src/eval.rs b/crates/nu-engine/src/eval.rs
index 4aa8465492..cfcc2ba85a 100644
--- a/crates/nu-engine/src/eval.rs
+++ b/crates/nu-engine/src/eval.rs
@@ -62,11 +62,13 @@ fn eval_call(context: &EvaluationContext, call: &Call, input: Value) -> Result<V
         }
 
         for named in decl.signature().named {
+            let var_id = named
+                .var_id
+                .expect("internal error: all custom parameters must have var_ids");
+
+            let mut found = false;
             for call_named in &call.named {
                 if call_named.0.item == named.long {
-                    let var_id = named
-                        .var_id
-                        .expect("internal error: all custom parameters must have var_ids");
                     if let Some(arg) = &call_named.1 {
                         let result = eval_expression(&state, arg)?;
 
@@ -80,8 +82,19 @@ fn eval_call(context: &EvaluationContext, call: &Call, input: Value) -> Result<V
                             },
                         )
                     }
+                    found = true;
                 }
             }
+
+            if !found && named.arg.is_none() {
+                state.add_var(
+                    var_id,
+                    Value::Bool {
+                        val: false,
+                        span: call.head,
+                    },
+                )
+            }
         }
         let engine_state = state.engine_state.borrow();
         let block = engine_state.get_block(block_id);
diff --git a/crates/nu-parser/src/parser.rs b/crates/nu-parser/src/parser.rs
index da71974f05..2f8f94597d 100644
--- a/crates/nu-parser/src/parser.rs
+++ b/crates/nu-parser/src/parser.rs
@@ -1664,6 +1664,7 @@ pub fn parse_shape_name(
         b"variable" => SyntaxShape::Variable,
         b"signature" => SyntaxShape::Signature,
         b"expr" => SyntaxShape::Expression,
+        b"bool" => SyntaxShape::Boolean,
         _ => return (SyntaxShape::Any, Some(ParseError::UnknownType(span))),
     };
 
@@ -1671,10 +1672,17 @@ pub fn parse_shape_name(
 }
 
 pub fn parse_type(_working_set: &StateWorkingSet, bytes: &[u8]) -> Type {
-    if bytes == b"int" {
-        Type::Int
-    } else {
-        Type::Unknown
+    match bytes {
+        b"int" => Type::Int,
+        b"bool" => Type::Bool,
+        b"string" => Type::String,
+        b"block" => Type::Block,
+        b"float" => Type::Float,
+        b"filesize" => Type::Filesize,
+        b"binary" => Type::Binary,
+        b"date" => Type::Date,
+
+        _ => Type::Unknown,
     }
 }
 
@@ -2153,8 +2161,11 @@ pub fn parse_signature_helper(
                                         *shape = syntax_shape;
                                     }
                                     Arg::Flag(Flag { arg, var_id, .. }) => {
-                                        working_set.set_variable_type(var_id.expect("internal error: all custom parameters must have var_ids"), syntax_shape.to_type());
-                                        *arg = Some(syntax_shape)
+                                        // Flags with a boolean type are just present/not-present switches
+                                        if syntax_shape != SyntaxShape::Boolean {
+                                            working_set.set_variable_type(var_id.expect("internal error: all custom parameters must have var_ids"), syntax_shape.to_type());
+                                            *arg = Some(syntax_shape)
+                                        }
                                     }
                                 }
                             }
@@ -2689,6 +2700,25 @@ pub fn parse_value(
                 error,
             )
         }
+        SyntaxShape::Boolean => {
+            // Redundant, though we catch bad boolean parses here
+            if bytes == b"$true" || bytes == b"$false" {
+                (
+                    Expression {
+                        expr: Expr::Bool(true),
+                        span,
+                        ty: Type::Bool,
+                        custom_completion: None,
+                    },
+                    None,
+                )
+            } else {
+                (
+                    garbage(span),
+                    Some(ParseError::Expected("bool".into(), span)),
+                )
+            }
+        }
         SyntaxShape::Any => {
             if bytes.starts_with(b"[") {
                 parse_value(working_set, span, &SyntaxShape::Table)
diff --git a/crates/nu-protocol/src/syntax_shape.rs b/crates/nu-protocol/src/syntax_shape.rs
index d1202528ac..7f058e18e1 100644
--- a/crates/nu-protocol/src/syntax_shape.rs
+++ b/crates/nu-protocol/src/syntax_shape.rs
@@ -73,6 +73,9 @@ pub enum SyntaxShape {
     /// A general expression, eg `1 + 2` or `foo --bar`
     Expression,
 
+    /// A boolean value
+    Boolean,
+
     /// A custom shape with custom completion logic
     Custom(Box<SyntaxShape>, String),
 }
@@ -102,6 +105,7 @@ impl SyntaxShape {
             SyntaxShape::Operator => Type::Unknown,
             SyntaxShape::Range => Type::Unknown,
             SyntaxShape::RowCondition => Type::Bool,
+            SyntaxShape::Boolean => Type::Bool,
             SyntaxShape::Signature => Type::Unknown,
             SyntaxShape::String => Type::String,
             SyntaxShape::Table => Type::List(Box::new(Type::Unknown)), // FIXME
diff --git a/src/tests.rs b/src/tests.rs
index 8e2ecd2742..13c1f8a55f 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -727,3 +727,35 @@ fn flag_param_value() -> TestResult {
 fn do_rest_args() -> TestResult {
     run_test(r#"(do { |...rest| $rest } 1 2).1 + 10"#, "12")
 }
+
+#[test]
+fn custom_switch1() -> TestResult {
+    run_test(
+        r#"def florb [ --dry-run: bool ] { if ($dry-run) { "foo" } else { "bar" } }; florb --dry-run"#,
+        "foo",
+    )
+}
+
+#[test]
+fn custom_switch2() -> TestResult {
+    run_test(
+        r#"def florb [ --dry-run: bool ] { if ($dry-run) { "foo" } else { "bar" } }; florb"#,
+        "bar",
+    )
+}
+
+#[test]
+fn custom_switch3() -> TestResult {
+    run_test(
+        r#"def florb [ --dry-run ] { if ($dry-run) { "foo" } else { "bar" } }; florb --dry-run"#,
+        "foo",
+    )
+}
+
+#[test]
+fn custom_switch4() -> TestResult {
+    run_test(
+        r#"def florb [ --dry-run ] { if ($dry-run) { "foo" } else { "bar" } }; florb"#,
+        "bar",
+    )
+}