use Protocol over ABC

Author: maxjakob

elasticsearch/helpers/vectorstore/_async/embedding_service.py

@@ -15,14 +15,12 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-from abc import ABC, abstractmethod
-from typing import List, Optional
+from typing import List, Optional, Protocol
 
 from elasticsearch import AsyncElasticsearch
 
 
-class AsyncEmbeddingService(ABC):
-    @abstractmethod
+class AsyncEmbeddingService(Protocol):
     async def embed_documents(self, texts: List[str]) -> List[List[float]]:
         """Generate embeddings for a list of documents.
 
@@ -33,7 +31,6 @@ async def embed_documents(self, texts: List[str]) -> List[List[float]]:
             A list of embeddings, one for each document in the input.
         """
 
-    @abstractmethod
     async def embed_query(self, query: str) -> List[float]:
         """Generate an embedding for a single query text.
 

elasticsearch/helpers/vectorstore/_async/strategies.py

@@ -15,16 +15,14 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Protocol, Tuple, Union, cast
 
 from elasticsearch import AsyncElasticsearch
 from elasticsearch.helpers.vectorstore._async._utils import model_must_be_deployed
 from elasticsearch.helpers.vectorstore._utils import DistanceMetric
 
 
-class AsyncRetrievalStrategy(ABC):
-    @abstractmethod
+class AsyncRetrievalStrategy(Protocol):
     def es_query(
         self,
         query: Optional[str],
@@ -50,7 +48,6 @@ def es_query(
             Dict: The Elasticsearch query body.
         """
 
-    @abstractmethod
     def es_mappings_settings(
         self,
         text_field: str,

elasticsearch/helpers/vectorstore/_sync/embedding_service.py

@@ -15,14 +15,12 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-from abc import ABC, abstractmethod
-from typing import List, Optional
+from typing import List, Optional, Protocol
 
 from elasticsearch import Elasticsearch
 
 
-class EmbeddingService(ABC):
-    @abstractmethod
+class EmbeddingService(Protocol):
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
         """Generate embeddings for a list of documents.
 
@@ -33,7 +31,6 @@ def embed_documents(self, texts: List[str]) -> List[List[float]]:
             A list of embeddings, one for each document in the input.
         """
 
-    @abstractmethod
     def embed_query(self, query: str) -> List[float]:
         """Generate an embedding for a single query text.
 

elasticsearch/helpers/vectorstore/_sync/strategies.py

@@ -15,16 +15,14 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Protocol, Tuple, Union, cast
 
 from elasticsearch import Elasticsearch
 from elasticsearch.helpers.vectorstore._sync._utils import model_must_be_deployed
 from elasticsearch.helpers.vectorstore._utils import DistanceMetric
 
 
-class RetrievalStrategy(ABC):
-    @abstractmethod
+class RetrievalStrategy(Protocol):
     def es_query(
         self,
         query: Optional[str],
@@ -50,7 +48,6 @@ def es_query(
             Dict: The Elasticsearch query body.
         """
 
-    @abstractmethod
     def es_mappings_settings(
         self,
         text_field: str,

test_elasticsearch/utils.py

@@ -215,16 +215,10 @@ def wipe_data_streams(client):
 
 
 def wipe_indices(client):
-    response = client.indices.get(index="_all")
-    index_names = response.keys()
-    for index_name in index_names:
-        if index_name.startswith("test_"):
-            client.indices.delete(index=index_name)
-    client.indices.refresh(index="_all")
-    # client.options(ignore_status=404).indices.delete(
-    #     index="*,-.ds-ilm-history-*",
-    #     expand_wildcards="all",
-    # )
+    client.options(ignore_status=404).indices.delete(
+        index="*,-.ds-ilm-history-*",
+        expand_wildcards="all",
+    )
 
 
 def wipe_searchable_snapshot_indices(client):

utils/run-unasync.py

@@ -24,12 +24,11 @@
 
 
 def cleanup(source_dir: Path, output_dir: Path, patterns: list[str]):
-    if patterns:
-        for file in glob("*.py", root_dir=source_dir):
-            path = Path(output_dir) / file
-            for pattern in patterns:
-                subprocess.check_call(["sed", "-i.bak", pattern, str(path)])
-            subprocess.check_call(["rm", f"{path}.bak"])
+    for file in glob("*.py", root_dir=source_dir):
+        path = Path(output_dir) / file
+        for pattern in patterns:
+            subprocess.check_call(["sed", "-i.bak", pattern, str(path)])
+        subprocess.check_call(["rm", f"{path}.bak"])
 
 
 def format_dir(dir: Path):
@@ -57,7 +56,8 @@ def run(
 
     unasync.unasync_files(filepaths, [rule])
 
-    cleanup(source_dir, output_dir, cleanup_patterns)
+    if cleanup_patterns:
+        cleanup(source_dir, output_dir, cleanup_patterns)
 
     if format:
         format_dir(source_dir)