Bug 1256573 - Add a new @imports primitive that allows to import modules into the decorated functions. r=nalexander
authorMike Hommey <mh+mozilla@glandium.org>
Sun, 27 Mar 2016 09:54:00 +0900
changeset 290990 96ef4c62378d5ddabb8c27c745bc928517273c20
parent 290989 f01a178847233cf15e9cddaaa97a5db6fc371cc0
child 290991 89b86fc355d7611a3a43872ee047bd32d651ddd1
push id19656
push usergwagner@mozilla.com
push dateMon, 04 Apr 2016 13:43:23 +0000
treeherderb2g-inbound@e99061fde28a [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersnalexander
bugs1256573
milestone48.0a1
Bug 1256573 - Add a new @imports primitive that allows to import modules into the decorated functions. r=nalexander Currently, we have @advanced, that gives the decorated functions access to all the builtins and consequently, to the import statement. That is a quite broad approach and doesn't allow to easily introspect what functions are importing which modules. This change introduces a new decorator that allows to import modules one by one into the decorated functions. Note: by the end of the change series, @advanced will be gone.
python/mozbuild/mozbuild/configure/__init__.py
python/mozbuild/mozbuild/test/configure/test_configure.py
--- a/python/mozbuild/mozbuild/configure/__init__.py
+++ b/python/mozbuild/mozbuild/configure/__init__.py
@@ -2,16 +2,17 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 from __future__ import absolute_import, print_function, unicode_literals
 
 import inspect
 import logging
 import os
+import re
 import sys
 import types
 from collections import OrderedDict
 from contextlib import contextmanager
 from functools import wraps
 from mozbuild.configure.options import (
     CommandLineHelper,
     ConflictingOptionError,
@@ -52,28 +53,29 @@ def forbidden_import(*args, **kwargs):
     raise ImportError('Importing modules is forbidden')
 
 
 class ConfigureSandbox(dict):
     """Represents a sandbox for executing Python code for build configuration.
     This is a different kind of sandboxing than the one used for moz.build
     processing.
 
-    The sandbox has 8 primitives:
+    The sandbox has 9 primitives:
     - option
     - depends
     - template
+    - imports
     - advanced
     - include
     - set_config
     - set_define
     - imply_option
 
     `option`, `include`, `set_config`, `set_define` and `imply_option` are
-    functions. `depends`, `template` and `advanced` are decorators.
+    functions. `depends`, `template`, `imports` and `advanced` are decorators.
 
     These primitives are declared as name_impl methods to this class and
     the mapping name -> name_impl is done automatically in __getitem__.
 
     Additional primitives should be frowned upon to keep the sandbox itself as
     simple as possible. Instead, helpers should be created within the sandbox
     with the existing primitives.
 
@@ -106,16 +108,18 @@ class ConfigureSandbox(dict):
         dict.__setitem__(self, '__builtins__', self.BUILTINS)
 
         self._paths = []
         self._templates = set()
         # Store the real function and its dependencies, behind each
         # DependsFunction generated from @depends.
         self._depends = {}
         self._seen = set()
+        # Store the @imports added to a given function.
+        self._imports = {}
 
         self._options = OrderedDict()
         # Store the raw values returned by @depends functions
         self._results = {}
         # Store values for each Option, as per returned by Option.get_value
         self._option_values = {}
         # Store raw option (as per command line or environment) for each Option
         self._raw_options = {}
@@ -423,19 +427,72 @@ class ConfigureSandbox(dict):
         self._templates.add(wrapper)
         return wrapper
 
     def advanced_impl(self, func):
         '''Implementation of @advanced.
         This function gives the decorated function access to the complete set
         of builtins, allowing the import keyword as an expected side effect.
         '''
-        func, glob = self._prepare_function(func)
-        glob.update(__builtins__=__builtins__)
-        return func
+        return self.imports_impl(_import='__builtin__', _as='__builtins__')(func)
+
+    RE_MODULE = re.compile('^[a-zA-Z0-9_\.]+$')
+
+    def imports_impl(self, _import, _from=None, _as=None):
+        '''Implementation of @imports.
+        This decorator imports the given _import from the given _from module
+        optionally under a different _as name.
+        The options correspond to the various forms for the import builtin.
+            @imports('sys')
+            @imports(_from='mozpack', _import='path', _as='mozpath')
+        '''
+        for value, required in (
+                (_import, True), (_from, False), (_as, False)):
+            if not isinstance(value, types.StringTypes) and not (
+                    required or value is None):
+                raise TypeError("Unexpected type: '%s'" % type(value))
+            if value is not None and not self.RE_MODULE.match(value):
+                raise ValueError("Invalid argument to @imports: '%s'" % value)
+
+        def decorator(func):
+            if func in self._prepared_functions:
+                raise ConfigureError(
+                    '@imports must appear after other decorators')
+            # For the imports to apply in the order they appear in the
+            # .configure file, we accumulate them in reverse order and apply
+            # them later.
+            imports = self._imports.setdefault(func, [])
+            imports.insert(0, (_from, _import, _as))
+            return func
+
+        return decorator
+
+    def _apply_imports(self, func, glob):
+        for _from, _import, _as in self._imports.get(func, ()):
+            # The special `__sandbox__` module gives access to the sandbox
+            # instance.
+            if _from is None and _import == '__sandbox__':
+                glob[_as or _import] = self
+                continue
+            # Special case for the open() builtin, because otherwise, using it
+            # fails with "IOError: file() constructor not accessible in
+            # restricted mode"
+            if _from == '__builtin__' and _import == 'open':
+                glob[_as or _import] = \
+                    lambda *args, **kwargs: open(*args, **kwargs)
+                continue
+            # Until this proves to be a performance problem, just construct an
+            # import statement and execute it.
+            import_line = ''
+            if _from:
+                import_line += 'from %s ' % _from
+            import_line += 'import %s' % _import
+            if _as:
+                import_line += ' as %s' % _as
+            exec(import_line, {}, glob)
 
     def _resolve_and_set(self, data, name, value):
         # Don't set anything when --help was on the command line
         if self._help:
             return
         name = self._resolve(name, need_help_dependency=False)
         if name is None:
             return
@@ -560,16 +617,17 @@ class ConfigureSandbox(dict):
 
         glob = SandboxedGlobal(func.func_globals)
         glob.update(
             __builtins__=self.BUILTINS,
             __file__=self._paths[-1] if self._paths else '',
             os=self.OS,
             log=self.log_impl,
         )
+        self._apply_imports(func, glob)
         func = wraps(func)(types.FunctionType(
             func.func_code,
             glob,
             func.__name__,
             func.func_defaults,
             func.func_closure
         ))
         self._prepared_functions.add(func)
--- a/python/mozbuild/mozbuild/test/configure/test_configure.py
+++ b/python/mozbuild/mozbuild/test/configure/test_configure.py
@@ -1,16 +1,18 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 from __future__ import absolute_import, print_function, unicode_literals
 
 from StringIO import StringIO
+import os
 import sys
+import textwrap
 import unittest
 
 from mozunit import main
 
 from mozbuild.configure.options import (
     InvalidOptionError,
     NegativeOptionValue,
     PositiveOptionValue,
@@ -228,16 +230,119 @@ class TestConfigure(unittest.TestCase):
     def test_advanced(self):
         config = self.get_config(['--with-advanced'])
         self.assertIn('ADVANCED', config)
         self.assertEquals(config['ADVANCED'], True)
 
         with self.assertRaises(ImportError):
             self.get_config(['--with-advanced=break'])
 
+    def test_imports(self):
+        config = {}
+        out = StringIO()
+        sandbox = ConfigureSandbox(config, {}, [], out, out)
+
+        with self.assertRaises(ImportError):
+            exec(textwrap.dedent('''
+                @template
+                def foo():
+                    import sys
+                foo()'''),
+                sandbox
+            )
+
+        exec(textwrap.dedent('''
+            @template
+            @imports('sys')
+            def foo():
+                return sys'''),
+            sandbox
+        )
+
+        self.assertIs(sandbox['foo'](), sys)
+
+        exec(textwrap.dedent('''
+            @template
+            @imports(_from='os', _import='path')
+            def foo():
+                return path'''),
+            sandbox
+        )
+
+        self.assertIs(sandbox['foo'](), os.path)
+
+        exec(textwrap.dedent('''
+            @template
+            @imports(_from='os', _import='path', _as='os_path')
+            def foo():
+                return os_path'''),
+            sandbox
+        )
+
+        self.assertIs(sandbox['foo'](), os.path)
+
+        exec(textwrap.dedent('''
+            @template
+            @imports('__builtin__')
+            def foo():
+                return __builtin__'''),
+            sandbox
+        )
+
+        import __builtin__
+        self.assertIs(sandbox['foo'](), __builtin__)
+
+        exec(textwrap.dedent('''
+            @template
+            @imports(_from='__builtin__', _import='open')
+            def foo():
+                return open('%s')''' % os.devnull),
+            sandbox
+        )
+
+        f = sandbox['foo']()
+        self.assertEquals(f.name, os.devnull)
+        f.close()
+
+        # This unlocks the sandbox
+        exec(textwrap.dedent('''
+            @template
+            @imports(_import='__builtin__', _as='__builtins__')
+            def foo():
+                import sys
+                return sys'''),
+            sandbox
+        )
+
+        self.assertIs(sandbox['foo'](), sys)
+
+        exec(textwrap.dedent('''
+            @template
+            @imports('__sandbox__')
+            def foo():
+                return __sandbox__'''),
+            sandbox
+        )
+
+        self.assertIs(sandbox['foo'](), sandbox)
+
+        exec(textwrap.dedent('''
+            @template
+            @imports(_import='__sandbox__', _as='s')
+            def foo():
+                return s'''),
+            sandbox
+        )
+
+        self.assertIs(sandbox['foo'](), sandbox)
+
+        # Nothing leaked from the function being executed
+        self.assertEquals(sandbox.keys(), ['__builtins__', 'foo'])
+        self.assertEquals(sandbox['__builtins__'], ConfigureSandbox.BUILTINS)
+
     def test_os_path(self):
         config = self.get_config(['--with-advanced=%s' % __file__])
         self.assertIn('IS_FILE', config)
         self.assertEquals(config['IS_FILE'], True)
 
         config = self.get_config(['--with-advanced=%s.no-exist' % __file__])
         self.assertIn('IS_FILE', config)
         self.assertEquals(config['IS_FILE'], False)