diff --git a/pyproject.toml b/pyproject.toml index 665fd5ce..6b60d162 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ "orjson", "questionary>=2.0,<3.0", "semble-grammars>=0.1.2", + "tqdm>=4.60", ] [project.optional-dependencies] diff --git a/src/semble/cli.py b/src/semble/cli.py index 1ba51b29..37574122 100644 --- a/src/semble/cli.py +++ b/src/semble/cli.py @@ -31,11 +31,12 @@ def _build_index(path: str, content: list[ContentType]) -> SembleIndex: - """Build an index from a local path or git URL.""" + """Build an index from a local path or git URL, showing a progress bar on a tty.""" + show_progress_bar = sys.stderr.isatty() # Only show the progress bar in a terminal return ( - SembleIndex.from_git(path, content=content) + SembleIndex.from_git(path, content=content, show_progress_bar=show_progress_bar) if is_git_url(path) - else SembleIndex.from_path(path, content=content) + else SembleIndex.from_path(path, content=content, show_progress_bar=show_progress_bar) ) @@ -114,17 +115,41 @@ def _load_index(path: str, content: list[ContentType]) -> SembleIndex: sys.exit(1) -def _run_search(path: str, query: str, top_k: int, content: list[ContentType], max_snippet_lines: int | None) -> None: +def _print_pretty(out: dict) -> None: + """Print a format_results() payload as human-readable text instead of JSON.""" + if "error" in out: + print(out["error"]) + return + for r in out["results"]: + print(f"{r['file_path']}:{r['start_line']}-{r['end_line']}") + if "content" in r: + print() + print(r["content"]) + print() + + +def _run_search( + path: str, query: str, top_k: int, content: list[ContentType], max_snippet_lines: int | None, pretty: bool +) -> None: """Handle the `search` subcommand.""" index = _load_index(path, content) results = index.search(query, top_k=top_k, max_snippet_lines=max_snippet_lines) out = format_results(query, results, max_snippet_lines) if results else {"error": "No results found."} - print(json.dumps(out)) + if pretty: + _print_pretty(out) + else: + print(json.dumps(out)) _maybe_save_index(index, path) def _run_find_related( - path: str, file_path: str, line: int, top_k: int, content: list[ContentType], max_snippet_lines: int | None + path: str, + file_path: str, + line: int, + top_k: int, + content: list[ContentType], + max_snippet_lines: int | None, + pretty: bool, ) -> None: """Handle the `find-related` subcommand.""" index = _load_index(path, content) @@ -139,7 +164,10 @@ def _run_find_related( if results else {"error": f"No related chunks found for {file_path}:{line}."} ) - print(json.dumps(out)) + if pretty: + _print_pretty(out) + else: + print(json.dumps(out)) _maybe_save_index(index, path) @@ -241,6 +269,7 @@ def _cli_main() -> None: metavar="N", help="Lines of source per result (default: full chunk). 10 = signature + body, 0 = no code.", ) + search_p.add_argument("--pretty", action="store_true", help="Human-readable text output instead of JSON.") _add_content_args(search_p) clear_p = sub.add_parser("clear", help="Clear the index cache.") @@ -262,6 +291,7 @@ def _cli_main() -> None: metavar="N", help="Lines of source per result (default: full chunk). 10 = signature + body, 0 = no code.", ) + related_p.add_argument("--pretty", action="store_true", help="Human-readable text output instead of JSON.") _add_content_args(related_p) sub.add_parser("savings", help="Show token savings and usage stats.") @@ -311,6 +341,7 @@ def _cli_main() -> None: args.top_k, _resolve_content(args.content, args.include_text_files), args.max_snippet_lines, + args.pretty, ) elif args.command == "find-related": _run_find_related( @@ -320,4 +351,5 @@ def _cli_main() -> None: args.top_k, _resolve_content(args.content, args.include_text_files), args.max_snippet_lines, + args.pretty, ) diff --git a/src/semble/index/create.py b/src/semble/index/create.py index 54e39028..43e066c0 100644 --- a/src/semble/index/create.py +++ b/src/semble/index/create.py @@ -1,10 +1,12 @@ import contextlib import logging +import sys from collections.abc import Sequence from pathlib import Path import numpy as np from model2vec.model import StaticModel +from tqdm import tqdm from vicinity.backends.basic import BasicArgs from semble.chunking import chunk_source @@ -72,6 +74,7 @@ def create_index_from_path( content: ContentType | Sequence[ContentType] = (ContentType.CODE,), display_root: Path | None = None, previous: PreviousIndex | None = None, + show_progress_bar: bool = False, ) -> tuple[BM25, SelectableBasicBackend, list[Chunk], dict[str, FileManifestEntry]]: """Create an index from a resolved directory, optionally reusing a previous index's unchanged files. @@ -80,6 +83,7 @@ def create_index_from_path( :param content: Content types to index. :param display_root: If set, chunk file paths are stored relative to this root. :param previous: A previously built index to reuse unchanged files' chunks/embeddings/postings from. + :param show_progress_bar: Show a progress bar on stderr while indexing. :raises ValueError: if no items were found, no index can be created. :return: A BM25 index, semantic index, list of chunks, and file manifest. """ @@ -98,7 +102,10 @@ def create_index_from_path( skipped_large: list[str] = [] - for file_path in walk_files(path, resolved_extensions): + files = list(walk_files(path, resolved_extensions)) + for file_path in tqdm( + files, desc="Indexing", unit="file", file=sys.stderr, leave=False, colour="green", disable=not show_progress_bar + ): language = detect_language(file_path) with contextlib.suppress(OSError): file_status = get_file_status(file_path, None) diff --git a/src/semble/index/index.py b/src/semble/index/index.py index cc199edd..9ec48b07 100644 --- a/src/semble/index/index.py +++ b/src/semble/index/index.py @@ -136,6 +136,7 @@ def from_path( content: ContentType | Sequence[ContentType] = _DEFAULT_CONTENT, include_text_files: bool | None = None, model_path: str | None = None, + show_progress_bar: bool = False, ) -> SembleIndex: """Create and index a SembleIndex from a directory. @@ -143,6 +144,7 @@ def from_path( :param content: Content types to index, e.g. ContentType.CODE or [ContentType.CODE, ContentType.DOCS]. :param include_text_files: Deprecated. Pass a content sequence directly instead. :param model_path: Path to the model to use. If None, the default model will be used. + :param show_progress_bar: Show a progress bar while indexing. :return: An indexed SembleIndex. Chunk file paths are relative to ``path``. :raises FileNotFoundError: If `path` does not exist. :raises NotADirectoryError: If `path` exists but is not a directory. @@ -167,6 +169,7 @@ def from_path( content=normalized, display_root=path, previous=previous, + show_progress_bar=show_progress_bar, ) return SembleIndex( @@ -181,6 +184,7 @@ def from_git( model_path: str | None = None, content: ContentType | Sequence[ContentType] = _DEFAULT_CONTENT, include_text_files: bool | None = None, + show_progress_bar: bool = False, ) -> SembleIndex: """Clone a git repository and index it. @@ -194,6 +198,7 @@ def from_git( :param model_path: Path to the model to use. If None, the default model will be used. :param content: Content types to index, e.g. (ContentType.CODE,) or (ContentType.CODE, ContentType.DOCS). :param include_text_files: Deprecated. Pass content=(ContentType.CODE, ContentType.DOCS, ...) instead. + :param show_progress_bar: Show a progress bar while indexing. :return: An indexed SembleIndex. Chunk file paths are repo-relative (e.g. ``src/foo.py``). :raises RuntimeError: If git is not on PATH, the clone fails, or times out. """ @@ -224,6 +229,7 @@ def from_git( model=model, content=normalized, display_root=resolved_path, + show_progress_bar=show_progress_bar, ) return SembleIndex( diff --git a/tests/test_cli.py b/tests/test_cli.py index 029c4079..8ccc614c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -34,6 +34,8 @@ def test_main_calls_asyncio_run(argv: list[str], monkeypatch: pytest.MonkeyPatch [ (["semble", "search", "query text", "/some/path"], ["query text", "0.9"]), (["semble", "search", "nothing", "/some/path", "--top-k", "3"], ["No results found"]), + (["semble", "search", "query text", "/some/path", "--pretty"], ["src/foo.py:1-1\n\ndef foo(): pass"]), + (["semble", "search", "nothing", "/some/path", "--pretty"], ["No results found."]), ], ) def test_cli_search( @@ -59,6 +61,7 @@ def test_cli_search( ("scenario", "expected_stdout", "expected_stderr", "expected_exit_code"), [ ("with_results", ["src/bar.py", "0.8"], None, None), + ("pretty", ["src/bar.py:1-1\n\nclass Bar: pass"], None, None), ("no_results", ["No related chunks found"], None, None), ("unknown_chunk", [], "No chunk found", 1), ], @@ -75,9 +78,11 @@ def test_cli_find_related( chunk = make_chunk("class Bar: pass", "src/bar.py") fake_index = MagicMock() fake_index.chunks = [] if scenario == "unknown_chunk" else [chunk] - fake_index.find_related.return_value = [SearchResult(chunk=chunk, score=0.8)] if scenario == "with_results" else [] + has_results = scenario in ("with_results", "pretty") + fake_index.find_related.return_value = [SearchResult(chunk=chunk, score=0.8)] if has_results else [] file_path = "unknown.py" if scenario == "unknown_chunk" else "src/bar.py" - monkeypatch.setattr(sys, "argv", ["semble", "find-related", file_path, "1", "/some/path"]) + argv = ["semble", "find-related", file_path, "1", "/some/path"] + (["--pretty"] if scenario == "pretty" else []) + monkeypatch.setattr(sys, "argv", argv) with patch("semble.cli.SembleIndex.from_path", return_value=fake_index): if expected_exit_code is None: _cli_main() diff --git a/uv.lock b/uv.lock index 7630799b..1e69d081 100644 --- a/uv.lock +++ b/uv.lock @@ -3140,6 +3140,7 @@ dependencies = [ { name = "pathspec" }, { name = "questionary" }, { name = "semble-grammars" }, + { name = "tqdm" }, { name = "vicinity" }, ] @@ -3186,6 +3187,7 @@ requires-dist = [ { name = "semble-grammars", specifier = ">=0.1.2" }, { name = "sentence-transformers", marker = "extra == 'benchmark'", specifier = ">=3.0" }, { name = "tiktoken", marker = "extra == 'benchmark'", specifier = ">=0.7" }, + { name = "tqdm", specifier = ">=4.60" }, { name = "vicinity", specifier = ">=0.4.4" }, ] provides-extras = ["mcp", "benchmark", "dev"]