feat: enhence glob for `--file` and rag docs (#1230)

1 files changed, 101 insertions, 28 deletions

diff --git a/src/utils/path.rs b/src/utils/path.rs
index 20a365e..67af0a4 100644
--- a/src/utils/path.rs
+++ b/src/utils/path.rs
@@ -36,16 +36,12 @@ pub async fn expand_glob_paths<T: AsRef<str>>(
 ) -> Result<IndexSet<String>> {
     let mut new_paths = IndexSet::new();
     for path in paths {
-        let (path_str, suffixes) = parse_glob(path.as_ref())?;
-        let suffixes = if suffixes.is_empty() {
-            None
-        } else {
-            Some(&suffixes)
-        };
+        let (path_str, suffixes, current_only) = parse_glob(path.as_ref())?;
         list_files(
             &mut new_paths,
             Path::new(&path_str),
-            suffixes,
+            suffixes.as_ref(),
+            current_only,
             bail_non_exist,
         )
         .await?;
@@ -90,30 +86,68 @@ pub fn resolve_home_dir(path: &str) -> String {
     path
 }
 
-fn parse_glob(path_str: &str) -> Result<(String, Vec<String>)> {
-    if let Some(start) = path_str.find("/**/*.").or_else(|| path_str.find(r"\**\*.")) {
-        let base_path = path_str[..start].to_string();
-        if let Some(curly_brace_end) = path_str[start..].find('}') {
+fn parse_glob(path_str: &str) -> Result<(String, Option<Vec<String>>, bool)> {
+    let glob_result =
+        if let Some(start) = path_str.find("/**/*.").or_else(|| path_str.find(r"\**\*.")) {
+            Some((start, 6, false))
+        } else if let Some(start) = path_str.find("**/*.").or_else(|| path_str.find(r"**\*.")) {
+            if start == 0 {
+                Some((start, 5, false))
+            } else {
+                None
+            }
+        } else if let Some(start) = path_str.find("/*.").or_else(|| path_str.find(r"\*.")) {
+            Some((start, 3, true))
+        } else if let Some(start) = path_str.find("*.") {
+            if start == 0 {
+                Some((start, 2, true))
+            } else {
+                None
+            }
+        } else {
+            None
+        };
+    if let Some((start, offset, current_only)) = glob_result {
+        let mut base_path = path_str[..start].to_string();
+        if base_path.is_empty() {
+            base_path = if path_str
+                .chars()
+                .next()
+                .map(|v| v == '/')
+                .unwrap_or_default()
+            {
+                "/"
+            } else {
+                "."
+            }
+            .into();
+        }
+
+        let extensions = if let Some(curly_brace_end) = path_str[start..].find('}') {
             let end = start + curly_brace_end;
-            let extensions_str = &path_str[start + 6..end + 1];
-            let extensions = if extensions_str.starts_with('{') && extensions_str.ends_with('}') {
+            let extensions_str = &path_str[start + offset..end + 1];
+            if extensions_str.starts_with('{') && extensions_str.ends_with('}') {
                 extensions_str[1..extensions_str.len() - 1]
                     .split(',')
                     .map(|s| s.to_string())
                     .collect::<Vec<String>>()
             } else {
                 bail!("Invalid path '{path_str}'");
-            };
-            Ok((base_path, extensions))
+            }
         } else {
-            let extensions_str = &path_str[start + 6..];
-            let extensions = vec![extensions_str.to_string()];
-            Ok((base_path, extensions))
-        }
+            let extensions_str = &path_str[start + offset..];
+            vec![extensions_str.to_string()]
+        };
+        let extensions = if extensions.is_empty() {
+            None
+        } else {
+            Some(extensions)
+        };
+        Ok((base_path, extensions, current_only))
     } else if path_str.ends_with("/**") || path_str.ends_with(r"\**") {
-        Ok((path_str[0..path_str.len() - 3].to_string(), vec![]))
+        Ok((path_str[0..path_str.len() - 3].to_string(), None, false))
     } else {
-        Ok((path_str.to_string(), vec![]))
+        Ok((path_str.to_string(), None, false))
     }
 }
 
@@ -122,6 +156,7 @@ async fn list_files(
     files: &mut IndexSet<String>,
     entry_path: &Path,
     suffixes: Option<&Vec<String>>,
+    current_only: bool,
     bail_non_exist: bool,
 ) -> Result<()> {
     if !entry_path.exists() {
@@ -136,7 +171,9 @@ async fn list_files(
         while let Some(entry) = reader.next_entry().await? {
             let path = entry.path();
             if path.is_dir() {
-                list_files(files, &path, suffixes, bail_non_exist).await?;
+                if !current_only {
+                    list_files(files, &path, suffixes, current_only, bail_non_exist).await?;
+                }
             } else {
                 add_file(files, suffixes, &path);
             }
@@ -174,23 +211,59 @@ mod tests {
 
     #[test]
     fn test_parse_glob() {
-        assert_eq!(parse_glob("dir").unwrap(), ("dir".into(), vec![]));
-        assert_eq!(parse_glob("dir/**").unwrap(), ("dir".into(), vec![]));
+        assert_eq!(parse_glob("dir").unwrap(), ("dir".into(), None, false));
+        assert_eq!(parse_glob("dir/**").unwrap(), ("dir".into(), None, false));
         assert_eq!(
             parse_glob("dir/file.md").unwrap(),
-            ("dir/file.md".into(), vec![])
+            ("dir/file.md".into(), None, false)
+        );
+        assert_eq!(
+            parse_glob("**/*.md").unwrap(),
+            (".".into(), Some(vec!["md".into()]), false)
+        );
+        assert_eq!(
+            parse_glob("/**/*.md").unwrap(),
+            ("/".into(), Some(vec!["md".into()]), false)
         );
         assert_eq!(
             parse_glob("dir/**/*.md").unwrap(),
-            ("dir".into(), vec!["md".into()])
+            ("dir".into(), Some(vec!["md".into()]), false)
         );
         assert_eq!(
             parse_glob("dir/**/*.{md,txt}").unwrap(),
-            ("dir".into(), vec!["md".into(), "txt".into()])
+            ("dir".into(), Some(vec!["md".into(), "txt".into()]), false)
         );
         assert_eq!(
             parse_glob("C:\\dir\\**\\*.{md,txt}").unwrap(),
-            ("C:\\dir".into(), vec!["md".into(), "txt".into()])
+            (
+                "C:\\dir".into(),
+                Some(vec!["md".into(), "txt".into()]),
+                false
+            )
+        );
+        assert_eq!(
+            parse_glob("*.md").unwrap(),
+            (".".into(), Some(vec!["md".into()]), true)
+        );
+        assert_eq!(
+            parse_glob("/*.md").unwrap(),
+            ("/".into(), Some(vec!["md".into()]), true)
+        );
+        assert_eq!(
+            parse_glob("dir/*.md").unwrap(),
+            ("dir".into(), Some(vec!["md".into()]), true)
+        );
+        assert_eq!(
+            parse_glob("dir/*.{md,txt}").unwrap(),
+            ("dir".into(), Some(vec!["md".into(), "txt".into()]), true)
+        );
+        assert_eq!(
+            parse_glob("C:\\dir\\*.{md,txt}").unwrap(),
+            (
+                "C:\\dir".into(),
+                Some(vec!["md".into(), "txt".into()]),
+                true
+            )
         );
     }
 }