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

12 files changed, 88 insertions, 10 deletions

diff --git a/Lib/importlib/_bootstrap.py b/Lib/importlib/_bootstrap.py
index beaa9b3..6b8c9ea 100644
--- a/Lib/importlib/_bootstrap.py
+++ b/Lib/importlib/_bootstrap.py
@@ -419,12 +419,13 @@ def _call_with_frames_removed(f, *args, **kwds):
 #     Python 3.4a4  3290 (changes to __qualname__ computation)
 #     Python 3.4a4  3300 (more changes to __qualname__ computation)
 #     Python 3.4rc2 3310 (alter __qualname__ computation)
+#     Python 3.5a0  3320 (matrix multiplication operator)
 #
 # MAGIC must change whenever the bytecode emitted by the compiler may no
 # longer be understood by older implementations of the eval loop (usually
 # due to the addition of new opcodes).
 
-MAGIC_NUMBER = (3310).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3320).to_bytes(2, 'little') + b'\r\n'
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 
 _PYCACHE = '__pycache__'
diff --git a/Lib/opcode.py b/Lib/opcode.py
index 0bd1ee6..bfd3c4d 100644
--- a/Lib/opcode.py
+++ b/Lib/opcode.py
@@ -70,6 +70,9 @@ def_op('UNARY_NOT', 12)
 
 def_op('UNARY_INVERT', 15)
 
+def_op('BINARY_MATRIX_MULTIPLY', 16)
+def_op('INPLACE_MATRIX_MULTIPLY', 17)
+
 def_op('BINARY_POWER', 19)
 def_op('BINARY_MULTIPLY', 20)
 
diff --git a/Lib/operator.py b/Lib/operator.py
index b60349f..856036d 100644
--- a/Lib/operator.py
+++ b/Lib/operator.py
@@ -105,6 +105,10 @@ def mul(a, b):
     "Same as a * b."
     return a * b
 
+def matmul(a, b):
+    "Same as a @ b."
+    return a @ b
+
 def neg(a):
     "Same as -a."
     return -a
@@ -326,6 +330,11 @@ def imul(a, b):
     a *= b
     return a
 
+def imatmul(a, b):
+    "Same as a @= b."
+    a @= b
+    return a
+
 def ior(a, b):
     "Same as a |= b."
     a |= b
@@ -383,6 +392,7 @@ __invert__ = invert
 __lshift__ = lshift
 __mod__ = mod
 __mul__ = mul
+__matmul__ = matmul
 __neg__ = neg
 __or__ = or_
 __pos__ = pos
@@ -403,6 +413,7 @@ __ifloordiv__ = ifloordiv
 __ilshift__ = ilshift
 __imod__ = imod
 __imul__ = imul
+__imatmul__ = imatmul
 __ior__ = ior
 __ipow__ = ipow
 __irshift__ = irshift
diff --git a/Lib/test/test_augassign.py b/Lib/test/test_augassign.py
index 9a59c58..19b7687 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 408f12c..ba7f2c4 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 8bb7d6a..e65edb2 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 bba8820..a7bad2d 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 ab58a98..1bd0391 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 b82358e..a809fd7 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 38611a7..8f74a06 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,
diff --git a/Lib/token.py b/Lib/token.py
index 7470c8c..bdfcec8 100644
--- a/Lib/token.py
+++ b/Lib/token.py
@@ -60,11 +60,12 @@ DOUBLESTAREQUAL = 46
 DOUBLESLASH = 47
 DOUBLESLASHEQUAL = 48
 AT = 49
-RARROW = 50
-ELLIPSIS = 51
-OP = 52
-ERRORTOKEN = 53
-N_TOKENS = 54
+ATEQUAL = 50
+RARROW = 51
+ELLIPSIS = 52
+OP = 53
+ERRORTOKEN = 54
+N_TOKENS = 55
 NT_OFFSET = 256
 #--end constants--
 
diff --git a/Lib/tokenize.py b/Lib/tokenize.py
index 98e9122..742abd1 100644
--- a/Lib/tokenize.py
+++ b/Lib/tokenize.py
@@ -91,7 +91,8 @@ EXACT_TOKEN_TYPES = {
     '**=': DOUBLESTAREQUAL,
     '//':  DOUBLESLASH,
     '//=': DOUBLESLASHEQUAL,
-    '@':   AT
+    '@':   AT,
+    '@=':  ATEQUAL,
 }
 
 class TokenInfo(collections.namedtuple('TokenInfo', 'type string start end line')):
@@ -150,7 +151,7 @@ String = group(StringPrefix + r"'[^\n'\\]*(?:\\.[^\n'\\]*)*'",
 # recognized as two instances of =).
 Operator = group(r"\*\*=?", r">>=?", r"<<=?", r"!=",
                  r"//=?", r"->",
-                 r"[+\-*/%&|^=<>]=?",
+                 r"[+\-*/%&@|^=<>]=?",
                  r"~")
 
 Bracket = '[][(){}]'