From cfd96f63f1aa2e362e4fbfd8462b9642ec5ae5aa Mon Sep 17 00:00:00 2001 From: Ludwik Trammer Date: Thu, 17 Oct 2024 10:16:42 +0200 Subject: [PATCH] Apply suggestions from code review Co-authored-by: Alan Konarski <129968242+akonarski-ds@users.noreply.github.com> --- .../src/ragbits/document_search/documents/sources.py | 5 +---- .../tests/unit/test_local_file_source.py | 2 ++ 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/packages/ragbits-document-search/src/ragbits/document_search/documents/sources.py b/packages/ragbits-document-search/src/ragbits/document_search/documents/sources.py index 73e50df8..a4a78f05 100644 --- a/packages/ragbits-document-search/src/ragbits/document_search/documents/sources.py +++ b/packages/ragbits-document-search/src/ragbits/document_search/documents/sources.py @@ -78,10 +78,7 @@ def list_sources(cls, path: Path, file_pattern: str = "*") -> list["LocalFileSou Returns: List of source objects. """ - sources = [] - for file_path in path.glob(file_pattern): - sources.append(cls(path=file_path)) - return sources + return [cls(path=file_path) for file_path in path.glob(file_pattern)] class GCSSource(Source): diff --git a/packages/ragbits-document-search/tests/unit/test_local_file_source.py b/packages/ragbits-document-search/tests/unit/test_local_file_source.py index a2cfe6cb..bde7549e 100644 --- a/packages/ragbits-document-search/tests/unit/test_local_file_source.py +++ b/packages/ragbits-document-search/tests/unit/test_local_file_source.py @@ -9,11 +9,13 @@ async def test_local_source_fetch(): source = LocalFileSource(path=TEST_FILE_PATH) path = await source.fetch() + assert path == TEST_FILE_PATH async def test_local_source_list_sources(): example_files = TEST_FILE_PATH.parent / "example_files" + sources = LocalFileSource.list_sources(example_files, file_pattern="*.md") assert len(sources) == 2