Remove max_count from storage.find

Author: natoverse

packages/graphrag/graphrag/storage/blob_pipeline_storage.py

@@ -98,13 +98,11 @@ def _container_exists(self) -> bool:
     def find(
         self,
         file_pattern: re.Pattern[str],
-        max_count=-1,
     ) -> Iterator[str]:
         """Find blobs in a container using a file pattern.
 
         Params:
             file_pattern: The file pattern to use.
-            max_count: The maximum number of blobs to return. If -1, all blobs are returned.
 
         Returns
         -------
@@ -138,8 +136,6 @@ def _blobname(blob_name: str) -> str:
                 if match:
                     yield _blobname(blob.name)
                     num_loaded += 1
-                    if max_count > 0 and num_loaded >= max_count:
-                        break
                 else:
                     num_filtered += 1
             logger.debug(

packages/graphrag/graphrag/storage/cosmosdb_pipeline_storage.py

@@ -120,13 +120,11 @@ def _delete_container(self) -> None:
     def find(
         self,
         file_pattern: re.Pattern[str],
-        max_count=-1,
     ) -> Iterator[str]:
         """Find documents in a Cosmos DB container using a file pattern regex.
 
         Params:
             file_pattern: The file pattern to use.
-            max_count: The maximum number of documents to return. If -1, all documents are returned.
 
         Returns
         -------
@@ -164,8 +162,6 @@ def find(
                 if match:
                     yield item["id"]
                     num_loaded += 1
-                    if max_count > 0 and num_loaded >= max_count:
-                        break
                 else:
                     num_filtered += 1
 

packages/graphrag/graphrag/storage/file_pipeline_storage.py

@@ -40,7 +40,6 @@ def __init__(self, **kwargs: Any) -> None:
     def find(
         self,
         file_pattern: re.Pattern[str],
-        max_count=-1,
     ) -> Iterator[str]:
         """Find files in the storage using a file pattern."""
         logger.info(
@@ -59,8 +58,6 @@ def find(
                     filename = filename[1:]
                 yield filename
                 num_loaded += 1
-                if max_count > 0 and num_loaded >= max_count:
-                    break
             else:
                 num_filtered += 1
         logger.debug(

packages/graphrag/graphrag/storage/pipeline_storage.py

@@ -17,7 +17,6 @@ class PipelineStorage(metaclass=ABCMeta):
     def find(
         self,
         file_pattern: re.Pattern[str],
-        max_count=-1,
     ) -> Iterator[str]:
         """Find files in the storage using a file pattern."""
 

tests/integration/storage/test_factory.py

@@ -113,7 +113,6 @@ def __init__(self, **kwargs):
         def find(
             self,
             file_pattern: re.Pattern[str],
-            max_count=-1,
         ) -> Iterator[str]:
             return iter([])