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 77861331 feat(stub): use distinct container type origins in type
schema and stub generation (#469)
77861331 is described below
commit 7786133167902a810dd4717fd0e7d39f4c6f7e99
Author: Junru Shao <[email protected]>
AuthorDate: Sat Feb 21 20:36:48 2026 -0800
feat(stub): use distinct container type origins in type schema and stub
generation (#469)
## Summary
- **Distinct container type origins**: TypeSchema now emits `Array`,
`List`, `Map`, `Dict` instead of collapsing all sequence types to `list`
and all mapping types to `dict`. This preserves the immutable
(copy-on-write) vs mutable (shared-reference) distinction through the
entire schema → annotation pipeline.
- **Correct stub annotations**: Generated stubs now use
`MutableSequence`/`MutableMapping` for `List`/`Dict` parameters and
`Sequence`/`Mapping` for `Array`/`Map`, matching the actual Python ABCs
each container implements.
- **New containers concept page**: Added `docs/concepts/containers.rst`
documenting all five container types, their mutability semantics, thread
safety, and usage guidance. Updated C++ and Python language guides with
List and Dict sections.
## Details
### Type schema changes
- `_TYPE_SCHEMA_ORIGIN_CONVERTER` in `type_info.pxi`: `ffi.Array` →
`Array`, `ffi.List` → `List`, `ffi.Map` → `Map`, `ffi.Dict` → `Dict`
- `TY_MAP_DEFAULTS` in `stub/consts.py`: `Array` → `Sequence`, `List` →
`MutableSequence`, `Map` → `Mapping`, `Dict` → `MutableMapping`
- Backward compat: raw `list`/`dict` origins still parse correctly in
`TypeSchema`
### Stub codegen improvements
- Collision avoidance: when a function name matches a type import
suffix, the type import gets an alias (e.g. `_StructuralKey`) to avoid
shadowing
- D105 (undocumented-magic-method) added to global ruff ignore list —
generated dunder stubs no longer need inline `# noqa` comments
- New `DeepCopy` and `__ffi_shallow_copy__` stubs appear in generated
code
### Documentation
- New `docs/concepts/containers.rst` concept page
- Updated `docs/concepts/any.rst` (List/Dict in type table)
- Updated `docs/guides/cpp_lang_guide.md` (List, Dict sections)
- Updated `docs/guides/python_lang_guide.md` (List, Dict, Tuple
sections)
- Updated `CLAUDE.md` and `.claude/skills/devtools/SKILL.md` (stub
generation workflow)
## Test plan
- [x] `uv run pytest -vvs tests/python` — 543 passed, 22 skipped, 1
xfailed
- [x] `uv run pre-commit run --all-files` — 27/27 hooks passed
- [x] `uv run ty check python/tvm_ffi/_ffi_api.py` — all checks passed
(no more `invalid-type-form` error)
- [x] CI: lint, clang-tidy, full test suite
- [x] Verify Sphinx doc build renders new containers page correctly
---
.claude/skills/devtools/SKILL.md | 12 +++
CLAUDE.md | 12 +++
docs/concepts/any.rst | 10 +-
docs/concepts/containers.rst | 208 ++++++++++++++++++++++++++++++++++++
docs/guides/cpp_lang_guide.md | 86 ++++++++++++++-
docs/guides/python_lang_guide.md | 73 ++++++++++++-
docs/index.rst | 1 +
pyproject.toml | 1 +
python/tvm_ffi/_ffi_api.py | 54 +++++-----
python/tvm_ffi/access_path.py | 3 +
python/tvm_ffi/cython/type_info.pxi | 30 +++---
python/tvm_ffi/structural.py | 2 +
python/tvm_ffi/stub/codegen.py | 20 +++-
python/tvm_ffi/stub/consts.py | 7 +-
python/tvm_ffi/testing/_ffi_api.py | 12 +--
python/tvm_ffi/testing/testing.py | 7 +-
tests/python/test_container.py | 12 +--
tests/python/test_metadata.py | 84 ++++++++-------
tests/python/test_stubgen.py | 90 ++++++++++++++++
19 files changed, 614 insertions(+), 110 deletions(-)
diff --git a/.claude/skills/devtools/SKILL.md b/.claude/skills/devtools/SKILL.md
index d97bae21..b1ef105f 100644
--- a/.claude/skills/devtools/SKILL.md
+++ b/.claude/skills/devtools/SKILL.md
@@ -173,6 +173,18 @@ Verify the install:
uv run python -c "import tvm_ffi; print(tvm_ffi.__version__)"
```
+### Update Python stubs
+
+After building (or after C++/Cython reflection changes), regenerate inline
type stubs:
+
+```bash
+uv run tvm-ffi-stubgen python
+```
+
+This updates inline stub blocks (between `tvm-ffi-stubgen(begin)` /
`tvm-ffi-stubgen(end)`
+markers) inside `.py` files with type annotations derived from the C++
reflection
+registry (field types, method signatures, global function schemas).
+
### Run Python tests
Install with test dependencies, then run pytest:
diff --git a/CLAUDE.md b/CLAUDE.md
index ad41b51c..b2b73e3b 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -54,6 +54,18 @@ uv pip install --force-reinstall --verbose -e .
C++/Cython changes always require re-running this command. Pure Python changes
are reflected immediately.
+### Update Python stubs
+
+After building (or after C++/Cython reflection changes), regenerate inline
type stubs:
+
+```bash
+uv run tvm-ffi-stubgen python
+```
+
+This updates inline stub blocks (between `tvm-ffi-stubgen(begin)` /
`tvm-ffi-stubgen(end)`
+markers) inside `.py` files with type annotations derived from the C++
reflection
+registry (field types, method signatures, global function schemas).
+
### C++-only build
```bash
diff --git a/docs/concepts/any.rst b/docs/concepts/any.rst
index 601e448d..95095109 100644
--- a/docs/concepts/any.rst
+++ b/docs/concepts/any.rst
@@ -69,7 +69,7 @@ Use :cpp:class:`~tvm::ffi::Any` for return values to transfer
ownership to the c
Container Storage
~~~~~~~~~~~~~~~~~
-:cpp:class:`~tvm::ffi::Any` can be stored in containers like
:cpp:class:`~tvm::ffi::Map` and :cpp:class:`~tvm::ffi::Array`:
+:cpp:class:`~tvm::ffi::Any` can be stored in containers like
:cpp:class:`~tvm::ffi::Array`, :cpp:class:`~tvm::ffi::List`,
:cpp:class:`~tvm::ffi::Map`, and :cpp:class:`~tvm::ffi::Dict` (see
:doc:`containers` for details):
.. code-block:: cpp
@@ -393,9 +393,15 @@ how reference counting works.
* - :cpp:class:`ModuleObj* <tvm::ffi::ModuleObj>`
- :cpp:enumerator:`kTVMFFIModule <TVMFFITypeIndex::kTVMFFIModule>` = 73
- :cpp:member:`~TVMFFIAny::v_obj`
+ * - :cpp:class:`ListObj* <tvm::ffi::ListObj>`
+ - :cpp:enumerator:`kTVMFFIList <TVMFFITypeIndex::kTVMFFIList>` = 75
+ - :cpp:member:`~TVMFFIAny::v_obj`
+ * - :cpp:class:`DictObj* <tvm::ffi::DictObj>`
+ - :cpp:enumerator:`kTVMFFIDict <TVMFFITypeIndex::kTVMFFIDict>` = 76
+ - :cpp:member:`~TVMFFIAny::v_obj`
-Heap-allocated objects - :cpp:class:`~tvm::ffi::String`,
:cpp:class:`~tvm::ffi::Function`, :cpp:class:`~tvm::ffi::Tensor`,
:cpp:class:`~tvm::ffi::Array`, :cpp:class:`~tvm::ffi::Map`, and custom types -
are
+Heap-allocated objects - :cpp:class:`~tvm::ffi::String`,
:cpp:class:`~tvm::ffi::Function`, :cpp:class:`~tvm::ffi::Tensor`,
:cpp:class:`~tvm::ffi::Array`, :cpp:class:`~tvm::ffi::List`,
:cpp:class:`~tvm::ffi::Map`, :cpp:class:`~tvm::ffi::Dict`, and custom types -
are
stored as pointers to reference-counted :cpp:class:`TVMFFIObject` headers:
.. code-block:: cpp
diff --git a/docs/concepts/containers.rst b/docs/concepts/containers.rst
new file mode 100644
index 00000000..45e35254
--- /dev/null
+++ b/docs/concepts/containers.rst
@@ -0,0 +1,208 @@
+.. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+.. http://www.apache.org/licenses/LICENSE-2.0
+
+.. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+Containers
+==========
+
+TVM-FFI provides five built-in container types for storing and exchanging
+collections of values across C++, Python, and Rust. They are all
+heap-allocated, reference-counted objects that can be stored in
+:cpp:class:`~tvm::ffi::Any` and passed through the FFI boundary.
+
+The containers split into two categories: **immutable** containers that use
+copy-on-write semantics, and **mutable** containers that use shared-reference
+semantics.
+
+Overview
+--------
+
+.. list-table::
+ :header-rows: 1
+ :widths: 14 20 20 16 30
+
+ * - Type
+ - C++ Class
+ - Python Class
+ - Mutability
+ - Semantics
+ * - Array
+ - :cpp:class:`Array\<T\> <tvm::ffi::Array>`
+ - :py:class:`tvm_ffi.Array`
+ - Immutable
+ - Homogeneous sequence with copy-on-write
+ * - List
+ - :cpp:class:`List\<T\> <tvm::ffi::List>`
+ - :py:class:`tvm_ffi.List`
+ - Mutable
+ - Homogeneous sequence with shared-reference
+ * - Tuple
+ - :cpp:class:`Tuple\<Ts...\> <tvm::ffi::Tuple>`
+ - (backed by :py:class:`tvm_ffi.Array`)
+ - Immutable
+ - Heterogeneous fixed-size sequence (backed by ArrayObj)
+ * - Map
+ - :cpp:class:`Map\<K, V\> <tvm::ffi::Map>`
+ - :py:class:`tvm_ffi.Map`
+ - Immutable
+ - Homogeneous key-value mapping with copy-on-write
+ * - Dict
+ - :cpp:class:`Dict\<K, V\> <tvm::ffi::Dict>`
+ - :py:class:`tvm_ffi.Dict`
+ - Mutable
+ - Homogeneous key-value mapping with shared-reference
+
+Immutable Containers (Copy-on-Write)
+-------------------------------------
+
+Array
+~~~~~
+
+``Array<T>`` is an immutable homogeneous sequence backed by
+:cpp:class:`~tvm::ffi::ArrayObj`. It implements **copy-on-write** semantics:
+when a mutation method is called in C++ (e.g. ``push_back``, ``Set``), the
+array checks whether the backing storage is uniquely owned. If it is shared
+with other handles, it copies the data first so that existing handles are
+unaffected.
+
+.. code-block:: cpp
+
+ ffi::Array<int> a = {1, 2, 3};
+ ffi::Array<int> b = a; // b shares the same ArrayObj
+ a.push_back(4); // copy-on-write: a gets a new backing storage
+ assert(a.size() == 4);
+ assert(b.size() == 3); // b is unchanged
+
+In Python, :py:class:`tvm_ffi.Array` implements ``collections.abc.Sequence``
+(read-only). When a Python ``list`` or ``tuple`` is passed to an FFI function,
+it is automatically converted to ``Array``.
+
+Tuple
+~~~~~
+
+``Tuple<T1, T2, ...>`` is an immutable heterogeneous fixed-size sequence. It is
+backed by the same :cpp:class:`~tvm::ffi::ArrayObj` as ``Array``, but provides
+compile-time type safety for each element position via C++ variadic templates.
+
+.. code-block:: cpp
+
+ ffi::Tuple<int, ffi::String, bool> t(42, "hello", true);
+ int x = t.get<0>(); // 42
+ ffi::String s = t.get<1>(); // "hello"
+
+In Python, ``Tuple`` does not have a separate class -- Python tuples passed
+through the FFI are converted to ``Array``.
+
+Map
+~~~
+
+``Map<K, V>`` is an immutable homogeneous key-value mapping backed by
+:cpp:class:`~tvm::ffi::MapObj`. It implements **copy-on-write** semantics
+(same principle as ``Array``). Insertion order is preserved.
+
+.. code-block:: cpp
+
+ ffi::Map<ffi::String, int> m = {{"Alice", 100}, {"Bob", 95}};
+ ffi::Map<ffi::String, int> m2 = m; // m2 shares the same MapObj
+ m.Set("Charlie", 88); // copy-on-write
+ assert(m.size() == 3);
+ assert(m2.size() == 2); // m2 is unchanged
+
+In Python, :py:class:`tvm_ffi.Map` implements ``collections.abc.Mapping``
+(read-only). When a Python ``dict`` is passed to an FFI function, it is
+automatically converted to ``Map``.
+
+Mutable Containers (Shared Reference)
+--------------------------------------
+
+List
+~~~~
+
+``List<T>`` is a mutable homogeneous sequence backed by
+:cpp:class:`~tvm::ffi::ListObj`. Unlike ``Array``, it does **not** use
+copy-on-write. Mutations happen directly on the underlying shared object, and
+**all handles** sharing the same ``ListObj`` see the mutations immediately.
+
+.. code-block:: cpp
+
+ ffi::List<int> a = {1, 2, 3};
+ ffi::List<int> b = a; // b shares the same ListObj
+ a.push_back(4); // in-place mutation
+ assert(a.size() == 4);
+ assert(b.size() == 4); // b sees the mutation
+
+In Python, :py:class:`tvm_ffi.List` implements
+``collections.abc.MutableSequence`` and supports ``append``, ``insert``,
+``__setitem__``, ``__delitem__``, ``pop``, ``reverse``, ``clear``, and
+``extend``.
+
+Dict
+~~~~
+
+``Dict<K, V>`` is a mutable homogeneous key-value mapping backed by
+:cpp:class:`~tvm::ffi::DictObj`. Like ``List``, mutations happen directly on
+the shared object with no copy-on-write.
+
+.. code-block:: cpp
+
+ ffi::Dict<ffi::String, int> d = {{"Alice", 100}};
+ ffi::Dict<ffi::String, int> d2 = d; // d2 shares the same DictObj
+ d.Set("Bob", 95); // in-place mutation
+ assert(d.size() == 2);
+ assert(d2.size() == 2); // d2 sees the mutation
+
+In Python, :py:class:`tvm_ffi.Dict` implements
+``collections.abc.MutableMapping`` and supports ``__setitem__``,
+``__delitem__``, ``pop``, ``clear``, and ``update``.
+
+When to Use Each Type
+---------------------
+
+.. list-table::
+ :header-rows: 1
+ :widths: 50 20
+
+ * - Use Case
+ - Container
+ * - Immutable snapshot of a sequence (e.g. function arguments)
+ - Array
+ * - Building up or modifying a sequence in-place
+ - List
+ * - Fixed heterogeneous collection (e.g. a multi-typed return value)
+ - Tuple
+ * - Immutable key-value lookup (e.g. configuration)
+ - Map
+ * - Mutable key-value store (e.g. accumulating results)
+ - Dict
+
+Thread Safety
+-------------
+
+**Immutable containers** (``Array``, ``Tuple``, ``Map``) can be safely shared
+across threads for **read-only** access. Copy-on-write mutations are
+thread-safe because they create a new backing object when the storage is
shared.
+
+**Mutable containers** (``List``, ``Dict``) are **NOT** thread-safe. If
+multiple threads need to read or write the same ``List`` or ``Dict``, external
+synchronization (e.g. a mutex) is required.
+
+Further Reading
+---------------
+
+- :doc:`../guides/cpp_lang_guide`: C++ examples for all container types
+- :doc:`../guides/python_lang_guide`: Python examples and conversion rules
+- :doc:`any`: How containers are stored in the type-erased
:cpp:class:`~tvm::ffi::Any` value
+- :doc:`object_and_class`: The object system underlying all container types
diff --git a/docs/guides/cpp_lang_guide.md b/docs/guides/cpp_lang_guide.md
index f8ab8e59..f222e8b7 100644
--- a/docs/guides/cpp_lang_guide.md
+++ b/docs/guides/cpp_lang_guide.md
@@ -486,7 +486,16 @@ will happen automatically.
## Container Types
To enable effective passing and storing of collections of values that are
compatible with tvm-ffi,
-we provide several built-in container types.
+we provide several built-in container types. See
[Containers](../concepts/containers.rst) for a
+conceptual overview.
+
+| Type | Header | Mutability | Semantics |
+| ------ | -------- | ------------ | ----------- |
+| `Array<T>` | `container/array.h` | Immutable (copy-on-write) | Homogeneous
sequence |
+| `List<T>` | `container/list.h` | Mutable (shared reference) | Homogeneous
sequence |
+| `Tuple<Ts...>` | `container/tuple.h` | Immutable (copy-on-write) |
Heterogeneous fixed-size sequence |
+| `Map<K,V>` | `container/map.h` | Immutable (copy-on-write) | Homogeneous
key-value mapping |
+| `Dict<K,V>` | `container/dict.h` | Mutable (shared reference) | Homogeneous
key-value mapping |
### Array
@@ -529,6 +538,44 @@ being passed into the Function.
**Performance note:** Repeatedly converting Any to `Array<T>` can incur
repeated
checking overhead at each element. Consider using `Array<Any>` to defer
checking or only run conversion once.
+### List
+
+`List<T>` provides a mutable sequence container with shared reference
semantics.
+Unlike `Array`, mutations happen directly on the underlying shared `ListObj` --
+there is no copy-on-write. All handles sharing the same `ListObj` see mutations
+immediately.
+
+```cpp
+#include <tvm/ffi/container/list.h>
+
+void ExampleList() {
+ namespace ffi = tvm::ffi;
+ ffi::List<int> numbers = {1, 2, 3};
+ // EXPECT_EQ is used here for demonstration purposes (testing framework)
+ EXPECT_EQ(numbers.size(), 3);
+ EXPECT_EQ(numbers[0], 1);
+
+ // Mutate in-place
+ numbers.push_back(4);
+ EXPECT_EQ(numbers.size(), 4);
+ numbers.Set(0, 10);
+ EXPECT_EQ(numbers[0], 10);
+
+ // Shared reference semantics: both handles see the mutation
+ ffi::List<int> alias = numbers;
+ alias.push_back(5);
+ EXPECT_EQ(numbers.size(), 5);
+}
+```
+
+Under the hood, `List` is backed by a reference-counted `ListObj` that stores
+a collection of Any values. Like `Array`, conversion from Any to `List<T>`
will result in
+runtime checks of elements.
+
+**When to use List vs Array:** Use `Array<T>` when you need an immutable
snapshot
+(e.g., passing a collection through FFI boundaries where the receiver should
not
+mutate). Use `List<T>` when you need to build up or modify a collection in
place.
+
### Tuple
`Tuple<Types...>` provides type-safe fixed-size collections.
@@ -585,6 +632,43 @@ being passed into the Function.
**Performance note:** Repeatedly converting Any to `Map<K, V>` can incur
repeated
checking overhead at each element. Consider using `Map<Any, Any>` to defer
checking or only run conversion once.
+### Dict
+
+`Dict<K, V>` provides a mutable key-value mapping with shared reference
semantics.
+Unlike `Map`, mutations happen directly on the underlying shared `DictObj`.
+All handles sharing the same `DictObj` see mutations immediately.
+
+```cpp
+#include <tvm/ffi/container/dict.h>
+
+void ExampleDict() {
+ namespace ffi = tvm::ffi;
+ ffi::Dict<ffi::String, int> scores = {{"Alice", 100}, {"Bob", 95}};
+ // EXPECT_EQ is used here for demonstration purposes (testing framework)
+ EXPECT_EQ(scores.size(), 2);
+ EXPECT_EQ(scores.at("Alice"), 100);
+
+ // Mutate in-place
+ scores.Set("Charlie", 88);
+ EXPECT_EQ(scores.size(), 3);
+
+ // Shared reference semantics: both handles see the mutation
+ ffi::Dict<ffi::String, int> alias = scores;
+ alias.Set("Dave", 92);
+ EXPECT_EQ(scores.size(), 4);
+
+ // Erase
+ scores.erase("Bob");
+ EXPECT_EQ(scores.size(), 3);
+}
+```
+
+Under the hood, `Dict` is backed by a reference-counted `DictObj` that stores
+a collection of Any key-value pairs. Like `Map`, the `Dict` preserves
insertion order.
+
+**When to use Dict vs Map:** Use `Map<K, V>` when you need an immutable
snapshot
+with copy-on-write semantics. Use `Dict<K, V>` when you need mutable in-place
operations.
+
### Optional
`Optional<T>` provides a safe way to handle values that may or may not exist.
diff --git a/docs/guides/python_lang_guide.md b/docs/guides/python_lang_guide.md
index 18741599..5d4c7909 100644
--- a/docs/guides/python_lang_guide.md
+++ b/docs/guides/python_lang_guide.md
@@ -112,19 +112,55 @@ assert tvm_ffi.get_global_func("example.add_one")(1) == 2
## Container Types
+TVM FFI provides five container types that split into **immutable**
(copy-on-write) and
+**mutable** (shared reference) variants. See the [Containers concept
page](../concepts/containers.rst) for a
+detailed overview.
+
+| Type | Mutability | Python ABC | Semantics |
+| ------ | ----------- | ------------ | ----------- |
+| {py:class}`tvm_ffi.Array` | Immutable | `Sequence[T]` | Homogeneous sequence
(copy-on-write) |
+| {py:class}`tvm_ffi.List` | Mutable | `MutableSequence[T]` | Homogeneous
sequence (shared reference) |
+| {py:class}`tvm_ffi.Map` | Immutable | `Mapping[K, V]` | Homogeneous mapping
(copy-on-write) |
+| {py:class}`tvm_ffi.Dict` | Mutable | `MutableMapping[K, V]` | Homogeneous
mapping (shared reference) |
+
+### Array (immutable sequence)
+
When an FFI function takes arguments from lists/tuples, they will be converted
into {py:class}`tvm_ffi.Array`.
+Arrays are **read-only** from Python -- they support indexing and iteration
but not mutation.
```python
import tvm_ffi
-# Lists become Arrays
+# Python lists become Arrays
arr = tvm_ffi.convert([1, 2, 3, 4])
assert isinstance(arr, tvm_ffi.Array)
assert len(arr) == 4
assert arr[0] == 1
+# arr[0] = 10 # TypeError: Array does not support item assignment
```
-Dictionaries will be converted to {py:class}`tvm_ffi.Map`
+### List (mutable sequence)
+
+{py:class}`tvm_ffi.List` is a mutable sequence with shared-reference semantics.
+All handles sharing the same underlying `ListObj` see mutations immediately.
+
+```python
+import tvm_ffi
+
+lst = tvm_ffi.List([1, 2, 3])
+assert len(lst) == 3
+
+# Mutate in-place
+lst.append(4)
+assert len(lst) == 4
+lst[0] = 10
+assert lst[0] == 10
+```
+
+### Map (immutable mapping)
+
+Dictionaries will be converted to {py:class}`tvm_ffi.Map`.
+Maps are **read-only** from Python -- they support key lookup and iteration
but not mutation.
```python
import tvm_ffi
@@ -134,10 +170,39 @@ assert isinstance(map_obj, tvm_ffi.Map)
assert len(map_obj) == 2
assert map_obj["a"] == 1
assert map_obj["b"] == 2
+# map_obj["c"] = 3 # TypeError: Map does not support item assignment
+```
+
+### Dict (mutable mapping)
+
+{py:class}`tvm_ffi.Dict` is a mutable mapping with shared-reference semantics.
+All handles sharing the same underlying `DictObj` see mutations immediately.
+
+```python
+import tvm_ffi
+
+d = tvm_ffi.Dict({"a": 1, "b": 2})
+d["c"] = 3
+assert len(d) == 3
+del d["a"]
+assert len(d) == 2
+```
+
+### Tuple conversion
+
+When a Python tuple is passed to an FFI function, it is converted to
{py:class}`tvm_ffi.Array`
+(the same backing type). The C++ `Tuple<Ts...>` template provides compile-time
+heterogeneous type safety, but on the Python side both lists and tuples map to
`Array`.
+
+```python
+import tvm_ffi
+
+t = tvm_ffi.convert((1, "hello", 3.14))
+assert isinstance(t, tvm_ffi.Array)
```
-When container values are returned from FFI functions, they are also stored in
these
-types respectively.
+When container values are returned from FFI functions, they are stored in the
+corresponding container type (Array, List, Map, or Dict).
## Inline Module
diff --git a/docs/index.rst b/docs/index.rst
index 73b1eea9..c00e6613 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -61,6 +61,7 @@ Table of Contents
concepts/abi_overview.rst
concepts/any.rst
+ concepts/containers.rst
concepts/object_and_class.rst
concepts/tensor.rst
concepts/func_module.rst
diff --git a/pyproject.toml b/pyproject.toml
index c1a8f76e..1c1f6339 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -197,6 +197,7 @@ select = [
ignore = [
"PLR2004", # pylint: magic-value-comparison
"ANN401", # flake8-annotations: any-type
+ "D105", # pydocstyle: undocumented-magic-method
"D203", # pydocstyle: incorrect-blank-line-before-class
"D213", # pydocstyle: multi-line-summary-second-line
]
diff --git a/python/tvm_ffi/_ffi_api.py b/python/tvm_ffi/_ffi_api.py
index b60ab808..97f19792 100644
--- a/python/tvm_ffi/_ffi_api.py
+++ b/python/tvm_ffi/_ffi_api.py
@@ -23,8 +23,8 @@ from __future__ import annotations
from .registry import init_ffi_api as _FFI_INIT_FUNC
from typing import TYPE_CHECKING
if TYPE_CHECKING:
- from collections.abc import Mapping, Sequence
- from tvm_ffi import Module
+ from collections.abc import Mapping, MutableMapping, MutableSequence,
Sequence
+ from tvm_ffi import Module, Object, StructuralKey as _StructuralKey
from tvm_ffi.access_path import AccessPath
from typing import Any, Callable
# isort: on
@@ -36,40 +36,41 @@ if TYPE_CHECKING:
_FFI_INIT_FUNC("ffi", __name__)
if TYPE_CHECKING:
def Array(*args: Any) -> Any: ...
+ def ArrayContains(_0: Sequence[Any], _1: Any, /) -> bool: ...
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 DeepCopy(_0: Any, /) -> Any: ...
def Dict(*args: Any) -> Any: ...
- def DictClear(_0: Any, /) -> None: ...
- def DictCount(_0: Any, _1: Any, /) -> int: ...
- def DictErase(_0: Any, _1: Any, /) -> None: ...
- def DictForwardIterFunctor(_0: Any, /) -> Callable[..., Any]: ...
- def DictGetItem(_0: Any, _1: Any, /) -> Any: ...
- def DictGetItemOrMissing(_0: Any, _1: Any, /) -> Any: ...
- def DictSetItem(_0: Any, _1: Any, _2: Any, /) -> None: ...
- def DictSize(_0: Any, /) -> int: ...
+ def DictClear(_0: MutableMapping[Any, Any], /) -> None: ...
+ def DictCount(_0: MutableMapping[Any, Any], _1: Any, /) -> int: ...
+ def DictErase(_0: MutableMapping[Any, Any], _1: Any, /) -> None: ...
+ def DictForwardIterFunctor(_0: MutableMapping[Any, Any], /) ->
Callable[..., Any]: ...
+ def DictGetItem(_0: MutableMapping[Any, Any], _1: Any, /) -> Any: ...
+ def DictGetItemOrMissing(_0: MutableMapping[Any, Any], _1: Any, /) -> Any:
...
+ def DictSetItem(_0: MutableMapping[Any, Any], _1: Any, _2: Any, /) ->
None: ...
+ def DictSize(_0: MutableMapping[Any, Any], /) -> int: ...
def FromJSONGraph(_0: Any, /) -> Any: ...
def FromJSONGraphString(_0: str, /) -> Any: ...
- def GetInvalidObject() -> Any: ...
def FunctionListGlobalNamesFunctor() -> Callable[..., Any]: ...
def FunctionRemoveGlobal(_0: str, /) -> bool: ...
def GetFirstStructuralMismatch(_0: Any, _1: Any, _2: bool, _3: bool, /) ->
tuple[AccessPath, AccessPath] | None: ...
def GetGlobalFuncMetadata(_0: str, /) -> str: ...
+ def GetInvalidObject() -> Object: ...
def GetRegisteredTypeKeys() -> Sequence[str]: ...
def List(*args: Any) -> Any: ...
- def ListAppend(_0: Sequence[Any], _1: Any, /) -> None: ...
- def ListClear(_0: Sequence[Any], /) -> None: ...
- def ListContains(_0: Sequence[Any], _1: Any, /) -> bool: ...
- def ListErase(_0: Sequence[Any], _1: int, /) -> None: ...
- def ListEraseRange(_0: Sequence[Any], _1: int, _2: int, /) -> None: ...
- def ListGetItem(_0: Sequence[Any], _1: int, /) -> Any: ...
- def ListInsert(_0: Sequence[Any], _1: int, _2: Any, /) -> None: ...
- def ListPop(_0: Sequence[Any], _1: int, /) -> Any: ...
- def ListReplaceSlice(_0: Sequence[Any], _1: int, _2: int, _3:
Sequence[Any], /) -> None: ...
- def ListReverse(_0: Sequence[Any], /) -> None: ...
- def ListSetItem(_0: Sequence[Any], _1: int, _2: Any, /) -> None: ...
- def ListSize(_0: Sequence[Any], /) -> int: ...
+ def ListAppend(_0: MutableSequence[Any], _1: Any, /) -> None: ...
+ def ListClear(_0: MutableSequence[Any], /) -> None: ...
+ def ListContains(_0: MutableSequence[Any], _1: Any, /) -> bool: ...
+ def ListErase(_0: MutableSequence[Any], _1: int, /) -> None: ...
+ def ListEraseRange(_0: MutableSequence[Any], _1: int, _2: int, /) -> None:
...
+ def ListGetItem(_0: MutableSequence[Any], _1: int, /) -> Any: ...
+ def ListInsert(_0: MutableSequence[Any], _1: int, _2: Any, /) -> None: ...
+ def ListPop(_0: MutableSequence[Any], _1: int, /) -> Any: ...
+ def ListReplaceSlice(_0: MutableSequence[Any], _1: int, _2: int, _3:
MutableSequence[Any], /) -> None: ...
+ def ListReverse(_0: MutableSequence[Any], /) -> None: ...
+ def ListSetItem(_0: MutableSequence[Any], _1: int, _2: Any, /) -> None: ...
+ def ListSize(_0: MutableSequence[Any], /) -> int: ...
def MakeObjectFromPackedArgs(*args: Any) -> Any: ...
def Map(*args: Any) -> Any: ...
def MapCount(_0: Mapping[Any, Any], _1: Any, /) -> int: ...
@@ -94,10 +95,10 @@ if TYPE_CHECKING:
def ReprPrint(_0: Any, /) -> str: ...
def Shape(*args: Any) -> Any: ...
def String(_0: str, /) -> str: ...
- def StructuralKey(_0: Any, /) -> Any: ...
- def StructuralKeyEqual(_0: Any, _1: Any, /) -> bool: ...
def StructuralEqual(_0: Any, _1: Any, _2: bool, _3: bool, /) -> bool: ...
def StructuralHash(_0: Any, _1: bool, _2: bool, /) -> int: ...
+ def StructuralKey(_0: Any, /) -> _StructuralKey: ...
+ def StructuralKeyEqual(_0: Any, _1: Any, /) -> bool: ...
def SystemLib(*args: Any) -> Any: ...
def ToJSONGraph(_0: Any, _1: Any, /) -> Any: ...
def ToJSONGraphString(_0: Any, _1: Any, /) -> str: ...
@@ -112,6 +113,7 @@ __all__ = [
"ArrayGetItem",
"ArraySize",
"Bytes",
+ "DeepCopy",
"Dict",
"DictClear",
"DictCount",
diff --git a/python/tvm_ffi/access_path.py b/python/tvm_ffi/access_path.py
index 6980e49a..243a5d4c 100644
--- a/python/tvm_ffi/access_path.py
+++ b/python/tvm_ffi/access_path.py
@@ -55,6 +55,8 @@ class AccessStep(Object):
# fmt: off
kind: int
key: Any
+ if TYPE_CHECKING:
+ def __ffi_shallow_copy__(self, /) -> Object: ...
# fmt: on
# tvm-ffi-stubgen(end)
@@ -88,6 +90,7 @@ class AccessPath(Object):
step: AccessStep | None
depth: int
if TYPE_CHECKING:
+ def __ffi_shallow_copy__(self, /) -> Object: ...
@staticmethod
def _root() -> AccessPath: ...
def _extend(self, _1: AccessStep, /) -> AccessPath: ...
diff --git a/python/tvm_ffi/cython/type_info.pxi
b/python/tvm_ffi/cython/type_info.pxi
index e2ecf9e2..050be909 100644
--- a/python/tvm_ffi/cython/type_info.pxi
+++ b/python/tvm_ffi/cython/type_info.pxi
@@ -72,10 +72,10 @@ _TYPE_SCHEMA_ORIGIN_CONVERTER = {
"Optional": "Optional",
"Tuple": "tuple",
"ffi.Function": "Callable",
- "ffi.Array": "list",
- "ffi.List": "list",
- "ffi.Map": "dict",
- "ffi.Dict": "dict",
+ "ffi.Array": "Array",
+ "ffi.List": "List",
+ "ffi.Map": "Map",
+ "ffi.Dict": "Dict",
"ffi.OpaquePyObject": "Any",
"ffi.Object": "Object",
"ffi.Tensor": "Tensor",
@@ -113,12 +113,12 @@ class TypeSchema:
assert len(args) >= 2, "Union must have at least two arguments"
elif origin == "Optional":
assert len(args) == 1, "Optional must have exactly one argument"
- elif origin == "list":
- assert len(args) in (0, 1), "list must have 0 or 1 argument"
+ elif origin in ("list", "Array", "List"):
+ assert len(args) in (0, 1), f"{origin} must have 0 or 1 argument"
if args == ():
self.args = (TypeSchema("Any"),)
- elif origin == "dict":
- assert len(args) in (0, 2), "dict must have 0 or 2 arguments"
+ elif origin in ("dict", "Map", "Dict"):
+ assert len(args) in (0, 2), f"{origin} must have 0 or 2 arguments"
if args == ():
self.args = (TypeSchema("Any"), TypeSchema("Any"))
elif origin == "tuple":
@@ -149,8 +149,8 @@ class TypeSchema:
----------
ty_map : Callable[[str], str], optional
A mapping function applied to the schema origin name before
- rendering (e.g. map ``"list" -> "Sequence"`` and
- ``"dict" -> "Mapping"``). If ``None``, the raw origin is used.
+ rendering (e.g. map ``"Array" -> "Array"`` and
+ ``"Map" -> "Map"``). If ``None``, the raw origin is used.
Returns
-------
@@ -173,12 +173,12 @@ class TypeSchema:
s = TypeSchema("Callable", (TypeSchema("int"), TypeSchema("str")))
assert s.repr() == "Callable[[str], int]"
- # Custom mapping to stdlib typing collections
- def _map(t: str) -> str:
- return {"list": "Sequence", "dict": "Mapping"}.get(t, t)
+ # Container types from C++ FFI schemas
+ s =
TypeSchema.from_json_str('{"type":"ffi.Map","args":[{"type":"str"},{"type":"int"}]}')
+ assert s.repr() == "Map[str, int]"
- s =
TypeSchema.from_json_str('{"type":"dict","args":[{"type":"str"},{"type":"int"}]}')
- assert s.repr(_map) == "Mapping[str, int]"
+ s =
TypeSchema.from_json_str('{"type":"ffi.Array","args":[{"type":"int"}]}')
+ assert s.repr() == "Array[int]"
"""
if ty_map is None:
diff --git a/python/tvm_ffi/structural.py b/python/tvm_ffi/structural.py
index 120268a1..798e264f 100644
--- a/python/tvm_ffi/structural.py
+++ b/python/tvm_ffi/structural.py
@@ -205,6 +205,8 @@ class StructuralKey(Object):
# fmt: off
key: Any
hash_i64: int
+ if TYPE_CHECKING:
+ def __ffi_shallow_copy__(self, /) -> Object: ...
# fmt: on
# tvm-ffi-stubgen(end)
diff --git a/python/tvm_ffi/stub/codegen.py b/python/tvm_ffi/stub/codegen.py
index 45e62318..9b9294ec 100644
--- a/python/tvm_ffi/stub/codegen.py
+++ b/python/tvm_ffi/stub/codegen.py
@@ -26,14 +26,22 @@ from .utils import FuncInfo, ImportItem, InitConfig,
ObjectInfo, Options
def _type_suffix_and_record(
- ty_map: dict[str, str], imports: list[ImportItem]
+ ty_map: dict[str, str],
+ imports: list[ImportItem],
+ func_names: set[str] | None = None,
) -> Callable[[str], str]:
def _run(name: str) -> str:
nonlocal ty_map, imports
name = ty_map.get(name, name)
+ suffix = name.rsplit(".", 1)[-1]
if "." in name:
- imports.append(ImportItem(name, type_checking_only=True,
alias=None))
- return name.rsplit(".", 1)[-1]
+ alias = None
+ if func_names and suffix in func_names:
+ alias = f"_{suffix}"
+ imports.append(ImportItem(name, type_checking_only=True,
alias=alias))
+ if alias:
+ return alias
+ return suffix
return _run
@@ -69,7 +77,8 @@ def generate_global_funcs(
),
]
)
- fn_ty_map = _type_suffix_and_record(ty_map, imports)
+ func_names = {f.schema.name.rsplit(".", 1)[-1] for f in global_funcs}
+ fn_ty_map = _type_suffix_and_record(ty_map, imports, func_names=func_names)
results: list[str] = [
"# fmt: off",
f'_FFI_INIT_FUNC("{prefix}", __name__)',
@@ -98,7 +107,8 @@ def generate_object(
"""
assert len(code.lines) >= 2
info = obj_info
- fn_ty_map = _type_suffix_and_record(ty_map, imports)
+ method_names = {m.schema.name.rsplit(".", 1)[-1] for m in info.methods}
+ fn_ty_map = _type_suffix_and_record(ty_map, imports,
func_names=method_names)
if info.methods:
imports.append(
ImportItem(
diff --git a/python/tvm_ffi/stub/consts.py b/python/tvm_ffi/stub/consts.py
index d8b17bcb..d7b13b51 100644
--- a/python/tvm_ffi/stub/consts.py
+++ b/python/tvm_ffi/stub/consts.py
@@ -53,9 +53,10 @@ DEFAULT_SOURCE_EXTS = {".py", ".pyi"}
TY_MAP_DEFAULTS = {
"Any": "typing.Any",
"Callable": "typing.Callable",
- "Mapping": "typing.Mapping",
- "list": "collections.abc.Sequence",
- "dict": "collections.abc.Mapping",
+ "Array": "collections.abc.Sequence",
+ "List": "collections.abc.MutableSequence",
+ "Map": "collections.abc.Mapping",
+ "Dict": "collections.abc.MutableMapping",
"Object": "ffi.Object",
"Tensor": "ffi.Tensor",
"dtype": "ffi.dtype",
diff --git a/python/tvm_ffi/testing/_ffi_api.py
b/python/tvm_ffi/testing/_ffi_api.py
index a83254f2..50c36598 100644
--- a/python/tvm_ffi/testing/_ffi_api.py
+++ b/python/tvm_ffi/testing/_ffi_api.py
@@ -23,7 +23,7 @@ from __future__ import annotations
from ..registry import init_ffi_api as _FFI_INIT_FUNC
from typing import TYPE_CHECKING
if TYPE_CHECKING:
- from collections.abc import Mapping, Sequence
+ from collections.abc import Mapping, MutableMapping, MutableSequence,
Sequence
from tvm_ffi import Device, Object, Tensor, dtype
from tvm_ffi.testing import TestIntPair
from typing import Any, Callable
@@ -55,17 +55,17 @@ if TYPE_CHECKING:
def schema_id_bool(_0: bool, /) -> bool: ...
def schema_id_bytes(_0: bytes, /) -> bytes: ...
def schema_id_device(_0: Device, /) -> Device: ...
- def schema_id_dict_str_int(_0: Mapping[str, int], /) -> Mapping[str, int]:
...
- def schema_id_dict_str_str(_0: Mapping[str, str], /) -> Mapping[str, str]:
...
+ def schema_id_dict_str_int(_0: MutableMapping[str, int], /) ->
MutableMapping[str, int]: ...
+ def schema_id_dict_str_str(_0: MutableMapping[str, str], /) ->
MutableMapping[str, str]: ...
def schema_id_dltensor(_0: Tensor, /) -> Tensor: ...
def schema_id_dtype(_0: dtype, /) -> dtype: ...
def schema_id_float(_0: float, /) -> float: ...
def schema_id_func(_0: Callable[..., Any], /) -> Callable[..., Any]: ...
def schema_id_func_typed(_0: Callable[[int, float, Callable[..., Any]],
None], /) -> Callable[[int, float, Callable[..., Any]], None]: ...
def schema_id_int(_0: int, /) -> int: ...
- def schema_id_list_int(_0: Sequence[int], /) -> Sequence[int]: ...
- def schema_id_list_obj(_0: Sequence[Object], /) -> Sequence[Object]: ...
- def schema_id_list_str(_0: Sequence[str], /) -> Sequence[str]: ...
+ def schema_id_list_int(_0: MutableSequence[int], /) ->
MutableSequence[int]: ...
+ def schema_id_list_obj(_0: MutableSequence[Object], /) ->
MutableSequence[Object]: ...
+ def schema_id_list_str(_0: MutableSequence[str], /) ->
MutableSequence[str]: ...
def schema_id_map(_0: Mapping[Any, Any], /) -> Mapping[Any, Any]: ...
def schema_id_map_str_int(_0: Mapping[str, int], /) -> Mapping[str, int]:
...
def schema_id_map_str_obj(_0: Mapping[str, Object], /) -> Mapping[str,
Object]: ...
diff --git a/python/tvm_ffi/testing/testing.py
b/python/tvm_ffi/testing/testing.py
index 08cc3c38..c5578910 100644
--- a/python/tvm_ffi/testing/testing.py
+++ b/python/tvm_ffi/testing/testing.py
@@ -16,7 +16,7 @@
# under the License.
"""Testing utilities."""
-# ruff: noqa: D102,D105
+# ruff: noqa: D102
# tvm-ffi-stubgen(begin): import-section
# fmt: off
# isort: off
@@ -48,6 +48,7 @@ class TestObjectBase(Object):
v_f64: float
v_str: str
if TYPE_CHECKING:
+ def __ffi_shallow_copy__(self, /) -> Object: ...
def add_i64(self, _1: int, /) -> int: ...
# fmt: on
# tvm-ffi-stubgen(end)
@@ -62,6 +63,7 @@ class TestIntPair(Object):
a: int
b: int
if TYPE_CHECKING:
+ def __ffi_shallow_copy__(self, /) -> Object: ...
@staticmethod
def __c_ffi_init__(_0: int, _1: int, /) -> Object: ...
def sum(self, /) -> int: ...
@@ -77,6 +79,8 @@ class TestObjectDerived(TestObjectBase):
# fmt: off
v_map: Mapping[Any, Any]
v_array: Sequence[Any]
+ if TYPE_CHECKING:
+ def __ffi_shallow_copy__(self, /) -> Object: ...
# fmt: on
# tvm-ffi-stubgen(end)
@@ -109,6 +113,7 @@ class _SchemaAllTypes:
v_variant: str | Sequence[int] | Mapping[str, int]
v_opt_arr_variant: Sequence[int | str] | None
if TYPE_CHECKING:
+ def __ffi_shallow_copy__(self, /) -> Object: ...
def add_int(self, _1: int, /) -> int: ...
def append_int(self, _1: Sequence[int], _2: int, /) -> Sequence[int]:
...
def maybe_concat(self, _1: str | None, _2: str | None, /) -> str |
None: ...
diff --git a/tests/python/test_container.py b/tests/python/test_container.py
index bd7e1d16..a4da236e 100644
--- a/tests/python/test_container.py
+++ b/tests/python/test_container.py
@@ -438,7 +438,7 @@ def test_seq_cross_conv_list_to_array_int() -> None:
def test_seq_cross_conv_array_to_list_int() -> None:
"""Array<int> passed to a function expecting List<int>."""
arr = tvm_ffi.Array([10, 20, 30])
- result = testing.schema_id_list_int(arr)
+ result = testing.schema_id_list_int(arr) # type:
ignore[invalid-argument-type]
assert isinstance(result, tvm_ffi.List)
assert list(result) == [10, 20, 30]
@@ -454,7 +454,7 @@ def test_seq_cross_conv_list_to_array_str() -> None:
def test_seq_cross_conv_array_to_list_str() -> None:
"""Array<String> passed to a function expecting List<String>."""
arr = tvm_ffi.Array(["a", "b", "c"])
- result = testing.schema_id_list_str(arr)
+ result = testing.schema_id_list_str(arr) # type:
ignore[invalid-argument-type]
assert isinstance(result, tvm_ffi.List)
assert list(result) == ["a", "b", "c"]
@@ -470,7 +470,7 @@ def test_seq_cross_conv_empty_list_to_array() -> None:
def test_seq_cross_conv_empty_array_to_list() -> None:
"""Empty Array passed to a function expecting List<int>."""
arr = tvm_ffi.Array([])
- result = testing.schema_id_list_int(arr)
+ result = testing.schema_id_list_int(arr) # type:
ignore[invalid-argument-type]
assert isinstance(result, tvm_ffi.List)
assert len(result) == 0
@@ -682,7 +682,7 @@ def test_map_cross_conv_dict_to_map_str_int() -> None:
def test_map_cross_conv_map_to_dict_str_int() -> None:
"""Map<String, int> passed to a function expecting Dict<String, int>."""
m = tvm_ffi.Map({"a": 1, "b": 2})
- result = testing.schema_id_dict_str_int(m)
+ result = testing.schema_id_dict_str_int(m) # type:
ignore[invalid-argument-type]
assert isinstance(result, tvm_ffi.Dict)
assert result["a"] == 1
assert result["b"] == 2
@@ -699,7 +699,7 @@ def test_map_cross_conv_dict_to_map_str_str() -> None:
def test_map_cross_conv_map_to_dict_str_str() -> None:
"""Map<String, String> passed to a function expecting Dict<String,
String>."""
m = tvm_ffi.Map({"x": "hello", "y": "world"})
- result = testing.schema_id_dict_str_str(m)
+ result = testing.schema_id_dict_str_str(m) # type:
ignore[invalid-argument-type]
assert isinstance(result, tvm_ffi.Dict)
assert result["x"] == "hello"
assert result["y"] == "world"
@@ -716,7 +716,7 @@ def test_map_cross_conv_empty_dict_to_map() -> None:
def test_map_cross_conv_empty_map_to_dict() -> None:
"""Empty Map passed to a function expecting Dict<String, int>."""
m = tvm_ffi.Map({})
- result = testing.schema_id_dict_str_int(m)
+ result = testing.schema_id_dict_str_int(m) # type:
ignore[invalid-argument-type]
assert isinstance(result, tvm_ffi.Dict)
assert len(result) == 0
diff --git a/tests/python/test_metadata.py b/tests/python/test_metadata.py
index ec0d0de7..24883c04 100644
--- a/tests/python/test_metadata.py
+++ b/tests/python/test_metadata.py
@@ -22,10 +22,12 @@ from tvm_ffi.core import TypeInfo, TypeSchema,
_lookup_type_attr
from tvm_ffi.testing import _SchemaAllTypes
-def _replace_list_dict(ty: str) -> str:
+def _replace_container_types(ty: str) -> str:
return {
- "list": "Sequence",
- "dict": "Mapping",
+ "Array": "Sequence",
+ "List": "MutableSequence",
+ "Map": "Mapping",
+ "Dict": "MutableMapping",
}.get(ty, ty)
@@ -52,23 +54,26 @@ def _replace_list_dict(ty: str) -> str:
("testing.schema_id_opt_int", "Callable[[int | None], int | None]"),
("testing.schema_id_opt_str", "Callable[[str | None], str | None]"),
("testing.schema_id_opt_obj", "Callable[[Object | None], Object |
None]"),
- ("testing.schema_id_arr_int", "Callable[[list[int]], list[int]]"),
- ("testing.schema_id_arr_str", "Callable[[list[str]], list[str]]"),
- ("testing.schema_id_arr_obj", "Callable[[list[Object]],
list[Object]]"),
- ("testing.schema_id_arr", "Callable[[list[Any]], list[Any]]"),
- ("testing.schema_id_map_str_int", "Callable[[dict[str, int]],
dict[str, int]]"),
- ("testing.schema_id_map_str_str", "Callable[[dict[str, str]],
dict[str, str]]"),
- ("testing.schema_id_map_str_obj", "Callable[[dict[str, Object]],
dict[str, Object]]"),
- ("testing.schema_id_map", "Callable[[dict[Any, Any]], dict[Any,
Any]]"),
- ("testing.schema_id_dict_str_int", "Callable[[dict[str, int]],
dict[str, int]]"),
- ("testing.schema_id_dict_str_str", "Callable[[dict[str, str]],
dict[str, str]]"),
+ ("testing.schema_id_arr_int", "Callable[[Array[int]], Array[int]]"),
+ ("testing.schema_id_arr_str", "Callable[[Array[str]], Array[str]]"),
+ ("testing.schema_id_arr_obj", "Callable[[Array[Object]],
Array[Object]]"),
+ ("testing.schema_id_arr", "Callable[[Array[Any]], Array[Any]]"),
+ ("testing.schema_id_map_str_int", "Callable[[Map[str, int]], Map[str,
int]]"),
+ ("testing.schema_id_map_str_str", "Callable[[Map[str, str]], Map[str,
str]]"),
+ ("testing.schema_id_map_str_obj", "Callable[[Map[str, Object]],
Map[str, Object]]"),
+ ("testing.schema_id_map", "Callable[[Map[Any, Any]], Map[Any, Any]]"),
+ ("testing.schema_id_dict_str_int", "Callable[[Dict[str, int]],
Dict[str, int]]"),
+ ("testing.schema_id_dict_str_str", "Callable[[Dict[str, str]],
Dict[str, str]]"),
("testing.schema_id_variant_int_str", "Callable[[int | str], int |
str]"),
("testing.schema_packed", "Callable[..., Any]"),
(
"testing.schema_arr_map_opt",
- "Callable[[list[int | None], dict[str, list[int]], str | None],
dict[str, list[int]]]",
+ "Callable[[Array[int | None], Map[str, Array[int]], str | None],
Map[str, Array[int]]]",
+ ),
+ (
+ "testing.schema_variant_mix",
+ "Callable[[int | str | Array[int]], int | str | Array[int]]",
),
- ("testing.schema_variant_mix", "Callable[[int | str | list[int]], int
| str | list[int]]"),
("testing.schema_no_args", "Callable[[], int]"),
("testing.schema_no_return", "Callable[[int], None]"),
("testing.schema_no_args_no_return", "Callable[[], None]"),
@@ -78,12 +83,11 @@ def test_schema_global_func(func_name: str, expected: str)
-> None:
metadata: dict[str, Any] = get_global_func_metadata(func_name)
actual: TypeSchema = TypeSchema.from_json_str(metadata["type_schema"])
assert str(actual) == expected, f"{func_name}: {actual}"
- assert actual.repr(_replace_list_dict) == expected.replace(
- "list",
- "Sequence",
- ).replace(
- "dict",
- "Mapping",
+ assert actual.repr(_replace_container_types) == (
+ expected.replace("Array", "Sequence")
+ .replace("List", "MutableSequence")
+ .replace("Map", "Mapping")
+ .replace("Dict", "MutableMapping")
)
@@ -99,12 +103,12 @@ def test_schema_global_func(func_name: str, expected: str)
-> None:
("v_bytes", "bytes"),
("v_opt_int", "int | None"),
("v_opt_str", "str | None"),
- ("v_arr_int", "list[int]"),
- ("v_arr_str", "list[str]"),
- ("v_map_str_int", "dict[str, int]"),
- ("v_map_str_arr_int", "dict[str, list[int]]"),
- ("v_variant", "str | list[int] | dict[str, int]"),
- ("v_opt_arr_variant", "list[int | str] | None"),
+ ("v_arr_int", "Array[int]"),
+ ("v_arr_str", "Array[str]"),
+ ("v_map_str_int", "Map[str, int]"),
+ ("v_map_str_arr_int", "Map[str, Array[int]]"),
+ ("v_variant", "str | Array[int] | Map[str, int]"),
+ ("v_opt_arr_variant", "Array[int | str] | None"),
],
)
def test_schema_field(field_name: str, expected: str) -> None:
@@ -113,12 +117,11 @@ def test_schema_field(field_name: str, expected: str) ->
None:
if field.name == field_name:
actual: TypeSchema =
TypeSchema.from_json_str(field.metadata["type_schema"])
assert str(actual) == expected, f"{field_name}: {actual}"
- assert actual.repr(_replace_list_dict) == expected.replace(
- "list",
- "Sequence",
- ).replace(
- "dict",
- "Mapping",
+ assert actual.repr(_replace_container_types) == (
+ expected.replace("Array", "Sequence")
+ .replace("List", "MutableSequence")
+ .replace("Map", "Mapping")
+ .replace("Dict", "MutableMapping")
)
break
else:
@@ -129,11 +132,11 @@ def test_schema_field(field_name: str, expected: str) ->
None:
"method_name,expected",
[
("add_int", "Callable[[testing.SchemaAllTypes, int], int]"),
- ("append_int", "Callable[[testing.SchemaAllTypes, list[int], int],
list[int]]"),
+ ("append_int", "Callable[[testing.SchemaAllTypes, Array[int], int],
Array[int]]"),
("maybe_concat", "Callable[[testing.SchemaAllTypes, str | None, str |
None], str | None]"),
(
"merge_map",
- "Callable[[testing.SchemaAllTypes, dict[str, list[int]], dict[str,
list[int]]], dict[str, list[int]]]",
+ "Callable[[testing.SchemaAllTypes, Map[str, Array[int]], Map[str,
Array[int]]], Map[str, Array[int]]]",
),
("make_with", "Callable[[int, float, str], testing.SchemaAllTypes]"),
],
@@ -144,12 +147,11 @@ def test_schema_member_method(method_name: str, expected:
str) -> None:
if method.name == method_name:
actual: TypeSchema =
TypeSchema.from_json_str(method.metadata["type_schema"])
assert str(actual) == expected, f"{method_name}: {actual}"
- assert actual.repr(_replace_list_dict) == expected.replace(
- "list",
- "Sequence",
- ).replace(
- "dict",
- "Mapping",
+ assert actual.repr(_replace_container_types) == (
+ expected.replace("Array", "Sequence")
+ .replace("List", "MutableSequence")
+ .replace("Map", "Mapping")
+ .replace("Dict", "MutableMapping")
)
break
else:
diff --git a/tests/python/test_stubgen.py b/tests/python/test_stubgen.py
index 0c2eb0cf..d615c33b 100644
--- a/tests/python/test_stubgen.py
+++ b/tests/python/test_stubgen.py
@@ -232,6 +232,67 @@ def test_objectinfo_gen_fields_and_methods() -> None:
]
+def test_type_schema_container_origins() -> None:
+ """Test that Array/List/Map/Dict origins are distinct and validated
correctly."""
+ # Array and List: 0 or 1 arg, default to (Any,)
+ for origin in ("Array", "List"):
+ s = TypeSchema(origin)
+ assert s.args == (TypeSchema("Any"),), f"{origin} should default to
(Any,)"
+ s = TypeSchema(origin, (TypeSchema("int"),))
+ assert s.repr() == f"{origin}[int]"
+
+ # Map and Dict: 0 or 2 args, default to (Any, Any)
+ for origin in ("Map", "Dict"):
+ s = TypeSchema(origin)
+ assert s.args == (TypeSchema("Any"), TypeSchema("Any")), (
+ f"{origin} should default to (Any, Any)"
+ )
+ s = TypeSchema(origin, (TypeSchema("str"), TypeSchema("float")))
+ assert s.repr() == f"{origin}[str, float]"
+
+ # from_json_str round-trip through _TYPE_SCHEMA_ORIGIN_CONVERTER
+ s =
TypeSchema.from_json_str('{"type":"ffi.Array","args":[{"type":"int"}]}')
+ assert s.origin == "Array"
+ assert s.repr() == "Array[int]"
+
+ s = TypeSchema.from_json_str('{"type":"ffi.List","args":[{"type":"str"}]}')
+ assert s.origin == "List"
+ assert s.repr() == "List[str]"
+
+ s =
TypeSchema.from_json_str('{"type":"ffi.Map","args":[{"type":"str"},{"type":"int"}]}')
+ assert s.origin == "Map"
+ assert s.repr() == "Map[str, int]"
+
+ s =
TypeSchema.from_json_str('{"type":"ffi.Dict","args":[{"type":"str"},{"type":"float"}]}')
+ assert s.origin == "Dict"
+ assert s.repr() == "Dict[str, float]"
+
+ # Backward compat: "list" and "dict" origins still work
+ s = TypeSchema("list", (TypeSchema("int"),))
+ assert s.repr() == "list[int]"
+ s = TypeSchema("dict", (TypeSchema("str"), TypeSchema("int")))
+ assert s.repr() == "dict[str, int]"
+
+
+def test_objectinfo_gen_fields_container_types() -> None:
+ """Test that ObjectInfo fields render distinct container annotations."""
+ info = ObjectInfo(
+ fields=[
+ NamedTypeSchema("arr", TypeSchema("Array", (TypeSchema("int"),))),
+ NamedTypeSchema("lst", TypeSchema("List", (TypeSchema("str"),))),
+ NamedTypeSchema("mp", TypeSchema("Map", (TypeSchema("str"),
TypeSchema("int")))),
+ NamedTypeSchema("dt", TypeSchema("Dict", (TypeSchema("str"),
TypeSchema("float")))),
+ ],
+ methods=[],
+ )
+ assert info.gen_fields(_type_suffix, indent=0) == [
+ "arr: Sequence[int]",
+ "lst: MutableSequence[str]",
+ "mp: Mapping[str, int]",
+ "dt: MutableMapping[str, float]",
+ ]
+
+
def test_generate_global_funcs_updates_block() -> None:
code = CodeBlock(
kind="global",
@@ -303,6 +364,35 @@ def
test_generate_global_funcs_respects_custom_import_from() -> None:
assert ImportItem("custom.mod.init_ffi_api", alias="_FFI_INIT_FUNC") in
imports
+def test_generate_global_funcs_aliases_colliding_type() -> None:
+ """When a function name matches a type name, the type import gets an
alias."""
+ code = CodeBlock(
+ kind="global",
+ param=("demo", "mockpkg"),
+ lineno_start=1,
+ lineno_end=2,
+ lines=[f"{C.STUB_BEGIN} global/demo@mockpkg", C.STUB_END],
+ )
+ # Function "demo.Foo" returns type "demo.Foo" — name collision
+ funcs = [
+ FuncInfo(
+ schema=NamedTypeSchema(
+ "demo.Foo",
+ TypeSchema("Callable", (TypeSchema("demo.Foo"),
TypeSchema("Any"))),
+ ),
+ is_member=False,
+ )
+ ]
+ ty_map = _default_ty_map()
+ ty_map["demo.Foo"] = "somepkg.Foo"
+ imports: list[ImportItem] = []
+ generate_global_funcs(code, funcs, ty_map, imports, Options(indent=4))
+ # The type import should use an alias to avoid shadowing the function
+ assert ImportItem("somepkg.Foo", type_checking_only=True, alias="_Foo") in
imports
+ # The function annotation should use the alias
+ assert any("-> _Foo:" in line for line in code.lines)
+
+
def test_generate_object_fields_only_block() -> None:
code = CodeBlock(
kind="object",