Gen

brentyi · Oct 30, 2024 · 6eb4a3f · 6eb4a3f
1 parent c5ca29b
commit 6eb4a3f
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/tests/test_py311_generated/test_custom_primitive_generated.py b/tests/test_py311_generated/test_custom_primitive_generated.py
@@ -14,9 +14,9 @@
 
 def test_custom_primitive_registry():
     """Test that we can use a custom primitive registry to parse a custom type."""
-    registry = tyro.constructors.PrimitiveConstructorRegistry()
+    primitive_registry = tyro.constructors.PrimitiveConstructorRegistry()
 
-    @registry.define_rule(
+    @primitive_registry.define_rule(
         matcher_fn=lambda type_info: type_info.type_origin is dict
         and get_args(type_info.type) == (str, Any)
     )
@@ -29,9 +29,8 @@ def json_dict_spec(
     def main(x: Dict[str, Any]) -> Dict[str, Any]:
         return x
 
-    assert tyro.cli(
-        main, args=["--x", '{"a": 1}'], primitive_constructor_registry=registry
-    ) == {"a": 1}
+    with primitive_registry:
+        assert tyro.cli(main, args=["--x", '{"a": 1}']) == {"a": 1}
 
 
 def test_custom_primitive_annotated():