From ac8b46b077c3c16e81cd478df7053122335fb3ed 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:19:21 +0200
Subject: [PATCH] Improve 'extend' statement

---
 smnp/ast/node/extend.py | 22 +++++++++++++++++++++-
 smnp/token/tokenizer.py |  1 +
 smnp/token/type.py      |  1 +
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/smnp/ast/node/extend.py b/smnp/ast/node/extend.py
index 270ec7b..5a23763 100644
--- a/smnp/ast/node/extend.py
+++ b/smnp/ast/node/extend.py
@@ -45,14 +45,34 @@ class Extend(Node):
         node.methods = methods
         return node
 
+
 def ExtendParser(input):
-    return Parser.allOf(
+
+    simpleExtend = Parser.allOf(
+        Parser.terminal(TokenType.EXTEND),
+        TypeParser,
+        Parser.terminal(TokenType.AS),
+        IdentifierLiteralParser,
+        Parser.terminal(TokenType.WITH),
+        Parser.wrap(FunctionDefinitionParser, lambda method: Block.withChildren([ method ], method.pos)),
+        createNode=lambda extend, type, _, variable, __, methods: Extend.withValues(extend.pos, type, variable, methods),
+        name="simple extend"
+    )
+
+    multiExtend = Parser.allOf(
         Parser.terminal(TokenType.EXTEND),
         TypeParser,
         Parser.terminal(TokenType.AS),
         IdentifierLiteralParser,
         MethodsDeclarationParser,
         createNode=lambda extend, type, _, variable, methods: Extend.withValues(extend.pos, type, variable, methods),
+        name="multiple extend"
+    )
+
+
+    return Parser.oneOf(
+        simpleExtend,
+        multiExtend,
         name="extend"
     )(input)
 
diff --git a/smnp/token/tokenizer.py b/smnp/token/tokenizer.py
index 2e23425..0caa830 100644
--- a/smnp/token/tokenizer.py
+++ b/smnp/token/tokenizer.py
@@ -51,6 +51,7 @@ tokenizers = (
     separated(defaultTokenizer(TokenType.EXTEND)),
     separated(defaultTokenizer(TokenType.IMPORT)),
     separated(defaultTokenizer(TokenType.FROM)),
+    separated(defaultTokenizer(TokenType.WITH)),
     separated(defaultTokenizer(TokenType.ELSE)),
     separated(defaultTokenizer(TokenType.AND)),
     separated(defaultTokenizer(TokenType.NOT)),
diff --git a/smnp/token/type.py b/smnp/token/type.py
index 35799fd..56b4e5f 100644
--- a/smnp/token/type.py
+++ b/smnp/token/type.py
@@ -36,6 +36,7 @@ class TokenType(Enum):
     EXTEND = 'extend'
     IMPORT = 'import'
     FROM = 'from'
+    WITH = 'with'
     ELSE = 'else'
     IF = 'if'
     AS = 'as'