This is an automated email from the ASF dual-hosted git repository.

tlopex 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 3518ef01 [STUBGEN][RUST] Mirror ffi.BigInt fields in generated Rust 
bindings (#816)
3518ef01 is described below

commit 3518ef0175f38ebb478521b1cad1f02540058ad2
Author: Linzhang Li <[email protected]>
AuthorDate: Mon Sep 28 00:08:58 2026 -0400

    [STUBGEN][RUST] Mirror ffi.BigInt fields in generated Rust bindings (#816)
    
    Stacked on #815 (which is stacked on #804). Only the commits after that
    branch belong here; the change is Python and Cython only.
    
    ## Summary
    
    A reflected class with a `BigInt` field describes it to Python as
    `{"type": "ffi.BigInt"}`, and the Cython `TypeSchema` folds that key
    into the
    `int` origin because a `BigInt` is a Python `int`. The Rust stub
    generator only
    saw the origin, so it mirrored the 16-byte cell as `i64` and the
    generated
    struct's size assertion failed to compile.
    
    `TypeSchema` now keeps the key a schema was parsed from when the fold
    loses
    information, in a new `type_key` attribute that takes no part in
    equality or
    `repr`. One table, `_TYPE_INDEX_TO_RETAINED_TYPE_KEY`, names such folds;
    today
    that is `ffi.BigInt` alone. `from_json_obj` and `from_type_index` both
    set it,
    `to_json` writes it back so a parse and serialize round trip is
    lossless, and
    `NamedTypeSchema` carries it for reflected fields.
    
    On the generator side, `render_rust_type` hands the leaf resolver the
    origin
    and the type key, and `_resolve` prefers a `ty-map` entry for the key
    over one
    for the origin. The default map already carries `ffi.BigInt ->
    tvm_ffi::BigInt`
    from #804, so `BigInt` and `Optional<BigInt>` fields now mirror as their
    own
    16-byte cells.
    
    Python-facing behaviour does not change: `.pyi` stubs and `repr` still
    show
    `int`. Lossless folds such as `ffi.String -> str` keep no key on
    purpose, so a
    user `ty-map` entry for `str` keeps winning.
    
    ## Testing
    
    - Python package rebuilt for the Cython change.
    - `tests/python/test_type_converter.py` and
    `tests/python/test_stubgen_rust.py`:
    **682 passed**, 2 pre-existing xfails; the dataclass, structural, BigInt
    and
      remaining stubgen suites: **1344 passed**.
    - New tests: complete-field mirrors for `BigInt` and `Optional<BigInt>`
    on a
      synthetic 64-byte class, including the size assertion; the same fields
    reflected from a C++ class compiled with `load_inline`, arriving as
    origin
    `int` with key `ffi.BigInt` and size 16; a user `ty-map` for `str` still
    applied to a reflected `ffi.String` field; `from_type_index` agreeing
    with the
    JSON path; the `to_json` round trip; the parsed-key rules (kept for the
    lossy fold, `None` otherwise, nested keys stay on the leaf); and the
    keys on
      a `Callable[[int], int]` schema in `test_big_int.py`.
    - `pre-commit` clean on the touched files (`ruff`, `ruff-format`, `ty`,
      `cython-lint`). No C++ or Rust changes.
---
 python/tvm_ffi/core.pyi                       |   2 +
 python/tvm_ffi/cython/type_info.pxi           |  43 +++++++---
 python/tvm_ffi/stub/rust_generator/codegen.py |  21 +++--
 python/tvm_ffi/stub/rust_generator/utils.py   |  18 +++--
 python/tvm_ffi/stub/utils.py                  |   1 +
 tests/python/test_big_int.py                  |   6 +-
 tests/python/test_stubgen_rust.py             | 111 +++++++++++++++++++++++++-
 tests/python/test_type_converter.py           |  42 ++++++++++
 8 files changed, 216 insertions(+), 28 deletions(-)

diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi
index f4629515..237eb39b 100644
--- a/python/tvm_ffi/core.pyi
+++ b/python/tvm_ffi/core.pyi
@@ -273,6 +273,7 @@ class TypeSchema:
     origin_type_index: int
     fallback: TypeSchema | None
     named_args: dict[str, tuple[TypeSchema, ...]] | None
+    type_key: str | None
 
     def __init__(
         self,
@@ -281,6 +282,7 @@ class TypeSchema:
         origin_type_index: int = ...,
         fallback: TypeSchema | None = None,
         named_args: dict[str, tuple[TypeSchema, ...]] | None = None,
+        type_key: str | None = None,
     ) -> None: ...
     @staticmethod
     def from_json_obj(obj: dict[str, Any]) -> TypeSchema: ...
diff --git a/python/tvm_ffi/cython/type_info.pxi 
b/python/tvm_ffi/cython/type_info.pxi
index 4e7fdc2f..37cb68e4 100644
--- a/python/tvm_ffi/cython/type_info.pxi
+++ b/python/tvm_ffi/cython/type_info.pxi
@@ -156,6 +156,11 @@ _TYPE_INDEX_TO_ORIGIN[kTVMFFISmallBytes] = "bytes"
 _TYPE_INDEX_TO_ORIGIN[kTVMFFIObjectRValueRef] = "Object"
 _TYPE_INDEX_TO_ORIGIN[kTVMFFIBigInt] = "int"
 
+# Folds that lose information: the folded origin is also a native type, so the
+# schema retains the key it was parsed from (``TypeSchema.type_key``).
+_TYPE_INDEX_TO_RETAINED_TYPE_KEY = {kTVMFFIBigInt: "ffi.BigInt"}
+_RETAINED_TYPE_KEYS = frozenset(_TYPE_INDEX_TO_RETAINED_TYPE_KEY.values())
+
 
 @dataclasses.dataclass(repr=False)
 class TypeSchema:
@@ -170,6 +175,9 @@ class TypeSchema:
     origin_type_index: int = 
dataclasses.field(default=_ORIGIN_TYPE_INDEX_UNKNOWN, repr=False)
     fallback: "TypeSchema | None" = None
     named_args: "dict[str, tuple[TypeSchema, ...]] | None" = None
+    #: The key the schema was parsed from when it was folded into ``origin``
+    #: (``"ffi.BigInt"`` for an ``int`` origin); ``None`` otherwise. Not part 
of equality.
+    type_key: str | None = dataclasses.field(default=None, repr=False, 
compare=False)
 
     def __post_init__(self):
         origin = self.origin
@@ -300,10 +308,15 @@ class TypeSchema:
             raise TypeError(
                 f"expected schema dict with 'type' key, got 
{type(obj).__name__}"
             )
-        origin = obj["type"]
-        if not isinstance(origin, str):
+        raw_origin = obj["type"]
+        if not isinstance(raw_origin, str):
             raise TypeError("schema type must be a string")
-        origin = _TYPE_SCHEMA_ORIGIN_CONVERTER.get(origin, origin)
+        origin = _TYPE_SCHEMA_ORIGIN_CONVERTER.get(raw_origin, raw_origin)
+        # Keep the key only when the fold loses it (`ffi.BigInt` -> `int`, 
where
+        # `int` is also a native type) for consumers that must tell the two 
apart,
+        # such as the Rust stub generator. Other folds (`ffi.String` -> `str`) 
keep
+        # None so a `ty-map` entry for the origin is not shadowed by the key's.
+        type_key = raw_origin if raw_origin in _RETAINED_TYPE_KEYS else None
         fallback = TypeSchema.from_json_obj(obj["fallback"]) if "fallback" in 
obj else None
         raw_named = obj.get("named_args")
         named_args = None
@@ -318,7 +331,7 @@ class TypeSchema:
                     raise TypeError("named_args lists must contain type 
schemas")
                 named_args[name] = tuple(TypeSchema.from_json_obj(value) for 
value in values)
         if "args" not in obj:
-            return TypeSchema(origin, fallback=fallback, named_args=named_args)
+            return TypeSchema(origin, fallback=fallback, 
named_args=named_args, type_key=type_key)
         raw_args = obj["args"]
         if (obj["type"] == "std::function" and isinstance(raw_args, (list, 
tuple))
                 and len(raw_args) == 2 and isinstance(raw_args[1], (list, 
tuple))):
@@ -326,14 +339,14 @@ class TypeSchema:
             ret = TypeSchema.from_json_obj(raw_args[0])
             params = tuple(TypeSchema.from_json_obj(value) for value in 
raw_args[1])
             return TypeSchema(origin, fallback=fallback,
-                              named_args={"return": (ret,), "params": params})
+                              named_args={"return": (ret,), "params": params}, 
type_key=type_key)
         if not isinstance(raw_args, (list, tuple)):
             raw_args = ()
         args = tuple(
             TypeSchema.from_json_obj(a) for a in raw_args
             if isinstance(a, dict)
         )
-        return TypeSchema(origin, args, fallback=fallback, 
named_args=named_args)
+        return TypeSchema(origin, args, fallback=fallback, 
named_args=named_args, type_key=type_key)
 
     @staticmethod
     def from_json_str(s: str) -> "TypeSchema":
@@ -362,7 +375,12 @@ class TypeSchema:
         origin = _TYPE_INDEX_TO_ORIGIN.get(type_index, None)
         if origin is None:
             origin = _type_index_to_key(type_index)
-        return TypeSchema(origin, args, origin_type_index=type_index)
+        return TypeSchema(
+            origin,
+            args,
+            origin_type_index=type_index,
+            type_key=_TYPE_INDEX_TO_RETAINED_TYPE_KEY.get(type_index),
+        )
 
     @staticmethod
     def from_annotation(annotation: object) -> "TypeSchema":
@@ -679,14 +697,19 @@ class TypeSchema:
             return f"{origin}[{args}]"
 
     def to_json(self) -> dict[str, Any]:
-        """Convert a TypeSchema to a JSON-compatible dict."""
+        """Convert a TypeSchema to a JSON-compatible dict.
+
+        A retained :attr:`type_key` is written back in place of the folded 
origin
+        (``ffi.BigInt`` rather than ``int``), so parsing the result restores 
the key.
+        """
+        origin = self.type_key or self.origin
         if self.args is not None and (self.args or self.origin == "tuple"):
             result = {
-                "type": self.origin,
+                "type": origin,
                 "args": [a.to_json() for a in self.args],
             }
         else:
-            result = {"type": self.origin}
+            result = {"type": origin}
         if self.named_args is not None:
             result["named_args"] = {
                 name: [value.to_json() for value in values]
diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py 
b/python/tvm_ffi/stub/rust_generator/codegen.py
index 7385dab7..4180b41a 100644
--- a/python/tvm_ffi/stub/rust_generator/codegen.py
+++ b/python/tvm_ffi/stub/rust_generator/codegen.py
@@ -119,9 +119,16 @@ class _ObjectRenderer:
 
     # --- name resolution ---------------------------------------------------
 
-    def _resolve(self, origin: str, imports: RustImports) -> str | None:
-        """Resolve a leaf origin to its in-scope Rust name (recording its 
``use``), or ``None``."""
-        mapped = self.ty_map.get(origin)
+    def _resolve(
+        self, origin: str, imports: RustImports, type_key: str | None = None
+    ) -> str | None:
+        """Resolve a leaf origin to its in-scope Rust name (recording its 
``use``), or ``None``.
+
+        A ``ty-map`` entry for the type key the schema was parsed from wins 
over the origin.
+        """
+        mapped = self.ty_map.get(type_key) if type_key is not None else None
+        if mapped is None:
+            mapped = self.ty_map.get(origin)
         if mapped is None:
             if "." not in origin or origin.startswith("ctypes."):
                 return None
@@ -129,8 +136,8 @@ class _ObjectRenderer:
             mapped = self._generated_type_path(origin)
         return imports.record(mapped)
 
-    def _ty_render(self, origin: str) -> str | None:
-        return self._resolve(origin, self.imports)
+    def _ty_render(self, origin: str, type_key: str | None = None) -> str | 
None:
+        return self._resolve(origin, self.imports, type_key)
 
     def _generated_type_path(self, type_key: str) -> str:
         """Spell a generated type key from this file.
@@ -230,7 +237,7 @@ class _ObjectRenderer:
             mirror = self._optional_mirror(field, imports)
         else:
             narrowed = C_RUST.RUST_SCALAR_BY_SIZE.get((field.origin, 
field.size))
-            mirror = narrowed or render_rust_type(field, lambda o: 
self._resolve(o, imports))
+            mirror = narrowed or render_rust_type(field, lambda o, k: 
self._resolve(o, imports, k))
         if mirror is None:
             return None
         if target in directives.nullable and not mirror.startswith("Option<"):
@@ -252,7 +259,7 @@ class _ObjectRenderer:
         (payload,) = field.args  # TypeSchema's post_init enforces exactly one 
argument.
         if payload.origin == "Any":
             return None
-        inner = render_rust_type(payload, lambda o: self._resolve(o, imports))
+        inner = render_rust_type(payload, lambda o, k: self._resolve(o, 
imports, k))
         if inner is None:
             return None
         any_backed = (
diff --git a/python/tvm_ffi/stub/rust_generator/utils.py 
b/python/tvm_ffi/stub/rust_generator/utils.py
index dcd68e3d..56fb181e 100644
--- a/python/tvm_ffi/stub/rust_generator/utils.py
+++ b/python/tvm_ffi/stub/rust_generator/utils.py
@@ -110,25 +110,31 @@ class RustImports:
         return probe.leaf
 
 
-def render_rust_type(schema: TypeSchema, ty_render: Callable[[str], str | 
None]) -> str | None:
-    """Render ``schema`` as a Rust value type via ``ty_render`` (leaf origin 
-> name), or ``None``."""
+def render_rust_type(
+    schema: TypeSchema, ty_render: Callable[[str, str | None], str | None]
+) -> str | None:
+    """Render ``schema`` as a Rust value type, or ``None``.
+
+    ``ty_render`` names a leaf from its origin and the type key it was parsed 
from (``None``
+    when they agree): the schema folds ``ffi.BigInt`` into ``int``.
+    """
     origin, args = schema.origin, schema.args
     if origin in C.RUST_UNSUPPORTED_ORIGINS:
         return None
     if origin == "Array":
         assert args  # TypeSchema's post_init fills a missing element type.
-        return _generic(ty_render("Array"), render_rust_type(args[0], 
ty_render))
+        return _generic(ty_render("Array", None), render_rust_type(args[0], 
ty_render))
     if origin == "Map":
         assert len(args) == 2  # TypeSchema's post_init fills a bare Map to 
(Any, Any).
         key = render_rust_type(args[0], ty_render)
         value = render_rust_type(args[1], ty_render)
-        return _generic(ty_render("Map"), key, value)
+        return _generic(ty_render("Map", None), key, value)
     if origin == "Optional":
         (payload,) = args  # TypeSchema's post_init enforces exactly one 
argument.
         return _generic("Option", render_rust_type(payload, ty_render))
     if origin == "Callable":
-        return ty_render("Callable")  # the crate's Function is type-erased
-    return ty_render(origin)
+        return ty_render("Callable", None)  # the crate's Function is 
type-erased
+    return ty_render(origin, schema.type_key)
 
 
 def _generic(base: str | None, *params: str | None) -> str | None:
diff --git a/python/tvm_ffi/stub/utils.py b/python/tvm_ffi/stub/utils.py
index e7f7eb94..848a85ae 100644
--- a/python/tvm_ffi/stub/utils.py
+++ b/python/tvm_ffi/stub/utils.py
@@ -137,6 +137,7 @@ class NamedTypeSchema(TypeSchema):
             origin_type_index=schema.origin_type_index,
             fallback=schema.fallback,
             named_args=schema.named_args,
+            type_key=schema.type_key,
         )
         self.name = name
         self.size = size
diff --git a/tests/python/test_big_int.py b/tests/python/test_big_int.py
index ae4a915f..cb19fe8a 100644
--- a/tests/python/test_big_int.py
+++ b/tests/python/test_big_int.py
@@ -143,7 +143,11 @@ def test_integer_type_schema(bigint_module: 
tvm_ffi.Module) -> None:
     metadata = bigint_module.get_function_metadata("identity")
     assert metadata is not None
     assert "ffi.BigInt" in metadata["type_schema"]
-    assert str(TypeSchema.from_json_str(metadata["type_schema"])) == 
"Callable[[int], int]"
+    schema = TypeSchema.from_json_str(metadata["type_schema"])
+    assert str(schema) == "Callable[[int], int]"
+    assert schema.named_args is not None
+    assert schema.named_args["return"][0].type_key == "ffi.BigInt"
+    assert schema.named_args["params"][0].type_key == "ffi.BigInt"
 
 
 def test_integer_callbacks(bigint_module: tvm_ffi.Module) -> None:
diff --git a/tests/python/test_stubgen_rust.py 
b/tests/python/test_stubgen_rust.py
index 789e9b0a..0ccdd166 100644
--- a/tests/python/test_stubgen_rust.py
+++ b/tests/python/test_stubgen_rust.py
@@ -23,8 +23,9 @@ from collections.abc import Container, Iterator
 from pathlib import Path
 
 import pytest
+import tvm_ffi.cpp
 import tvm_ffi.stub.cli as stub_cli
-import tvm_ffi.testing  # noqa: F401  (loads the `testing.*` fixture types)
+import tvm_ffi.testing  # loads the `testing.*` fixture types
 from tvm_ffi.core import TypeSchema
 from tvm_ffi.stub import consts as C
 from tvm_ffi.stub.cli import _stage_1, _stage_3
@@ -179,8 +180,9 @@ def _render_type(schema: TypeSchema) -> tuple[str | None, 
RustImports]:
     imports = RustImports()
     ty_map = RC.RUST_TY_MAP_DEFAULTS
 
-    def ty_render(origin: str) -> str | None:
-        return imports.record(ty_map[origin]) if origin in ty_map else None
+    def ty_render(origin: str, type_key: str | None) -> str | None:
+        leaf = type_key if type_key in ty_map else origin
+        return imports.record(ty_map[leaf]) if leaf in ty_map else None
 
     return render_rust_type(schema, ty_render), imports
 
@@ -188,7 +190,7 @@ def _render_type(schema: TypeSchema) -> tuple[str | None, 
RustImports]:
 def test_render_rust_type_value_positions() -> None:
     assert _render_type(TypeSchema("int"))[0] == "i64"
     assert _render_type(TypeSchema("str"))[0] == "String"
-    assert _render_type(TypeSchema("ffi.BigInt"))[0] == "BigInt"
+    assert _render_type(TypeSchema.from_json_str('{"type":"ffi.BigInt"}'))[0] 
== "BigInt"
     assert _render_type(TypeSchema("Any"))[0] == "Any"
     assert _render_type(TypeSchema("Callable", (TypeSchema("int"),)))[0] == 
"Function"
     assert _render_type(TypeSchema("Optional", (TypeSchema("str"),)))[0] == 
"Option<String>"
@@ -882,6 +884,107 @@ def test_complete_optional_field_mirrors() -> None:
     assert "tvm_ffi::Optional" in _uses(imports)
 
 
+def test_complete_big_int_field_mirrors() -> None:
+    """The schema folds `ffi.BigInt` into `int`; the parsed type key keeps the 
cell apart."""
+    big = TypeSchema.from_json_str('{"type":"ffi.BigInt"}')
+    maybe = 
TypeSchema.from_json_str('{"type":"Optional","args":[{"type":"ffi.BigInt"}]}')
+    assert (big.origin, big.type_key) == ("int", "ffi.BigInt")
+    info = _info(
+        "demo.Big",
+        (
+            _field("value", big, 24, 16),
+            _field("maybe", maybe, 40, 16),
+            _field("count", "int", 56, 8),
+        ),
+        total_size=64,
+        is_final=True,
+    )
+    text, imports = _render(info)
+    assert (
+        "    base: Object,\n"
+        "    pub value: BigInt,\n"
+        "    pub maybe: Optional<BigInt>,\n"
+        "    pub count: i64,\n"
+        "}"
+    ) in text
+    assert "    pub fn new(value: BigInt, maybe: Optional<BigInt>, count: i64) 
-> Self {" in text
+    assert "assert!(::core::mem::size_of::<BigObj>() == 64);" in text
+    assert {"tvm_ffi::BigInt", "tvm_ffi::Optional"} <= _uses(imports)
+
+
+def test_ty_map_origin_override_survives_lossless_fold() -> None:
+    """A `ty-map` for `str` still wins for a reflected `ffi.String`: the 
lossless fold keeps no key."""
+    name = TypeSchema.from_json_str('{"type":"ffi.String"}')
+    assert (name.origin, name.type_key) == ("str", None)
+    info = _info("demo.Named", (_field("name", name, 24, 16),), total_size=40, 
is_final=True)
+    ty_map = dict(RUST.default_ty_map())
+    ty_map["str"] = "crate::MyStr"
+    imports = RustImports()
+    assert info.type_key is not None
+    block = _object_block(info.type_key)
+    generate_rust_object(block, ty_map, imports, Options(), info, ALL_DECLARED)
+    text = "\n".join(block.lines[1:-1])
+    assert "    pub name: MyStr,\n" in text
+    assert "crate::MyStr" in _uses(imports)
+
+
[email protected](scope="module")
+def big_int_holder() -> ObjectInfo:
+    """Reflection of a C++ class with `BigInt` fields, compiled and registered 
at test time."""
+    tvm_ffi.cpp.load_inline(
+        name="test_stubgen_rust_big_int",
+        cpp_sources=r"""
+            #include <tvm/ffi/big_int.h>
+            #include <tvm/ffi/optional.h>
+            #include <tvm/ffi/reflection/registry.h>
+
+            namespace stubgen_test {
+            using namespace tvm::ffi;
+
+            class BigIntHolder : public Object {
+             public:
+              BigInt value;
+              Optional<BigInt> maybe;
+              int64_t count = 0;
+              static constexpr bool _type_mutable = true;
+              TVM_FFI_DECLARE_OBJECT_INFO_FINAL("stubgen_test.BigIntHolder", 
BigIntHolder, Object);
+            };
+
+            TVM_FFI_STATIC_INIT_BLOCK() {
+              namespace refl = tvm::ffi::reflection;
+              refl::ObjectDef<BigIntHolder>()
+                  .def_rw("value", &BigIntHolder::value)
+                  .def_rw("maybe", &BigIntHolder::maybe)
+                  .def_rw("count", &BigIntHolder::count);
+            }
+            }  // namespace stubgen_test
+
+            int64_t touch() { return 0; }
+        """,
+        functions=["touch"],
+    )
+    return object_info_from_type_key("stubgen_test.BigIntHolder")
+
+
+def test_registry_big_int_fields_are_mirrored(big_int_holder: ObjectInfo) -> 
None:
+    """Reflected `BigInt` and `Optional<BigInt>` fields arrive as `int` plus 
the parsed type key."""
+    value, maybe, count = sorted(big_int_holder.fields, key=lambda f: f.offset 
or 0)
+    assert (value.origin, value.type_key, value.size) == ("int", "ffi.BigInt", 
16)
+    assert (maybe.args[0].origin, maybe.args[0].type_key, maybe.size) == 
("int", "ffi.BigInt", 16)
+    assert (count.origin, count.type_key, count.size) == ("int", None, 8)
+    text, imports = _render(big_int_holder)
+    assert (
+        "pub struct BigIntHolderObj {\n"
+        "    base: Object,\n"
+        "    pub value: BigInt,\n"
+        "    pub maybe: Optional<BigInt>,\n"
+        "    pub count: i64,\n"
+        "}"
+    ) in text
+    assert "assert!(::core::mem::size_of::<BigIntHolderObj>() == 64);" in text
+    assert {"tvm_ffi::BigInt", "tvm_ffi::Optional"} <= _uses(imports)
+
+
 @pytest.mark.parametrize(
     "field",
     [
diff --git a/tests/python/test_type_converter.py 
b/tests/python/test_type_converter.py
index fd54313f..32ff9b76 100644
--- a/tests/python/test_type_converter.py
+++ b/tests/python/test_type_converter.py
@@ -4662,3 +4662,45 @@ class TestSmallStringOptimization:
         b = tvm_ffi.core.Bytes(b"x" * 200)
         cany = A(bytes).convert(b)
         assert cany.type_index == _BYTES_TYPE_INDEX
+
+
+# ---------------------------------------------------------------------------
+# Category 64: parsed type key
+# ---------------------------------------------------------------------------
+class TestParsedTypeKey:
+    """``from_json_*`` keeps the raw key when it folds the origin into another 
name."""
+
+    def test_folded_key_is_kept(self) -> None:
+        s = TypeSchema.from_json_str('{"type":"ffi.BigInt"}')
+        assert (s.origin, s.type_key) == ("int", "ffi.BigInt")
+        assert s == TypeSchema("int")  # the key takes no part in equality
+        assert repr(s) == "int"
+
+    def test_unfolded_key_is_none(self) -> None:
+        assert TypeSchema.from_json_str('{"type":"int"}').type_key is None
+        assert TypeSchema("int").type_key is None
+
+    def test_nested_key_stays_on_the_leaf(self) -> None:
+        s = 
TypeSchema.from_json_str('{"type":"Optional","args":[{"type":"ffi.BigInt"}]}')
+        assert s.type_key is None
+        assert s.args[0].type_key == "ffi.BigInt"
+
+    def test_lossless_fold_keeps_no_key(self) -> None:
+        # `ffi.String` -> `str` loses nothing (`str` only ever means a 
string), so no key is
+        # kept and a `ty-map` entry for the origin is not shadowed by one for 
the key.
+        s = TypeSchema.from_json_str('{"type":"ffi.String"}')
+        assert (s.origin, s.type_key) == ("str", None)
+
+    def test_from_type_index_agrees_with_json(self) -> None:
+        big_int_index = _object_type_key_to_index("ffi.BigInt")
+        assert big_int_index is not None
+        s = TypeSchema.from_type_index(big_int_index)
+        assert (s.origin, s.type_key) == ("int", "ffi.BigInt")
+        assert TypeSchema.from_type_index(1).type_key is None  # kTVMFFIInt
+
+    def test_to_json_round_trips_the_key(self) -> None:
+        obj = {"type": "Optional", "args": [{"type": "ffi.BigInt"}]}
+        s = TypeSchema.from_json_obj(obj)
+        assert s.to_json() == obj
+        assert TypeSchema.from_json_obj(s.to_json()).args[0].type_key == 
"ffi.BigInt"
+        assert TypeSchema("int").to_json() == {"type": "int"}

Reply via email to