Allow class definitions and support decorators, bases and global __me…

…taclass__.

docs/usage/basic_usage.rst

@@ -95,6 +95,9 @@ Necessary setup
 `RestrictedPython` requires some predefined names in globals in order to work
 properly.
 
+To use classes in Python 3
+    ``__metaclass__`` must be set. Set it to ``type`` to use no custom metaclass.
+
 To use ``for`` statements and comprehensions
     ``_iter_unpack_sequence_`` must point to :func:`RestrictedPython.Guards.guarded_iter_unpack_sequence`.
 

src/RestrictedPython/transformer.py

@@ -29,6 +29,7 @@
 
 import ast
 import contextlib
+import textwrap
 
 
 # For AugAssign the operator must be converted to a string.
@@ -1346,7 +1347,22 @@ def visit_Nonlocal(self, node):
     def visit_ClassDef(self, node):
         """Check the name of a class definition."""
         self.check_name(node, node.name)
-        return self.node_contents_visit(node)
+        node = self.node_contents_visit(node)
+        if IS_PY2:
+            new_class_node = node
+        else:
+            if any(keyword.arg == 'metaclass' for keyword in node.keywords):
+                self.error(
+                    node, 'The keyword argument "metaclass" is not allowed.')
+            CLASS_DEF = textwrap.dedent('''\
+                class {0.name}(metaclass=__metaclass__):
+                    pass
+            '''.format(node))
+            new_class_node = ast.parse(CLASS_DEF).body[0]
+            new_class_node.body = node.body
+            new_class_node.bases = node.bases
+            new_class_node.decorator_list = node.decorator_list
+        return new_class_node
 
     def visit_Module(self, node):
         """Add the print_collector (only if print is used) at the top."""

tests/__init__.py

@@ -19,6 +19,8 @@ def _exec(source, glb=None):
             glb = {}
         exec(code, glb)
         return glb
+    # The next line can be dropped after the old implementation was dropped.
+    _exec.compile_func = compile_func
     return _exec
 
 

tests/test_Guards.py

@@ -33,6 +33,7 @@ def test_Guards__safe_builtins__2(e_exec):
     restricted_globals = dict(
         __builtins__=safe_builtins, b=None,
         __name__='restricted_module',
+        __metaclass__=type,
         _write_=lambda x: x,
         _getattr_=getattr)
 

tests/test_print_function.py

@@ -305,6 +305,7 @@ def test_print_function_no_new_scope():
     code, errors = compiler(NO_PRINT_SCOPES)[:2]
     glb = {
         '_print_': PrintCollector,
+        '__metaclass__': type,
         '_getattr_': None,
         '_getiter_': lambda ob: ob
     }

tests/transformer/test_transformer.py

@@ -3,6 +3,7 @@
 from RestrictedPython._compat import IS_PY3
 from RestrictedPython.Guards import guarded_iter_unpack_sequence
 from RestrictedPython.Guards import guarded_unpack_sequence
+from RestrictedPython.Guards import safe_builtins
 from tests import c_exec
 from tests import e_eval
 from tests import e_exec
@@ -1304,6 +1305,83 @@ def test_transformer__RestrictingNodeTransformer__visit_ClassDef__2(c_exec):
         'because it starts with "_"',)
 
 
+IMPLICIT_METACLASS = '''
+class Meta:
+    pass
+
+b = Meta().foo
+'''
+
+
+@pytest.mark.parametrize(*e_exec)
+def test_transformer__RestrictingNodeTransformer__visit_ClassDef__3(e_exec):
+    """It applies the global __metaclass__ to all generated classes if present.
+    """
+    def _metaclass(name, bases, dict):
+        ob = type(name, bases, dict)
+        ob.foo = 2411
+        return ob
+
+    restricted_globals = dict(
+        __metaclass__=_metaclass, b=None, _getattr_=getattr)
+
+    e_exec(IMPLICIT_METACLASS, restricted_globals)
+
+    assert restricted_globals['b'] == 2411
+
+
+EXPLICIT_METACLASS = '''
+class WithMeta(metaclass=MyMetaClass):
+    pass
+'''
+
+
+@pytest.mark.skipif(IS_PY2, reason="No valid syntax in Python 2.")
+@pytest.mark.parametrize(*c_exec)
+def test_transformer__RestrictingNodeTransformer__visit_ClassDef__4(c_exec):
+    """It does not allow to pass a metaclass to class definitions."""
+
+    result = c_exec(EXPLICIT_METACLASS)
+
+    assert result.errors == (
+        'Line 2: The keyword argument "metaclass" is not allowed.',)
+    assert result.code is None
+
+
+DECORATED_CLASS = '''\
+def wrap(cls):
+    cls.wrap_att = 23
+    return cls
+
+class Base:
+    base_att = 42
+
+@wrap
+class Combined(Base):
+    class_att = 2342
+
+comb = Combined()
+'''
+
+
+@pytest.mark.parametrize(*e_exec)
+def test_transformer__RestrictingNodeTransformer__visit_ClassDef__5(e_exec):
+    """It preserves base classes and decorators for classes."""
+
+    restricted_globals = dict(
+        comb=None, _getattr_=getattr, _write_=lambda x: x, __metaclass__=type,
+        __name__='restricted_module', __builtins__=safe_builtins)
+
+    e_exec(DECORATED_CLASS, restricted_globals)
+
+    comb = restricted_globals['comb']
+    assert comb.class_att == 2342
+    assert comb.base_att == 42
+    if e_exec.compile_func is RestrictedPython.compile.compile_restricted_exec:
+        # Class decorators are only supported by the new implementation.
+        assert comb.wrap_att == 23
+
+
 @pytest.mark.parametrize(*e_exec)
 def test_transformer__RestrictingNodeTransformer__test_ternary_if(
         e_exec, mocker):