gemini-code-assist[bot] commented on code in PR #18639:
URL: https://github.com/apache/tvm/pull/18639#discussion_r2663381252
##########
tests/python/testing/test_type_annotation_checker.py:
##########
@@ -187,5 +187,43 @@ def func(_: type_annotation):
func(case)
[email protected](
+ ["type_annotation", "expected_key", "expected_subtypes"],
+ [
+ pytest.param(Union[str, int], "union", [str, int], id="Union[str,
int]"),
+ pytest.param(List[str], "list", [str], id="List[str]"),
+ pytest.param(Dict[str, int], "dict", [str, int], id="Dict[str, int]"),
+ pytest.param(Tuple[str, int], "tuple", (str, int), id="Tuple[str,
int]"),
+ pytest.param(
+ Union[List[str], Dict[str, int]],
+ "union",
+ None, # We'll check length instead
+ id="Union[List[str], Dict[str, int]]",
+ ),
+ ],
+)
+def test_subscripted_generics(type_annotation, expected_key,
expected_subtypes):
+ """Test that _dispatcher correctly handles subscripted generics in Python
3.14+.
+
+ In Python 3.14, Union and other generic types have a different internal
representation.
+ This test ensures that the dispatcher correctly identifies these types.
+ """
+ from tvm.tir.schedule._type_checker import _dispatcher
+
+ key, subtypes = _dispatcher(type_annotation)
+ assert key == expected_key, f"Expected '{expected_key}' but got '{key}'"
+
+ if expected_subtypes is not None:
+ if isinstance(expected_subtypes, tuple):
+ assert tuple(subtypes) == expected_subtypes, (
+ f"Expected {expected_subtypes} but got {subtypes}"
+ )
+ else:
+ assert subtypes == expected_subtypes, f"Expected
{expected_subtypes} but got {subtypes}"
+ else:
+ # For nested Union test, just check that we got 2 subtypes
+ assert len(subtypes) == 2, f"Expected 2 subtypes but got
{len(subtypes)}"
Review Comment:
The test for the nested `Union` is a bit weak as it only checks the number
of subtypes. You can make it more robust by checking for the exact subtypes.
This also simplifies the test logic by removing the special case for
`expected_subtypes is None`.
```python
pytest.param(
Union[List[str], Dict[str, int]],
"union",
[List[str], Dict[str, int]],
id="Union[List[str], Dict[str, int]]",
),
],
)
def test_subscripted_generics(type_annotation, expected_key,
expected_subtypes):
"""Test that _dispatcher correctly handles subscripted generics in
Python 3.14+.
In Python 3.14, Union and other generic types have a different internal
representation.
This test ensures that the dispatcher correctly identifies these types.
"""
from tvm.tir.schedule._type_checker import _dispatcher
key, subtypes = _dispatcher(type_annotation)
assert key == expected_key, f"Expected '{expected_key}' but got '{key}'"
if isinstance(expected_subtypes, tuple):
assert tuple(subtypes) == expected_subtypes, (
f"Expected {expected_subtypes} but got {subtypes}"
)
else:
assert subtypes == expected_subtypes, f"Expected {expected_subtypes}
but got {subtypes}"
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
