Make augmented assignment operators variadic

NEWS.rst

@@ -11,6 +11,11 @@ Removals
 * Literal keywords are no longer parsed differently in calls to functions
   with certain names.
 
+New Features
+------------------------------
+* All augmented assignment operators (except `%=` and `^=`) now allow
+  more than two arguments.
+
 Bug Fixes
 ------------------------------
 * Statements in the second argument of `assert` are now executed.

hy/compiler.py

@@ -1265,19 +1265,21 @@ class HyASTCompiler(object):
         return ret + asty.Compare(
             expr, left=exprs[0], ops=ops, comparators=exprs[1:])
 
-    m_ops = {"+": ast.Add,
-             "/": ast.Div,
-             "//": ast.FloorDiv,
-             "*": ast.Mult,
-             "-": ast.Sub,
-             "%": ast.Mod,
-             "**": ast.Pow,
-             "<<": ast.LShift,
-             ">>": ast.RShift,
-             "|": ast.BitOr,
-             "^": ast.BitXor,
-             "&": ast.BitAnd,
-             "@": ast.MatMult}
+    # The second element of each tuple below is an aggregation operator
+    # that's used for augmented assignment with three or more arguments.
+    m_ops = {"+": (ast.Add, "+"),
+             "/": (ast.Div, "*"),
+             "//": (ast.FloorDiv, "*"),
+             "*": (ast.Mult, "*"),
+             "-": (ast.Sub, "+"),
+             "%": (ast.Mod, None),
+             "**": (ast.Pow, "**"),
+             "<<": (ast.LShift, "+"),
+             ">>": (ast.RShift, "+"),
+             "|": (ast.BitOr, "|"),
+             "^": (ast.BitXor, None),
+             "&": (ast.BitAnd, "&"),
+             "@": (ast.MatMult, "@")}
 
     @special(["+", "*", "|"], [many(FORM)])
     @special(["-", "/", "&", "@"], [oneplus(FORM)])
@@ -1302,7 +1304,7 @@ class HyASTCompiler(object):
                 # Return the argument unchanged.
                 return self.compile(args[0])
 
-        op = self.m_ops[root]
+        op = self.m_ops[root][0]
         right_associative = root == "**"
         ret = self.compile(args[-1 if right_associative else 0])
         for child in args[-2 if right_associative else 1 ::
@@ -1318,11 +1320,16 @@ class HyASTCompiler(object):
 
     a_ops = {x + "=": v for x, v in m_ops.items()}
 
-    @special(list(a_ops.keys()), [FORM, FORM])
-    def compile_augassign_expression(self, expr, root, target, value):
-        op = self.a_ops[root]
+    @special([x for x, (_, v) in a_ops.items() if v is not None], [FORM, oneplus(FORM)])
+    @special([x for x, (_, v) in a_ops.items() if v is None], [FORM, times(1, 1, FORM)])
+    def compile_augassign_expression(self, expr, root, target, values):
+        if len(values) > 1:
+            return self.compile(mkexpr(root, [target],
+                mkexpr(self.a_ops[root][1], rest=values)).replace(expr))
+
+        op = self.a_ops[root][0]
         target = self._storeize(target, self.compile(target))
-        ret = self.compile(value)
+        ret = self.compile(values[0])
         return ret + asty.AugAssign(
             expr, target=target, value=ret.force_expr, op=op())
 

tests/native_tests/operators.hy

@@ -303,3 +303,43 @@
   (assert (= (f "hello" 1) "e"))
   (assert (= (f [[1 2 3] [4 5 6] [7 8 9]] 1 2) 6))
   (assert (= (f {"x" {"y" {"z" 12}}} "x" "y" "z") 12)))
+
+
+(defn test-augassign []
+    (setv b 2  c 3  d 4)
+    (defmacro same-as [expr1 expr2 expected-value]
+      `(do
+        (setv a 4)
+        ~expr1
+        (setv expr1-value a)
+        (setv a 4)
+        ~expr2
+        (assert (= expr1-value a ~expected-value))))
+    (same-as (+= a b c d) (+= a (+ b c d)) 13)
+    (same-as (-= a b c d) (-= a (+ b c d)) -5)
+    (same-as (*= a b c d) (*= a (* b c d)) 96)
+    (same-as (**= a b c) (**= a (** b c)) 65,536)
+    (same-as (/= a b c d) (/= a (* b c d)) (/ 1 6))
+    (same-as (//= a b c d) (//= a (* b c d)) 0)
+    (same-as (<<= a b c d) (<<= a (+ b c d)) 0b10_00000_00000)
+    (same-as (>>= a b c d) (>>= a (+ b c d)) 0)
+    (same-as (&= a b c d) (&= a (& b c d)) 0)
+    (same-as (|= a b c d) (|= a (| b c d)) 0b111)
+
+    (defclass C [object]
+      (defn __init__ [self content] (setv self.content content))
+      (defn __matmul__ [self other] (C (+ self.content other.content))))
+    (setv  a (C "a")  b (C "b")  c (C "c")  d (C "d"))
+    (@= a b c d)
+    (assert (= a.content "abcd"))
+    (setv a (C "a"))
+    (@= a (@ b c d))
+    (assert (= a.content "abcd"))
+
+    (setv a 15)
+    (%= a 9)
+    (assert (= a 6))
+
+    (setv a 0b1100)
+    (^= a 0b1010)
+    (assert (= a 0b0110)))