Add ability to compose transforms of any dimension

trananso · trananso · commit 2dfd5b0a8a94 · 2024-04-21T18:28:14.000-04:00
diff --git a/lib/matplotlib/tests/test_transforms.py b/lib/matplotlib/tests/test_transforms.py
@@ -853,7 +853,7 @@ def test_str_transform():
                 BlendedAffine(
                     IdentityTransform(),
                     IdentityTransform())),
-            CompositeAffine2D(
+            CompositeAffine(
                 Affine2D().scale(1.0),
                 Affine2D().scale(1.0))),
         PolarTransform(
@@ -876,7 +876,7 @@ def test_str_transform():
                     (0.5, 0.5),
                     TransformedBbox(
                         Bbox(x0=0.0, y0=0.0, x1=6.283185307179586, y1=1.0),
-                        CompositeAffine2D(
+                        CompositeAffine(
                             Affine2D().scale(1.0),
                             Affine2D().scale(1.0))),
                     LockableBbox(
diff --git a/lib/matplotlib/transforms.py b/lib/matplotlib/transforms.py
@@ -2466,7 +2466,7 @@ def _invalidate_internal(self, level, invalidating_node):
         super()._invalidate_internal(level, invalidating_node)
 
     def __eq__(self, other):
-        if isinstance(other, (CompositeGenericTransform, CompositeAffine2D)):
+        if isinstance(other, (CompositeGenericTransform, CompositeAffine)):
             return self is other or (self._a == other._a
                                      and self._b == other._b)
         else:
@@ -2526,7 +2526,7 @@ def inverted(self):
             self._b.inverted(), self._a.inverted())
 
 
-class CompositeAffine2D(Affine2DBase):
+class CompositeAffine(AffineImmutable):
     """
     A composite transform formed by applying transform *a* then transform *b*.
 
@@ -2536,7 +2536,7 @@ class CompositeAffine2D(Affine2DBase):
     def __init__(self, a, b, **kwargs):
         """
         Create a new composite transform that is the result of
-        applying `Affine2DBase` *a* then `Affine2DBase` *b*.
+        applying `AffineImmutable` *a* then `AffineImmutable` *b*.
 
         You will generally not call this constructor directly but write ``a +
         b`` instead, which will automatically choose the best kind of composite
@@ -2547,10 +2547,8 @@ def __init__(self, a, b, **kwargs):
         if a.output_dims != b.input_dims:
             raise ValueError("The output dimension of 'a' must be equal to "
                              "the input dimensions of 'b'")
-        self.input_dims = a.input_dims
-        self.output_dims = b.output_dims
+        super().__init__(dims=a.output_dims, **kwargs)
 
-        super().__init__(**kwargs)
         self._a = a
         self._b = b
         self.set_children(a, b)
@@ -2579,6 +2577,11 @@ def get_matrix(self):
         return self._mtx
 
 
+@_api.deprecated("3.9", alternative="CompositeAffine")
+class CompositeAffine2D(CompositeAffine):
+    pass
+
+
 def composite_transform_factory(a, b):
     """
     Create a new composite transform that is the result of applying
@@ -2602,7 +2605,7 @@ def composite_transform_factory(a, b):
     elif isinstance(b, IdentityTransform):
         return a
     elif isinstance(a, Affine2D) and isinstance(b, Affine2D):
-        return CompositeAffine2D(a, b)
+        return CompositeAffine(a, b)
     return CompositeGenericTransform(a, b)
 
 
diff --git a/lib/matplotlib/transforms.pyi b/lib/matplotlib/transforms.pyi
@@ -282,11 +282,14 @@ class CompositeGenericTransform(Transform):
     pass_through: bool
     def __init__(self, a: Transform, b: Transform, **kwargs) -> None: ...
 
-class CompositeAffine2D(Affine2DBase):
-    def __init__(self, a: Affine2DBase, b: Affine2DBase, **kwargs) -> None: ...
+class CompositeAffine(AffineImmutable):
+    def __init__(self, a: AffineImmutable, b: AffineImmutable, **kwargs) -> None: ...
     @property
     def depth(self) -> int: ...
 
+class CompositeAffine2D(CompositeAffine):
+    pass
+
 def composite_transform_factory(a: Transform, b: Transform) -> Transform: ...
 
 class BboxTransform(AffineImmutable):
