feat: add `-c/--code` to generate only code (#327)

1 files changed, 18 insertions, 0 deletions

diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index b59b7ad..da98cc6 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -10,10 +10,16 @@ pub use self::prompt_input::*;
 pub use self::render_prompt::render_prompt;
 pub use self::tiktoken::cl100k_base_singleton;
 
+use fancy_regex::Regex;
+use lazy_static::lazy_static;
 use sha2::{Digest, Sha256};
 use std::env;
 use std::process::Command;
 
+lazy_static! {
+    static ref CODE_BLOCK_RE: Regex = Regex::new(r"(?ms)```\w*(.*?)```").unwrap();
+}
+
 pub fn now() -> String {
     let now = chrono::Local::now();
     now.to_rfc3339_opts(chrono::SecondsFormat::Secs, false)
@@ -150,6 +156,18 @@ pub fn run_command(eval_str: &str) -> anyhow::Result<i32> {
     Ok(status.code().unwrap_or_default())
 }
 
+pub fn extract_block(input: &str) -> String {
+    let output: String = CODE_BLOCK_RE
+        .captures_iter(input)
+        .filter_map(|m| {
+            m.ok()
+                .and_then(|cap| cap.get(1))
+                .map(|m| String::from(m.as_str()))
+        })
+        .collect();
+    output.trim().to_string()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;