pytholite/compiler: support range(constants) in for loops

diff --git a/migen/pytholite/compiler.py b/migen/pytholite/compiler.py
index 6d5b328034e6c53e08a801c6890e83b500d6dd85..116940918f14482a98e5dee9f65730d97107f376 100644 (file)
--- a/migen/pytholite/compiler.py
+++ b/migen/pytholite/compiler.py
@@ -128,7 +128,7 @@ class _Compiler:
                        target = statement.target.id
                        if target in self.symdict:
                                raise NotImplementedError("For loop target must use an available name")
-                       it = ast.literal_eval(statement.iter)
+                       it = self.visit_iterator(statement.iter)
                        last_exit_states = []
                        for iteration in it:
                                self.symdict[target] = iteration
@@ -143,6 +143,19 @@ class _Compiler:
                        raise NotImplementedError
                return states, exit_states
        
+       def visit_iterator(self, node):
+               if isinstance(node, ast.List):
+                       return ast.literal_eval(node)
+               elif isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+                       funcname = node.func.id
+                       args = map(ast.literal_eval, node.args)
+                       if funcname == "range":
+                               return range(*args)
+                       else:
+                               raise NotImplementedError
+               else:
+                       raise NotImplementedError
+       
        def visit_assign(self, node):
                if isinstance(node.targets[0], ast.Name):
                        self.targetname = node.targets[0].id