Protect extended slice subscript via _getitem_.

src/RestrictedPython/transformer.py

@@ -226,6 +226,51 @@ def gen_none_node(self):
         else:
             return ast.Name(id='None', ctx=ast.Load())
 
+    def transform_slice(self, slice_):
+        """Transforms slices into function parameters.
+
+        ast.Slice nodes are only allowed within a ast.Subscript node.
+        To use a slice as an argument of ast.Call it has to be converted.
+        Conversion is done by calling the 'slice' function from builtins
+        """
+
+        if isinstance(slice_, ast.Index):
+            return slice_.value
+
+        elif isinstance(slice_, ast.Slice):
+            # Create a python slice object.
+            args = []
+
+            if slice_.lower:
+                args.append(slice_.lower)
+            else:
+                args.append(self.gen_none_node())
+
+            if slice_.upper:
+                args.append(slice_.upper)
+            else:
+                args.append(self.gen_none_node())
+
+            if slice_.step:
+                args.append(slice_.step)
+            else:
+                args.append(self.gen_none_node())
+
+            return ast.Call(
+                func=ast.Name('slice', ast.Load()),
+                args=args,
+                keywords=[])
+
+        elif isinstance(slice_, ast.ExtSlice):
+            dims = ast.Tuple([], ast.Load())
+            for item in slice_.dims:
+                dims.elts.append(self.transform_slice(item))
+            return dims
+
+        else:
+            raise Exception("Unknown slice type: {0}".format(slice_))
+
+
     # Special Functions for an ast.NodeTransformer
 
     def generic_visit(self, node):
@@ -606,51 +651,18 @@ def visit_Subscript(self, node):
         'foo[ab:]' becomes '_getitem_(foo, slice(ab, None, None))'
         'foo[a:b]' becomes '_getitem_(foo, slice(a, b, None))'
         'foo[a:b:c]' becomes '_getitem_(foo, slice(a, b, c))'
+        'foo[a, b:c] becomes '_getitem_(foo, (a, slice(b, c, None)))'
         """
         node = self.generic_visit(node)
+
         if isinstance(node.ctx, ast.Load):
-            if isinstance(node.slice, ast.Slice):
-                # Create a python slice object.
-                args = []
-
-                if node.slice.lower:
-                    args.append(node.slice.lower)
-                else:
-                    args.append(self.gen_none_node())
-
-                if node.slice.upper:
-                    args.append(node.slice.upper)
-                else:
-                    args.append(self.gen_none_node())
-
-                if node.slice.step:
-                    args.append(node.slice.step)
-                else:
-                    args.append(self.gen_none_node())
-
-                slice_ob = ast.Call(
-                    func=ast.Name('slice', ast.Load()),
-                    args=args,
-                    keywords=[])
-
-                # Feed this slice object into _getitem_
-                new_node = ast.Call(
-                    func=ast.Name('_getitem_', ast.Load()),
-                    args=[node.value, slice_ob],
-                    keywords=[])
-
-                copy_locations(new_node, node)
-                return new_node
-
-            if isinstance(node.slice, ast.Index):
-                new_node = ast.Call(
-                    func=ast.Name('_getitem_', ast.Load()),
-                    args=[node.value, node.slice.value],
-                    keywords=[])
-                copy_locations(new_node, node)
-                return new_node
+            new_node = ast.Call(
+                func=ast.Name('_getitem_', ast.Load()),
+                args=[node.value, self.transform_slice(node.slice)],
+                keywords=[])
 
-        return node
+            copy_locations(new_node, node)
+            return new_node
 
     def visit_Index(self, node):
         """

tests/test_transformer.py

@@ -306,6 +306,9 @@ def slice_subscript_no_step(a):
 
 def slice_subscript_with_step(a):
     return a[1:2:3]
+
+def extended_slice_subscript(a):
+    return a[0, :1, 1:, 1:2, 1:2:3]
 """
 
 
@@ -349,3 +352,19 @@ def test_transformer__RestrictingNodeTransformer__visit_Subscript(compile, mocke
     assert ref == ret
     _getitem_.assert_called_once_with(*ref)
     _getitem_.reset_mock()
+
+    ret = glb['extended_slice_subscript'](value)
+    ref = (
+        value,
+        (
+            0,
+            slice(None, 1, None),
+            slice(1, None, None),
+            slice(1, 2, None),
+            slice(1, 2, 3)
+        )
+    )
+
+    assert ref == ret
+    _getitem_.assert_called_once_with(*ref)
+    _getitem_.reset_mock()