Fix ValueSpec.from_annotation on Tuple[<type>, ...].

PiperOrigin-RevId: 549462506
google · Jul 19, 2023 · 9e2a772 · 9e2a772
1 parent 627990e
commit 9e2a772
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/pyglove/core/typing/annotation_conversion.py b/pyglove/core/typing/annotation_conversion.py
@@ -102,10 +102,17 @@ def _value_spec_from_type_annotation(
             args[0], True)) if args else vs.List(vs.Any())
   # Handling tuple.
   elif origin in (tuple, typing.Tuple):
-    if args:
-      return vs.Tuple([_value_spec_from_annotation(arg, True) for arg in args])
-    else:
+    if not args:
       return vs.Tuple(vs.Any())
+    else:
+      if args[-1] is ...:
+        if len(args) != 2:
+          raise TypeError(
+              f'Tuple with ellipsis should have exact 2 type arguments. '
+              f'Encountered: {annotation}.')
+        return vs.Tuple(_value_spec_from_type_annotation(args[0], False))
+      return vs.Tuple([_value_spec_from_type_annotation(arg, False)
+                       for arg in args])
   # Handle sequence.
   elif origin in (collections.abc.Sequence,):
     elem = _value_spec_from_annotation(args[0], True) if args else vs.Any()

diff --git a/pyglove/core/typing/annotation_conversion_test.py b/pyglove/core/typing/annotation_conversion_test.py
@@ -167,6 +167,11 @@ def test_tuple(self):
         ValueSpec.from_annotation(typing.Tuple, True), vs.Tuple(vs.Any()))
     self.assertEqual(
         ValueSpec.from_annotation(tuple[int], True), vs.Tuple([vs.Int()]))
+    self.assertEqual(
+        ValueSpec.from_annotation(tuple[int, ...], True), vs.Tuple(vs.Int()))
+    with self.assertRaisesRegex(
+        TypeError, 'Tuple with ellipsis should have exact 2 type arguments'):
+      ValueSpec.from_annotation(tuple[...], True)
 
   def test_sequence(self):
     self.assertEqual(