feat: use ripgrep for scanning if available

Author: CyanVoxel

src/tagstudio/core/utils/refresh_dir.py

@@ -1,3 +1,4 @@
+import shutil
 from collections.abc import Iterator
 from dataclasses import dataclass, field
 from datetime import datetime as dt
@@ -10,6 +11,7 @@
 from tagstudio.core.library.alchemy.library import Library
 from tagstudio.core.library.alchemy.models import Entry
 from tagstudio.core.library.ignore import PATH_GLOB_FLAGS, Ignore
+from tagstudio.qt.helpers.silent_popen import silent_run
 
 logger = structlog.get_logger(__name__)
 
@@ -41,19 +43,120 @@ def save_new_files(self):
 
         yield
 
-    def refresh_dir(self, library_dir: Path) -> Iterator[int]:
-        """Scan a directory for files, and add those relative filenames to internal variables."""
+    def refresh_dir(self, library_dir: Path, force_internal_tools: bool = False) -> Iterator[int]:
+        """Scan a directory for files, and add those relative filenames to internal variables.
+
+        Args:
+            library_dir (Path): The library directory.
+            force_internal_tools (bool): Option to force the use of internal tools for scanning
+                (i.e. wcmatch) instead of using tools found on the system (i.e. ripgrep).
+        """
         if self.library.library_dir is None:
             raise ValueError("No library directory set.")
 
+        ignore_patterns = Ignore.get_patterns(library_dir)
+
+        if force_internal_tools:
+            return self.__wc_add(library_dir, ignore_patterns)
+
+        dir_list: list[str] | None = self.__get_dir_list(library_dir, ignore_patterns)
+
+        # Use ripgrep if it was found and working, else fallback to wcmatch.
+        if dir_list is not None:
+            return self.__rg_add(library_dir, dir_list)
+        else:
+            return self.__wc_add(library_dir, ignore_patterns)
+
+    def __get_dir_list(self, library_dir: Path, ignore_patterns: list[str]) -> list[str] | None:
+        """Use ripgrep to return a list of matched directories and files.
+
+        Return `None` if ripgrep not found on system.
+        """
+        rg_path = shutil.which("rg")
+        # Use ripgrep if found on system
+        if rg_path is not None:
+            logger.info("[Refresh: Using ripgrep for scanning]")
+
+            compiled_ignore_path = library_dir / ".TagStudio" / ".compiled_ignore"
+
+            # Write compiled ignore patterns (built-in + user) to a temp file to pass to ripgrep
+            with open(compiled_ignore_path, "w") as pattern_file:
+                pattern_file.write("\n".join(ignore_patterns))
+
+            result = silent_run(
+                " ".join(
+                    [
+                        "rg",
+                        "--files",
+                        "--follow",
+                        "--hidden",
+                        "--ignore-file",
+                        f'"{str(compiled_ignore_path)}"',
+                    ]
+                ),
+                cwd=library_dir,
+                capture_output=True,
+                text=True,
+                shell=True,
+            )
+            compiled_ignore_path.unlink()
+
+            if result.stderr:
+                logger.error(result.stderr)
+
+            return result.stdout.splitlines()  # pyright: ignore [reportReturnType]
+
+        logger.warning("[Refresh: ripgrep not found on system]")
+        return None
+
+    def __rg_add(self, library_dir: Path, dir_list: list[str]) -> Iterator[int]:
         start_time_total = time()
         start_time_loop = time()
-
+        dir_file_count = 0
         self.files_not_in_library = []
+
+        for r in dir_list:
+            f = pathlib.Path(r)
+
+            end_time_loop = time()
+            # Yield output every 1/30 of a second
+            if (end_time_loop - start_time_loop) > 0.034:
+                yield dir_file_count
+                start_time_loop = time()
+
+            # Skip if the file/path is already mapped in the Library
+            if f in self.library.included_files:
+                dir_file_count += 1
+                continue
+
+            # Ignore if the file is a directory
+            if f.is_dir():
+                continue
+
+            dir_file_count += 1
+            self.library.included_files.add(f)
+
+            if not self.library.has_path_entry(f):
+                self.files_not_in_library.append(f)
+
+        end_time_total = time()
+        yield dir_file_count
+        logger.info(
+            "[Refresh]: Directory scan time",
+            path=library_dir,
+            duration=(end_time_total - start_time_total),
+            files_scanned=dir_file_count,
+            tool_used="ripgrep (system)",
+        )
+
+    def __wc_add(self, library_dir: Path, ignore_patterns: list[str]) -> Iterator[int]:
+        start_time_total = time()
+        start_time_loop = time()
         dir_file_count = 0
+        self.files_not_in_library = []
+
+        logger.info("[Refresh]: Falling back to wcmatch for scanning")
 
-        ignore_patterns = Ignore.get_patterns(library_dir)
-        logger.info(ignore_patterns)
         for f in pathlib.Path(str(library_dir)).glob(
             "***/*", flags=PATH_GLOB_FLAGS, exclude=ignore_patterns
         ):
@@ -76,16 +179,16 @@ def refresh_dir(self, library_dir: Path) -> Iterator[int]:
             self.library.included_files.add(f)
 
             relative_path = f.relative_to(library_dir)
-            # TODO - load these in batch somehow
+
             if not self.library.has_path_entry(relative_path):
                 self.files_not_in_library.append(relative_path)
 
         end_time_total = time()
         yield dir_file_count
         logger.info(
-            "Directory scan time",
+            "[Refresh]: Directory scan time",
             path=library_dir,
             duration=(end_time_total - start_time_total),
             files_scanned=dir_file_count,
-            ignore_patterns=ignore_patterns,
+            tool_used="wcmatch (internal)",
         )

src/tagstudio/qt/helpers/silent_popen.py

@@ -86,7 +86,7 @@ def silent_Popen(  # noqa: N802
     )
 
 
-def silent_run(  # noqa: N802
+def silent_run(
     args,
     bufsize=-1,
     executable=None,

tests/macros/test_refresh_dir.py

@@ -20,7 +20,7 @@ def test_refresh_new_files(library, exclude_mode):
     (library.library_dir / "FOO.MD").touch()
 
     # When
-    assert len(list(registry.refresh_dir(library.library_dir))) == 1
+    assert len(list(registry.refresh_dir(library.library_dir, force_internal_tools=True))) == 1
 
     # Then
     assert registry.files_not_in_library == [Path("FOO.MD")]