This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 5bc7fcd feat: add array __contains__ support (#377)
5bc7fcd is described below
commit 5bc7fcdebd0fae2d3650a5b18ae69154c1c92d70
Author: Guan-Ming (Wesley) Chiu <[email protected]>
AuthorDate: Fri Jan 2 16:52:24 2026 +0800
feat: add array __contains__ support (#377)
## Why
Python arrays lack native in operator support, requiring manual
iteration to check if a value exists.
## How
- Add `ffi.ArrayContains` C++ FFI function
- Implement `__contains__` method in Python Array class
- Add tests and type stub for Python
---
python/tvm_ffi/_ffi_api.py | 2 ++
python/tvm_ffi/container.py | 4 ++++
src/ffi/container.cc | 6 ++++++
tests/cpp/test_array.cc | 20 ++++++++++++++++++++
tests/python/test_container.py | 21 +++++++++++++++++++++
5 files changed, 53 insertions(+)
diff --git a/python/tvm_ffi/_ffi_api.py b/python/tvm_ffi/_ffi_api.py
index 49fbf68..9678fd0 100644
--- a/python/tvm_ffi/_ffi_api.py
+++ b/python/tvm_ffi/_ffi_api.py
@@ -38,6 +38,7 @@ if TYPE_CHECKING:
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 @@ if TYPE_CHECKING:
__all__ = [
# tvm-ffi-stubgen(begin): __all__
"Array",
+ "ArrayContains",
"ArrayGetItem",
"ArraySize",
"Bytes",
diff --git a/python/tvm_ffi/container.py b/python/tvm_ffi/container.py
index bda4c57..d2aff29 100644
--- a/python/tvm_ffi/container.py
+++ b/python/tvm_ffi/container.py
@@ -188,6 +188,10 @@ class Array(core.Object, Sequence[T]):
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))
diff --git a/src/ffi/container.cc b/src/ffi/container.cc
index 57eda37..dd1d004 100644
--- a/src/ffi/container.cc
+++ b/src/ffi/container.cc
@@ -70,6 +70,12 @@ 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;
+ return std::any_of(n->begin(), n->end(),
+ [&](const Any& elem) { return eq(elem, value);
});
+ })
.def_packed("ffi.Map",
[](ffi::PackedArgs args, Any* ret) {
TVM_FFI_ICHECK_EQ(args.size() % 2, 0);
diff --git a/tests/cpp/test_array.cc b/tests/cpp/test_array.cc
index 0204aa1..a86cc29 100644
--- a/tests/cpp/test_array.cc
+++ b/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,23 @@ TEST(Array, Upcast) {
static_assert(details::type_contains_v<Any, Array<float>>);
}
+TEST(Array, Contains) {
+ Function f = Function::GetGlobalRequired("ffi.ArrayContains");
+
+ Array<int> arr = {1, 2, 3, 4, 5};
+ EXPECT_TRUE(f(arr, 3).cast<bool>());
+ EXPECT_TRUE(f(arr, 1).cast<bool>());
+ EXPECT_TRUE(f(arr, 5).cast<bool>());
+ EXPECT_FALSE(f(arr, 10).cast<bool>());
+ EXPECT_FALSE(f(arr, 0).cast<bool>());
+
+ Array<int> empty_arr;
+ EXPECT_FALSE(f(empty_arr, 1).cast<bool>());
+
+ Array<String> str_arr = {String("hello"), String("world")};
+ EXPECT_TRUE(f(str_arr, String("hello")).cast<bool>());
+ EXPECT_TRUE(f(str_arr, String("world")).cast<bool>());
+ EXPECT_FALSE(f(str_arr, String("foo")).cast<bool>());
+}
+
} // namespace
diff --git a/tests/python/test_container.py b/tests/python/test_container.py
index e15f00b..8650e08 100644
--- a/tests/python/test_container.py
+++ b/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
+
+
[email protected](
+ "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
