Adjust IndexList and IndexModel

pinecone-io · Jul 18, 2024 · afcd31b · afcd31b
1 parent fb4c78e
commit afcd31b
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/pinecone/models/index_list.py b/pinecone/models/index_list.py
@@ -5,19 +5,20 @@
 class IndexList:
     def __init__(self, index_list: OpenAPIIndexList):
         self.index_list = index_list
+        self.indexes = [IndexModel(i) for i in self.index_list.indexes]
         self.current = 0
 
     def names(self):
-        return [i["name"] for i in self.index_list.indexes]
+        return [i.name for i in self.indexes]
 
     def __getitem__(self, key):
-        return IndexModel(self.index_list.indexes[key])
+        return IndexModel(self.indexes[key])
 
     def __len__(self):
-        return len(self.index_list.indexes)
+        return len(self.indexes)
 
     def __iter__(self):
-        return iter(self.index_list.indexes)
+        return iter(self.indexes)
 
     def __str__(self):
         return str(self.index_list)

diff --git a/pinecone/models/index_model.py b/pinecone/models/index_model.py
@@ -14,3 +14,6 @@ def __repr__(self):
 
     def __getattr__(self, attr):
         return getattr(self.index, attr)
+
+    def __getitem__(self, key):
+        return self.__getattr__(key)
diff --git a/tests/unit/test_control.py b/tests/unit/test_control.py
@@ -2,7 +2,7 @@
 import re
 from unittest.mock import patch, MagicMock
 from pinecone import ConfigBuilder, Pinecone, PodSpec, ServerlessSpec
-from pinecone.core.openapi.control.models import IndexList, IndexModel
+from pinecone.core.openapi.control.models import IndexList, IndexModel, DeletionProtection
 from pinecone.core.openapi.control.api.manage_indexes_api import ManageIndexesApi
 from pinecone.core.openapi.shared.configuration import Configuration as OpenApiConfiguration
 
@@ -20,6 +20,7 @@ def index_list_response():
                 host="asdf",
                 status={"ready": True},
                 spec={},
+                deletion_protection=DeletionProtection("enabled"),
                 _check_type=False,
             ),
             IndexModel(
@@ -29,6 +30,7 @@ def index_list_response():
                 host="asdf",
                 status={"ready": True},
                 spec={},
+                deletion_protection=DeletionProtection("enabled"),
                 _check_type=False,
             ),
             IndexModel(
@@ -38,6 +40,7 @@ def index_list_response():
                 host="asdf",
                 status={"ready": True},
                 spec={},
+                deletion_protection=DeletionProtection("disabled"),
                 _check_type=False,
             ),
         ]