Commits

Victor Stinner committed 26e4023

log removed nodes

  • Participants
  • Parent commits b316361

Comments (0)

Files changed (5)

File astoptimizer/__init__.py

 from astoptimizer.config import Config
 from astoptimizer.optimizer import Optimizer
 
-def optimize_ast(tree, config=None):
+def optimize_ast(tree, config=None, filename=None):
     """
     Optimize an AST with the specified config (or with default config).
     Return the optimized AST
     """
     optimizer = Optimizer(config)
-    return optimizer.optimize(tree)
+    return optimizer.optimize(tree, filename)
 
 def optimize_code(code, filename=None, mode=None, flags=0, dont_inherit=False, optimize=-1, config=None):
     """
     (or default config).
     """
     tree = parse_ast(code, filename, mode, flags, dont_inherit, optimize)
-    tree = optimize_ast(tree, config)
+    tree = optimize_ast(tree, config, filename)
     return compile_ast(tree, filename, mode, flags, dont_inherit, optimize)
 
 def patch_compile(config=None):
         tree = parse_ast(source, filename, mode, flags, dont_inherit, optimize)
         if flags & ast.PyCF_ONLY_AST:
             return tree
-        tree = optimize_ast(tree, config)
+        tree = optimize_ast(tree, config, filename)
         code = compile_ast(tree, filename, mode, flags, dont_inherit, optimize)
         return code
 

File astoptimizer/config.py

         return self._call(config, func, args)
 
 
+def display_info(message):
+    sys.stderr.write("%s\n" % message)
+    sys.stderr.flush()
+
 def display_warning(message):
     sys.stderr.write("WARNING: %s\n" % message)
     sys.stderr.flush()
         # Example: "if 0: print('debug')" => "pass"
         self.remove_dead_code = True
 
-        # Callback to display a message, prototype:
+        # Callback to display a information message, prototype:
+        #
+        # def info(message):
+        #     ...
+        self.info_func = display_info
+
+        # Callback to display a warning message, prototype:
         #
         # def warning(message):
         #     ...

File astoptimizer/optimizer.py

         self.namespace = None
         self.root = None
         self.is_conditional = None
+        self.filename = None
 
     def _optimize(self, tree, namespace):
         self.namespace = namespace
         self.is_conditional = False
         return self.visit(tree)
 
-    def optimize(self, tree):
+    def optimize(self, tree, filename=None):
         namespace = Namespace()
+        self.filename = filename
         return self._optimize(tree, namespace)
 
+    def info(self, message):
+        self.config.info_func(message)
+
+    def log_node_removal(self, node):
+        if self.filename:
+            filename = self.filename
+        else:
+            filename = "<string>"
+        where = "%s:%s" % (filename, node.lineno)
+        text = ast.dump(node)
+        self.info("Remove dead code at %s: %s" % (where, text))
+
     def warning(self, message):
         self.config.warning_func(message)
 
             return
         if truncate == len(node_list) - 1:
             return
+        for node in node_list[truncate+1:]:
+            self.log_node_removal(node)
         del node_list[truncate+1:]
 
     def optimize_node_list(self, node_list):
 
     def visit_list(self, node_list, conditional=False):
         new_node_list = []
-        for node in node_list:
+        for  node in node_list:
             new_node = self.visit(node, conditional=conditional)
             if new_node is None:
                 continue
 
     def fullvisit_FunctionDef(self, node):
         optimizer = FunctionOptimizer(self.config)
+        optimizer.filename = self.filename
         func_namespace = self.namespace.copy()
         return optimizer._optimize(node, func_namespace)
 
             return self.generic_visit(node)
         else:
             optimizer = Optimizer(self.config)
+            optimizer.filename = self.filename
             class_namespace = self.namespace.copy()
             return optimizer._optimize(node, class_namespace)
 

File astoptimizer/tests.py

     def __init__(self, *args):
         TestCase.__init__(self, *args)
         self.maxDiff = 4096
+        self.infos = []
         self.warnings = []
 
     def setUp(self):
         if self.warnings:
             raise Exception("WARNINGS: %s" % self.warnings)
 
+    def log_info(self, message):
+        self.infos.append(message)
+
     def log_warning(self, message):
         self.warnings.append(message)
 
     def create_config(self, *features):
         config = Config(*features)
+        config.info_func = self.log_info
         config.warning_func = self.log_warning
         return config
 
                 raise
         return tree, optimizer
 
-    def _check_warnings(self, message):
-        emitted = self.warnings[:]
-        del self.warnings[:]
+    def _check_messages(self, messages, message):
+        emitted = messages[:]
+        del messages[:]
         if message:
             if isinstance(message, (list, tuple)):
                 expected = message
         elif emitted:
             raise Exception("WARNINGS: %s" % emitted)
 
-    def check(self, code, expected, config=None, warning=None):
+    def _check_infos(self, message):
+        self._check_messages(self.infos, message)
+
+    def _check_warnings(self, message):
+        self._check_messages(self.warnings, message)
+
+    def check(self, code, expected, config=None, info=None, warning=None):
         tree = parse_ast(code)
         tree, optimizer = self._optimize_ast(tree, config)
         text = ast.dump(tree)
         self.assertEqual(text, expected)
+        self._check_infos(info)
         self._check_warnings(warning)
 
-    def check_not_optimized(self, code, config=None, catch_syntaxerror=False, warning=None):
+    def check_not_optimized(self, code,
+                            config=None, catch_syntaxerror=False,
+                            info=None, warning=None):
         old_tree = parse_ast(code)
         old = ast.dump(old_tree)
         old = re.sub(r"UnaryOp\(op=USub\(\), operand=Num\(n=([0-9]+(?:\.[0-9]*)?)\)\)", lambda regs: "Num(n=-%s)" % regs.group(1), old)
         new_tree, optimizer = self._optimize_ast(old_tree, config, catch_syntaxerror=catch_syntaxerror)
         new = ast.dump(new_tree)
         self.assertEqual(new, old)
+        self._check_infos(info)
         self._check_warnings(warning)
 
-    def check_qualnames(self, code, expected, warning=None):
+    def check_qualnames(self, code, expected,
+                        info=None, warning=None):
         tree = parse_ast(code)
         tree, optimizer = self._optimize_ast(tree)
         self.assertEqual(optimizer.namespace.qualnames, expected)
+        self._check_infos(info)
         self._check_warnings(warning)
 
 
     def test_remove_dead_code(self):
         # return, raise
         self.check('def f():\n return 1\n return 2',
-                   self.text_ast('def f():\n return 1'))
+                   self.text_ast('def f():\n return 1'),
+                   info='Remove dead code at <string>:3: Return(value=Num(n=2))')
         self.check('def f():\n raise\n return 2',
-                   self.text_ast('def f():\n raise'))
+                   self.text_ast('def f():\n raise'),
+                   info='Remove dead code at <string>:3: Return(value=Num(n=2))')
         self.check('def f():\n g()\n return 1\n h()\n return 2',
-                   self.text_ast('def f():\n g()\n return 1'))
+                   self.text_ast('def f():\n g()\n return 1'),
+                   info=["Remove dead code at <string>:4: Expr(value=Call(func=Name(id='h', ctx=Load()), args=[], keywords=[], starargs=None, kwargs=None))",
+                        "Remove dead code at <string>:5: Return(value=Num(n=2))"])
         self.check('def f():\n g()\n raise ValueError("error")\n h()\n return 2',
-                   self.text_ast('def f():\n g()\n raise ValueError("error")'))
+                   self.text_ast('def f():\n g()\n raise ValueError("error")'),
+                   info=["Remove dead code at <string>:4: Expr(value=Call(func=Name(id='h', ctx=Load()), args=[], keywords=[], starargs=None, kwargs=None))",
+                         "Remove dead code at <string>:5: Return(value=Num(n=2))"])
 
         # global, nonlocal
         self.check_not_optimized('''
         parts = line.split('``')
         before = parts[0]
         after = parts[2]
+        after2 = None
+        info = None
 
         if (sys.version_info < (2, 7)
         and ('.bit_length()' in before
             before = 'if a:\n if b:\n  print("true")'
         if after == 'if a and b: print("true")':
             after = 'if a and b:\n print("true")'
+        if 'return 1; return 2' in before:
+            info = 'Remove dead code at <string>:1: Return(value=Num(n=2))'
 
         if 'if DEBUG:' in before:
             config = self.create_default_config()
             config = None
         if len(parts) >= 5 and ' or ' in parts[3]:
             after2 = parts[4]
-        else:
-            after2 = None
         try:
-            self.check(before, self.text_ast(after), config)
+            self.check(before, self.text_ast(after), config, info=info)
         except (AssertionError, SyntaxError):
             if after2 is not None:
                 self.check(before, self.text_ast(after2), config)
 
 def main():
     parser = OptionParser(usage="%prog [options] code [line2 line3 ...]")
+    parser.add_option("-f", "--file", metavar="SCRIPT",
+        help="execute specified script",
+        type="str", default=None)
     parser.add_option("-k",
         help="Don't remove dead code",
         action="store_true", default=False)
     parser.add_option("-n",
         help='Disable all configuration features',
         action="store_true", default=False)
-    parser.add_option("-f", "--features",
+    parser.add_option("-e", "--enable", metavar="FEATURES",
         help='Enable configuration features FEATURES (default: builtin_funcs and pythonbin)',
         action="store", type="str", default="builtin_funcs,pythonbin")
     parser.add_option("-a", "--all",
         help="Maximum length of frozenset, set and tuple",
         action="store_true", default=False)
     options, args = parser.parse_args()
-    if not args:
+    if not options.file and not args:
         parser.print_help()
         exit(1)
 
-    if args == ['-']:
-        code_str = sys.stdin.read()
+    if options.file:
+        filename = options.file
+        with open(filename, 'r') as fp:
+            code_str = fp.read()
     else:
+        filename = None
         code_str = '\n'.join(args)
 
     config = Config()
     if not options.n:
-        for feature in options.features.split(','):
+        for feature in options.enable.split(','):
             config.enable(feature) #'builtin_funcs', 'pythonenv')
     if options.x:
         config.use_experimental_vars = True
         dump_bytecode(code)
     print("")
 
-    tree = optimize_ast(tree, config)
+    tree = optimize_ast(tree, config, filename)
     print("Optimised AST:")
     print(ast.dump(tree))
     print("")