This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 622e96b7 [FEAT] Add named type schema groups for function signatures
(#820)
622e96b7 is described below
commit 622e96b766a3ae3024967e7ad07904e1b674824e
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Sep 27 13:36:50 2026 -0400
[FEAT] Add named type schema groups for function signatures (#820)
- Add recursive `named_args` lists to type schemas while retaining
positional `args` and `fallback`.
- Describe typed functions with named `return` and `params` groups,
including an explicit empty parameter list.
- Update C++ and Python producers, reflection consumers, binding
generation, and type-schema documentation.
---
.github/actions/build-wheel-for-publish/action.yml | 2 +-
.github/workflows/ci_mainline_only.yml | 4 +-
.github/workflows/ci_test.yml | 8 +-
.github/workflows/torch_c_dlpack.yml | 6 +-
docs/guides/cpp_lang_guide.md | 15 ++-
include/tvm/ffi/extra/stl.h | 7 +-
include/tvm/ffi/function.h | 4 +-
include/tvm/ffi/function_details.h | 9 +-
python/tvm_ffi/core.pyi | 2 +
python/tvm_ffi/cython/type_info.pxi | 100 +++++++++++---
python/tvm_ffi/dataclasses/py_class.py | 4 +-
python/tvm_ffi/stub/python_generator/utils.py | 19 ++-
python/tvm_ffi/stub/utils.py | 1 +
tests/cpp/test_metadata.cc | 148 +++++++++++++--------
tests/cpp/test_reflection.cc | 7 +-
tests/python/test_stubgen.py | 42 +++++-
tests/python/test_type_converter.py | 22 ++-
17 files changed, 289 insertions(+), 111 deletions(-)
diff --git a/.github/actions/build-wheel-for-publish/action.yml
b/.github/actions/build-wheel-for-publish/action.yml
index f51808c5..fded3153 100644
--- a/.github/actions/build-wheel-for-publish/action.yml
+++ b/.github/actions/build-wheel-for-publish/action.yml
@@ -46,7 +46,7 @@ runs:
using: "composite"
steps:
- name: Set up uv
- uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 #
v8.2.0
+ uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 #
v10.2.0
- name: Check out source
uses: actions/checkout@v5
diff --git a/.github/workflows/ci_mainline_only.yml
b/.github/workflows/ci_mainline_only.yml
index bfb0ac68..c919e9da 100644
--- a/.github/workflows/ci_mainline_only.yml
+++ b/.github/workflows/ci_mainline_only.yml
@@ -50,7 +50,7 @@ jobs:
fetch-depth: 0
fetch-tags: true
- name: Set up uv
- uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 #
v8.2.0
+ uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 #
v10.2.0
- uses: ./.github/actions/detect-env-vars
id: env_vars
- name: Run clang-tidy
@@ -108,7 +108,7 @@ jobs:
fetch-tags: true
- name: Set up uv
- uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 #
v8.2.0
+ uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 #
v10.2.0
with:
python-version: ${{ matrix.python_version }}
activate-environment: true
diff --git a/.github/workflows/ci_test.yml b/.github/workflows/ci_test.yml
index 2db75128..a4d9bf36 100644
--- a/.github/workflows/ci_test.yml
+++ b/.github/workflows/ci_test.yml
@@ -72,7 +72,7 @@ jobs:
fetch-depth: 0
fetch-tags: true
- name: Set up uv
- uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 #
v8.2.0
+ uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 #
v10.2.0
- name: Set up Python environment
run: uv sync --group dev --no-install-project
- name: Run pre-commit
@@ -89,7 +89,7 @@ jobs:
fetch-depth: 0
fetch-tags: true
- name: Set up uv
- uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 #
v8.2.0
+ uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 #
v10.2.0
with:
python-version: 3.13
- uses: ./.github/actions/detect-env-vars
@@ -112,7 +112,7 @@ jobs:
fetch-depth: 0
fetch-tags: true
- name: Set up uv
- uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 #
v8.2.0
+ uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 #
v10.2.0
with:
python-version: 3.13
- name: Generate docs
@@ -183,7 +183,7 @@ jobs:
# Run Python tests
- name: Setup Python ${{ matrix.python_version }}
- uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 #
v8.2.0
+ uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 #
v10.2.0
with:
python-version: ${{ matrix.python_version }}
activate-environment: true
diff --git a/.github/workflows/torch_c_dlpack.yml
b/.github/workflows/torch_c_dlpack.yml
index 5aacba94..cd4afbd3 100644
--- a/.github/workflows/torch_c_dlpack.yml
+++ b/.github/workflows/torch_c_dlpack.yml
@@ -83,7 +83,7 @@ jobs:
path: tvm-ffi
- uses: Jimver/cuda-toolkit@3d45d157f327c09c04b50ee6ccdea2d9d017ec76 #
v0.2.35
id: cuda-toolkit
- - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 #
v8.2.0
+ - uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 #
v10.2.0
- name: build wheels
env:
CUDA_HOME: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}
@@ -109,7 +109,7 @@ jobs:
fetch-depth: 0
submodules: recursive
path: tvm-ffi
- - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 #
v8.2.0
+ - uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 #
v10.2.0
- name: build torch libs and wheels
working-directory: ./tvm-ffi
run: |
@@ -128,7 +128,7 @@ jobs:
fetch-depth: 0
submodules: recursive
path: tvm-ffi
- - uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 #
v8.2.0
+ - uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 #
v10.2.0
with:
python-version: "3.12"
activate-environment: true
diff --git a/docs/guides/cpp_lang_guide.md b/docs/guides/cpp_lang_guide.md
index d964421c..fe88ea5c 100644
--- a/docs/guides/cpp_lang_guide.md
+++ b/docs/guides/cpp_lang_guide.md
@@ -330,7 +330,20 @@ signature validation without calling the function.
The metadata contains:
-- **type_schema**: JSON string describing function signature (return type and
argument types)
+- **type_schema**: JSON string describing the function signature.
+
+A type schema has a nonempty `type` and may have positional `args`, named
+`named_args`, and `fallback`. `named_args` maps each group name to a list of
+complete type schemas; the list may be empty, and types in it may have their
+own `args`, `named_args`, or `fallback`. A function uses a one-element `return`
+list and an ordered `params` list:
+
+```json
+{"type":"ffi.Function","named_args":{"return":[{"type":"int"}],"params":[{"type":"int"}]}}
+```
+
+An empty `params` list represents a known zero-argument function. A bare
+`{"type":"ffi.Function"}` leaves the signature unspecified.
```cpp
ffi::Module mod = ffi::Module::LoadFromFile("path/to/export_lib.so");
diff --git a/include/tvm/ffi/extra/stl.h b/include/tvm/ffi/extra/stl.h
index 3f3e4213..e30c1d1b 100644
--- a/include/tvm/ffi/extra/stl.h
+++ b/include/tvm/ffi/extra/stl.h
@@ -453,6 +453,7 @@ struct TypeTraits<std::variant<Args...>> : public
TypeTraitsBase {
os << R"({"type":"std::variant","args":[)";
const char* sep = "";
((os << sep << details::TypeSchema<Args>::v(), sep = ", "), ...);
+ (void)sep;
os << "]}";
return std::move(os).str();
}
@@ -641,10 +642,12 @@ struct TypeTraits<std::function<Ret(Args...)>> :
TypeTraitsBase {
TVM_FFI_INLINE static std::string TypeSchema() {
std::ostringstream os;
- os << R"({"type":"std::function","args":[)" <<
details::TypeSchema<Ret>::v() << ",[";
+ os << R"({"type":"std::function","named_args":{"return":[)" <<
details::TypeSchema<Ret>::v()
+ << R"(],"params":[)";
const char* sep = "";
((os << sep << details::TypeSchema<Args>::v(), sep = ", "), ...);
- os << "]]}";
+ (void)sep;
+ os << "]}}";
return std::move(os).str();
}
};
diff --git a/include/tvm/ffi/function.h b/include/tvm/ffi/function.h
index 30aed159..a7a02959 100644
--- a/include/tvm/ffi/function.h
+++ b/include/tvm/ffi/function.h
@@ -891,8 +891,8 @@ class TypedFunction<R(Args...)> {
/*! \return Whether the packed function is not nullptr */
bool operator!=(std::nullptr_t null) const { return packed_ != nullptr; }
/*!
- * \brief Get the type schema of `TypedFunction<R(Args...)>` in json format.
- * \return The type schema of the function in json format.
+ * \brief Get the JSON type schema of `TypedFunction<R(Args...)>`.
+ * \return An `ffi.Function` schema with named `return` and `params` type
lists.
*/
static std::string TypeSchema() { return details::FuncFunctorImpl<R,
Args...>::TypeSchema(); }
diff --git a/include/tvm/ffi/function_details.h
b/include/tvm/ffi/function_details.h
index cf33a79c..b2f41ce0 100644
--- a/include/tvm/ffi/function_details.h
+++ b/include/tvm/ffi/function_details.h
@@ -99,10 +99,13 @@ struct FuncFunctorImpl {
}
TVM_FFI_INLINE static std::string TypeSchema() {
std::ostringstream oss;
- oss << R"({"type":")" << StaticTypeKey::kTVMFFIFunction << R"(","args":[)";
+ oss << R"({"type":")" << StaticTypeKey::kTVMFFIFunction <<
R"(","named_args":{"return":[)";
oss << details::TypeSchema<R>::v();
- ((oss << "," << details::TypeSchema<Args>::v()), ...);
- oss << "]}";
+ oss << R"(],"params":[)";
+ const char* sep = "";
+ ((oss << sep << details::TypeSchema<Args>::v(), sep = ","), ...);
+ (void)sep;
+ oss << "]}}";
return oss.str();
}
};
diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi
index de5ddac4..f4629515 100644
--- a/python/tvm_ffi/core.pyi
+++ b/python/tvm_ffi/core.pyi
@@ -272,6 +272,7 @@ class TypeSchema:
args: tuple[TypeSchema, ...] = ()
origin_type_index: int
fallback: TypeSchema | None
+ named_args: dict[str, tuple[TypeSchema, ...]] | None
def __init__(
self,
@@ -279,6 +280,7 @@ class TypeSchema:
args: tuple[TypeSchema, ...] = (),
origin_type_index: int = ...,
fallback: TypeSchema | None = None,
+ named_args: dict[str, tuple[TypeSchema, ...]] | 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 cc6e0006..4e7fdc2f 100644
--- a/python/tvm_ffi/cython/type_info.pxi
+++ b/python/tvm_ffi/cython/type_info.pxi
@@ -161,14 +161,15 @@ _TYPE_INDEX_TO_ORIGIN[kTVMFFIBigInt] = "int"
class TypeSchema:
"""Type schema that describes a TVM FFI type.
- The schema is expressed using a compact JSON-compatible structure
- and can be rendered as a Python typing string with
- :py:meth:`repr`.
+ ``args`` holds positional type arguments. ``named_args`` maps each name
+ to an ordered list of complete type schemas. A typed function uses
+ ``return`` and ``params`` as named groups.
"""
origin: str
args: tuple["TypeSchema", ...] | None = None
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
def __post_init__(self):
origin = self.origin
@@ -176,9 +177,39 @@ class TypeSchema:
if args is not None and not isinstance(args, tuple):
args = tuple(args)
self.args = args
+ if not isinstance(origin, str) or not origin:
+ raise ValueError("schema type must be a nonempty string")
if origin != "tuple" and args is None:
args = ()
self.args = args
+ if args is not None and any(not isinstance(a, TypeSchema) for a in
args):
+ raise TypeError("schema arguments must be TypeSchema instances")
+ named_args = self.named_args
+ if named_args is not None:
+ if not isinstance(named_args, dict):
+ raise TypeError("named_args must be a map of names to type
lists")
+ normalized = {}
+ for name, values in named_args.items():
+ if not isinstance(name, str) or not name:
+ raise ValueError("named_args keys must be nonempty
strings")
+ if not isinstance(values, (list, tuple)):
+ raise TypeError("named_args values must be lists of type
schemas")
+ if any(not isinstance(value, TypeSchema) for value in values):
+ raise TypeError("named_args values must contain type
schemas")
+ normalized[name] = tuple(values)
+ self.named_args = named_args = normalized
+ if origin == "Callable":
+ if args and named_args is None:
+ # Legacy flat schemas retain return-only unknown parameters.
+ self.named_args = named_args = {"return": (args[0],)}
+ if len(args) > 1:
+ named_args["params"] = args[1:]
+ self.args = args = ()
+ if named_args is not None:
+ if args or set(named_args) - {"return", "params"}:
+ raise ValueError("Callable uses only named return and
params groups")
+ if "return" not in named_args or len(named_args["return"]) !=
1:
+ raise ValueError("Callable return group must contain
exactly one type")
if origin == "Union":
if len(args) < 2:
raise ValueError("Union must have at least two arguments")
@@ -270,18 +301,39 @@ class TypeSchema:
f"expected schema dict with 'type' key, got
{type(obj).__name__}"
)
origin = obj["type"]
+ if not isinstance(origin, str):
+ raise TypeError("schema type must be a string")
origin = _TYPE_SCHEMA_ORIGIN_CONVERTER.get(origin, origin)
fallback = TypeSchema.from_json_obj(obj["fallback"]) if "fallback" in
obj else None
+ raw_named = obj.get("named_args")
+ named_args = None
+ if "named_args" in obj:
+ if not isinstance(raw_named, dict):
+ raise TypeError("named_args must be a map of names to type
lists")
+ named_args = {}
+ for name, values in raw_named.items():
+ if not isinstance(name, str) or not name or not
isinstance(values, list):
+ raise TypeError("named_args must map names to type lists")
+ if any(not isinstance(value, dict) for value in values):
+ 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)
+ return TypeSchema(origin, fallback=fallback, named_args=named_args)
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))):
+ # Legacy std::function encoded parameters as a bare JSON array.
+ 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})
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)
+ return TypeSchema(origin, args, fallback=fallback,
named_args=named_args)
@staticmethod
def from_json_str(s: str) -> "TypeSchema":
@@ -420,13 +472,15 @@ class TypeSchema:
params, ret = targs
ret_schema = TypeSchema.from_annotation(ret)
if isinstance(params, list):
- # Callable[[P1, P2], R] → (R, P1, P2)
+ # Callable[[P1, P2], R] → named return and params groups.
param_schemas = tuple(
TypeSchema.from_annotation(p) for p in params
)
- return TypeSchema("Callable", (ret_schema,) +
param_schemas)
+ return TypeSchema(
+ "Callable", named_args={"return": (ret_schema,),
"params": param_schemas}
+ )
# Callable[..., R]
- return TypeSchema("Callable", (ret_schema,))
+ return TypeSchema("Callable", named_args={"return":
(ret_schema,)})
return TypeSchema("Callable")
# --- Parameterised CObject subclasses (Array[int], Dict[str, V], …)
---
@@ -547,8 +601,9 @@ class TypeSchema:
s =
TypeSchema.from_json_str('{"type":"Optional","args":[{"type":"int"}]}')
assert s.repr() == "int | None"
- # Callable where the first arg is return type, remaining are
parameters
- s = TypeSchema("Callable", (TypeSchema("int"), TypeSchema("str")))
+ # Named return and parameter groups
+ s = TypeSchema("Callable", named_args={"return":
(TypeSchema("int"),),
+ "params":
(TypeSchema("str"),)})
assert s.repr() == "Callable[[str], int]"
# Container types from C++ FFI schemas
@@ -594,6 +649,19 @@ class TypeSchema:
origin = self.origin if ty_map is None else ty_map(self.origin)
schema_args = self.args
+ if origin == "Callable":
+ if not self.named_args:
+ return "Callable[..., Any]"
+ ret = self.named_args["return"][0]._repr_impl(
+ ty_map, input_mode, expanded_convert_types
+ )
+ if "params" not in self.named_args:
+ return f"Callable[..., {ret}]"
+ params = ", ".join(
+ arg._repr_impl(ty_map, input_mode, expanded_convert_types)
+ for arg in self.named_args["params"]
+ )
+ return f"Callable[[{params}], {ret}]"
args = [
i._repr_impl(ty_map, input_mode, expanded_convert_types)
for i in (() if schema_args is None else schema_args)
@@ -602,13 +670,6 @@ class TypeSchema:
return " | ".join(args)
elif origin == "Optional":
return args[0] + " | None"
- elif origin == "Callable":
- if not args:
- return "Callable[..., Any]"
- else:
- ret = args[0]
- args = ", ".join(args[1:])
- return f"Callable[[{args}], {ret}]"
elif origin == "tuple" and schema_args == ():
return "tuple[()]"
elif not args:
@@ -626,6 +687,11 @@ class TypeSchema:
}
else:
result = {"type": self.origin}
+ if self.named_args is not None:
+ result["named_args"] = {
+ name: [value.to_json() for value in values]
+ for name, values in self.named_args.items()
+ }
if self.fallback is not None:
result["fallback"] = self.fallback.to_json()
return result
diff --git a/python/tvm_ffi/dataclasses/py_class.py
b/python/tvm_ffi/dataclasses/py_class.py
index 1c581e75..4329ba42 100644
--- a/python/tvm_ffi/dataclasses/py_class.py
+++ b/python/tvm_ffi/dataclasses/py_class.py
@@ -289,7 +289,9 @@ def _method_type_schema_json(
ret_annotation = hints.get("return", Any)
ret_schema = TypeSchema.from_annotation(ret_annotation)
- schema = TypeSchema("Callable", (ret_schema, *arg_schemas))
+ schema = TypeSchema(
+ "Callable", named_args={"return": (ret_schema,), "params":
tuple(arg_schemas)}
+ )
return json.dumps({"type_schema": schema.to_json()})
diff --git a/python/tvm_ffi/stub/python_generator/utils.py
b/python/tvm_ffi/stub/python_generator/utils.py
index bf19b367..e4d6dfc9 100644
--- a/python/tvm_ffi/stub/python_generator/utils.py
+++ b/python/tvm_ffi/stub/python_generator/utils.py
@@ -113,12 +113,16 @@ def render_func_signature(
buf.write(f"def {func_name}(")
if func.schema.origin != "Callable":
raise ValueError(f"Expected Callable type schema, but got:
{func.schema}")
- if not func.schema.args:
+ if not func.schema.named_args:
ty_map("Any")
buf.write("*args: Any) -> Any: ...")
return buf.getvalue()
- arg_ret = func.schema.args[0]
- arg_args = func.schema.args[1:]
+ arg_ret = func.schema.named_args["return"][0]
+ if "params" not in func.schema.named_args:
+ ty_map("Any")
+ buf.write(f"*args: Any) -> {arg_ret.output_repr(ty_map)}: ...")
+ return buf.getvalue()
+ arg_args = func.schema.named_args["params"]
for i, arg in enumerate(arg_args):
if func.is_member and i == 0:
buf.write("self, ")
@@ -192,12 +196,15 @@ def _render_ffi_init_from_method(
input_ty_map = ty_map
indent_str = " " * indent
schema = method.schema
- if schema.origin != "Callable" or not schema.args:
+ if schema.origin != "Callable" or not schema.named_args:
+ ty_map("Any")
+ return f"{indent_str}def __ffi_init__(self, *args: Any) -> None: ..."
+ if "params" not in schema.named_args:
ty_map("Any")
return f"{indent_str}def __ffi_init__(self, *args: Any) -> None: ..."
- # schema.args[0] is return type, schema.args[1:] are param types.
+ # Function parameters are an ordered list in named_args["params"].
parts: list[str] = []
- for i, arg in enumerate(schema.args[1:]):
+ for i, arg in enumerate(schema.named_args["params"]):
parts.append(f"_{i}: {arg.input_repr(input_ty_map)}")
if parts:
params = ", ".join(parts)
diff --git a/python/tvm_ffi/stub/utils.py b/python/tvm_ffi/stub/utils.py
index f317f9c4..e7f7eb94 100644
--- a/python/tvm_ffi/stub/utils.py
+++ b/python/tvm_ffi/stub/utils.py
@@ -136,6 +136,7 @@ class NamedTypeSchema(TypeSchema):
args=schema.args,
origin_type_index=schema.origin_type_index,
fallback=schema.fallback,
+ named_args=schema.named_args,
)
self.name = name
self.size = size
diff --git a/tests/cpp/test_metadata.cc b/tests/cpp/test_metadata.cc
index ddb1abda..af701416 100644
--- a/tests/cpp/test_metadata.cc
+++ b/tests/cpp/test_metadata.cc
@@ -20,9 +20,11 @@
#include <tvm/ffi/c_api.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/container/tuple.h>
#include <tvm/ffi/container/variant.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/extra/json.h>
+#include <tvm/ffi/extra/stl.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/object.h>
#include <tvm/ffi/optional.h>
@@ -71,70 +73,85 @@ TEST(Schema, GlobalFuncTypeSchema) {
return ParseMetadataToSchema(metadata);
};
// Simple IDs
- EXPECT_EQ(fetch("testing.schema_id_int"),
-
R"({"type":"ffi.Function","args":[{"type":"int"},{"type":"int"}]})");
- EXPECT_EQ(fetch("testing.schema_id_float"),
-
R"({"type":"ffi.Function","args":[{"type":"float"},{"type":"float"}]})");
- EXPECT_EQ(fetch("testing.schema_id_bool"),
-
R"({"type":"ffi.Function","args":[{"type":"bool"},{"type":"bool"}]})");
- EXPECT_EQ(fetch("testing.schema_id_device"),
-
R"({"type":"ffi.Function","args":[{"type":"Device"},{"type":"Device"}]})");
- EXPECT_EQ(fetch("testing.schema_id_dtype"),
-
R"({"type":"ffi.Function","args":[{"type":"DataType"},{"type":"DataType"}]})");
- EXPECT_EQ(fetch("testing.schema_id_string"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.String"},{"type":"ffi.String"}]})");
- EXPECT_EQ(fetch("testing.schema_id_bytes"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Bytes"},{"type":"ffi.Bytes"}]})");
- EXPECT_EQ(fetch("testing.schema_id_func"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Function"},{"type":"ffi.Function"}]})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_int"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"int"}],"params":[{"type":"int"}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_float"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"float"}],"params":[{"type":"float"}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_bool"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"bool"}],"params":[{"type":"bool"}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_device"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"Device"}],"params":[{"type":"Device"}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_dtype"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"DataType"}],"params":[{"type":"DataType"}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_string"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.String"}],"params":[{"type":"ffi.String"}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_bytes"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Bytes"}],"params":[{"type":"ffi.Bytes"}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_func"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Function"}],"params":[{"type":"ffi.Function"}]}})");
EXPECT_EQ(
fetch("testing.schema_id_func_typed"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Function","args":[{"type":"None"},{"type":"int"},{"type":"float"},{"type":"ffi.Function"}]},{"type":"ffi.Function","args":[{"type":"None"},{"type":"int"},{"type":"float"},{"type":"ffi.Function"}]}]})");
-
- EXPECT_EQ(fetch("testing.schema_id_any"),
-
R"({"type":"ffi.Function","args":[{"type":"Any"},{"type":"Any"}]})");
- EXPECT_EQ(fetch("testing.schema_id_object"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Object"},{"type":"ffi.Object"}]})");
- EXPECT_EQ(fetch("testing.schema_id_dltensor"),
-
R"({"type":"ffi.Function","args":[{"type":"DLTensor*"},{"type":"DLTensor*"}]})");
- EXPECT_EQ(fetch("testing.schema_id_tensor"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Tensor"},{"type":"ffi.Tensor"}]})");
- EXPECT_EQ(fetch("testing.schema_tensor_view_input"),
-
R"({"type":"ffi.Function","args":[{"type":"None"},{"type":"DLTensor*"}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Function","named_args":{"return":[{"type":"None"}],"params":[{"type":"int"},{"type":"float"},{"type":"ffi.Function"}]}}],"params":[{"type":"ffi.Function","named_args":{"return":[{"type":"None"}],"params":[{"type":"int"},{"type":"float"},{"type":"ffi.Function"}]}}]}})");
+
+ EXPECT_EQ(
+ fetch("testing.schema_id_any"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"Any"}],"params":[{"type":"Any"}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_object"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Object"}],"params":[{"type":"ffi.Object"}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_dltensor"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"DLTensor*"}],"params":[{"type":"DLTensor*"}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_tensor"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Tensor"}],"params":[{"type":"ffi.Tensor"}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_tensor_view_input"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"None"}],"params":[{"type":"DLTensor*"}]}})");
EXPECT_EQ(
fetch("testing.schema_id_opt_int"),
-
R"({"type":"ffi.Function","args":[{"type":"Optional","args":[{"type":"int"}]},{"type":"Optional","args":[{"type":"int"}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"Optional","args":[{"type":"int"}]}],"params":[{"type":"Optional","args":[{"type":"int"}]}]}})");
EXPECT_EQ(
fetch("testing.schema_id_opt_str"),
-
R"({"type":"ffi.Function","args":[{"type":"Optional","args":[{"type":"ffi.String"}]},{"type":"Optional","args":[{"type":"ffi.String"}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"Optional","args":[{"type":"ffi.String"}]}],"params":[{"type":"Optional","args":[{"type":"ffi.String"}]}]}})");
EXPECT_EQ(
fetch("testing.schema_id_opt_obj"),
-
R"({"type":"ffi.Function","args":[{"type":"Optional","args":[{"type":"ffi.Object"}]},{"type":"Optional","args":[{"type":"ffi.Object"}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"Optional","args":[{"type":"ffi.Object"}]}],"params":[{"type":"Optional","args":[{"type":"ffi.Object"}]}]}})");
EXPECT_EQ(
fetch("testing.schema_id_arr_int"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Array","args":[{"type":"int"}]},{"type":"ffi.Array","args":[{"type":"int"}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Array","args":[{"type":"int"}]}],"params":[{"type":"ffi.Array","args":[{"type":"int"}]}]}})");
EXPECT_EQ(
fetch("testing.schema_id_arr_str"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Array","args":[{"type":"ffi.String"}]},{"type":"ffi.Array","args":[{"type":"ffi.String"}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Array","args":[{"type":"ffi.String"}]}],"params":[{"type":"ffi.Array","args":[{"type":"ffi.String"}]}]}})");
EXPECT_EQ(
fetch("testing.schema_id_arr_obj"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Array","args":[{"type":"ffi.Object"}]},{"type":"ffi.Array","args":[{"type":"ffi.Object"}]}]})");
- EXPECT_EQ(fetch("testing.schema_id_arr"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Array"},{"type":"ffi.Array"}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Array","args":[{"type":"ffi.Object"}]}],"params":[{"type":"ffi.Array","args":[{"type":"ffi.Object"}]}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_arr"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Array"}],"params":[{"type":"ffi.Array"}]}})");
EXPECT_EQ(
fetch("testing.schema_id_map_str_int"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"int"}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"int"}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"int"}]}],"params":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"int"}]}]}})");
EXPECT_EQ(
fetch("testing.schema_id_map_str_str"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.String"}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.String"}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.String"}]}],"params":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.String"}]}]}})");
EXPECT_EQ(
fetch("testing.schema_id_map_str_obj"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Object"}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Object"}]}]})");
- EXPECT_EQ(fetch("testing.schema_id_map"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Map"},{"type":"ffi.Map"}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Object"}]}],"params":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Object"}]}]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_id_map"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Map"}],"params":[{"type":"ffi.Map"}]}})");
EXPECT_EQ(
fetch("testing.schema_id_variant_int_str"),
-
R"({"type":"ffi.Function","args":[{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"}]},{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"}]}],"params":[{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"}]}]}})");
// Packed function registered via def_packed: schema is plain ffi.Function
EXPECT_EQ(fetch("testing.schema_packed"), R"({"type":"ffi.Function"})");
@@ -142,18 +159,31 @@ TEST(Schema, GlobalFuncTypeSchema) {
// Mixed containers and optionals
EXPECT_EQ(
fetch("testing.schema_arr_map_opt"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"ffi.Array","args":[{"type":"Optional","args":[{"type":"int"}]}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"Optional","args":[{"type":"ffi.String"}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]}],"params":[{"type":"ffi.Array","args":[{"type":"Optional","args":[{"type":"int"}]}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"Optional","args":[{"type":"ffi.String"}]}]}})");
EXPECT_EQ(
fetch("testing.schema_variant_mix"),
-
R"({"type":"ffi.Function","args":[{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]}],"params":[{"type":"Variant","args":[{"type":"int"},{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]}]}})");
// No-arg and no-return combinations
- EXPECT_EQ(fetch("testing.schema_no_args"),
R"({"type":"ffi.Function","args":[{"type":"int"}]})");
- EXPECT_EQ(fetch("testing.schema_no_return"),
-
R"({"type":"ffi.Function","args":[{"type":"None"},{"type":"int"}]})");
+ EXPECT_EQ(fetch("testing.schema_no_args"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"int"}],"params":[]}})");
+ EXPECT_EQ(
+ fetch("testing.schema_no_return"),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"None"}],"params":[{"type":"int"}]}})");
EXPECT_EQ(fetch("testing.schema_no_args_no_return"),
- R"({"type":"ffi.Function","args":[{"type":"None"}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"None"}],"params":[]}})");
+}
+
+TEST(Schema, TypedFunctionGroups) {
+ EXPECT_EQ(
+ (TypedFunction<Tuple<int, String>()>::TypeSchema()),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"Tuple","args":[{"type":"int"},{"type":"ffi.String"}]}],"params":[]}})");
+ EXPECT_EQ(
+ (TypeTraits<std::function<int(int)>>::TypeSchema()),
+
R"({"type":"std::function","named_args":{"return":[{"type":"int"}],"params":[{"type":"int"}]}})");
+ EXPECT_EQ((TypeTraits<std::function<int()>>::TypeSchema()),
+
R"({"type":"std::function","named_args":{"return":[{"type":"int"}],"params":[]}})");
}
TEST(Schema, FieldTypeSchemas) {
@@ -199,34 +229,36 @@ TEST(Schema, MethodTypeSchemas) {
// Instance methods
EXPECT_EQ(
method_schema("add_int"),
-
R"({"type":"ffi.Function","args":[{"type":"int"},{"type":"testing.SchemaAllTypes"},{"type":"int"}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"int"}],"params":[{"type":"testing.SchemaAllTypes"},{"type":"int"}]}})");
EXPECT_EQ(
method_schema("append_int"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Array","args":[{"type":"int"}]},{"type":"testing.SchemaAllTypes"},{"type":"ffi.Array","args":[{"type":"int"}]},{"type":"int"}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Array","args":[{"type":"int"}]}],"params":[{"type":"testing.SchemaAllTypes"},{"type":"ffi.Array","args":[{"type":"int"}]},{"type":"int"}]}})");
EXPECT_EQ(
method_schema("maybe_concat"),
-
R"({"type":"ffi.Function","args":[{"type":"Optional","args":[{"type":"ffi.String"}]},{"type":"testing.SchemaAllTypes"},{"type":"Optional","args":[{"type":"ffi.String"}]},{"type":"Optional","args":[{"type":"ffi.String"}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"Optional","args":[{"type":"ffi.String"}]}],"params":[{"type":"testing.SchemaAllTypes"},{"type":"Optional","args":[{"type":"ffi.String"}]},{"type":"Optional","args":[{"type":"ffi.String"}]}]}})");
EXPECT_EQ(
method_schema("merge_map"),
-
R"({"type":"ffi.Function","args":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"testing.SchemaAllTypes"},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]}],"params":[{"type":"testing.SchemaAllTypes"},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]},{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"ffi.Array","args":[{"type":"int"}]}]}]}})");
// Static method make_with: return type is the object type itself.
// Build expected JSON as ffi.Function with return type = type_key and args
= (int, float, str)
EXPECT_EQ(
method_schema("make_with"),
-
R"({"type":"ffi.Function","args":[{"type":"testing.SchemaAllTypes"},{"type":"int"},{"type":"float"},{"type":"ffi.String"}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"testing.SchemaAllTypes"}],"params":[{"type":"int"},{"type":"float"},{"type":"ffi.String"}]}})");
}
TEST(Schema, DLLExportedFuncMetadata) {
// Minimal sanity check that DLL export metadata mechanism works.
-
EXPECT_EQ(ParseMetadataToSchema(CallMetadataFunc(__tvm_ffi__metadata_testing_dll_schema_id_int)),
-
R"({"type":"ffi.Function","args":[{"type":"int"},{"type":"int"}]})");
+ EXPECT_EQ(
+
ParseMetadataToSchema(CallMetadataFunc(__tvm_ffi__metadata_testing_dll_schema_id_int)),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"int"}],"params":[{"type":"int"}]}})");
}
TEST(Schema, DLLExportedFuncDocumentation) {
- EXPECT_EQ(ParseMetadataToSchema(
-
CallMetadataFunc(__tvm_ffi__metadata_testing_dll_test_add_with_docstring)),
-
R"({"type":"ffi.Function","args":[{"type":"int"},{"type":"int"},{"type":"int"}]})");
+ EXPECT_EQ(
+ ParseMetadataToSchema(
+
CallMetadataFunc(__tvm_ffi__metadata_testing_dll_test_add_with_docstring)),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"int"}],"params":[{"type":"int"},{"type":"int"}]}})");
String doc =
CallMetadataFunc(__tvm_ffi__doc_testing_dll_test_add_with_docstring);
std::string doc_str(doc);
EXPECT_EQ(doc_str,
diff --git a/tests/cpp/test_reflection.cc b/tests/cpp/test_reflection.cc
index f28a0574..84a80b6a 100644
--- a/tests/cpp/test_reflection.cc
+++ b/tests/cpp/test_reflection.cc
@@ -226,8 +226,9 @@ TEST(Reflection, ObjectPtrMethod) {
const TVMFFIMethodInfo* info =
reflection::GetMethodInfo("test.ObjectPtrHolder", "identity");
Map<String, Any> metadata =
json::Parse(String(info->metadata)).cast<Map<String, Any>>();
- EXPECT_EQ(metadata["type_schema"].cast<String>(),
-
R"({"type":"ffi.Function","args":[{"type":"test.Int"},{"type":"test.Int"}]})");
+ EXPECT_EQ(
+ metadata["type_schema"].cast<String>(),
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"test.Int"}],"params":[{"type":"test.Int"}]}})");
const TVMFFIMethodInfo* optional_info =
reflection::GetMethodInfo("test.ObjectPtrHolder", "optional_identity");
@@ -235,7 +236,7 @@ TEST(Reflection, ObjectPtrMethod) {
json::Parse(String(optional_info->metadata)).cast<Map<String, Any>>();
EXPECT_EQ(
optional_metadata["type_schema"].cast<String>(),
-
R"({"type":"ffi.Function","args":[{"type":"Optional","args":[{"type":"test.Int"}]},{"type":"Optional","args":[{"type":"test.Int"}]}]})");
+
R"({"type":"ffi.Function","named_args":{"return":[{"type":"Optional","args":[{"type":"test.Int"}]}],"params":[{"type":"Optional","args":[{"type":"test.Int"}]}]}})");
}
TEST(Reflection, FieldInfo) {
diff --git a/tests/python/test_stubgen.py b/tests/python/test_stubgen.py
index 712835cf..0de76e91 100644
--- a/tests/python/test_stubgen.py
+++ b/tests/python/test_stubgen.py
@@ -196,6 +196,28 @@ def test_funcinfo_gen_variants() -> None:
assert render_func_signature(func, ty_map, indent=2) == " def
no_args(*args: Any) -> Any: ..."
assert called == ["Any"]
+ return_only = FuncInfo.from_schema("demo.unknown", TypeSchema("Callable",
(TypeSchema("int"),)))
+ assert (
+ render_func_signature(return_only, _identity_ty_map, indent=0)
+ == "def unknown(*args: Any) -> int: ..."
+ )
+
+ grouped = TypeSchema.from_json_obj(
+ {
+ "type": "ffi.Function",
+ "named_args": {
+ "return": [{"type": "Tuple", "args": [{"type": "int"},
{"type": "str"}]}],
+ "params": [],
+ },
+ }
+ )
+ copied = NamedTypeSchema("demo.grouped", grouped)
+ assert copied.named_args == grouped.named_args
+ assert (
+ render_func_signature(FuncInfo.from_schema("demo.grouped", copied),
_identity_ty_map, 0)
+ == "def grouped() -> tuple[int, str]: ..."
+ )
+
schema_member = NamedTypeSchema(
"pkg.Class.method",
TypeSchema(
@@ -236,7 +258,12 @@ def test_objectinfo_gen_fields_and_methods() -> None:
],
methods=[
FuncInfo(
- schema=NamedTypeSchema("demo.static", TypeSchema("Callable",
(TypeSchema("int"),))),
+ schema=NamedTypeSchema(
+ "demo.static",
+ TypeSchema(
+ "Callable", named_args={"return":
(TypeSchema("int"),), "params": ()}
+ ),
+ ),
is_member=False,
),
FuncInfo(
@@ -274,7 +301,7 @@ def test_objectinfo_gen_overloaded_static_methods() -> None:
),
FuncInfo.from_schema(
"demo.Factory.reset",
- TypeSchema("Callable", (TypeSchema("None"),)),
+ TypeSchema("Callable", named_args={"return":
(TypeSchema("None"),), "params": ()}),
is_member=False,
),
FuncInfo.from_schema(
@@ -462,8 +489,9 @@ def test_py_class_method_metadata_renders_stub_signature()
-> None:
describe_schema = methods["describe"].schema
assert describe_schema.origin == "Callable"
- assert [arg.origin for arg in describe_schema.args] == [
- "str",
+ assert describe_schema.named_args is not None
+ assert [arg.origin for arg in describe_schema.named_args["return"]] ==
["str"]
+ assert [arg.origin for arg in describe_schema.named_args["params"]] == [
MethodMetadata.__tvm_ffi_type_info__.type_key, # ty:
ignore[unresolved-attribute]
"List",
"str",
@@ -522,7 +550,9 @@ def
test_generate_global_funcs_imports_enum_from_dataclasses() -> None:
FuncInfo(
schema=NamedTypeSchema(
"demo.get_enum",
- TypeSchema("Callable", (TypeSchema("ffi.Enum"),)),
+ TypeSchema(
+ "Callable", named_args={"return":
(TypeSchema("ffi.Enum"),), "params": ()}
+ ),
),
is_member=False,
)
@@ -661,7 +691,7 @@ def test_generate_object_with_methods() -> None:
),
FuncInfo.from_schema(
"demo.IntPair.sum",
- TypeSchema("Callable", (TypeSchema("int"),)),
+ TypeSchema("Callable", named_args={"return":
(TypeSchema("int"),), "params": ()}),
is_member=True,
),
],
diff --git a/tests/python/test_type_converter.py
b/tests/python/test_type_converter.py
index 16c6c7ab..fd54313f 100644
--- a/tests/python/test_type_converter.py
+++ b/tests/python/test_type_converter.py
@@ -84,6 +84,21 @@ def _make_str_enum_type() -> typing.Any:
A = TypeSchema.from_annotation
+def test_named_args_round_trip() -> None:
+ payload = TypeSchema("custom.Payload", fallback=TypeSchema("int"))
+ function = TypeSchema(
+ "Callable",
+ named_args={"return": (A(tuple[int, str]),), "params": (payload,)},
+ )
+ schema = TypeSchema("custom.Wrapper", named_args={"members": (function,),
"empty": ()})
+ assert TypeSchema.from_json_obj(schema.to_json()) == schema
+ assert function.input_repr() == "Callable[[int], tuple[int, str]]"
+ assert A(Callable[[], int]).to_json()["named_args"]["params"] == []
+ assert "params" not in A(Callable[..., int]).to_json()["named_args"]
+ with pytest.raises(TypeError):
+ TypeSchema.from_json_obj({"type": "int", "named_args": {"member":
{"type": "str"}}})
+
+
# ---------------------------------------------------------------------------
# Category 1: POD type exact match (check_value)
# ---------------------------------------------------------------------------
@@ -3488,9 +3503,10 @@ class TestSTLOriginParsing:
def test_std_function(self) -> None:
"""std::function maps to Callable."""
s = TypeSchema.from_json_str(
- '{"type":"std::function","args":[{"type":"int"},{"type":"str"}]}'
+ '{"type":"std::function","args":[{"type":"int"},[{"type":"str"}]]}'
)
assert s.origin == "Callable"
+ assert s.to_json()["named_args"]["params"] == [{"type": "str"}]
def test_object_rvalue_ref_origin(self) -> None:
"""ObjectRValueRef maps to Object."""
@@ -4213,7 +4229,9 @@ class TestFromAnnotationCallable:
def test_no_params(self) -> None:
"""Callable[[], int]."""
- assert A(Callable[[], int]) == S("Callable", S("int"))
+ assert A(Callable[[], int]) == TypeSchema(
+ "Callable", named_args={"return": (S("int"),), "params": ()}
+ )
class TestFromAnnotationList:
