From dbcabd8f44d32dea8a698474427ba54f6900a91d Mon Sep 17 00:00:00 2001
From: Qinren Zhou <zhouqinren.zqr@alibaba-inc.com>
Date: Wed, 4 Mar 2026 19:16:18 +0800
Subject: [PATCH] minor: add function overload in python api (#197)

---
 python/zvec/model/collection.py | 34 ++++++++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/python/zvec/model/collection.py b/python/zvec/model/collection.py
index ec4d630..eee571d 100644
--- a/python/zvec/model/collection.py
+++ b/python/zvec/model/collection.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import Optional, Union
+from typing import Optional, Union, overload
 
 from _zvec import _Collection
 
@@ -230,6 +230,14 @@ class Collection:
         self._schema = CollectionSchema._from_core(self._obj.Schema())
 
     # ========== Collection DDL Methods ==========
+    @overload
+    def insert(self, docs: Doc) -> Status:
+        pass
+
+    @overload
+    def insert(self, docs: list[Doc]) -> list[Status]:
+        pass
+
     def insert(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:
         """Insert new documents into the collection.
 
@@ -249,6 +257,14 @@ class Collection:
         )
         return results[0] if is_single else results
 
+    @overload
+    def upsert(self, docs: Doc) -> Status:
+        pass
+
+    @overload
+    def upsert(self, docs: list[Doc]) -> list[Status]:
+        pass
+
     def upsert(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:
         """Insert new documents or update existing ones by ID.
 
@@ -266,6 +282,14 @@ class Collection:
         )
         return results[0] if is_single else results
 
+    @overload
+    def update(self, docs: Doc) -> Status:
+        pass
+
+    @overload
+    def update(self, docs: list[Doc]) -> list[Status]:
+        pass
+
     def update(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:
         """Update existing documents by ID.
 
@@ -285,6 +309,14 @@ class Collection:
         )
         return results[0] if is_single else results
 
+    @overload
+    def delete(self, ids: str) -> Status:
+        pass
+
+    @overload
+    def delete(self, ids: list[str]) -> list[Status]:
+        pass
+
     def delete(self, ids: Union[str, list[str]]) -> Union[Status, list[Status]]:
         """Delete documents by ID.