diff --git a/latticejson/madx.lark b/latticejson/madx.lark index 4679f99..bbaa152 100644 --- a/latticejson/madx.lark +++ b/latticejson/madx.lark @@ -12,17 +12,19 @@ string : ESCAPED_STRING word : /[\w\.]+/ start : (_statement ";")* -_statement : element | lattice | command | assignment +_statement : element | lattice | sequence | command | assignment element : word ":" [word] ("," attribute)* ","? attribute : word ("=" | ":=") (expr | string) - lattice : word ":" "LINE"i "=" arrangement arrangement : [int "*"] [/-/] "(" object ("," object)* ")" ?object : ref_name | arrangement ref_name : [int "*"] [/-/] word +sequence : word ":" "SEQUENCE"i ("," attribute)* ";" seq_elements "ENDSEQUENCE"i +seq_elements: (element ";")* + command : word ("," (word | string | attribute))* // Arithmeitc Expressions diff --git a/latticejson/parse.py b/latticejson/parse.py index b97245c..cdbee02 100644 --- a/latticejson/parse.py +++ b/latticejson/parse.py @@ -90,7 +90,8 @@ def transform(self, tree): string = lambda self, item: item[1:-1] def element(self, name, type_, *attributes): - self.elements[name.lower()] = type_.lower(), dict(attributes) + self.elements[name.upper()] = type_.upper(), dict(attributes) + return name def attribute(self, name, value): return name.lower(), value @@ -150,7 +151,12 @@ def command(self, *items): @v_args(inline=True) class MADXTransformer(ArithmeticTransformer, AbstractLatticeFileTransformer): - pass + def sequence(self, name, *items): + *attributes, elements = items + self.lattices[name.upper()] = elements + + def seq_elements(self, *elements): + return list(elements) @v_args(inline=True)