From b7192ea52b384ecf86546020bc75c6bb3dc502e9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bart=C5=82omiej=20Pluta?= <the.orleander@gmail.com>
Date: Fri, 12 Jul 2019 14:12:55 +0200
Subject: [PATCH] Add optional 'as' operator to loop

---
 smnp/ast/node/factor.py | 40 +++++++++++++++++++++++++++++++++++++++-
 smnp/ast/parser.py      | 11 +++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/smnp/ast/node/factor.py b/smnp/ast/node/factor.py
index 25e231d..f1bd8a5 100644
--- a/smnp/ast/node/factor.py
+++ b/smnp/ast/node/factor.py
@@ -1,4 +1,7 @@
 from smnp.ast.node.chain import ChainParser
+from smnp.ast.node.iterable import abstractIterableParser
+from smnp.ast.node.model import Node
+from smnp.ast.node.none import NoneNode
 from smnp.ast.node.operator import BinaryOperator, Operator, UnaryOperator
 from smnp.ast.node.valuable import Valuable
 from smnp.ast.parser import Parser
@@ -14,12 +17,36 @@ class NotOperator(UnaryOperator):
 
 
 class Loop(BinaryOperator):
+    def __init__(self, pos):
+        super().__init__(pos)
+        self.children.append(NoneNode())
+
+    @property
+    def parameters(self):
+        return self[3]
+
+    @parameters.setter
+    def parameters(self, value):
+        self[3] = value
+
+    @classmethod
+    def loop(cls, left, parameters, operator, right):
+        node = cls(left.pos)
+        node.left = left
+        node.parameters = parameters
+        node.operator = operator
+        node.right = right
+        return node
+
+
+class LoopParameters(Node):
     pass
 
 
 def FactorParser(input):
     from smnp.ast.node.expression import ExpressionParser
     from smnp.ast.node.statement import StatementParser
+    from smnp.ast.node.identifier import IdentifierLiteralParser
 
     powerFactor = Parser.leftAssociativeOperatorParser(
         ChainParser,
@@ -50,11 +77,22 @@ def FactorParser(input):
         name="not"
     )
 
+    loopParameters = Parser.allOf(
+        Parser.terminal(TokenType.AS),
+        Parser.oneOf(
+            Parser.wrap(IdentifierLiteralParser, lambda id: LoopParameters.withChildren([id], id.pos)),
+            abstractIterableParser(LoopParameters, TokenType.OPEN_PAREN, TokenType.CLOSE_PAREN, IdentifierLiteralParser)
+        ),
+        createNode=lambda asKeyword, parameters: parameters,
+        name="loop parameters"
+    )
+
     loopFactor = Parser.allOf(
         factorParser,
+        Parser.optional(loopParameters),
         Parser.terminal(TokenType.DASH, createNode=Operator.withValue),
         StatementParser,
-        createNode=Loop.withValues,
+        createNode=Loop.loop,
         name="dash-loop"
     )
 
diff --git a/smnp/ast/parser.py b/smnp/ast/parser.py
index 202d75e..dbe49b7 100644
--- a/smnp/ast/parser.py
+++ b/smnp/ast/parser.py
@@ -207,3 +207,14 @@ class Parser:
                     return ParseResult.OK(createNode(results, pos) if len(results) > 0 else NoneNode())
 
         return Parser(parse, name, parsers=[parser])
+
+    @staticmethod
+    def wrap(parser, createNode):
+        def parse(input):
+            result = parser(input)
+            if result.result:
+                return ParseResult.OK(createNode(result.node))
+
+            return result
+
+        return parse
\ No newline at end of file