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

6 files changed, 138 insertions, 42 deletions

diff --git a/Lib/compileall.py b/Lib/compileall.py
index 17cc61d..aefdb89 100644
--- a/Lib/compileall.py
+++ b/Lib/compileall.py
@@ -19,8 +19,8 @@ import struct
 
 __all__ = ["compile_dir","compile_file","compile_path"]
 
-def compile_dir(dir, maxlevels=10, ddir=None,
-                force=False, rx=None, quiet=False, legacy=False):
+def compile_dir(dir, maxlevels=10, ddir=None, force=False, rx=None,
+                quiet=False, legacy=False, optimize=-1):
     """Byte-compile all modules in the given directory tree.
 
     Arguments (only dir is required):
@@ -32,6 +32,7 @@ def compile_dir(dir, maxlevels=10, ddir=None,
     force:     if True, force compilation, even if timestamps are up-to-date
     quiet:     if True, be quiet during compilation
     legacy:    if True, produce legacy pyc paths instead of PEP 3147 paths
+    optimize:  optimization level or -1 for level of the interpreter
     """
     if not quiet:
         print('Listing', dir, '...')
@@ -51,7 +52,8 @@ def compile_dir(dir, maxlevels=10, ddir=None,
         else:
             dfile = None
         if not os.path.isdir(fullname):
-            if not compile_file(fullname, ddir, force, rx, quiet, legacy):
+            if not compile_file(fullname, ddir, force, rx, quiet,
+                                legacy, optimize):
                 success = 0
         elif (maxlevels > 0 and name != os.curdir and name != os.pardir and
               os.path.isdir(fullname) and not os.path.islink(fullname)):
@@ -61,7 +63,7 @@ def compile_dir(dir, maxlevels=10, ddir=None,
     return success
 
 def compile_file(fullname, ddir=None, force=0, rx=None, quiet=False,
-                 legacy=False):
+                 legacy=False, optimize=-1):
     """Byte-compile file.
     fullname:  the file to byte-compile
     ddir:      if given, purported directory name (this is the
@@ -69,6 +71,7 @@ def compile_file(fullname, ddir=None, force=0, rx=None, quiet=False,
     force:     if True, force compilation, even if timestamps are up-to-date
     quiet:     if True, be quiet during compilation
     legacy:    if True, produce legacy pyc paths instead of PEP 3147 paths
+    optimize:  optimization level or -1 for level of the interpreter
     """
     success = 1
     name = os.path.basename(fullname)
@@ -84,7 +87,11 @@ def compile_file(fullname, ddir=None, force=0, rx=None, quiet=False,
         if legacy:
             cfile = fullname + ('c' if __debug__ else 'o')
         else:
-            cfile = imp.cache_from_source(fullname)
+            if optimize >= 0:
+                cfile = imp.cache_from_source(fullname,
+                                              debug_override=not optimize)
+            else:
+                cfile = imp.cache_from_source(fullname)
             cache_dir = os.path.dirname(cfile)
         head, tail = name[:-3], name[-3:]
         if tail == '.py':
@@ -101,7 +108,8 @@ def compile_file(fullname, ddir=None, force=0, rx=None, quiet=False,
             if not quiet:
                 print('Compiling', fullname, '...')
             try:
-                ok = py_compile.compile(fullname, cfile, dfile, True)
+                ok = py_compile.compile(fullname, cfile, dfile, True,
+                                        optimize=optimize)
             except py_compile.PyCompileError as err:
                 if quiet:
                     print('*** Error compiling', fullname, '...')
@@ -126,7 +134,7 @@ def compile_file(fullname, ddir=None, force=0, rx=None, quiet=False,
     return success
 
 def compile_path(skip_curdir=1, maxlevels=0, force=False, quiet=False,
-                 legacy=False):
+                 legacy=False, optimize=-1):
     """Byte-compile all module on sys.path.
 
     Arguments (all optional):
@@ -136,6 +144,7 @@ def compile_path(skip_curdir=1, maxlevels=0, force=False, quiet=False,
     force: as for compile_dir() (default False)
     quiet: as for compile_dir() (default False)
     legacy: as for compile_dir() (default False)
+    optimize: as for compile_dir() (default -1)
     """
     success = 1
     for dir in sys.path:
@@ -144,7 +153,7 @@ def compile_path(skip_curdir=1, maxlevels=0, force=False, quiet=False,
         else:
             success = success and compile_dir(dir, maxlevels, None,
                                               force, quiet=quiet,
-                                              legacy=legacy)
+                                              legacy=legacy, optimize=optimize)
     return success
 
 
diff --git a/Lib/py_compile.py b/Lib/py_compile.py
index d241434..e0f98cb 100644
--- a/Lib/py_compile.py
+++ b/Lib/py_compile.py
@@ -72,7 +72,7 @@ def wr_long(f, x):
                    (x >> 16) & 0xff,
                    (x >> 24) & 0xff]))
 
-def compile(file, cfile=None, dfile=None, doraise=False):
+def compile(file, cfile=None, dfile=None, doraise=False, optimize=-1):
     """Byte-compile one Python source file to Python bytecode.
 
     :param file: The source file name.
@@ -86,6 +86,10 @@ def compile(file, cfile=None, dfile=None, doraise=False):
         will be printed, and the function will return to the caller. If an
         exception occurs and this flag is set to True, a PyCompileError
         exception will be raised.
+    :param optimize: The optimization level for the compiler.  Valid values
+        are -1, 0, 1 and 2.  A value of -1 means to use the optimization
+        level of the current interpreter, as given by -O command line options.
+
     :return: Path to the resulting byte compiled file.
 
     Note that it isn't necessary to byte-compile Python modules for
@@ -111,7 +115,8 @@ def compile(file, cfile=None, dfile=None, doraise=False):
             timestamp = int(os.stat(file).st_mtime)
         codestring = f.read()
     try:
-        codeobject = builtins.compile(codestring, dfile or file,'exec')
+        codeobject = builtins.compile(codestring, dfile or file, 'exec',
+                                      optimize=optimize)
     except Exception as err:
         py_exc = PyCompileError(err.__class__, err, dfile or file)
         if doraise:
@@ -120,7 +125,10 @@ def compile(file, cfile=None, dfile=None, doraise=False):
             sys.stderr.write(py_exc.msg + '\n')
             return
     if cfile is None:
-        cfile = imp.cache_from_source(file)
+        if optimize >= 0:
+            cfile = imp.cache_from_source(file, debug_override=not optimize)
+        else:
+            cfile = imp.cache_from_source(file)
     try:
         os.makedirs(os.path.dirname(cfile))
     except OSError as error:
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:
diff --git a/Lib/zipfile.py b/Lib/zipfile.py
index bfe41b7..35bba73 100644
--- a/Lib/zipfile.py
+++ b/Lib/zipfile.py
@@ -1295,6 +1295,12 @@ class ZipFile:
 class PyZipFile(ZipFile):
     """Class to create ZIP archives with Python library files and packages."""
 
+    def __init__(self, file, mode="r", compression=ZIP_STORED,
+                 allowZip64=False, optimize=-1):
+        ZipFile.__init__(self, file, mode=mode, compression=compression,
+                         allowZip64=allowZip64)
+        self._optimize = optimize
+
     def writepy(self, pathname, basename=""):
         """Add all files from "pathname" to the ZIP archive.
 
@@ -1367,44 +1373,63 @@ class PyZipFile(ZipFile):
         archive name, compiling if necessary.  For example, given
         /python/lib/string, return (/python/lib/string.pyc, string).
         """
+        def _compile(file, optimize=-1):
+            import py_compile
+            if self.debug:
+                print("Compiling", file)
+            try:
+                py_compile.compile(file, doraise=True, optimize=optimize)
+            except py_compile.PyCompileError as error:
+                print(err.msg)
+                return False
+            return True
+
         file_py  = pathname + ".py"
         file_pyc = pathname + ".pyc"
         file_pyo = pathname + ".pyo"
         pycache_pyc = imp.cache_from_source(file_py, True)
         pycache_pyo = imp.cache_from_source(file_py, False)
-        if (os.path.isfile(file_pyo) and
-            os.stat(file_pyo).st_mtime >= os.stat(file_py).st_mtime):
-            # Use .pyo file.
-            arcname = fname = file_pyo
-        elif (os.path.isfile(file_pyc) and
-              os.stat(file_pyc).st_mtime >= os.stat(file_py).st_mtime):
-            # Use .pyc file.
-            arcname = fname = file_pyc
-        elif (os.path.isfile(pycache_pyc) and
-              os.stat(pycache_pyc).st_mtime >= os.stat(file_py).st_mtime):
-            # Use the __pycache__/*.pyc file, but write it to the legacy pyc
-            # file name in the archive.
-            fname = pycache_pyc
-            arcname = file_pyc
-        elif (os.path.isfile(pycache_pyo) and
-              os.stat(pycache_pyo).st_mtime >= os.stat(file_py).st_mtime):
-            # Use the __pycache__/*.pyo file, but write it to the legacy pyo
-            # file name in the archive.
-            fname = pycache_pyo
-            arcname = file_pyo
+        if self._optimize == -1:
+            # legacy mode: use whatever file is present
+            if (os.path.isfile(file_pyo) and
+                os.stat(file_pyo).st_mtime >= os.stat(file_py).st_mtime):
+                # Use .pyo file.
+                arcname = fname = file_pyo
+            elif (os.path.isfile(file_pyc) and
+                  os.stat(file_pyc).st_mtime >= os.stat(file_py).st_mtime):
+                # Use .pyc file.
+                arcname = fname = file_pyc
+            elif (os.path.isfile(pycache_pyc) and
+                  os.stat(pycache_pyc).st_mtime >= os.stat(file_py).st_mtime):
+                # Use the __pycache__/*.pyc file, but write it to the legacy pyc
+                # file name in the archive.
+                fname = pycache_pyc
+                arcname = file_pyc
+            elif (os.path.isfile(pycache_pyo) and
+                  os.stat(pycache_pyo).st_mtime >= os.stat(file_py).st_mtime):
+                # Use the __pycache__/*.pyo file, but write it to the legacy pyo
+                # file name in the archive.
+                fname = pycache_pyo
+                arcname = file_pyo
+            else:
+                # Compile py into PEP 3147 pyc file.
+                if _compile(file_py):
+                    fname = (pycache_pyc if __debug__ else pycache_pyo)
+                    arcname = (file_pyc if __debug__ else file_pyo)
+                else:
+                    fname = arcname = file_py
         else:
-            # Compile py into PEP 3147 pyc file.
-            import py_compile
-            if self.debug:
-                print("Compiling", file_py)
-            try:
-                py_compile.compile(file_py, doraise=True)
-            except py_compile.PyCompileError as error:
-                print(err.msg)
-                fname = file_py
+            # new mode: use given optimization level
+            if self._optimize == 0:
+                fname = pycache_pyc
+                arcname = file_pyc
             else:
-                fname = (pycache_pyc if __debug__ else pycache_pyo)
-                arcname = (file_pyc if __debug__ else file_pyo)
+                fname = pycache_pyo
+                arcname = file_pyo
+            if not (os.path.isfile(fname) and
+                    os.stat(fname).st_mtime >= os.stat(file_py).st_mtime):
+                if not _compile(file_py, optimize=self._optimize):
+                    fname = arcname = file_py
         archivename = os.path.split(arcname)[1]
         if basename:
             archivename = "%s/%s" % (basename, archivename)