feat: add array __contains__ support

python/tvm_ffi/_ffi_api.py

@@ -38,6 +38,7 @@
     def Array(*args: Any) -> Any: ...
     def ArrayGetItem(_0: Sequence[Any], _1: int, /) -> Any: ...
     def ArraySize(_0: Sequence[Any], /) -> int: ...
+    def ArrayContains(_0: Sequence[Any], _1: Any, /) -> bool: ...
     def Bytes(_0: bytes, /) -> bytes: ...
     def FromJSONGraph(_0: Any, /) -> Any: ...
     def FromJSONGraphString(_0: str, /) -> Any: ...
@@ -81,6 +82,7 @@ def ToJSONGraphString(_0: Any, _1: Any, /) -> str: ...
 __all__ = [
     # tvm-ffi-stubgen(begin): __all__
     "Array",
+    "ArrayContains",
     "ArrayGetItem",
     "ArraySize",
     "Bytes",

python/tvm_ffi/container.py

@@ -188,6 +188,10 @@ def __repr__(self) -> str:
             return type(self).__name__ + "(chandle=None)"
         return "[" + ", ".join([x.__repr__() for x in self]) + "]"
 
+    def __contains__(self, value: object) -> bool:
+        """Check if the array contains a value."""
+        return _ffi_api.ArrayContains(self, value)
+
     def __add__(self, other: Iterable[T]) -> Array[T]:
         """Concatenate two arrays."""
         return type(self)(itertools.chain(self, other))

src/ffi/container.cc

@@ -70,6 +70,16 @@ TVM_FFI_STATIC_INIT_BLOCK() {
       .def("ffi.ArrayGetItem", [](const ffi::ArrayObj* n, int64_t i) -> Any { return n->at(i); })
       .def("ffi.ArraySize",
            [](const ffi::ArrayObj* n) -> int64_t { return static_cast<int64_t>(n->size()); })
+      .def("ffi.ArrayContains",
+           [](const ffi::ArrayObj* n, const Any& value) -> bool {
+             AnyEqual eq;
+             for (const Any& elem : *n) {
+               if (eq(elem, value)) {
+                 return true;
+               }
+             }
+             return false;
+           })
       .def_packed("ffi.Map",
                   [](ffi::PackedArgs args, Any* ret) {
                     TVM_FFI_ICHECK_EQ(args.size() % 2, 0);

tests/cpp/test_array.cc

@@ -19,6 +19,7 @@
 #include <gtest/gtest.h>
 #include <tvm/ffi/container/array.h>
 #include <tvm/ffi/function.h>
+#include <tvm/ffi/string.h>
 
 #include "./testing_object.h"
 
@@ -292,4 +293,51 @@ TEST(Array, Upcast) {
   static_assert(details::type_contains_v<Any, Array<float>>);
 }
 
+TEST(Array, Contains) {
+  Array<int> arr = {1, 2, 3, 4, 5};
+  AnyEqual eq;
+
+  // Test element is present
+  bool found = false;
+  for (const auto& elem : *arr.GetArrayObj()) {
+    if (eq(elem, Any(3))) {
+      found = true;
+      break;
+    }
+  }
+  EXPECT_TRUE(found);
+
+  // Test element is not present
+  found = false;
+  for (const auto& elem : *arr.GetArrayObj()) {
+    if (eq(elem, Any(10))) {
+      found = true;
+      break;
+    }
+  }
+  EXPECT_FALSE(found);
+
+  // Test empty array
+  Array<int> empty_arr;
+  found = false;
+  for (const auto& elem : *empty_arr.GetArrayObj()) {
+    if (eq(elem, Any(1))) {
+      found = true;
+      break;
+    }
+  }
+  EXPECT_FALSE(found);
+
+  // Test with strings
+  Array<String> str_arr = {String("hello"), String("world")};
+  found = false;
+  for (const auto& elem : *str_arr.GetArrayObj()) {
+    if (eq(elem, Any(String("world")))) {
+      found = true;
+      break;
+    }
+  }
+  EXPECT_TRUE(found);
+}
+
 }  // namespace

tests/python/test_container.py

@@ -14,6 +14,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from __future__ import annotations
+
 import pickle
 import sys
 from typing import Any
@@ -206,3 +208,22 @@ def test_large_map_get() -> None:
     amap = tvm_ffi.convert({k: k**2 for k in range(100)})
     assert amap.get(101) is None
     assert amap.get(3) == 9
+
+
+@pytest.mark.parametrize(
+    "arr, value, expected",
+    [
+        ([1, 2, 3, 4, 5], 3, True),
+        ([1, 2, 3, 4, 5], 1, True),
+        ([1, 2, 3, 4, 5], 5, True),
+        ([1, 2, 3, 4, 5], 10, False),
+        ([1, 2, 3, 4, 5], 0, False),
+        ([], 1, False),
+        (["hello", "world"], "hello", True),
+        (["hello", "world"], "world", True),
+        (["hello", "world"], "foo", False),
+    ],
+)
+def test_array_contains(arr: list[Any], value: Any, expected: bool) -> None:
+    a = tvm_ffi.convert(arr)
+    assert (value in a) == expected