Refactorize visitors

Author: boriel

src/api/optimize.py

@@ -1,8 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-import types
-
 from typing import NamedTuple
 from typing import Optional
 from typing import Set
@@ -35,6 +33,7 @@ class GenericVisitor(NodeVisitor):
     """ A slightly different visitor, that just traverses an AST, but does not return
     a translation of it. Used to examine the AST or do transformations
     """
+    node_type = ToVisit
 
     @property
     def O_LEVEL(self):
@@ -56,25 +55,9 @@ def TYPE(type_):
         return gl.SYMBOL_TABLE.basic_types[type_]
 
     def visit(self, node):
-        stack = [ToVisit(node)]
-        last_result = None
-
-        while stack:
-            try:
-                last = stack[-1]
-                if isinstance(last, types.GeneratorType):
-                    stack.append(last.send(last_result))
-                    last_result = None
-                elif isinstance(last, ToVisit):
-                    stack.append(self._visit(stack.pop()))
-                else:
-                    last_result = stack.pop()
-            except StopIteration:
-                stack.pop()
-
-        return last_result
-
-    def _visit(self, node):
+        return super().visit(ToVisit(node))
+
+    def _visit(self, node: ToVisit):
         if node.obj is None:
             return None
 

src/ast/ast.py

@@ -9,7 +9,7 @@
 #                    the GNU General License
 # ----------------------------------------------------------------------
 
-from typing import Callable, Any
+from typing import Any, Callable, Type
 
 import types
 from .tree import Tree
@@ -27,6 +27,8 @@ def token(self):
 
 
 class NodeVisitor:
+    node_type: Type = Ast
+
     def visit(self, node):
         stack = [node]
         last_result = None
@@ -37,7 +39,7 @@ def visit(self, node):
                 if isinstance(last, types.GeneratorType):
                     stack.append(last.send(last_result))
                     last_result = None
-                elif isinstance(last, Ast):
+                elif isinstance(last, self.node_type):
                     stack.append(self._visit(stack.pop()))
                 else:
                     last_result = stack.pop()
@@ -47,9 +49,7 @@ def visit(self, node):
         return last_result
 
     def _visit(self, node):
-        methname = 'visit_' + node.token
-        meth = getattr(self, methname, self.generic_visit)
-
+        meth = getattr(self, f"visit_{node.token}", self.generic_visit)
         return meth(node)
 
     @staticmethod