This is an automated email from the ASF dual-hosted git repository.
tqchen 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 88e066ea [PY] Support pickling FFI String and Bytes (#643)
88e066ea is described below
commit 88e066ea2ed100da3c51e081fd4c036a33075fe4
Author: Yixin Dong <[email protected]>
AuthorDate: Tue Jun 23 15:44:33 2026 -0400
[PY] Support pickling FFI String and Bytes (#643)
## Summary
- Add pickle reductions for `tvm_ffi.core.String` and
`tvm_ffi.core.Bytes` so only the Python payload is serialized.
- Ensure unpickled values drop the hidden `_tvm_ffi_cached_object`
backing object.
- Add regression coverage for cached FFI String/Bytes returned from
`testing.echo`.
## Notes
- This is a clean branch based on the latest `apache/main` and
supersedes #642, which became dirty after main advanced.
- `Bytes.__slots__` is intentionally unchanged because Python does not
support non-empty `__slots__` on `bytes` subclasses.
## Test plan
- `uv pip install --reinstall --verbose . && .venv/bin/pytest -vvs
tests/python/test_string.py`
-
`PYTHONPATH=/raid/user_data/yixind/miniforge3/lib/python3.12/site-packages
pre-commit run --files python/tvm_ffi/cython/string.pxi
tests/python/test_string.py`
---
python/tvm_ffi/cython/string.pxi | 6 ++++++
tests/python/test_string.py | 20 +++++++++++++++++++-
2 files changed, 25 insertions(+), 1 deletion(-)
diff --git a/python/tvm_ffi/cython/string.pxi b/python/tvm_ffi/cython/string.pxi
index 399c8199..edb361d9 100644
--- a/python/tvm_ffi/cython/string.pxi
+++ b/python/tvm_ffi/cython/string.pxi
@@ -57,6 +57,9 @@ class String(str, PyNativeObject):
val._tvm_ffi_cached_object = None
return val
+ def __reduce_ex__(self, protocol):
+ return (str, (str(self),))
+
# pylint: disable=no-self-argument
def __from_tvm_ffi_object__(cls, obj: Any) -> "String":
"""Construct a ``String`` from an FFI object (internal)."""
@@ -80,6 +83,9 @@ class Bytes(bytes, PyNativeObject):
val._tvm_ffi_cached_object = None
return val
+ def __reduce_ex__(self, protocol):
+ return (bytes, (bytes(self),))
+
# pylint: disable=no-self-argument
def __from_tvm_ffi_object__(cls, obj: Any) -> "Bytes":
"""Construct ``Bytes`` from an FFI object (internal)."""
diff --git a/tests/python/test_string.py b/tests/python/test_string.py
index 499270bb..c8dfdb46 100644
--- a/tests/python/test_string.py
+++ b/tests/python/test_string.py
@@ -18,6 +18,7 @@
import pickle
import tvm_ffi
+import tvm_ffi.testing
def test_string() -> None:
@@ -33,6 +34,15 @@ def test_string() -> None:
s4 = pickle.loads(pickle.dumps(s))
assert s4 == "hello"
+ assert type(s4) is str
+
+ cached = fecho("x" * 200)
+ assert isinstance(cached, tvm_ffi.core.String)
+ assert cached._tvm_ffi_cached_object is not None
+
+ cached_roundtrip = pickle.loads(pickle.dumps(cached))
+ assert cached_roundtrip == cached
+ assert type(cached_roundtrip) is str
def test_bytes() -> None:
@@ -52,7 +62,15 @@ def test_bytes() -> None:
b5 = pickle.loads(pickle.dumps(b))
assert b5 == b"hello"
- assert isinstance(b5, tvm_ffi.core.Bytes)
+ assert type(b5) is bytes
+
+ cached = fecho(b"x" * 200)
+ assert isinstance(cached, tvm_ffi.core.Bytes)
+ assert cached._tvm_ffi_cached_object is not None
+
+ cached_roundtrip = pickle.loads(pickle.dumps(cached))
+ assert cached_roundtrip == cached
+ assert type(cached_roundtrip) is bytes
def test_string_find_substr() -> None:
