fix: implement date and time in prompts (#298)

reubeno · Dec 26, 2024 · eacb489 · eacb489
1 parent 0b59daf
commit eacb489
Show file tree

Hide file tree

Showing 6 changed files with 155 additions and 8 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/brush-core/Cargo.toml b/brush-core/Cargo.toml
@@ -22,6 +22,7 @@ async-trait = "0.1.83"
 brush-parser = { version = "^0.2.11", path = "../brush-parser" }
 cached = "0.54.0"
 cfg-if = "1.0.0"
+chrono = "0.4.39"
 clap = { version = "4.5.21", features = ["derive", "wrap_help"] }
 fancy-regex = "0.14.0"
 futures = "0.3.31"

diff --git a/brush-core/benches/shell.rs b/brush-core/benches/shell.rs
@@ -76,6 +76,13 @@ mod unix {
         c.bench_function("expand_one_string", |b| {
             b.iter(|| black_box(expand_one_string()));
         });
+        c.bench_function("for_loop", |b| {
+            b.to_async(tokio()).iter(|| {
+                black_box(run_one_command(
+                    "for ((i = 0; i < 100000; i++)); do :; done",
+                ))
+            });
+        });
     }
 }
 

diff --git a/brush-core/src/prompt.rs b/brush-core/src/prompt.rs
@@ -52,7 +52,9 @@ pub(crate) fn format_prompt_piece(
             tilde_replaced,
             basename,
         } => format_current_working_directory(shell, tilde_replaced, basename),
-        brush_parser::prompt::PromptPiece::Date(_) => return error::unimp("prompt: date"),
+        brush_parser::prompt::PromptPiece::Date(format) => {
+            format_date(&chrono::Local::now(), &format)
+        }
         brush_parser::prompt::PromptPiece::DollarOrPound => {
             if users::is_root() {
                 "#".to_owned()
@@ -77,9 +79,7 @@ pub(crate) fn format_prompt_piece(
             hn
         }
         brush_parser::prompt::PromptPiece::Newline => "\n".to_owned(),
-        brush_parser::prompt::PromptPiece::NumberOfManagedJobs => {
-            return error::unimp("prompt: number of managed jobs")
-        }
+        brush_parser::prompt::PromptPiece::NumberOfManagedJobs => shell.jobs.jobs.len().to_string(),
         brush_parser::prompt::PromptPiece::ShellBaseName => {
             if let Some(shell_name) = &shell.shell_name {
                 Path::new(shell_name)
@@ -100,7 +100,9 @@ pub(crate) fn format_prompt_piece(
         brush_parser::prompt::PromptPiece::TerminalDeviceBaseName => {
             return error::unimp("prompt: terminal device base name")
         }
-        brush_parser::prompt::PromptPiece::Time(_) => return error::unimp("prompt: time"),
+        brush_parser::prompt::PromptPiece::Time(time_fmt) => {
+            format_time(&chrono::Local::now(), &time_fmt)
+        }
     };
 
     Ok(formatted)
@@ -125,3 +127,101 @@ fn format_current_working_directory(shell: &Shell, tilde_replaced: bool, basenam
 
     working_dir_str
 }
+
+fn format_time<Tz: chrono::TimeZone>(
+    datetime: &chrono::DateTime<Tz>,
+    format: &brush_parser::prompt::PromptTimeFormat,
+) -> String
+where
+    Tz::Offset: std::fmt::Display,
+{
+    let formatted = match format {
+        brush_parser::prompt::PromptTimeFormat::TwelveHourAM => datetime.format("%I:%M %p"),
+        brush_parser::prompt::PromptTimeFormat::TwelveHourHHMMSS => datetime.format("%I:%M:%S"),
+        brush_parser::prompt::PromptTimeFormat::TwentyFourHourHHMMSS => datetime.format("%H:%M:%S"),
+    };
+
+    formatted.to_string()
+}
+
+fn format_date<Tz: chrono::TimeZone>(
+    datetime: &chrono::DateTime<Tz>,
+    format: &brush_parser::prompt::PromptDateFormat,
+) -> String
+where
+    Tz::Offset: std::fmt::Display,
+{
+    match format {
+        brush_parser::prompt::PromptDateFormat::WeekdayMonthDate => {
+            datetime.format("%a %b %d").to_string()
+        }
+        brush_parser::prompt::PromptDateFormat::Custom(fmt) => {
+            let fmt_items = chrono::format::StrftimeItems::new(fmt);
+            datetime.format_with_items(fmt_items).to_string()
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_format_time() {
+        // Create a well-known test date/time.
+        let dt = chrono::DateTime::parse_from_rfc3339("2024-12-25T13:34:56.789Z").unwrap();
+
+        assert_eq!(
+            format_time(&dt, &brush_parser::prompt::PromptTimeFormat::TwelveHourAM),
+            "01:34 PM"
+        );
+
+        assert_eq!(
+            format_time(
+                &dt,
+                &brush_parser::prompt::PromptTimeFormat::TwentyFourHourHHMMSS
+            ),
+            "13:34:56"
+        );
+
+        assert_eq!(
+            format_time(
+                &dt,
+                &brush_parser::prompt::PromptTimeFormat::TwelveHourHHMMSS
+            ),
+            "01:34:56"
+        );
+    }
+
+    #[test]
+    fn test_format_date() {
+        // Create a well-known test date/time.
+        let dt = chrono::DateTime::parse_from_rfc3339("2024-12-25T12:34:56.789Z").unwrap();
+
+        assert_eq!(
+            format_date(
+                &dt,
+                &brush_parser::prompt::PromptDateFormat::WeekdayMonthDate
+            ),
+            "Wed Dec 25"
+        );
+
+        assert_eq!(
+            format_date(
+                &dt,
+                &brush_parser::prompt::PromptDateFormat::Custom(String::from("%Y-%m-%d"))
+            ),
+            "2024-12-25"
+        );
+
+        assert_eq!(
+            format_date(
+                &dt,
+                &brush_parser::prompt::PromptDateFormat::Custom(String::from(
+                    "%Y-%m-%d %H:%M:%S.%f"
+                ))
+            ),
+            "2024-12-25 12:34:56.789000000"
+        );
+    }
+}
diff --git a/brush-parser/src/prompt.rs b/brush-parser/src/prompt.rs
@@ -123,7 +123,7 @@ peg::parser! {
             s:$((!special_sequence() [c])+) { PromptPiece::Literal(s.to_owned()) }
 
         rule date_format() -> String =
-            s:$(!"}" [c]+) { s.to_owned() }
+            s:$([c if c != '}']*) { s.to_owned() }
 
         rule octal_number() -> u32 =
             s:$(['0'..='9']*<3,3>) {? u32::from_str_radix(s, 8).or(Err("invalid octal number")) }

diff --git a/brush-shell/tests/cases/prompt.yaml b/brush-shell/tests/cases/prompt.yaml
@@ -54,3 +54,39 @@ cases:
 
       prompt='\V'
       [[ "${prompt@P}" == ^\d+\.\d+\.\d+$ ]] && echo "Release is correct"
+
+  - name: "Simple date"
+    stdin: |
+      prompt='\d'
+      [[ "${prompt@P}" == $(date +'%a %b %d') ]] && echo '\d date matches'
+
+  - name: "Date format with plain string"
+    stdin: |
+      prompt='\D{something}'
+      echo "Date format with plain string: '${prompt@P}'"
+
+  - name: "Date format with year"
+    stdin: |
+      prompt='\d{%Y}'
+      echo "Date format with year: '${prompt@P}'"
+
+  - name: "Time format: @"
+    stdin: |
+      prompt='\@'
+      expanded=${prompt@P}
+      roundtripped=$(date --date="${expanded}" +'%I:%M %p')
+      [[ ${expanded} == ${roundtripped} ]] && echo "Time matches"
+
+  - name: "Time format: T"
+    stdin: |
+      prompt='\T'
+      expanded=${prompt@P}
+      roundtripped=$(date --date="${expanded}" +'%I:%M:%S')
+      [[ ${expanded} == ${roundtripped} ]] && echo "Time matches"
+
+  - name: "Time format: t"
+    stdin: |
+      prompt='\t'
+      expanded=${prompt@P}
+      roundtripped=$(date --date="${expanded}" +'%H:%M:%S')
+      [[ ${expanded} == ${roundtripped} ]] && echo "Time matches"