From a14edebafe5667f193eecb04011ae6715ca1391b Mon Sep 17 00:00:00 2001 From: Matteo Rosati Date: Thu, 23 Apr 2026 21:49:46 +0200 Subject: [PATCH] add colors! --- chromy/chroma_functions.py | 5 +++-- chromy/cli.py | 3 ++- chromy/handlers/count_collection.py | 1 + chromy/handlers/create_collection.py | 3 ++- chromy/handlers/delete_collection.py | 5 +++-- chromy/handlers/import_data.py | 4 +++- pyproject.toml | 1 + tests/test_cli.py | 10 ++++++---- tests/test_handlers.py | 2 +- uv.lock | 2 ++ 10 files changed, 24 insertions(+), 12 deletions(-) diff --git a/chromy/chroma_functions.py b/chromy/chroma_functions.py index d2d6528..731f035 100644 --- a/chromy/chroma_functions.py +++ b/chromy/chroma_functions.py @@ -54,10 +54,11 @@ def delete_data(collection_name: str, where: dict[str, str]) -> int: return int(result.get("deleted", 0)) -def count_collection(collection_name: str) -> int: +def count_collection(collection_name: str) -> str: _, collection = _get_client_and_collection(collection_name) + count = collection.count() - return collection.count() + return f"The '{collection_name}' collection contains [bold green]{count}[/] records." def add_data( diff --git a/chromy/cli.py b/chromy/cli.py index d294822..12f0095 100644 --- a/chromy/cli.py +++ b/chromy/cli.py @@ -3,6 +3,7 @@ from __future__ import annotations from typing import Annotated, Callable import typer +from rich import print from chromadb.errors import InternalError, NotFoundError from chromy.handlers.import_data import handle_import @@ -27,7 +28,7 @@ def _run(handler: ExitCodeHandler) -> None: def _fail(message: str) -> None: - typer.echo(message) + print("[bold red]Error[/]:", message) raise typer.Exit(1) diff --git a/chromy/handlers/count_collection.py b/chromy/handlers/count_collection.py index 559c92b..53da70f 100644 --- a/chromy/handlers/count_collection.py +++ b/chromy/handlers/count_collection.py @@ -1,5 +1,6 @@ from __future__ import annotations +from rich import print from chromy.chroma_functions import count_collection diff --git a/chromy/handlers/create_collection.py b/chromy/handlers/create_collection.py index b42723d..2113046 100644 --- a/chromy/handlers/create_collection.py +++ b/chromy/handlers/create_collection.py @@ -1,9 +1,10 @@ from __future__ import annotations +from rich import print from chromy.chroma_functions import create_collection def handle_create_collection(collection: str) -> int: collection_name = create_collection(collection) - print(f"Created collection '{collection_name}'.") + print(f"[bold green]Created[/]: collection '{collection_name}'.") return 0 diff --git a/chromy/handlers/delete_collection.py b/chromy/handlers/delete_collection.py index 0aa5f1c..af1a41a 100644 --- a/chromy/handlers/delete_collection.py +++ b/chromy/handlers/delete_collection.py @@ -1,5 +1,6 @@ from __future__ import annotations +from rich import print from chromy.chroma_functions import delete_collection, delete_data @@ -22,7 +23,7 @@ def _parse_where_clause(where_clause: str) -> dict[str, str]: def handle_delete_collection(collection: str) -> int: delete_collection(collection) - print(f"Deleted collection '{collection}'.") + print(f"[bold green]Deleted[/] collection '{collection}'.") return 0 @@ -31,7 +32,7 @@ def handle_delete_records(collection: str, where_clause: str) -> int: deleted = delete_data(collection, where) condition, value = next(iter(where.items())) print( - f"Deleted {deleted} record(s) from collection '{collection}' " + f"[bold green]Deleted[/] {deleted} record(s) from collection '{collection}' " f"where {condition}={value}." ) return 0 diff --git a/chromy/handlers/import_data.py b/chromy/handlers/import_data.py index 99b2285..8558b8a 100644 --- a/chromy/handlers/import_data.py +++ b/chromy/handlers/import_data.py @@ -2,6 +2,7 @@ from __future__ import annotations import os from pathlib import Path +from rich import print from chromy.utilities import ingest_file @@ -25,5 +26,6 @@ def _get_absolute_path(file: str) -> str: def handle_import(collection: str, file: str) -> int: records_added = ingest_file(collection, _get_absolute_path(file)) - print(f"Added {records_added} records to collection '{collection}'.") + print( + f"[bold green]Added[/] {records_added} records to collection '{collection}'.") return 0 diff --git a/pyproject.toml b/pyproject.toml index d0e1196..c102d50 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ dependencies = [ "openai>=2.32.0", "pymupdf4llm>=1.27.2.2", "python-dotenv>=1.2.2", + "rich>=15.0.0", "semchunk>=4.0.0", "tiktoken>=0.12.0", "transformers>=5.5.4", diff --git a/tests/test_cli.py b/tests/test_cli.py index 73b6fef..2d64187 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -46,7 +46,7 @@ class CliTests(unittest.TestCase): create_collection.assert_called_once_with("notes") self.assertEqual(result.exit_code, 0) - self.assertEqual(result.stdout, "Created collection 'notes'.\n") + self.assertEqual(result.stdout, "Created: collection 'notes'.\n") def test_create_collection_with_same_name(self) -> None: with patch( @@ -58,7 +58,8 @@ class CliTests(unittest.TestCase): create_collection.assert_called_once_with("notes") self.assertEqual(result.exit_code, 1) - self.assertEqual(result.stdout, "Collection 'notes' already exists.\n") + self.assertEqual( + result.stdout, "Error: Collection 'notes' already exists.\n") def test_delete_collection(self) -> None: with patch( @@ -79,7 +80,8 @@ class CliTests(unittest.TestCase): delete_collection.assert_called_once_with("notes") self.assertEqual(result.exit_code, 1) - self.assertEqual(result.stdout, "Collection 'notes' does not exist.\n") + self.assertEqual( + result.stdout, "Error: Collection 'notes' does not exist.\n") def test_count(self) -> None: with patch( @@ -147,7 +149,7 @@ class CliTests(unittest.TestCase): self.assertEqual(result.exit_code, 1) self.assertEqual( result.stdout, - "Invalid --where value. Expected =.\n", + "Error: Invalid --where value. Expected =.\n", ) def test_delete_requires_where_option(self) -> None: diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 144947c..897ab0d 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -61,7 +61,7 @@ class HandlerTests(unittest.TestCase): create_collection.assert_called_once_with("notes") self.assertEqual(exit_code, 0) - self.assertEqual(output, "Created collection 'notes'.\n") + self.assertEqual(output, "Created: collection 'notes'.\n") def test_delete_collection_uses_typed_input(self) -> None: with patch("chromy.handlers.delete_collection.delete_collection") as delete: diff --git a/uv.lock b/uv.lock index f0535fd..5d617f5 100644 --- a/uv.lock +++ b/uv.lock @@ -261,6 +261,7 @@ dependencies = [ { name = "openai" }, { name = "pymupdf4llm" }, { name = "python-dotenv" }, + { name = "rich" }, { name = "semchunk" }, { name = "tiktoken" }, { name = "transformers" }, @@ -281,6 +282,7 @@ requires-dist = [ { name = "openai", specifier = ">=2.32.0" }, { name = "pymupdf4llm", specifier = ">=1.27.2.2" }, { name = "python-dotenv", specifier = ">=1.2.2" }, + { name = "rich", specifier = ">=15.0.0" }, { name = "semchunk", specifier = ">=4.0.0" }, { name = "tiktoken", specifier = ">=0.12.0" }, { name = "transformers", specifier = ">=5.5.4" },