support "paths.with.123.inside"

PiperOrigin-RevId: 732093550
google-research · Feb 28, 2025 · cb65a07 · cb65a07
1 parent afafe04
commit cb65a07
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 1 deletion.
diff --git a/kauldron/kontext/path_grammar.lark b/kauldron/kontext/path_grammar.lark
@@ -27,7 +27,7 @@ tensor_slice_key.2:  tensor_axis_key ","
 !slice_key: [number] ":" [number]
           | [number] ":" [number] ":" [number]
 
-_identifier: STAR | DOUBLE_STAR | IDENTIFIER
+_identifier: STAR | DOUBLE_STAR | IDENTIFIER | DEC_NUMBER
 
 NONE: "None"
 BOOLEAN: "True" | "False"

diff --git a/kauldron/kontext/path_parser.py b/kauldron/kontext/path_parser.py
@@ -116,6 +116,10 @@ def BOOLEAN(args: str) -> bool:
   def STRING(args: str) -> str:
     return ast.literal_eval(args)
 
+  @staticmethod
+  def DEC_NUMBER(args: str) -> Any:
+    return int(args)
+
   @staticmethod
   def tuple_key(args: list[Any]) -> tuple[Any, ...]:
     return tuple(args)

diff --git a/kauldron/kontext/path_parser_test.py b/kauldron/kontext/path_parser_test.py
@@ -27,6 +27,16 @@ def test_parse():
   ]
 
 
+def test_parse_integer_key():
+  parts = path_parser.parse_parts('a.123.b[10]')
+  assert parts == [
+      'a',
+      123,
+      'b',
+      10,
+  ]
+
+
 def test_parse_interms_example():
   parts = path_parser.parse_parts('interms.model.__call__[0]')
   assert parts == ['interms', 'model', '__call__', 0]