From 11d585df1580415ffefee551ba9bc75135eb3799 Mon Sep 17 00:00:00 2001
From: ScR4PP <95632351+JarcauCristian@users.noreply.github.com>
Date: Mon, 18 Nov 2024 11:43:19 +0200
Subject: [PATCH] Update mage_to_python_utils.py

---
 mage_to_cwl/mage_to_python_utils.py | 27 +++++++++++++--------------
 1 file changed, 13 insertions(+), 14 deletions(-)

diff --git a/mage_to_cwl/mage_to_python_utils.py b/mage_to_cwl/mage_to_python_utils.py
index a4443ff..1f76bea 100644
--- a/mage_to_cwl/mage_to_python_utils.py
+++ b/mage_to_cwl/mage_to_python_utils.py
@@ -62,18 +62,18 @@ def visit_Module(self, node):
         node.body = [self.visit(stmt) for stmt in node.body if stmt is not None]
         return node
 
-    def remove_unused_code(self, code):
-        try:
-            tree = ast.parse(code)
-
-            transformer = UnusedCodeRemover()
-            sanitized_tree = transformer.visit(tree)
-
-            sanitized_code = ast.unparse(sanitized_tree)
-            return sanitized_code
-        except Exception as e:
-            print(f"Error while removing unused code: {e}")
-            return code
+def remove_unused_code(code):
+     try:
+        tree = ast.parse(code)
+
+        transformer = RemoveUnusedCode()
+        sanitized_tree = transformer.visit(tree)
+
+        sanitized_code = ast.unparse(sanitized_tree)
+        return sanitized_code
+    except Exception as e:
+        print(f"Error while removing unused code: {e}")
+        return code
     
 
 class MageToPythonTransformer(ast.NodeTransformer):
@@ -234,8 +234,7 @@ def remove_imports_with_word(code_string: str, word: str, block_name: str, previ
         cleaned_tree.body.insert(1, pandas_import)
 
     cleaned_code = ast.unparse(cleaned_tree)
-    cleaner = RemoveUnusedCode()
-    cleaned_code = cleaner.remove_unused_code(cleaned_code)
+    cleaned_code = remove_unused_code(cleaned_code)
     return cleaned_code, transformer.env_vars