Add extended mechanism

src/zfitwrapper/ModelParser.py

@@ -14,6 +14,7 @@ class TokenType(enum.Enum):
     T_LPAR = 5
     T_RPAR = 6
     T_END = 7
+    T_EXT = 8
 
 
 class Token(NamedTuple):
@@ -22,9 +23,15 @@ class Token(NamedTuple):
 
 
 class Node:
-    def __init__(self, token_type: TokenType, value: Optional[str] = None):
+    def __init__(
+        self,
+        token_type: TokenType,
+        value: Optional[str] = None,
+        ext_param: Optional[str] = None,
+    ):
         self.token_type = token_type
         self.value = value
+        self.ext_param = ext_param  # For extended models, stores the parameter
         self.children = []  # type: List[Node]
 
 
@@ -38,10 +45,7 @@ KnownMappings = {
 
 
 class ModelParser:
-    def __init__(
-        self,
-        modelstring: str,
-    ) -> None:
+    def __init__(self, modelstring: str) -> None:
         self.modelstring = modelstring.replace(" ", "")  # type: str
         self._nodeTokens = []  # type: List[Node]
 
@@ -80,6 +84,8 @@ class ModelParser:
     def __parenthesesParse(tokens: List[Node]) -> Node:
         if tokens[0].token_type == TokenType.T_MODEL:
             return tokens.pop(0)
+        if tokens[0].token_type == TokenType.T_EXT:
+            return tokens.pop(0)
 
         ModelParser.__match(tokens, TokenType.T_LPAR)
         expression = ModelParser.__baseParse(tokens)
@@ -88,30 +94,25 @@ class ModelParser:
         return expression
 
     def _lexical_analysis(self) -> List[Node]:
-        for expr in re.split(r"(\W+)", self.modelstring):
-            if expr:
-                if re.match(r"\d+", expr):
-                    raise ValueError("[ModelParser]: No explicit numerical values possible")
-
-                elif re.match(r"\w+", expr):
-                    token = Node(TokenType.T_MODEL, value=expr)
-                    self._nodeTokens.append(token)
-
-                elif expr in KnownMappings.keys():
-                    token_type = KnownMappings[expr]
-                    token = Node(token_type.type, value=expr)
-                    self._nodeTokens.append(token)
-
-                else:
-                    for symbol in re.split("(\(|\))", expr):
-                        if symbol:
-                            if symbol in KnownMappings.keys():
-                                token_type = KnownMappings[symbol]
-                                token = Node(token_type.type, value=symbol)
-                                self._nodeTokens.append(token)
-
-                            else:
-                                raise Exception("Invalid token: {}".format(expr))
+        token_pattern = r"(\w+@\w+|\w+|[+*()@])"
+        tokens = re.findall(token_pattern, self.modelstring)
+
+        for expr in tokens:
+            if re.match(r"\w+@\w+", expr):
+                param_name, model_name = expr.split("@")
+                token = Node(TokenType.T_EXT, value=model_name, ext_param=param_name)
+                self._nodeTokens.append(token)
+            elif re.match(r"\w+", expr):
+                token = Node(TokenType.T_MODEL, value=expr)
+                self._nodeTokens.append(token)
+            elif expr in KnownMappings.keys():
+                token_type = KnownMappings[expr]
+                token = Node(token_type.type, value=expr)
+                self._nodeTokens.append(token)
+            elif expr == "@":
+                continue  # '@' is handled as part of the extended model token
+            else:
+                raise Exception("Invalid token: {}".format(expr))
 
         self._nodeTokens.append(Node(TokenType.T_END))