Port py_compile over to importlib

2 files changed, 31 insertions, 46 deletions

diff --git a/Lib/importlib/_bootstrap.py b/Lib/importlib/_bootstrap.py
index 344ef84..fa3a243 100644
--- a/Lib/importlib/_bootstrap.py
+++ b/Lib/importlib/_bootstrap.py
@@ -689,6 +689,15 @@ def _compile_bytecode(data, name=None, bytecode_path=None, source_path=None):
         raise ImportError("Non-code object in {!r}".format(bytecode_path),
                           name=name, path=bytecode_path)
 
+def _code_to_bytecode(code, mtime=0, source_size=0):
+    """Compile a code object into bytecode for writing out to a byte-compiled
+    file."""
+    data = bytearray(_MAGIC_BYTES)
+    data.extend(_w_long(mtime))
+    data.extend(_w_long(source_size))
+    data.extend(marshal.dumps(code))
+    return data
+
 
 # Loaders #####################################################################
 
@@ -951,13 +960,13 @@ class SourceLoader(_LoaderBasics):
             raise ImportError("Failed to decode source file",
                               name=fullname) from exc
 
-    def source_to_code(self, data, path):
+    def source_to_code(self, data, path, *, _optimize=-1):
         """Return the code object compiled from source.
 
         The 'data' argument can be any object type that compile() supports.
         """
         return _call_with_frames_removed(compile, data, path, 'exec',
-                                        dont_inherit=True)
+                                        dont_inherit=True, optimize=_optimize)
 
     def get_code(self, fullname):
         """Concrete implementation of InspectLoader.get_code.
@@ -1000,11 +1009,9 @@ class SourceLoader(_LoaderBasics):
         code_object = self.source_to_code(source_bytes, source_path)
         _verbose_message('code object from {}', source_path)
         if (not sys.dont_write_bytecode and bytecode_path is not None and
-            source_mtime is not None):
-            data = bytearray(_MAGIC_BYTES)
-            data.extend(_w_long(source_mtime))
-            data.extend(_w_long(len(source_bytes)))
-            data.extend(marshal.dumps(code_object))
+                source_mtime is not None):
+            data = _code_to_bytecode(code_object, source_mtime,
+                    len(source_bytes))
             try:
                 self._cache_bytecode(source_path, bytecode_path, data)
                 _verbose_message('wrote {!r}', bytecode_path)
diff --git a/Lib/py_compile.py b/Lib/py_compile.py
index 3cb46d0..d255a2d 100644
--- a/Lib/py_compile.py
+++ b/Lib/py_compile.py
@@ -3,17 +3,13 @@
 This module has intimate knowledge of the format of .pyc files.
 """
 
-import builtins
-import errno
 import imp
-import marshal
+import importlib._bootstrap
+import importlib.machinery
 import os
 import sys
-import tokenize
 import traceback
 
-MAGIC = imp.get_magic()
-
 __all__ = ["compile", "main", "PyCompileError"]
 
 
@@ -65,13 +61,6 @@ class PyCompileError(Exception):
         return self.msg
 
 
-def wr_long(f, x):
-    """Internal; write a 32-bit int to a file in little-endian order."""
-    f.write(bytes([x         & 0xff,
-                   (x >> 8)  & 0xff,
-                   (x >> 16) & 0xff,
-                   (x >> 24) & 0xff]))
-
 def compile(file, cfile=None, dfile=None, doraise=False, optimize=-1):
     """Byte-compile one Python source file to Python bytecode.
 
@@ -108,17 +97,16 @@ def compile(file, cfile=None, dfile=None, doraise=False, optimize=-1):
     byte-compile all installed files (or all files in selected
     directories).
     """
-    with tokenize.open(file) as f:
-        try:
-            st = os.fstat(f.fileno())
-        except AttributeError:
-            st = os.stat(file)
-        timestamp = int(st.st_mtime)
-        size = st.st_size & 0xFFFFFFFF
-        codestring = f.read()
+    if cfile is None:
+        if optimize >= 0:
+            cfile = imp.cache_from_source(file, debug_override=not optimize)
+        else:
+            cfile = imp.cache_from_source(file)
+    loader = importlib.machinery.SourceFileLoader('<py_compile>', file)
+    source_bytes = loader.get_data(file)
     try:
-        codeobject = builtins.compile(codestring, dfile or file, 'exec',
-                                      optimize=optimize)
+        code = loader.source_to_code(source_bytes, dfile or file,
+                _optimize=optimize)
     except Exception as err:
         py_exc = PyCompileError(err.__class__, err, dfile or file)
         if doraise:
@@ -126,26 +114,16 @@ def compile(file, cfile=None, dfile=None, doraise=False, optimize=-1):
         else:
             sys.stderr.write(py_exc.msg + '\n')
             return
-    if cfile is None:
-        if optimize >= 0:
-            cfile = imp.cache_from_source(file, debug_override=not optimize)
-        else:
-            cfile = imp.cache_from_source(file)
     try:
         dirname = os.path.dirname(cfile)
         if dirname:
             os.makedirs(dirname)
-    except OSError as error:
-        if error.errno != errno.EEXIST:
-            raise
-    with open(cfile, 'wb') as fc:
-        fc.write(b'\0\0\0\0')
-        wr_long(fc, timestamp)
-        wr_long(fc, size)
-        marshal.dump(codeobject, fc)
-        fc.flush()
-        fc.seek(0, 0)
-        fc.write(MAGIC)
+    except FileExistsError:
+        pass
+    source_stats = loader.path_stats(file)
+    bytecode = importlib._bootstrap._code_to_bytecode(code,
+        source_stats['mtime'], len(source_bytes))
+    loader._cache_bytecode(file, cfile, bytecode)
     return cfile
 
 def main(args=None):