[pypy-commit] pypy py3.3: Implement docstring stripping for compile(..., optimize=2).

[kvas]

Author: Vasily Kuznetsov <kvas.it at gmail.com>
Branch: py3.3
Changeset: r72561:17f769cdad0a
Date: 2014-07-27 13:27 +0200
http://bitbucket.org/pypy/pypy/changeset/17f769cdad0a/

Log:	Implement docstring stripping for compile(..., optimize=2).

diff --git a/pypy/interpreter/astcompiler/codegen.py b/pypy/interpreter/astcompiler/codegen.py
--- a/pypy/interpreter/astcompiler/codegen.py
+++ b/pypy/interpreter/astcompiler/codegen.py
@@ -244,7 +244,7 @@
         self.emit_op_arg(op, self.add_name(container, identifier))
 
     def possible_docstring(self, node):
-        if isinstance(node, ast.Expr):
+        if isinstance(node, ast.Expr) and self.compile_info.optimize < 2:
             expr_value = node.value
             if isinstance(expr_value, ast.Str):
                 return expr_value
diff --git a/pypy/module/__builtin__/test/test_compile.py b/pypy/module/__builtin__/test/test_compile.py
--- a/pypy/module/__builtin__/test/test_compile.py
+++ b/pypy/module/__builtin__/test/test_compile.py
@@ -1,8 +1,7 @@
 class AppTestCompile:
 
-    # TODO: This test still fails for now because the docstrings are not
-    #       removed with optimize=2.
-    def untest_compile(self):
+    def test_compile(self):
+        """Clone of the part of the original test that was failing."""
         import ast
 
         codestr = '''def f():
@@ -37,7 +36,7 @@
                 assert rv == (debugval, docstring)
 
     def test_assert_remove(self):
-        """Test just removal of the asserts with optimize=1."""
+        """Test removal of the asserts with optimize=1."""
         import ast
 
         code = """def f():
@@ -50,9 +49,41 @@
             exec(compiled, ns)
             ns['f']()
 
+    def test_docstring_remove(self):
+        """Test removal of docstrings with optimize=2."""
+        import ast
+        import marshal
 
-# TODO: Remove docstrings with optimize=2.
+        code = """
+'module_doc'
+
+def f():
+    'func_doc'
+
+class C:
+    'class_doc'
+"""
+        tree = ast.parse(code)
+        for to_compile in [code, tree]:
+            compiled = compile(to_compile, "<test>", "exec", optimize=2)
+
+            # check that the docstrings are really gone
+            marshalled = str(marshal.dumps(compiled))
+            assert 'module_doc' not in marshalled
+            assert 'func_doc' not in marshalled
+            assert 'class_doc' not in marshalled
+
+            # try to execute the bytecode and see what we get
+            ns = {}
+            exec(compiled, ns)
+            assert '__doc__' not in ns
+            assert ns['f'].__doc__ is None
+            assert ns['C'].__doc__ is None
+
+
 # TODO: Check the value of __debug__ inside of the compiled block!
 #       According to the documentation, it should follow the optimize flag.
+#       However, cpython3.3 behaves the same way as PyPy (__debug__ follows
+#       -O, -OO flags of the interpreter).
 # TODO: It would also be good to test that with the assert is not removed and
 #       is executed when -O flag is set but optimize=0.