Allow a set of datatype methods like __init__ for class methods

docs/CHANGES.rst

@@ -4,6 +4,20 @@ Changes
 4.0b4 (unreleased)
 ------------------
 
+- Allow the following magic methods to be defined on classes.
+  (`#104 <https://github.com/zopefoundation/RestrictedPython/issues/104>`_)
+  They cannot be called directly but by the built-in way to use them (e. g.
+  class instantiation, or comparison):
+
+  + ``__init__``
+  + ``__contains__``
+  + ``__lt__``
+  + ``__le__``
+  + ``__eq__``
+  + ``__ne__``
+  + ``__gt__``
+  + ``__ge__``
+
 - Imports like `from a import *` (so called star imports) are now forbidden as
   they allow to import names starting with an underscore which could override
   protected build-ins.

setup.cfg

@@ -36,6 +36,8 @@ not_skip =
 
 [flake8]
 ignore =
+    # W503 line break before binary operator: is no longer requested by PEP-8
+    W503
 no-accept-encodings = True
 
 [coverage:run]

src/RestrictedPython/transformer.py

@@ -53,6 +53,26 @@
     IOPERATOR_TO_STR[ast.MatMult] = '@='
 
 
+# For creation allowed magic method names. See also
+# https://docs.python.org/3/reference/datamodel.html#special-method-names
+ALLOWED_FUNC_NAMES = frozenset([
+    '__init__',
+    '__contains__',
+    '__lt__',
+    '__le__',
+    '__eq__',
+    '__ne__',
+    '__gt__',
+    '__ge__',
+])
+
+
+FORBIDDEN_FUNC_NAMES = frozenset([
+    'print',
+    'printed',
+])
+
+
 # When new ast nodes are generated they have no 'lineno' and 'col_offset'.
 # This function copies these two fields from the incoming node
 def copy_locations(new_node, old_node):
@@ -363,26 +383,30 @@ def transform_slice(self, slice_):
         else:
             raise Exception("Unknown slice type: {0}".format(slice_))
 
-    def check_name(self, node, name):
+    def check_name(self, node, name, allow_magic_methods=False):
+        """Check names if they are allowed.
+
+        If ``allow_magic_methods is True`` names in `ALLOWED_FUNC_NAMES`
+        are additionally allowed although their names start with `_`.
+
+        """
         if name is None:
             return
 
-        if name.startswith('_') and name != '_':
+        if (name.startswith('_')
+                and name != '_'
+                and not (allow_magic_methods
+                         and name in ALLOWED_FUNC_NAMES
+                         and node.col_offset != 0)):
             self.error(
                 node,
                 '"{name}" is an invalid variable name because it '
                 'starts with "_"'.format(name=name))
-
         elif name.endswith('__roles__'):
             self.error(node, '"%s" is an invalid variable name because '
                        'it ends with "__roles__".' % name)
-
-        elif name == "printed":
-            self.error(node, '"printed" is a reserved name.')
-
-        elif name == 'print':
-            # Assignments to 'print' would lead to funny results.
-            self.error(node, '"print" is a reserved name.')
+        elif name in FORBIDDEN_FUNC_NAMES:
+            self.error(node, '"{name}" is a reserved name.'.format(name=name))
 
     def check_function_argument_names(self, node):
         # In python3 arguments are always identifiers.
@@ -1234,7 +1258,7 @@ def visit_withitem(self, node):
 
     def visit_FunctionDef(self, node):
         """Allow function definitions (`def`) with some restrictions."""
-        self.check_name(node, node.name)
+        self.check_name(node, node.name, allow_magic_methods=True)
         self.check_function_argument_names(node)
 
         with self.print_info.new_print_scope():

tests/transformer/test_call.py

@@ -102,3 +102,21 @@ def test_RestrictingNodeTransformer__visit_Call__2(e_exec, mocker):
     assert ref == ret
     _apply_.assert_called_once_with(glb['foo'], *ref[0], **ref[1])
     _apply_.reset_mock()
+
+
+@pytest.mark.parametrize(*c_exec)
+def test_visit_Call__private_function(c_exec):
+    """Calling private functions is forbidden."""
+    result = c_exec('__init__(1)')
+    assert result.errors == (
+        'Line 1: "__init__" is an invalid variable name because it starts with "_"',  # NOQA: E501
+    )
+
+
+@pytest.mark.parametrize(*c_exec)
+def test_visit_Call__private_method(c_exec):
+    """Calling private methods is forbidden."""
+    result = c_exec('Int.__init__(1)')
+    assert result.errors == (
+        'Line 1: "__init__" is an invalid attribute name because it starts with "_".',  # NOQA: E501
+    )

tests/transformer/test_classdef.py

@@ -108,3 +108,92 @@ def test_RestrictingNodeTransformer__visit_ClassDef__5(e_exec):
     assert comb.class_att == 2342
     assert comb.base_att == 42
     assert comb.wrap_att == 23
+
+
+CONSTRUCTOR_TEST = """\
+class Test(object):
+    def __init__(self, input):
+        self.input = input
+
+t = Test(42)
+"""
+
+
+@pytest.mark.parametrize(*e_exec)
+def test_RestrictingNodeTransformer__visit_ClassDef__6(e_exec):
+    """It allows to define an ``__init__`` method."""
+    restricted_globals = dict(
+        t=None,
+        _write_=lambda x: x,
+        __metaclass__=type,
+    )
+
+    e_exec(CONSTRUCTOR_TEST, restricted_globals)
+    t = restricted_globals['t']
+    assert t.input == 42
+
+
+COMPARE_TEST = """\
+class Test(object):
+
+    def __init__(self, value):
+        self.value = value
+
+    def __eq__(self, other):
+        return self.value == other.value
+
+a = Test(42)
+b = Test(42)
+c = Test(43)
+
+result1 = (a == b)
+result2 = (b == c)
+"""
+
+
+@pytest.mark.parametrize(*e_exec)
+def test_RestrictingNodeTransformer__visit_ClassDef__7(e_exec):
+    """It allows to define an ``__eq__`` method."""
+    restricted_globals = dict(
+        result1=None,
+        result2=None,
+        _getattr_=getattr,
+        _write_=lambda x: x,
+        __metaclass__=type,
+    )
+
+    e_exec(COMPARE_TEST, restricted_globals)
+    assert restricted_globals['result1'] is True
+    assert restricted_globals['result2'] is False
+
+
+CONTAINER_TEST = """\
+class Test(object):
+
+    def __init__(self, values):
+        self.values = values
+
+    def __contains__(self, value):
+        return value in self.values
+
+a = Test([1, 2, 3])
+
+result1 = (1 in a)
+result2 = (4 not in a)
+"""
+
+
+@pytest.mark.parametrize(*e_exec)
+def test_RestrictingNodeTransformer__visit_ClassDef__8(e_exec):
+    """It allows to define a ``__contains__`` method."""
+    restricted_globals = dict(
+        result1=None,
+        result2=None,
+        _getattr_=getattr,
+        _write_=lambda x: x,
+        __metaclass__=type,
+    )
+
+    e_exec(CONTAINER_TEST, restricted_globals)
+    assert restricted_globals['result1'] is True
+    assert restricted_globals['result2'] is True

tests/transformer/test_functiondef.py

@@ -136,3 +136,25 @@ def test_RestrictingNodeTransformer__visit_FunctionDef__8(
         mocker.call((1, 2)),
         mocker.call((3, 4))])
     _getiter_.reset_mock()
+
+
+BLACKLISTED_FUNC_NAMES_CALL_TEST = """
+def __init__(test):
+    test
+
+__init__(1)
+"""
+
+
+@pytest.mark.parametrize(*c_exec)
+def test_RestrictingNodeTransformer__module_func_def_name_call(c_exec):
+    """It forbids definition and usage of magic methods as functions ...
+
+    ... at module level.
+    """
+    result = c_exec(BLACKLISTED_FUNC_NAMES_CALL_TEST)
+    # assert result.errors == ('Line 1: ')
+    assert result.errors == (
+        'Line 2: "__init__" is an invalid variable name because it starts with "_"',  # NOQA: E501
+        'Line 5: "__init__" is an invalid variable name because it starts with "_"',  # NOQA: E501
+    )