# HG changeset patch
# User paulb@jeremy
# Date 1152919318 -7200
# Node ID e51ec15e7979dd0cf2d5aad029df0a623482da47
# Parent  fbac48c3637bd33245001772412885c7873d3028
Added Subscript node support (including augmented subscript assignment) along
with support for slice/subscript deletion flags and special methods.


diff -r fbac48c3637b -r e51ec15e7979 simplify.py
--- a/simplify.py	Sat Jul 15 00:47:24 2006 +0200
+++ b/simplify.py	Sat Jul 15 01:21:58 2006 +0200
@@ -323,6 +323,23 @@
             result.code = [store_expr, store_lower, store_upper, get_incremented, store,
                 ReleaseTemp(index="expr"), ReleaseTemp(index="lower"), ReleaseTemp(index="upper"), ReleaseTemp()]
 
+        # Complicated augassign using subscripts: lvalue[subs] += expr
+
+        elif isinstance(augassign.node, compiler.ast.Subscript):
+
+            # <lvalue>, <subs> -> <lvalue>.__setitem__(<subs>, <lvalue>.__getitem__(<subs>).__xxx__(expr))
+
+            subscript = augassign.node
+            store_expr = StoreTemp(index="expr", expr=self.dispatch(subscript.expr))
+            subs = self._visitSubscriptSubs(subscript.subs)
+            store_subs = StoreTemp(index="subs", expr=subs)
+            node_subscript = self._visitSubscript(LoadTemp(index="expr"), LoadTemp(index="subs"), "OP_APPLY")
+            get_incremented = StoreTemp(
+                expr=Invoke(expr=LoadAttr(expr=node_subscript, name=self.augassign_methods[augassign.op]), args=[expr])
+                )
+            store = self._visitSubscript(LoadTemp(index="expr"), LoadTemp(index="subs"), "OP_ASSIGN", LoadTemp())
+            result.code = [store_expr, store_subs, get_incremented, store, ReleaseTemp(index="expr"), ReleaseTemp(index="subs"), ReleaseTemp()]
+
         else:
             raise NotImplementedError, augassign.node.__class__
 
@@ -372,6 +389,9 @@
         elif flags == "OP_APPLY":
             args = []
             result = Invoke(expr=LoadAttr(expr=expr, name="__getslice__"))
+        elif flags == "OP_DELETE":
+            args = []
+            result = Invoke(expr=LoadAttr(expr=expr, name="__delslice__"))
         else:
             raise NotImplementedError, flags
 
@@ -387,6 +407,37 @@
         value = self._visitAssNameOrAttr(slice, in_sequence)
         return self._visitSlice(self.dispatch(slice.expr), self.dispatch_or_none(slice.lower), self.dispatch_or_none(slice.upper), slice.flags, value)
 
+    def _visitSubscript(self, expr, subs, flags, value=None):
+        if flags == "OP_ASSIGN":
+            args = [value]
+            result = Invoke(expr=LoadAttr(expr=expr, name="__setitem__"))
+        elif flags == "OP_APPLY":
+            args = []
+            result = Invoke(expr=LoadAttr(expr=expr, name="__getitem__"))
+        elif flags == "OP_DELETE":
+            args = []
+            result = Invoke(expr=LoadAttr(expr=expr, name="__delitem__"))
+        else:
+            raise NotImplementedError, flags
+
+        # Add the dimensions.
+
+        args.insert(0, subs)
+
+        result.args = args
+        return result
+
+    def _visitSubscriptSubs(self, subs):
+        if len(subs) == 1:
+            return self.dispatch(subs[0])
+        else:
+            return Invoke(expr=LoadName(name="Tuple"), args=self.dispatches(subs))
+
+    def visitSubscript(self, subscript, in_sequence=0):
+        value = self._visitAssNameOrAttr(subscript, in_sequence)
+        subs = self._visitSubscriptSubs(subscript.subs)
+        return self._visitSubscript(self.dispatch(subscript.expr), subs, subscript.flags, value)
+
     # Invocation and subprogram transformations.
 
     def _visitFunction(self, function, subprogram):
diff -r fbac48c3637b -r e51ec15e7979 tests/augassign.py
--- a/tests/augassign.py	Sat Jul 15 00:47:24 2006 +0200
+++ b/tests/augassign.py	Sat Jul 15 01:21:58 2006 +0200
@@ -3,3 +3,5 @@
 a.b().c += d
 a[b:c] += d
 a.b[c:] += d
+a[b] += c
+a[b,c] += d
diff -r fbac48c3637b -r e51ec15e7979 tests/subscript.py
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/tests/subscript.py	Sat Jul 15 01:21:58 2006 +0200
@@ -0,0 +1,3 @@
+a[x]
+p[q] = s
+p[q,r] = s