diff --git a/CHANGES.rst b/CHANGES.rst index 60c7eb4..dfdd7bd 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -9,6 +9,10 @@ Changes providing an ``__import__`` implementation hands the security boundary to the import policy of the calling application. +- Allow defining ``__enter__`` and ``__exit__`` on a class in restricted code, + so a class can be used as a context manager in a ``with`` statement. See + `issue 230 `_. + 8.5 (2026-08-19) ---------------- diff --git a/src/RestrictedPython/transformer.py b/src/RestrictedPython/transformer.py index e129b63..b8e8d2a 100644 --- a/src/RestrictedPython/transformer.py +++ b/src/RestrictedPython/transformer.py @@ -55,6 +55,8 @@ '__ne__', '__gt__', '__ge__', + '__enter__', + '__exit__', ]) diff --git a/tests/test_Guards.py b/tests/test_Guards.py index 5895d49..f0a6543 100644 --- a/tests/test_Guards.py +++ b/tests/test_Guards.py @@ -55,6 +55,37 @@ def display(self): assert restricted_globals['result'] == '2411' +def test_Guards__safe_builtins__3(): + """It allows classes to define `__enter__` and `__exit__` so they can be + used as context managers in a `with` statement. + """ + + context_manager_code = ''' +class MyContextManager: + def __enter__(self): + events.append('enter') + return self + + def __exit__(self, exc_type, exc_value, exc_traceback): + events.append('exit') + return False + +events = [] +with MyContextManager(): + events.append('inside') +result = events''' + + restricted_globals = dict( + result=None, + __name__='restricted_module', + __metaclass__=type, + _write_=_write_, + _getattr_=getattr) + + restricted_exec(context_manager_code, restricted_globals) + assert restricted_globals['result'] == ['enter', 'inside', 'exit'] + + def test_Guards__guarded_setattr__1(): """It allows use setattr and delattr when _guarded_writes is True. """