Add matmul to DataFrame, Series

Author: bnaul

doc/source/whatsnew/v0.23.0.txt

@@ -324,6 +324,7 @@ Other Enhancements
 - ``IntervalIndex.astype`` now supports conversions between subtypes when passed an ``IntervalDtype`` (:issue:`19197`)
 - :class:`IntervalIndex` and its associated constructor methods (``from_arrays``, ``from_breaks``, ``from_tuples``) have gained a ``dtype`` parameter (:issue:`19262`)
 - Added :func:`SeriesGroupBy.is_monotonic_increasing` and :func:`SeriesGroupBy.is_monotonic_decreasing` (:issue:`17015`)
+- :class:`DataFrame` and :class:`Series` now support matrix multiplication (```@```) operator (:issue:`10259`) for Python>=3.5
 
 .. _whatsnew_0230.api_breaking:
 

pandas/core/frame.py

@@ -830,7 +830,8 @@ def __len__(self):
 
     def dot(self, other):
         """
-        Matrix multiplication with DataFrame or Series objects
+        Matrix multiplication with DataFrame or Series objects.  Can also be
+        called using `self @ other` in Python >= 3.5.
 
         Parameters
         ----------
@@ -872,6 +873,13 @@ def dot(self, other):
         else:  # pragma: no cover
             raise TypeError('unsupported type: %s' % type(other))
 
+    def __matmul__(self, other):
+        """ Matrix multiplication using binary `@` operator in Python>=3.5 """
+        try:
+            return self.dot(other)
+        except TypeError:
+            return NotImplemented
+
     # ----------------------------------------------------------------------
     # IO methods (to / from other formats)
 

pandas/core/series.py

@@ -1559,7 +1559,7 @@ def autocorr(self, lag=1):
     def dot(self, other):
         """
         Matrix multiplication with DataFrame or inner-product with Series
-        objects
+        objects. Can also be called using `self @ other` in Python >= 3.5.
 
         Parameters
         ----------
@@ -1598,6 +1598,13 @@ def dot(self, other):
         else:  # pragma: no cover
             raise TypeError('unsupported type: %s' % type(other))
 
+    def __matmul__(self, other):
+        """ Matrix multiplication using binary `@` operator in Python>=3.5 """
+        try:
+            return self.dot(other)
+        except TypeError:
+            return NotImplemented
+
     @Substitution(klass='Series')
     @Appender(base._shared_docs['searchsorted'])
     @deprecate_kwarg(old_arg_name='v', new_arg_name='value')

pandas/tests/frame/test_analytics.py

@@ -2075,41 +2075,42 @@ def test_clip_with_na_args(self):
                               self.frame)
 
     # Matrix-like
-
+    @pytest.mark.parametrize('dot_fn', [DataFrame.dot, DataFrame.__matmul__])
     def test_dot(self):
+        # __matmul__ test is for GH #10259
         a = DataFrame(np.random.randn(3, 4), index=['a', 'b', 'c'],
                       columns=['p', 'q', 'r', 's'])
         b = DataFrame(np.random.randn(4, 2), index=['p', 'q', 'r', 's'],
                       columns=['one', 'two'])
 
-        result = a.dot(b)
+        result = dot_fn(a, b)
         expected = DataFrame(np.dot(a.values, b.values),
                              index=['a', 'b', 'c'],
                              columns=['one', 'two'])
         # Check alignment
         b1 = b.reindex(index=reversed(b.index))
-        result = a.dot(b)
+        result = dot_fn(a, b)
         tm.assert_frame_equal(result, expected)
 
         # Check series argument
-        result = a.dot(b['one'])
+        result = dot_fn(a, b['one'])
         tm.assert_series_equal(result, expected['one'], check_names=False)
         assert result.name is None
 
-        result = a.dot(b1['one'])
+        result = dot_fn(a, b1['one'])
         tm.assert_series_equal(result, expected['one'], check_names=False)
         assert result.name is None
 
         # can pass correct-length arrays
         row = a.iloc[0].values
 
-        result = a.dot(row)
-        exp = a.dot(a.iloc[0])
+        result = dot_fn(a, row)
+        exp = dot_fn(a, a.iloc[0])
         tm.assert_series_equal(result, exp)
 
         with tm.assert_raises_regex(ValueError,
                                     'Dot product shape mismatch'):
-            a.dot(row[:-1])
+            dot_fn(a, row[:-1])
 
         a = np.random.rand(1, 5)
         b = np.random.rand(5, 1)
@@ -2119,14 +2120,19 @@ def test_dot(self):
         B = DataFrame(b)  # noqa
 
         # it works
-        result = A.dot(b)
+        result = dot_fn(A, b)
 
         # unaligned
         df = DataFrame(randn(3, 4), index=[1, 2, 3], columns=lrange(4))
         df2 = DataFrame(randn(5, 3), index=lrange(5), columns=[1, 2, 3])
 
         with tm.assert_raises_regex(ValueError, 'aligned'):
-            df.dot(df2)
+            dot_fn(df, df2)
+
+        # invalid type
+        with tm.assert_raises_regex(TypeError, 'unsupported'):
+            x = 1
+            dot_fn(df, 1)
 
 
 @pytest.fixture

pandas/tests/series/test_analytics.py

@@ -895,31 +895,34 @@ def test_count(self):
         ts.iloc[[0, 3, 5]] = nan
         assert_series_equal(ts.count(level=1), right - 1)
 
-    def test_dot(self):
+    @pytest.mark.parametrize('dot_fn', [Series.dot, Series.__matmul__])
+    def test_dot(self, dot_fn):
+        # __matmul__ test is for GH #10259
         a = Series(np.random.randn(4), index=['p', 'q', 'r', 's'])
         b = DataFrame(np.random.randn(3, 4), index=['1', '2', '3'],
                       columns=['p', 'q', 'r', 's']).T
 
-        result = a.dot(b)
+        result = dot_fn(a, b)
         expected = Series(np.dot(a.values, b.values), index=['1', '2', '3'])
         assert_series_equal(result, expected)
 
         # Check index alignment
         b2 = b.reindex(index=reversed(b.index))
-        result = a.dot(b)
+        result = dot_fn(a, b)
         assert_series_equal(result, expected)
 
         # Check ndarray argument
-        result = a.dot(b.values)
+        result = dot_fn(a, b.values)
         assert np.all(result == expected.values)
-        assert_almost_equal(a.dot(b['2'].values), expected['2'])
+        assert_almost_equal(dot_fn(a, b['2'].values), expected['2'])
 
         # Check series argument
-        assert_almost_equal(a.dot(b['1']), expected['1'])
-        assert_almost_equal(a.dot(b2['1']), expected['1'])
+        assert_almost_equal(dot_fn(a, b['1']), expected['1'])
+        assert_almost_equal(dot_fn(a, b2['1']), expected['1'])
 
         pytest.raises(Exception, a.dot, a.values[:3])
         pytest.raises(ValueError, a.dot, b.T)
+        pytest.raises(TypeError, a.dot, 1)
 
     def test_value_counts_nunique(self):