Fix SF bug #1448804 and ad a test to ensure that all subscript operations continue to be handled correctly

2 files changed, 101 insertions, 24 deletions

diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py
index 93a2fb5..1d47f91 100644
--- a/Lib/test/test_compile.py
+++ b/Lib/test/test_compile.py
@@ -284,6 +284,78 @@ if 1:
         f1, f2 = f()
         self.assertNotEqual(id(f1.func_code), id(f2.func_code))
 
+    def test_subscripts(self):
+        # SF bug 1448804
+        # Class to make testing subscript results easy
+        class str_map(object):
+            def __init__(self):
+                self.data = {}
+            def __getitem__(self, key):
+                return self.data[str(key)]
+            def __setitem__(self, key, value):
+                self.data[str(key)] = value
+            def __delitem__(self, key):
+                del self.data[str(key)]
+            def __contains__(self, key):
+                return str(key) in self.data
+        d = str_map()
+        # Index
+        d[1] = 1
+        self.assertEqual(d[1], 1)
+        d[1] += 1
+        self.assertEqual(d[1], 2)
+        del d[1]
+        self.assertEqual(1 in d, False)
+        # Tuple of indices
+        d[1, 1] = 1
+        self.assertEqual(d[1, 1], 1)
+        d[1, 1] += 1
+        self.assertEqual(d[1, 1], 2)
+        del d[1, 1]
+        self.assertEqual((1, 1) in d, False)
+        # Simple slice
+        d[1:2] = 1
+        self.assertEqual(d[1:2], 1)
+        d[1:2] += 1
+        self.assertEqual(d[1:2], 2)
+        del d[1:2]
+        self.assertEqual(slice(1, 2) in d, False)
+        # Tuple of simple slices
+        d[1:2, 1:2] = 1
+        self.assertEqual(d[1:2, 1:2], 1)
+        d[1:2, 1:2] += 1
+        self.assertEqual(d[1:2, 1:2], 2)
+        del d[1:2, 1:2]
+        self.assertEqual((slice(1, 2), slice(1, 2)) in d, False)
+        # Extended slice
+        d[1:2:3] = 1
+        self.assertEqual(d[1:2:3], 1)
+        d[1:2:3] += 1
+        self.assertEqual(d[1:2:3], 2)
+        del d[1:2:3]
+        self.assertEqual(slice(1, 2, 3) in d, False)
+        # Tuple of extended slices
+        d[1:2:3, 1:2:3] = 1
+        self.assertEqual(d[1:2:3, 1:2:3], 1)
+        d[1:2:3, 1:2:3] += 1
+        self.assertEqual(d[1:2:3, 1:2:3], 2)
+        del d[1:2:3, 1:2:3]
+        self.assertEqual((slice(1, 2, 3), slice(1, 2, 3)) in d, False)
+        # Ellipsis
+        d[...] = 1
+        self.assertEqual(d[...], 1)
+        d[...] += 1
+        self.assertEqual(d[...], 2)
+        del d[...]
+        self.assertEqual(Ellipsis in d, False)
+        # Tuple of Ellipses
+        d[..., ...] = 1
+        self.assertEqual(d[..., ...], 1)
+        d[..., ...] += 1
+        self.assertEqual(d[..., ...], 2)
+        del d[..., ...]
+        self.assertEqual((Ellipsis, Ellipsis) in d, False)
+
 def test_main():
     test_support.run_unittest(TestSpecifics)
 
diff --git a/Python/compile.c b/Python/compile.c
index e3b3df8..baf3989 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -3853,42 +3853,47 @@ compiler_visit_nested_slice(struct compiler *c, slice_ty s,
 static int
 compiler_visit_slice(struct compiler *c, slice_ty s, expr_context_ty ctx)
 {
+	char * kindname = NULL;
 	switch (s->kind) {
+	case Index_kind:
+		kindname = "index";
+		if (ctx != AugStore) {
+			VISIT(c, expr, s->v.Index.value);
+		}
+		break;
 	case Ellipsis_kind:
-		ADDOP_O(c, LOAD_CONST, Py_Ellipsis, consts);
+		kindname = "ellipsis";
+		if (ctx != AugStore) {
+			ADDOP_O(c, LOAD_CONST, Py_Ellipsis, consts);
+		}
 		break;
 	case Slice_kind:
+		kindname = "slice";
 		if (!s->v.Slice.step) 
 			return compiler_simple_slice(c, s, ctx);
-		if (!compiler_slice(c, s, ctx))
-			return 0;
-		if (ctx == AugLoad) {
-			ADDOP_I(c, DUP_TOPX, 2);
-		}
-		else if (ctx == AugStore) {
-			ADDOP(c, ROT_THREE);
-		}
-		return compiler_handle_subscr(c, "slice", ctx);
-	case ExtSlice_kind: {
-		int i, n = asdl_seq_LEN(s->v.ExtSlice.dims);
-		for (i = 0; i < n; i++) {
-			slice_ty sub = asdl_seq_GET(s->v.ExtSlice.dims, i);
-			if (!compiler_visit_nested_slice(c, sub, ctx))
+		if (ctx != AugStore) {
+			if (!compiler_slice(c, s, ctx))
 				return 0;
 		}
-		ADDOP_I(c, BUILD_TUPLE, n);
-		return compiler_handle_subscr(c, "extended slice", ctx);
-	}
-	case Index_kind:
-		if (ctx != AugStore)
-			VISIT(c, expr, s->v.Index.value);
-		return compiler_handle_subscr(c, "index", ctx);
+		break;
+	case ExtSlice_kind:
+		kindname = "extended slice";
+		if (ctx != AugStore) {
+			int i, n = asdl_seq_LEN(s->v.ExtSlice.dims);
+			for (i = 0; i < n; i++) {
+				slice_ty sub = asdl_seq_GET(s->v.ExtSlice.dims, i);
+				if (!compiler_visit_nested_slice(c, sub, ctx))
+					return 0;
+			}
+			ADDOP_I(c, BUILD_TUPLE, n);
+		}
+		break;
 	default:
 		PyErr_Format(PyExc_SystemError,
-			     "invalid slice %d", s->kind);
+			     "invalid subscript kind %d", s->kind);
 		return 0;
 	}
-	return 1;
+	return compiler_handle_subscr(c, kindname, ctx);
 }
 
 /* do depth-first search of basic block graph, starting with block.