Add __debug__ load support

DinoV · DinoV · commit 977a6993aa12 · 2021-01-28T11:14:38.000-08:00
Summary: Loading __debug__ turns into an ast Constant node in the AST optimizer

Test Plan: ./python -m test.test_compiler
diff --git a/compiler/optimizer.py b/compiler/optimizer.py
@@ -60,6 +60,10 @@ class Py37Limits:
 
 
 class AstOptimizer(ASTRewriter):
+    def __init__(self, optimize = False):
+        super().__init__()
+        self.optimize = optimize
+
     def visitUnaryOp(self, node: ast.UnaryOp) -> ast.expr:
         op = self.visit(node.operand)
         if is_const(op):
@@ -180,3 +184,9 @@ def visitCompare(self, node: ast.Compare) -> ast.expr:
                 comparators[-1] = new_iter
 
         return self.update_node(node, left=left, comparators=comparators)
+
+    def visitName(self, node: ast.Name):
+        if node.id == "__debug__":
+            return copy_location(Constant(not self.optimize), node)
+
+        return self.generic_visit(node)
diff --git a/test_compiler/common.py b/test_compiler/common.py
@@ -58,6 +58,10 @@ def dump_graph(self, graph):
         graph.dump(io)
         return io.getvalue()
 
+    def graph_to_instrs(self, graph):
+        for block in graph.getBlocks():
+            yield from block.getInstructions()
+
     def assertNotInGraph(self, graph, opname, argval=_UNSPECIFIED):
         for block in graph.getBlocks():
             for instr in block.getInstructions():
diff --git a/test_compiler/test_py37.py b/test_compiler/test_py37.py
@@ -154,6 +154,30 @@ def test_compile_opt_enabled(self):
         graph = self.to_graph('x = -1', CodeGenerator)
         self.assertInGraph(graph, 'UNARY_NEGATIVE')
 
+    def test_opt_debug(self):
+        graph = self.to_graph('if not __debug__:\n    x = 42', Python37CodeGenerator)
+        self.assertNotInGraph(graph, 'STORE_NAME')
+
+        graph = self.to_graph('if not __debug__:\n    x = 42', CodeGenerator)
+        self.assertInGraph(graph, 'STORE_NAME')
+
+    def test_opt_debug_del(self):
+        code = 'def f(): del __debug__'
+        outer_graph = self.to_graph(code, Python37CodeGenerator)
+        for outer_instr in self.graph_to_instrs(outer_graph):
+            if outer_instr.opname == "LOAD_CONST" and isinstance(outer_instr.oparg, CodeGenerator):
+                graph = outer_instr.oparg.graph
+                self.assertInGraph(graph, 'LOAD_CONST', True)
+                self.assertNotInGraph(graph, 'DELETE_FAST', '__debug__')
+
+        outer_graph = self.to_graph(code, CodeGenerator)
+        for outer_instr in self.graph_to_instrs(outer_graph):
+            if outer_instr.opname == "LOAD_CONST" and isinstance(outer_instr.oparg, CodeGenerator):
+                graph = outer_instr.oparg.graph
+                self.assertNotInGraph(graph, 'LOAD_CONST', True)
+                self.assertInGraph(graph, 'DELETE_FAST', '__debug__')
+
+
     def test_ast_optimizer(self):
         cases = [
             ("+1", "1"),
