From 915d05aa1ec63a279b242f9f77244ede06c83059 Mon Sep 17 00:00:00 2001 From: Aryaz Eghbali Date: Thu, 18 Nov 2021 19:46:46 +0100 Subject: [PATCH] Fixed issue with __future__ imports --- src/lexecutor/CodeRewriter.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/lexecutor/CodeRewriter.py b/src/lexecutor/CodeRewriter.py index acd1d30..d504833 100644 --- a/src/lexecutor/CodeRewriter.py +++ b/src/lexecutor/CodeRewriter.py @@ -81,12 +81,21 @@ def visit_SimpleStatementLine(self, node): # add import of our runtime library to the file def leave_Module(self, node, updated_node): + imports_index = -1 + for i in range(len(updated_node.body)): + if isinstance(updated_node.body[i].body, tuple): + if isinstance(updated_node.body[i].body[0], (cst.Import, cst.ImportFrom)): + imports_index = i + else: + break + else: + break import_n = self.__create_import("_n_") import_a = self.__create_import("_a_") import_c = self.__create_import("_c_") import_b = self.__create_import("_b_") - new_body = [import_n, import_a, import_c, - import_b]+list(updated_node.body) + new_body = list(updated_node.body[:imports_index+1]) + [import_n, import_a, import_c, + import_b] + list(updated_node.body[imports_index+1:]) return updated_node.with_changes(body=new_body) # rewrite Call nodes to intercept function calls