Add an "optimize" parameter to compile() to control the optimization level, and provide an interface to it in py_compile, compileall and PyZipFile.

3 files changed, 54 insertions, 0 deletions

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index 7b73949..1469e36 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -6,6 +6,7 @@ import sys
 import warnings
 import collections
 import io
+import ast
 import types
 import builtins
 import random
@@ -285,6 +286,34 @@ class BuiltinTest(unittest.TestCase):
         self.assertRaises(TypeError, compile, chr(0), 'f', 'exec')
         self.assertRaises(ValueError, compile, str('a = 1'), 'f', 'bad')
 
+        # test the optimize argument
+
+        codestr = '''def f():
+        """doc"""
+        try:
+            assert False
+        except AssertionError:
+            return (True, f.__doc__)
+        else:
+            return (False, f.__doc__)
+        '''
+        def f(): """doc"""
+        values = [(-1, __debug__, f.__doc__),
+                  (0, True, 'doc'),
+                  (1, False, 'doc'),
+                  (2, False, None)]
+        for optval, debugval, docstring in values:
+            # test both direct compilation and compilation via AST
+            codeobjs = []
+            codeobjs.append(compile(codestr, "<test>", "exec", optimize=optval))
+            tree = ast.parse(codestr)
+            codeobjs.append(compile(tree, "<test>", "exec", optimize=optval))
+            for code in codeobjs:
+                ns = {}
+                exec(code, ns)
+                rv = ns['f']()
+                self.assertEqual(rv, (debugval, docstring))
+
     def test_delattr(self):
         sys.spam = 1
         delattr(sys, 'spam')
diff --git a/Lib/test/test_compileall.py b/Lib/test/test_compileall.py
index 1955006..4246b2f 100644
--- a/Lib/test/test_compileall.py
+++ b/Lib/test/test_compileall.py
@@ -88,6 +88,15 @@ class CompileallTests(unittest.TestCase):
         compileall.compile_file(data_file)
         self.assertFalse(os.path.exists(os.path.join(data_dir, '__pycache__')))
 
+    def test_optimize(self):
+        # make sure compiling with different optimization settings than the
+        # interpreter's creates the correct file names
+        optimize = 1 if __debug__ else 0
+        compileall.compile_dir(self.directory, quiet=True, optimize=optimize)
+        cached = imp.cache_from_source(self.source_path,
+                                       debug_override=not optimize)
+        self.assertTrue(os.path.isfile(cached))
+
 
 class EncodingTest(unittest.TestCase):
     """Issue 6716: compileall should escape source code when printing errors
diff --git a/Lib/test/test_zipfile.py b/Lib/test/test_zipfile.py
index 7f93b68..a0367e1 100644
--- a/Lib/test/test_zipfile.py
+++ b/Lib/test/test_zipfile.py
@@ -654,6 +654,22 @@ class PyZipFileTests(unittest.TestCase):
             self.assertTrue('email/mime/text.pyo' in names or
                             'email/mime/text.pyc' in names)
 
+    def test_write_with_optimization(self):
+        import email
+        packagedir = os.path.dirname(email.__file__)
+        # use .pyc if running test in optimization mode,
+        # use .pyo if running test in debug mode
+        optlevel = 1 if __debug__ else 0
+        ext = '.pyo' if optlevel == 1 else '.pyc'
+
+        with TemporaryFile() as t, \
+                 zipfile.PyZipFile(t, "w", optimize=optlevel) as zipfp:
+            zipfp.writepy(packagedir)
+
+            names = zipfp.namelist()
+            self.assertIn('email/__init__' + ext, names)
+            self.assertIn('email/mime/text' + ext, names)
+
     def test_write_python_directory(self):
         os.mkdir(TESTFN2)
         try: