Repository: rasterio/affine

Issue:
Add support for e.g. numpy.meshgrid on the right hand side of `@`
Rasterio's warp tests use the product of `Affine` and `numpy.meshgrid()` and I've seen this in Rasterio user code.

```python
def test_geoloc_warp_array(path_rgb_byte_tif, tmp_path):
    """Warp an array using external geolocation arrays."""
    with rasterio.open(path_rgb_byte_tif) as src:
        xs, ys = src.transform * np.meshgrid(
            np.arange(src.width), np.arange(src.height)
        )
```

Hints from issue comments:

Solution patch:
diff --git a/src/affine.py b/src/affine.py
index cbe823c..1a8318c 100644
--- a/src/affine.py
+++ b/src/affine.py
@@ -555,7 +555,7 @@ class Affine:
 
         Apply the transform using matrix multiplication, creating
         a resulting object of the same type.  A transform may be applied
-        to another transform or vector array.
+        to another transform, a vector, vector array, or shape.
 
         Parameters
         ----------
@@ -563,7 +563,7 @@ class Affine:
 
         Returns
         -------
-        Affine or a tuple of two or three floats
+        Affine or a tuple of two or three items
         """
         sa, sb, sc, sd, se, sf = self[:6]
         if isinstance(other, Affine):
@@ -576,23 +576,29 @@ class Affine:
                 sd * ob + se * oe,
                 sd * oc + se * of + sf,
             )
-        # vector of 2 or 3 values
+        # vector of 2 or 3 items
         try:
-            other = tuple(map(float, other))
+            num_items = len(other)
         except (TypeError, ValueError):
             return NotImplemented
-        num_values = len(other)
-        if num_values == 2:
+        if num_items == 2:
             vx, vy = other
-        elif num_values == 3:
+        elif num_items == 3:
             vx, vy, vw = other
-            if vw != 1.0:
-                raise ValueError("third value must be 1.0")
+            vw_eq_one = vw == 1.0
+            try:
+                is_eq_one = bool(vw_eq_one)
+                msg = "third value must be 1.0"
+            except ValueError:
+                is_eq_one = (vw_eq_one).all()
+                msg = "third values must all be 1.0"
+            if not is_eq_one:
+                raise ValueError(msg)
         else:
-            raise TypeError("expected vector of 2 or 3 values")
+            raise TypeError("expected vector of 2 or 3 items")
         px = vx * sa + vy * sb + sc
         py = vx * sd + vy * se + sf
-        if num_values == 2:
+        if num_items == 2:
             return (px, py)
         return (px, py, vw)
 
@@ -617,7 +623,7 @@ class Affine:
 
         Returns
         -------
-        Affine or a tuple of two floats
+        Affine or a tuple of two items
         """
         # TODO: consider enabling this for 3.1
         # warnings.warn(


Test patch:
diff --git a/tests/test_numpy.py b/tests/test_numpy.py
index 74bc49f..c7bfc29 100644
--- a/tests/test_numpy.py
+++ b/tests/test_numpy.py
@@ -2,7 +2,7 @@
 
 import pytest
 
-from affine import Affine, identity
+from affine import Affine
 
 try:
     import numpy as np
@@ -68,18 +68,23 @@ def test_linalg():
     testing.assert_allclose(np.linalg.inv(ar), expected_inv)
 
 
-def test_matmul():
+def test_matmul_3x3_array():
     A = Affine(2, 0, 3, 0, 3, 2)
     Ar = np.array(A)
 
     # matrix @ matrix = matrix
-    res = A @ identity
+    res = A @ Affine.identity()
     assert isinstance(res, Affine)
     testing.assert_equal(res, Ar)
     res = Ar @ np.eye(3)
     assert isinstance(res, np.ndarray)
     testing.assert_equal(res, Ar)
 
+
+def test_matmul_vector():
+    A = Affine(2, 0, 3, 0, 3, 2)
+    Ar = np.array(A)
+
     # matrix @ vector = vector
     v = (2, 3, 1)
     vr = np.array(v)
@@ -93,3 +98,63 @@ def test_matmul():
     res = Ar @ vr
     assert isinstance(res, np.ndarray)
     testing.assert_equal(res, expected_p)
+
+
+def test_matmul_items():
+    A = Affine(2, 0, 3, 0, 3, 2) @ Affine.rotation(45)
+    shape = (4, 5)
+    vx, vy = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
+    expected_px = np.array(
+        [
+            [3.0, 4.4142137, 5.8284273, 7.2426405, 8.656855],
+            [1.5857865, 3.0, 4.4142137, 5.8284273, 7.2426405],
+            [0.17157288, 1.5857865, 3.0, 4.4142137, 5.8284273],
+            [-1.2426407, 0.17157288, 1.5857865, 3.0, 4.4142137],
+        ]
+    )
+    expected_py = np.array(
+        [
+            [2.0, 4.1213202, 6.2426405, 8.363961, 10.485281],
+            [4.1213202, 6.2426405, 8.363961, 10.485281, 12.606602],
+            [6.2426405, 8.363961, 10.485281, 12.606602, 14.727922],
+            [8.363961, 10.485281, 12.606602, 14.727922, 16.849243],
+        ]
+    )
+    px, py = A @ (vx, vy)
+    testing.assert_allclose(px, expected_px)
+    testing.assert_allclose(py, expected_py)
+
+    px, py, pw = A @ (vx, vy, 1)
+    testing.assert_allclose(px, expected_px)
+    testing.assert_allclose(py, expected_py)
+    assert pw == 1
+
+    px, py, pw = A @ (vx, vy, np.ones(shape))
+    testing.assert_allclose(px, expected_px)
+    testing.assert_allclose(py, expected_py)
+    testing.assert_array_equal(pw, np.ones(shape))
+
+    # see GH-125 with mul `*` op
+    px, py = A * (vx, vy)
+    testing.assert_allclose(px, expected_px)
+    testing.assert_allclose(py, expected_py)
+
+
+def test_matmul_item_errors():
+    shape = (4, 5)
+    vx, vy = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
+    with pytest.raises(ValueError, match="third value must be 1.0"):
+        Affine.identity() @ (vx, vy, 2)
+    with pytest.raises(ValueError, match="third values must all be 1.0"):
+        Affine.identity() @ (vx, vy, np.zeros(shape))
+    with pytest.raises(TypeError, match="2 or 3 items"):
+        Affine.identity() @ (vx, vy, np.ones(shape), np.ones(shape))
+
+
+def test_mul_item_errors():
+    shape = (4, 5)
+    vx, vy = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
+    with pytest.raises(TypeError):
+        Affine.identity() * (vx, vy, 1)
+    with pytest.raises(TypeError):
+        Affine.identity() * (vx, vy, np.zeros(shape))
diff --git a/tests/test_transform.py b/tests/test_transform.py
index 441e717..e75b9db 100644
--- a/tests/test_transform.py
+++ b/tests/test_transform.py
@@ -654,5 +654,5 @@ def test_matmul_invalid_vector():
 
 @pytest.mark.parametrize("vec", [(2.0,), (2.0, 2.0, 1.0, 1.0)])
 def test_matmul_invalid_vector_2(vec):
-    with pytest.raises(TypeError, match="2 or 3 values"):
+    with pytest.raises(TypeError, match="2 or 3 items"):
         Affine.identity() @ vec