From 058299e500af86adbc065eced542e108ae325458 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Thu, 9 May 2024 22:59:06 +0000
Subject: [PATCH] refactor: refine detect code block

---
 src/config/role.rs     |  2 ++
 src/render/markdown.rs | 12 +++++++++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/config/role.rs b/src/config/role.rs
index b96e704..39b777c 100644
--- a/src/config/role.rs
+++ b/src/config/role.rs
@@ -90,9 +90,11 @@ APPLY MARKDOWN formatting when possible."#
 ### INPUT:
 async sleep in js
 ### OUTPUT:
+```javascript
 async function timeout(ms) {
   return new Promise(resolve => setTimeout(resolve, ms));
 }
+```
 "#
             .into(),
             temperature: None,
diff --git a/src/render/markdown.rs b/src/render/markdown.rs
index 3f68ea2..400cfe2 100644
--- a/src/render/markdown.rs
+++ b/src/render/markdown.rs
@@ -280,13 +280,14 @@ fn blend_fg_color(fg: SyntectColor, bg: SyntectColor) -> SyntectColor {
 }
 
 fn detect_code_block(line: &str) -> Option<String> {
+    let line = line.trim_start();
     if !line.starts_with("```") {
         return None;
     }
     let lang = line
         .chars()
         .skip(3)
-        .take_while(|v| v.is_alphanumeric())
+        .take_while(|v| !v.is_whitespace())
         .collect();
     Some(lang)
 }
@@ -380,4 +381,13 @@ std::error::Error>> {
         let output = render.render(TEXT);
         assert_eq!(TEXT_WRAP_ALL, output);
     }
+
+    #[test]
+    fn test_detect_code_block() {
+        assert_eq!(detect_code_block("```rust"), Some("rust".into()));
+        assert_eq!(detect_code_block("```c++"), Some("c++".into()));
+        assert_eq!(detect_code_block("  ```rust"), Some("rust".into()));
+        assert_eq!(detect_code_block("```"), Some("".into()));
+        assert_eq!(detect_code_block("``rust"), None);
+    }
 }