Enable extending existing types to add custom methods

smnp/newast/node/extend.py

@@ -0,0 +1,69 @@
+from smnp.newast.node.block import BlockNode
+from smnp.newast.node.function import FunctionDefinitionNode
+from smnp.newast.node.identifier import IdentifierNode
+from smnp.newast.node.none import NoneNode
+from smnp.newast.node.statement import StatementNode
+from smnp.newast.node.type import TypeNode
+from smnp.newast.parser import Parser
+from smnp.token.type import TokenType
+
+
+class ExtendNode(StatementNode):
+    def __init__(self, pos):
+        super().__init__(pos)
+        self.children = [NoneNode(), NoneNode(), NoneNode()]
+
+    @property
+    def type(self):
+        return self[0]
+
+    @type.setter
+    def type(self, value):
+        self[0] = value
+
+    @property
+    def variable(self):
+        return self[1]
+
+    @variable.setter
+    def variable(self, value):
+        self[1] = value
+
+    @property
+    def methods(self):
+        return self[2]
+
+    @methods.setter
+    def methods(self, value):
+        self[2] = value
+
+    @classmethod
+    def _parse(cls, input):
+        def createNode(extend, type, variable, methods):
+            node = ExtendNode(extend.pos)
+            node.type = type
+            node.variable = variable
+            node.methods = methods
+            return node
+
+        return Parser.allOf(
+            Parser.terminalParser(TokenType.EXTEND),
+            TypeNode.parse,
+            IdentifierNode.identifierParser(),
+            cls._methodsDeclarationsParser(),
+            createNode=createNode
+        )(input)
+
+    @classmethod
+    def _methodsDeclarationsParser(cls):
+        def createNode(openBracket, items, closeBracket):
+            node = BlockNode(openBracket.pos)
+            node.children = items
+            return node
+
+        return Parser.loop(
+            Parser.terminalParser(TokenType.OPEN_BRACKET),
+            FunctionDefinitionNode.parse,
+            Parser.terminalParser(TokenType.CLOSE_BRACKET),
+            createNode=createNode
+        )

smnp/newast/node/function.py

@@ -16,7 +16,7 @@ class ArgumentsDeclarationNode(Node):
         raise RuntimeError("This class is not supposed to be automatically called")
 
 
-class FunctionDefinition(StatementNode):
+class FunctionDefinitionNode(StatementNode):
     def __init__(self, pos):
         super().__init__(pos)
         self.children = [NoneNode(), NoneNode(), NoneNode()]
@@ -48,7 +48,7 @@ class FunctionDefinition(StatementNode):
     @classmethod
     def _parse(cls, input):
         def createNode(function, name, arguments, body):
-            node = FunctionDefinition(function.pos)
+            node = FunctionDefinitionNode(function.pos)
             node.name = name
             node.arguments = arguments
             node.body = body

smnp/newast/node/program.py

@@ -1,5 +1,7 @@
 from smnp.error.syntax import SyntaxException
 from smnp.newast.node.expression import ExpressionNode
+from smnp.newast.node.extend import ExtendNode
+from smnp.newast.node.function import FunctionDefinitionNode
 from smnp.newast.node.model import Node, ParseResult
 from smnp.newast.node.statement import StatementNode
 from smnp.newast.parser import Parser
@@ -13,9 +15,11 @@ class Program(Node):
     def _parse(cls, input):
         def parseToken(input):
             return Parser.oneOf(
+                FunctionDefinitionNode.parse,
+                ExtendNode.parse,
                 ExpressionNode.parse,
                 StatementNode.parse,
-                exception = SyntaxException("Unknown statement")
+                exception = SyntaxException(f"Unknown statement: {input.current().pos}")
             )(input)
 
         root = Program()

smnp/newast/node/statement.py

@@ -7,13 +7,11 @@ class StatementNode(Node):
     @classmethod
     def _parse(cls, input):
         from smnp.newast.node.block import BlockNode
-        from smnp.newast.node.function import FunctionDefinition
         from smnp.newast.node.expression import ExpressionNode
         from smnp.newast.node.ret import ReturnNode
 
         return Parser.oneOf(
             BlockNode.parse,
-            FunctionDefinition.parse,
             ReturnNode.parse,
             ExpressionNode.parse
         )(input)