PEP 465: a dedicated infix operator for matrix multiplication (closes #21176)

7 files changed, 63 insertions, 2 deletions

diff --git a/Lib/test/test_augassign.py b/Lib/test/test_augassign.py
index 9a59c58ec0..19b76874be 100644
--- a/Lib/test/test_augassign.py
+++ b/Lib/test/test_augassign.py
@@ -136,6 +136,14 @@ class AugAssignTest(unittest.TestCase):
                 output.append("__imul__ called")
                 return self
 
+            def __matmul__(self, val):
+                output.append("__matmul__ called")
+            def __rmatmul__(self, val):
+                output.append("__rmatmul__ called")
+            def __imatmul__(self, val):
+                output.append("__imatmul__ called")
+                return self
+
             def __div__(self, val):
                 output.append("__div__ called")
             def __rdiv__(self, val):
@@ -233,6 +241,10 @@ class AugAssignTest(unittest.TestCase):
         1 * x
         x *= 1
 
+        x @ 1
+        1 @ x
+        x @= 1
+
         x / 1
         1 / x
         x /= 1
@@ -279,6 +291,9 @@ __isub__ called
 __mul__ called
 __rmul__ called
 __imul__ called
+__matmul__ called
+__rmatmul__ called
+__imatmul__ called
 __truediv__ called
 __rtruediv__ called
 __itruediv__ called
diff --git a/Lib/test/test_capi.py b/Lib/test/test_capi.py
index 408f12c324..ba7f2c4c6d 100644
--- a/Lib/test/test_capi.py
+++ b/Lib/test/test_capi.py
@@ -150,6 +150,23 @@ class CAPITest(unittest.TestCase):
         self.assertEqual(_testcapi.docstring_with_signature_and_extra_newlines.__text_signature__,
             "($module, /, parameter)")
 
+    def test_c_type_with_matrix_multiplication(self):
+        M = _testcapi.matmulType
+        m1 = M()
+        m2 = M()
+        self.assertEqual(m1 @ m2, ("matmul", m1, m2))
+        self.assertEqual(m1 @ 42, ("matmul", m1, 42))
+        self.assertEqual(42 @ m1, ("matmul", 42, m1))
+        o = m1
+        o @= m2
+        self.assertEqual(o, ("imatmul", m1, m2))
+        o = m1
+        o @= 42
+        self.assertEqual(o, ("imatmul", m1, 42))
+        o = 42
+        o @= m1
+        self.assertEqual(o, ("matmul", 42, m1))
+
 
 @unittest.skipUnless(threading, 'Threading required for this test.')
 class TestPendingCalls(unittest.TestCase):
diff --git a/Lib/test/test_descr.py b/Lib/test/test_descr.py
index 8bb7d6a474..e65edb2f8c 100644
--- a/Lib/test/test_descr.py
+++ b/Lib/test/test_descr.py
@@ -4160,6 +4160,7 @@ order (MRO) for bases """
                 ('__add__',      'x + y',                   'x += y'),
                 ('__sub__',      'x - y',                   'x -= y'),
                 ('__mul__',      'x * y',                   'x *= y'),
+                ('__matmul__',   'x @ y',                   'x @= y'),
                 ('__truediv__',  'operator.truediv(x, y)',  None),
                 ('__floordiv__', 'operator.floordiv(x, y)', None),
                 ('__div__',      'x / y',                   'x /= y'),
diff --git a/Lib/test/test_grammar.py b/Lib/test/test_grammar.py
index bba8820ab1..a7bad2d2d9 100644
--- a/Lib/test/test_grammar.py
+++ b/Lib/test/test_grammar.py
@@ -985,6 +985,20 @@ class GrammarTests(unittest.TestCase):
         self.assertFalse((False is 2) is 3)
         self.assertFalse(False is 2 is 3)
 
+    def test_matrix_mul(self):
+        # This is not intended to be a comprehensive test, rather just to be few
+        # samples of the @ operator in test_grammar.py.
+        class M:
+            def __matmul__(self, o):
+                return 4
+            def __imatmul__(self, o):
+                self.other = o
+                return self
+        m = M()
+        self.assertEqual(m @ m, 4)
+        m @= 42
+        self.assertEqual(m.other, 42)
+
 
 def test_main():
     run_unittest(TokenTests, GrammarTests)
diff --git a/Lib/test/test_operator.py b/Lib/test/test_operator.py
index ab58a98365..1bd0391ee2 100644
--- a/Lib/test/test_operator.py
+++ b/Lib/test/test_operator.py
@@ -203,6 +203,15 @@ class OperatorTestCase:
         self.assertRaises(TypeError, operator.mul, None, None)
         self.assertTrue(operator.mul(5, 2) == 10)
 
+    def test_matmul(self):
+        operator = self.module
+        self.assertRaises(TypeError, operator.matmul)
+        self.assertRaises(TypeError, operator.matmul, 42, 42)
+        class M:
+            def __matmul__(self, other):
+                return other - 1
+        self.assertEqual(M() @ 42, 41)
+
     def test_neg(self):
         operator = self.module
         self.assertRaises(TypeError, operator.neg)
@@ -416,6 +425,7 @@ class OperatorTestCase:
             def __ilshift__  (self, other): return "ilshift"
             def __imod__     (self, other): return "imod"
             def __imul__     (self, other): return "imul"
+            def __imatmul__  (self, other): return "imatmul"
             def __ior__      (self, other): return "ior"
             def __ipow__     (self, other): return "ipow"
             def __irshift__  (self, other): return "irshift"
@@ -430,6 +440,7 @@ class OperatorTestCase:
         self.assertEqual(operator.ilshift  (c, 5), "ilshift")
         self.assertEqual(operator.imod     (c, 5), "imod")
         self.assertEqual(operator.imul     (c, 5), "imul")
+        self.assertEqual(operator.imatmul  (c, 5), "imatmul")
         self.assertEqual(operator.ior      (c, 5), "ior")
         self.assertEqual(operator.ipow     (c, 5), "ipow")
         self.assertEqual(operator.irshift  (c, 5), "irshift")
diff --git a/Lib/test/test_sys.py b/Lib/test/test_sys.py
index b82358ea78..a809fd737f 100644
--- a/Lib/test/test_sys.py
+++ b/Lib/test/test_sys.py
@@ -952,7 +952,7 @@ class SizeofTest(unittest.TestCase):
         check(int, s)
         # (PyTypeObject + PyNumberMethods + PyMappingMethods +
         #  PySequenceMethods + PyBufferProcs + 4P)
-        s = vsize('P2n15Pl4Pn9Pn11PIP') + struct.calcsize('34P 3P 10P 2P 4P')
+        s = vsize('P2n17Pl4Pn9Pn11PIP') + struct.calcsize('34P 3P 10P 2P 4P')
         # Separate block for PyDictKeysObject with 4 entries
         s += struct.calcsize("2nPn") + 4*struct.calcsize("n2P")
         # class
diff --git a/Lib/test/test_tokenize.py b/Lib/test/test_tokenize.py
index 38611a79ee..8f74a06df5 100644
--- a/Lib/test/test_tokenize.py
+++ b/Lib/test/test_tokenize.py
@@ -464,7 +464,7 @@ Additive
 
 Multiplicative
 
-    >>> dump_tokens("x = 1//1*1/5*12%0x12")
+    >>> dump_tokens("x = 1//1*1/5*12%0x12@42")
     ENCODING   'utf-8'       (0, 0) (0, 0)
     NAME       'x'           (1, 0) (1, 1)
     OP         '='           (1, 2) (1, 3)
@@ -479,6 +479,8 @@ Multiplicative
     NUMBER     '12'          (1, 13) (1, 15)
     OP         '%'           (1, 15) (1, 16)
     NUMBER     '0x12'        (1, 16) (1, 20)
+    OP         '@'           (1, 20) (1, 21)
+    NUMBER     '42'          (1, 21) (1, 23)
 
 Unary
 
@@ -1154,6 +1156,7 @@ class TestTokenize(TestCase):
         self.assertExactTypeEqual('//', token.DOUBLESLASH)
         self.assertExactTypeEqual('//=', token.DOUBLESLASHEQUAL)
         self.assertExactTypeEqual('@', token.AT)
+        self.assertExactTypeEqual('@=', token.ATEQUAL)
 
         self.assertExactTypeEqual('a**2+b**2==c**2',
                                   NAME, token.DOUBLESTAR, NUMBER,