From 948f8500beaba057b9095baed2a1566bcf7d3ec8 Mon Sep 17 00:00:00 2001 From: Matteo Rosati Date: Fri, 24 Apr 2026 18:20:22 +0200 Subject: [PATCH] types cleanup --- chromy/handlers/import_data.py | 7 ++++--- chromy/utilities.py | 16 +++++++--------- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/chromy/handlers/import_data.py b/chromy/handlers/import_data.py index 8558b8a..9bee641 100644 --- a/chromy/handlers/import_data.py +++ b/chromy/handlers/import_data.py @@ -1,8 +1,10 @@ from __future__ import annotations + import os from pathlib import Path from rich import print + from chromy.utilities import ingest_file @@ -21,11 +23,10 @@ def _get_absolute_path(file: str) -> str: raise FileNotFoundError() file_path = Path(file) - return str(file_path.resolve(file_path)) + return str(file_path.resolve()) def handle_import(collection: str, file: str) -> int: records_added = ingest_file(collection, _get_absolute_path(file)) - print( - f"[bold green]Added[/] {records_added} records to collection '{collection}'.") + print(f"[bold green]Added[/] {records_added} records to collection '{collection}'.") return 0 diff --git a/chromy/utilities.py b/chromy/utilities.py index 117491b..bd02b82 100644 --- a/chromy/utilities.py +++ b/chromy/utilities.py @@ -1,12 +1,11 @@ from __future__ import annotations -from rich.text import Text -from rich.rule import Rule -from rich.console import Console - from collections.abc import Mapping, Sequence from chromadb import QueryResult +from rich.console import Console +from rich.rule import Rule +from rich.text import Text from chromy.chroma_functions import add_data, query_data from chromy.chunk_functions import chunk_file @@ -31,7 +30,7 @@ def run_query(collection_name: str, query_text: str) -> QueryResult: return query_data(collection_name, [query_text]) -def format_query_result(result: QueryResult) -> list[str]: +def format_query_result(result: QueryResult) -> list[Rule | Text]: ids = result.get("ids", [[]]) documents = result.get("documents", [[]]) distances = result.get("distances", [[]]) @@ -43,12 +42,11 @@ def format_query_result(result: QueryResult) -> list[str]: first_metadatas = metadatas[0] if metadatas else [] if not first_ids: - return ["No results found."] + return [Text.from_markup("[yellow]No results found.[/]")] - lines = [Rule(title="Query results")] + lines: list[Rule | Text] = [Rule(title="Query results")] for index, document_id in enumerate(first_ids, start=1): - # lines.append(f"{index}.\tid: {document_id}") lines.append( Text.from_markup(f"[bold]{index}[/].\t[green]id[/]\t\t{document_id}") ) @@ -72,7 +70,7 @@ def format_query_result(result: QueryResult) -> list[str]: if i < len(first_documents): lines.append(Text.from_markup("\n[bold green]Retrieved contents[/]\n")) - lines.append(first_documents[i]) + lines.append(Text(first_documents[i])) # Print a separator between documents lines.append(Rule())