From 3cb9147f22966b76549db38aab521e081c0c6ef6 Mon Sep 17 00:00:00 2001
From: JT <547158+jntrnr@users.noreply.github.com>
Date: Mon, 29 Aug 2022 14:35:55 +1200
Subject: [PATCH] Make `$` on variable names optional (#6434)

* DRAFT: make var dollar optional

* couple fixes

* fix some tests + cleanup
---
 .../nu-command/src/dataframe/eager/append.rs  |  4 +--
 .../src/dataframe/eager/drop_nulls.rs         |  6 ++---
 crates/nu-command/src/dataframe/eager/take.rs |  4 +--
 crates/nu-parser/src/parser.rs                | 26 +++++++++++++++++++
 crates/nu-protocol/src/engine/engine_state.rs | 12 +++++++--
 5 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/crates/nu-command/src/dataframe/eager/append.rs b/crates/nu-command/src/dataframe/eager/append.rs
index 9d8c71a3f..8b0f21ea3 100644
--- a/crates/nu-command/src/dataframe/eager/append.rs
+++ b/crates/nu-command/src/dataframe/eager/append.rs
@@ -32,7 +32,7 @@ impl Command for AppendDF {
         vec![
             Example {
                 description: "Appends a dataframe as new columns",
-                example: r#"let a = ([[a b]; [1 2] [3 4]] | into df);
+                example: r#"let a = ([['a' 'b']; [1 2] [3 4]] | into df);
     $a | append $a"#,
                 result: Some(
                     NuDataFrame::try_from_columns(vec![
@@ -59,7 +59,7 @@ impl Command for AppendDF {
             },
             Example {
                 description: "Appends a dataframe merging at the end of columns",
-                example: r#"let a = ([[a b]; [1 2] [3 4]] | into df);
+                example: r#"let a = ([['a' 'b']; [1 2] [3 4]] | into df);
     $a | append $a --col"#,
                 result: Some(
                     NuDataFrame::try_from_columns(vec![
diff --git a/crates/nu-command/src/dataframe/eager/drop_nulls.rs b/crates/nu-command/src/dataframe/eager/drop_nulls.rs
index 27bc629cd..f7989960d 100644
--- a/crates/nu-command/src/dataframe/eager/drop_nulls.rs
+++ b/crates/nu-command/src/dataframe/eager/drop_nulls.rs
@@ -36,9 +36,9 @@ impl Command for DropNulls {
         vec![
             Example {
                 description: "drop null values in dataframe",
-                example: r#"let df = ([[a b]; [1 2] [3 0] [1 2]] | into df);
-    let res = ($df.b / $df.b);
-    let a = ($df | with-column $res --name res);
+                example: r#"let my_df = ([[a b]; [1 2] [3 0] [1 2]] | into df);
+    let res = ($my_df.b / $my_df.b);
+    let a = ($my_df | with-column $res --name 'res');
     $a | drop-nulls"#,
                 result: Some(
                     NuDataFrame::try_from_columns(vec![
diff --git a/crates/nu-command/src/dataframe/eager/take.rs b/crates/nu-command/src/dataframe/eager/take.rs
index 68e3d058d..6ee189c9c 100644
--- a/crates/nu-command/src/dataframe/eager/take.rs
+++ b/crates/nu-command/src/dataframe/eager/take.rs
@@ -38,9 +38,9 @@ impl Command for TakeDF {
         vec![
             Example {
                 description: "Takes selected rows from dataframe",
-                example: r#"let df = ([[a b]; [4 1] [5 2] [4 3]] | into df);
+                example: r#"let my_df = ([[a b]; [4 1] [5 2] [4 3]] | into df);
     let indices = ([0 2] | into df);
-    $df | take $indices"#,
+    $my_df | take $indices"#,
                 result: Some(
                     NuDataFrame::try_from_columns(vec![
                         Column::new(
diff --git a/crates/nu-parser/src/parser.rs b/crates/nu-parser/src/parser.rs
index 784454c53..d4a36e815 100644
--- a/crates/nu-parser/src/parser.rs
+++ b/crates/nu-parser/src/parser.rs
@@ -74,6 +74,10 @@ pub fn is_math_expression_like(
         return true;
     }
 
+    if bytes == b"nu" {
+        return false;
+    }
+
     let b = bytes[0];
 
     if b == b'('
@@ -115,6 +119,11 @@ pub fn is_math_expression_like(
         return true;
     }
 
+    let parsed_variable = parse_variable(working_set, span);
+    if parsed_variable.0.is_some() && parsed_variable.1.is_none() {
+        return true;
+    }
+
     false
 }
 
@@ -4008,6 +4017,23 @@ pub fn parse_value(
         return parse_variable_expr(working_set, span);
     }
 
+    let parsed_variable = parse_variable(working_set, span);
+    if parsed_variable.0.is_some() && parsed_variable.1.is_none() {
+        let var_id = parsed_variable
+            .0
+            .expect("internal error: already checked var id exists");
+        return (
+            Expression {
+                expr: Expr::Var(var_id),
+                span,
+                custom_completion: None,
+                ty: working_set.get_variable(var_id).ty.clone(),
+            },
+            None,
+        );
+    }
+    let bytes = working_set.get_span_contents(span);
+
     // Check for reserved keyword values
     match bytes {
         b"true" => {
diff --git a/crates/nu-protocol/src/engine/engine_state.rs b/crates/nu-protocol/src/engine/engine_state.rs
index 827f29f6d..6fbd0e528 100644
--- a/crates/nu-protocol/src/engine/engine_state.rs
+++ b/crates/nu-protocol/src/engine/engine_state.rs
@@ -1508,13 +1508,21 @@ impl<'a> StateWorkingSet<'a> {
     pub fn find_variable(&self, name: &[u8]) -> Option<VarId> {
         let mut removed_overlays = vec![];
 
+        let name = if name.starts_with(&[b'$']) {
+            name.to_vec()
+        } else {
+            let mut new_name = name.to_vec();
+            new_name.insert(0, b'$');
+            new_name
+        };
+
         for scope_frame in self.delta.scope.iter().rev() {
             for overlay_frame in scope_frame
                 .active_overlays(&mut removed_overlays)
                 .iter()
                 .rev()
             {
-                if let Some(var_id) = overlay_frame.vars.get(name) {
+                if let Some(var_id) = overlay_frame.vars.get(&name) {
                     return Some(*var_id);
                 }
             }
@@ -1526,7 +1534,7 @@ impl<'a> StateWorkingSet<'a> {
             .iter()
             .rev()
         {
-            if let Some(var_id) = overlay_frame.vars.get(name) {
+            if let Some(var_id) = overlay_frame.vars.get(&name) {
                 return Some(*var_id);
             }
         }