python · vstinner · Jun 23, 2017
diff --git a/Doc/library/sys.rst b/Doc/library/sys.rst
@@ -413,6 +413,13 @@ always available.
    .. versionadded:: 3.7
 
 
+.. function:: get_code_transformers()
+
+   Return the list of code transformers.; see :func:`set_code_transformers`.
+
+   .. versionadded:: 3.6
+
+
 .. function:: getcheckinterval()
 
    Return the interpreter's "check interval"; see :func:`setcheckinterval`.
@@ -719,6 +726,9 @@ always available.
    ``cache_tag`` is set to ``None``, it indicates that module caching should
    be disabled.
 
+   *optim_tag* is the optimized tag used by the import machinery in the
+   filenames of cached modules.
+
    :data:`sys.implementation` may contain additional attributes specific to
    the Python implementation.  These non-standard attributes must start with
    an underscore, and are not described here.  Regardless of its contents,
@@ -728,6 +738,9 @@ always available.
 
    .. versionadded:: 3.3
 
+   .. versionchanged:: 3.6
+      Added *optim_tag* attribute.
+
 
 .. data:: int_info
 
@@ -978,6 +991,13 @@ always available.
    implement a dynamic prompt.
 
 
+.. function:: set_code_transformers(transformers)
+
+   Set code transformers: *transformers* is a sequence of code transformers.
+
+   .. versionadded:: 3.6
+
+
 .. function:: setcheckinterval(interval)
 
    Set the interpreter's "check interval".  This integer value determines how often

diff --git a/Include/code.h b/Include/code.h
@@ -149,6 +149,8 @@ PyAPI_FUNC(int) _PyCode_SetExtra(PyObject *code, Py_ssize_t index,
                                  void *extra);
 #endif
 
+PyAPI_DATA(PyTypeObject) _PyPeepholeOptimizer_Type;
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/Include/compile.h b/Include/compile.h
@@ -22,6 +22,7 @@ PyAPI_FUNC(PyCodeObject *) PyNode_Compile(struct _node *, const char *);
 #define PyCF_DONT_IMPLY_DEDENT 0x0200
 #define PyCF_ONLY_AST 0x0400
 #define PyCF_IGNORE_COOKIE 0x0800
+#define PyCF_TRANSFORMED_AST 0x1000
 
 #ifndef Py_LIMITED_API
 typedef struct {

diff --git a/Include/pydebug.h b/Include/pydebug.h
@@ -23,6 +23,7 @@ PyAPI_DATA(int) Py_NoUserSiteDirectory;
 PyAPI_DATA(int) Py_UnbufferedStdioFlag;
 PyAPI_DATA(int) Py_HashRandomizationFlag;
 PyAPI_DATA(int) Py_IsolatedFlag;
+PyAPI_DATA(PyObject*) _Py_OptimTag;
 
 #ifdef MS_WINDOWS
 PyAPI_DATA(int) Py_LegacyWindowsStdioFlag;

diff --git a/Include/pystate.h b/Include/pystate.h
@@ -52,11 +52,11 @@ typedef struct _is {
 
     int64_t id;
 
-    PyObject *modules;
+    PyObject *sysdict;   /* sys.__dict__ */
+    PyObject *modules;   /* sys.modules list */
     PyObject *modules_by_index;
-    PyObject *sysdict;
-    PyObject *builtins;
-    PyObject *importlib;
+    PyObject *builtins;  /* dictionary of the builtins module */
+    PyObject *importlib; /* importlib module */
 
     PyObject *codec_search_path;
     PyObject *codec_search_cache;

diff --git a/Include/sysmodule.h b/Include/sysmodule.h
@@ -37,6 +37,26 @@ PyAPI_FUNC(PyObject *) PySys_GetXOptions(void);
 PyAPI_FUNC(size_t) _PySys_GetSizeOf(PyObject *);
 #endif
 
+#ifndef Py_LIMITED_API
+typedef struct {
+    /* The code transformer object */
+    PyObject *transformer;
+
+    /* name of the transformer */
+    PyObject *name;
+
+    /* ast_transformer() method of transformer, or NULL */
+    PyObject *ast_transformer;
+
+    /* code_transformer() method of transformer, or NULL */
+    PyObject *code_transformer;
+} _PySys_CodeTransformer;
+
+PyAPI_FUNC(int) _PySys_GetCodeTransformers(
+    _PySys_CodeTransformer **transformers,
+    Py_ssize_t *ntransformer);
+#endif
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/Lib/distutils/tests/test_bdist_dumb.py b/Lib/distutils/tests/test_bdist_dumb.py
@@ -87,7 +87,9 @@ def test_simple_built(self):
         contents = sorted(os.path.basename(fn) for fn in contents)
         wanted = ['foo-0.1-py%s.%s.egg-info' % sys.version_info[:2], 'foo.py']
         if not sys.dont_write_bytecode:
-            wanted.append('foo.%s.pyc' % sys.implementation.cache_tag)
+            fn = 'foo.%s.%s-0.pyc' % (sys.implementation.cache_tag,
+                                      sys.implementation.optim_tag)
+            wanted.append(fn)
         self.assertEqual(contents, sorted(wanted))
 
 def test_suite():

diff --git a/Lib/distutils/tests/test_build_py.py b/Lib/distutils/tests/test_build_py.py
@@ -62,8 +62,9 @@ def test_package_data(self):
             self.assertFalse(os.path.exists(pycache_dir))
         else:
             pyc_files = os.listdir(pycache_dir)
-            self.assertIn("__init__.%s.pyc" % sys.implementation.cache_tag,
-                          pyc_files)
+            fn = '__init__.%s.%s-0.pyc' % (sys.implementation.cache_tag,
+                                           sys.implementation.optim_tag)
+            self.assertIn(fn, pyc_files)
 
     def test_empty_package_dir(self):
         # See bugs #1668596/#1720897
@@ -102,8 +103,10 @@ def test_byte_compile(self):
         found = os.listdir(cmd.build_lib)
         self.assertEqual(sorted(found), ['__pycache__', 'boiledeggs.py'])
         found = os.listdir(os.path.join(cmd.build_lib, '__pycache__'))
+        fn = 'boiledeggs.%s.%s-0.pyc' % (sys.implementation.cache_tag,
+                                         sys.implementation.optim_tag)
         self.assertEqual(found,
-                         ['boiledeggs.%s.pyc' % sys.implementation.cache_tag])
+                         [fn])
 
     @unittest.skipIf(sys.dont_write_bytecode, 'byte-compile disabled')
     def test_byte_compile_optimized(self):
@@ -120,7 +123,9 @@ def test_byte_compile_optimized(self):
         found = os.listdir(cmd.build_lib)
         self.assertEqual(sorted(found), ['__pycache__', 'boiledeggs.py'])
         found = os.listdir(os.path.join(cmd.build_lib, '__pycache__'))
-        expect = 'boiledeggs.{}.opt-1.pyc'.format(sys.implementation.cache_tag)
+        expect = ('boiledeggs.%s.%s-1.pyc'
+                  % (sys.implementation.cache_tag,
+                     sys.implementation.optim_tag))
         self.assertEqual(sorted(found), [expect])
 
     def test_dir_in_package_data(self):

diff --git a/Lib/distutils/tests/test_install.py b/Lib/distutils/tests/test_install.py
@@ -191,7 +191,9 @@ def test_record(self):
             f.close()
 
         found = [os.path.basename(line) for line in content.splitlines()]
-        expected = ['hello.py', 'hello.%s.pyc' % sys.implementation.cache_tag,
+        fn_pyc = 'hello.%s.%s-0.pyc' % (sys.implementation.cache_tag,
+                                        sys.implementation.optim_tag)
+        expected = ['hello.py', fn_pyc,
                     'sayhi',
                     'UNKNOWN-0.0.0-py%s.%s.egg-info' % sys.version_info[:2]]
         self.assertEqual(found, expected)

diff --git a/Lib/importlib/_bootstrap_external.py b/Lib/importlib/_bootstrap_external.py
@@ -253,7 +253,6 @@ def _write_atomic(path, data, mode=0o666):
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 
 _PYCACHE = '__pycache__'
-_OPT = 'opt-'
 
 SOURCE_SUFFIXES = ['.py']  # _setup() adds .pyw as needed.
 
@@ -294,15 +293,15 @@ def cache_from_source(path, debug_override=None, *, optimization=None):
         raise NotImplementedError('sys.implementation.cache_tag is None')
     almost_filename = ''.join([(base if base else rest), sep, tag])
     if optimization is None:
-        if sys.flags.optimize == 0:
-            optimization = ''
-        else:
-            optimization = sys.flags.optimize
+        optimization = sys.flags.optimize
+    elif optimization == '':
+        optimization = 0
     optimization = str(optimization)
     if optimization != '':
         if not optimization.isalnum():
             raise ValueError('{!r} is not alphanumeric'.format(optimization))
-        almost_filename = '{}.{}{}'.format(almost_filename, _OPT, optimization)
+        part = sys.implementation.optim_tag + '-'
+        almost_filename = '{}.{}{}'.format(almost_filename, part, optimization)
     return _path_join(head, _PYCACHE, almost_filename + BYTECODE_SUFFIXES[0])
 
 
@@ -329,10 +328,12 @@ def source_from_cache(path):
                          '{!r}'.format(pycache_filename))
     elif dot_count == 3:
         optimization = pycache_filename.rsplit('.', 2)[-2]
-        if not optimization.startswith(_OPT):
+        optim_tag = sys.implementation.optim_tag + '-'
+        if optimization.startswith(optim_tag):
+            opt_level = optimization[len(optim_tag):]
+        else:
             raise ValueError("optimization portion of filename does not start "
-                             "with {!r}".format(_OPT))
-        opt_level = optimization[len(_OPT):]
+                             "with {!r}".format(optim_tag))
         if not opt_level.isalnum():
             raise ValueError("optimization level {!r} is not an alphanumeric "
                              "value".format(optimization))
@@ -683,6 +684,13 @@ def load_module(self, fullname):
         return _bootstrap._load_module_shim(self, fullname)
 
 
+def _transformers_tag():
+    transformers = sys.get_code_transformers()
+    if not transformers:
+        return 'noopt'
+    return '-'.join(transformer.name for transformer in transformers)
+
+
 class SourceLoader(_LoaderBasics):
 
     def path_mtime(self, path):
@@ -778,6 +786,15 @@ def get_code(self, fullname):
                         return _compile_bytecode(bytes_data, name=fullname,
                                                  bytecode_path=bytecode_path,
                                                  source_path=source_path)
+
+        optim_tag= sys.implementation.optim_tag
+        transformers_tag = _transformers_tag()
+        if optim_tag != transformers_tag:
+            # AST transformers are missing, the code cannot be compiled
+            raise ImportError("missing AST transformers for %r: "
+                              "optim_tag=%r, transformers tag=%r"
+                              % (source_path, optim_tag, transformers_tag))
+
         source_bytes = self.get_data(source_path)
         code_object = self.source_to_code(source_bytes, source_path)
         _bootstrap._verbose_message('code object from {}', source_path)

diff --git a/Lib/test/test_code_transformer.py b/Lib/test/test_code_transformer.py
@@ -0,0 +1,137 @@
+"""
+Test the implemenentation of the PEP 511: code transformers.
+"""
+import unittest
+import ast
+import os.path
+import sys
+import types
+from test.support.script_helper import assert_python_ok, assert_python_failure
+
+
+class BytecodeTransformer:
+    name = "bytecode"
+
+    def __init__(self):
+        self.call = None
+
+    def code_transformer(self, code, context):
+        self.call = (code, context)
+        consts = ['Ni! Ni! Ni!' if isinstance(const, str) else const
+                  for const in code.co_consts]
+        return types.CodeType(code.co_argcount,
+                              code.co_kwonlyargcount,
+                              code.co_nlocals,
+                              code.co_stacksize,
+                              code.co_flags,
+                              code.co_code,
+                              tuple(consts),
+                              code.co_names,
+                              code.co_varnames,
+                              code.co_filename,
+                              code.co_name,
+                              code.co_firstlineno,
+                              code.co_lnotab,
+                              code.co_freevars,
+                              code.co_cellvars)
+
+
+class KnightsWhoSayNi(ast.NodeTransformer):
+    def visit_Str(self, node):
+        node.s = 'Ni! Ni! Ni!'
+        return node
+
+
+class ASTTransformer:
+    name = "ast"
+
+    def __init__(self):
+        self.transformer = KnightsWhoSayNi()
+        self.call = None
+
+    def ast_transformer(self, tree, context):
+        self.call = (tree, context)
+        self.transformer.visit(tree)
+        return tree
+
+
+class CodeTransformerTests(unittest.TestCase):
+    def setUp(self):
+        transformers = sys.get_code_transformers()
+        self.addCleanup(sys.set_code_transformers, transformers)
+
+        sys.set_code_transformers([])
+
+    def test_bytecode(self):
+        sys.set_code_transformers([BytecodeTransformer()])
+        code = compile('print("Hello World")', '<string>', 'exec')
+        self.assertEqual(code.co_consts, ('Ni! Ni! Ni!', None))
+
+    def test_bytecode_call(self):
+        expr = 'x + 1'
+        filename = "test.py"
+        expected = compile(expr, filename, "exec")
+
+        transformer = BytecodeTransformer()
+        sys.set_code_transformers([transformer])
+        code = compile(expr, filename, "exec")
+
+        code, context = transformer.call
+        self.assertEqual(code, expected)
+        self.assertEqual(context.filename, filename)
+
+    def test_ast(self):
+        expected = ast.parse('print("Ni! Ni! Ni!")')
+
+        sys.set_code_transformers([ASTTransformer()])
+        tree = compile('print("Hello World")', 'string', 'exec',
+                       flags=ast.PyCF_TRANSFORMED_AST)
+        self.assertEqual(ast.dump(tree), ast.dump(expected))
+
+    def test_ast_call(self):
+        code = '1 + 1'
+        expected_ast = ast.dump(ast.parse(code))
+
+        filename = "test.py"
+        transformer = ASTTransformer()
+        sys.set_code_transformers([transformer])
+        compile(code, filename, "exec")
+
+        tree, context = transformer.call
+        self.assertEqual(ast.dump(tree), expected_ast)
+        self.assertEqual(context.filename, filename)
+
+
+class DefaultTransformerTests(unittest.TestCase):
+    def check_default(self, expected, *args):
+        code = ('import sys; ',
+                'transformers = list(map(type, sys.get_code_transformers()))',
+                'print(transformers, sys.implementation.optim_tag)')
+        res = assert_python_ok(*args, '-c', '\n'.join(code))
+        self.assertEqual(res.out.rstrip(), expected)
+
+    def test_default(self):
+        self.check_default(b"[<class 'PeepholeOptimizer'>] opt")
+
+    def test_opt(self):
+        self.check_default(b"[<class 'PeepholeOptimizer'>] opt", '-o', 'opt')
+
+    def test_noopt(self):
+        self.check_default(b"[] noopt", '-o', 'noopt')
+
+
+class OptimTagTests(unittest.TestCase):
+    """Test -o command line option."""
+
+    def test_invalid_optim_tags(self):
+        invalid_tags = ['a.b', 'a%sb' % os.path.sep]
+        if os.path.altsep:
+            invalid_tags.append('a%sb' % os.path.altsep)
+        for optim_tag in invalid_tags:
+            code = 'import sys; print(sys.implementation.optim_tag)'
+            res = assert_python_failure('-o', optim_tag, '-c', code)
+            self.assertIn(b'invalid optimization tag', res.err.rstrip())
+
+
+if __name__ == "__main__":
+    unittest.main()
-Original file line number
+Diff line change
@@ Expand Up @@
                                      void *extra);
     #endif
+    PyAPI_DATA(PyTypeObject) _PyPeepholeOptimizer_Type;
     #ifdef __cplusplus
     }
     #endif
@@ Expand Down @@