diff --git a/autocog/sta/ast.py b/autocog/sta/ast.py index 001910d..057d111 100644 --- a/autocog/sta/ast.py +++ b/autocog/sta/ast.py @@ -175,10 +175,13 @@ def gvlbl(self): class EnumType(ASTNode): kind: str + arguments: List[Argument] = [] source: Union[Path,List[Expression]] def gvtree(self): yield from super().gvtree() + for (i,arg) in enumerate(self.arguments): + yield ("argument",i,arg) if isinstance(self.source, Path): yield ("path",None,self.source) else: diff --git a/autocog/sta/frontend.py b/autocog/sta/frontend.py index c01a93c..5b0290d 100644 --- a/autocog/sta/frontend.py +++ b/autocog/sta/frontend.py @@ -410,18 +410,42 @@ def visit_string_expr(self, node, visited_children): return visited_children[0] def visit_repeat_def(self, node, visited_children): - assert len(visited_children) == 5 - return EnumType(kind='repeat', source=visited_children[3]) + assert len(visited_children) == 7 + arguments = visited_children[2] + if 'children' in arguments: + assert len(arguments['children']) == 1 + arguments = arguments['children'][0] + assert arguments['kind'] == 'param_list' + arguments = arguments['children'] + else: + arguments = [] + return EnumType(kind='repeat', source=visited_children[5], arguments=arguments) def visit_select_def(self, node, visited_children): - assert len(visited_children) == 5 - return EnumType(kind='select', source=visited_children[3]) + assert len(visited_children) == 7 + arguments = visited_children[2] + if 'children' in arguments: + assert len(arguments['children']) == 1 + arguments = arguments['children'][0] + assert arguments['kind'] == 'param_list' + arguments = arguments['children'] + else: + arguments = [] + return EnumType(kind='select', source=visited_children[5], arguments=arguments) def visit_enum_def(self, node, visited_children): - assert len(visited_children) == 5 - source = visited_children[3] + assert len(visited_children) == 7 + arguments = visited_children[2] + if 'children' in arguments: + assert len(arguments['children']) == 1 + arguments = arguments['children'][0] + assert arguments['kind'] == 'param_list' + arguments = arguments['children'] + else: + arguments = [] + source = visited_children[5] assert source['kind'] == 'string_expr_list' - return EnumType(kind='enum', source=source['children']) + return EnumType(kind='enum', source=source['children'], arguments=arguments) def visit_string_expr_list__(self, node, visited_children): if len(visited_children) == 1: diff --git a/autocog/sta/grammar.py b/autocog/sta/grammar.py index 48c3933..ad1eb44 100644 --- a/autocog/sta/grammar.py +++ b/autocog/sta/grammar.py @@ -83,9 +83,9 @@ array_slice_cont = array_slice__? array_slice__ = WS COLON WS expression -repeat_def = REPEAT WS LPAR path_expr RPAR -select_def = SELECT WS LPAR path_expr RPAR -enum_def = ENUM WS LPAR string_expr_list RPAR +repeat_def = REPEAT WS type_ref_param? WS LPAR path_expr RPAR +select_def = SELECT WS type_ref_param? WS LPAR path_expr RPAR +enum_def = ENUM WS type_ref_param? WS LPAR string_expr_list RPAR type_ref = identifier WS type_ref_param? type_ref_param = LT WS param_list WS GT