Commits

Victor Stinner committed fcfe360

Simplify Optimizer class hierarchy

Comments (0)

Files changed (2)

astoptimizer/__init__.py

 
 from astoptimizer.compatibility import ast
 from astoptimizer.config import Config
-from astoptimizer.optimizer import Optimizer, FunctionOptimizer
+from astoptimizer.optimizer import Optimizer
 
 def optimize_ast(tree, config=None):
     """
     Optimize an AST with the specified config (or with default config).
     Return the optimized AST
     """
-    if isinstance(tree, ast.FunctionDef):
-        optimizer = FunctionOptimizer(config)
-    else:
-        optimizer = Optimizer(config)
+    optimizer = Optimizer(config)
     return optimizer.optimize(tree)
 
 if sys.version_info >= (2,6):

astoptimizer/optimizer.py

     and isinstance(node_list[0], ast.Pass)):
         del node_list[0]
 
-class BaseOptimizer(ast.NodeTransformer):
+class Optimizer(ast.NodeTransformer):
     def __init__(self, config):
         ast.NodeTransformer.__init__(self)
         self.config = config
         self.is_conditional = False
         return self.visit(tree)
 
+    def optimize(self, tree):
+        namespace = Namespace()
+        return self._optimize(tree, namespace)
+
     def warning(self, message):
         self.config.warning_func(message)
 
                 self.unassign(arg)
 
     def fullvisit_FunctionDef(self, node):
-        if node is self.root:
-            return self.generic_visit(node)
-        else:
-            optimizer = FunctionOptimizer(self.config)
-            func_namespace = self.namespace.copy()
-            return optimizer.optimize(node, func_namespace)
+        optimizer = FunctionOptimizer(self.config)
+        func_namespace = self.namespace.copy()
+        return optimizer._optimize(node, func_namespace)
 
     def fullvisit_ClassDef(self, node):
         if node is self.root:
             return optimizer._optimize(node, class_namespace)
 
 
-class FunctionOptimizer(BaseOptimizer):
+class FunctionOptimizer(Optimizer):
     def __init__(self, config):
-        BaseOptimizer.__init__(self, config)
+        Optimizer.__init__(self, config)
         self.is_generator = None
         self.seen_yield = None
 
+    def fullvisit_FunctionDef(self, node):
+        if node is self.root:
+            return self.generic_visit(node)
+        else:
+            return Optimizer.fullvisit_FunctionDef(self, node)
+
     def can_drop(self, node_list):
         if not self.is_generator:
             return True
         seen = self.seen_yield
         try:
             self.seen_yield = None
-            BaseOptimizer.remove_dead_code(self, node_list)
+            Optimizer.remove_dead_code(self, node_list)
         finally:
             self.seen_yield = seen
 
     def visit_Yield(self, node):
         self.seen_yield = True
 
-    def optimize(self, tree, namespace):
+    def _optimize(self, tree, namespace):
         self.is_generator = ast_contains(tree, ast.Yield)
         self.seen_yield = False
-        tree = self._optimize(tree, namespace)
-        return tree
+        return Optimizer._optimize(self, tree, namespace)
 
-
-class Optimizer(BaseOptimizer):
-    def optimize(self, tree):
-        assert not isinstance(tree, ast.FunctionDef)
-        namespace = Namespace()
-        return self._optimize(tree, namespace)
-