Fix set_thetalim((min, max)).

anntzer · anntzer · commit 2431ecb83621 · 2021-04-16T17:19:45.000+02:00
Trying `set_xlim()` and undoing the changes if the range is wrong is
perhaps not the most elegant, but much easier than parsing the various
ways `*args, **kwargs` can be passed in.  Also this should be marginally
faster if the inputs are valid (and slower if they are invalid).
diff --git a/lib/matplotlib/projections/polar.py b/lib/matplotlib/projections/polar.py
@@ -1041,29 +1041,20 @@ def set_thetalim(self, *args, **kwargs):
         wrapped in to the range :math:`[0, 2\pi]` (in radians), so for example
         it is possible to do ``set_thetalim(-np.pi / 2, np.pi / 2)`` to have
         an axes symmetric around 0. A ValueError is raised if the absolute
-        angle difference is larger than :math:`2\pi`.
+        angle difference is larger than a full circle.
         """
-        thetamin = None
-        thetamax = None
-        left = None
-        right = None
-
-        if len(args) == 2:
-            if args[0] is not None and args[1] is not None:
-                left, right = args
-                if abs(right - left) > 2 * np.pi:
-                    raise ValueError('The angle range must be <= 2 pi')
-
+        orig_lim = self.get_xlim()  # in radians
         if 'thetamin' in kwargs:
-            thetamin = np.deg2rad(kwargs.pop('thetamin'))
+            kwargs['xmin'] = np.deg2rad(kwargs.pop('thetamin'))
         if 'thetamax' in kwargs:
-            thetamax = np.deg2rad(kwargs.pop('thetamax'))
-
-        if thetamin is not None and thetamax is not None:
-            if abs(thetamax - thetamin) > 2 * np.pi:
-                raise ValueError('The angle range must be <= 360 degrees')
-        return tuple(np.rad2deg(self.set_xlim(left=left, right=right,
-                                              xmin=thetamin, xmax=thetamax)))
+            kwargs['xmax'] = np.deg2rad(kwargs.pop('thetamax'))
+        new_min, new_max = self.set_xlim(*args, **kwargs)
+        # Parsing all permutations of *args, **kwargs is tricky; it is simpler
+        # to let set_xlim() do it and then validate the limits.
+        if abs(new_max - new_min) > 2 * np.pi:
+            self.set_xlim(orig_lim)  # un-accept the change
+            raise ValueError("The angle range must be less than a full circle")
+        return tuple(np.rad2deg((new_min, new_max)))
 
     def set_theta_offset(self, offset):
         """
diff --git a/lib/matplotlib/tests/test_polar.py b/lib/matplotlib/tests/test_polar.py
@@ -343,8 +343,17 @@ def test_thetalim_valid_invalid():
     ax = plt.subplot(projection='polar')
     ax.set_thetalim(0, 2 * np.pi)  # doesn't raise.
     ax.set_thetalim(thetamin=800, thetamax=440)  # doesn't raise.
-    with pytest.raises(ValueError, match='The angle range must be <= 2 pi'):
+    with pytest.raises(ValueError,
+                       match='angle range must be less than a full circle'):
         ax.set_thetalim(0, 3 * np.pi)
     with pytest.raises(ValueError,
-                       match='The angle range must be <= 360 degrees'):
+                       match='angle range must be less than a full circle'):
         ax.set_thetalim(thetamin=800, thetamax=400)
+
+
+def test_thetalim_args():
+    ax = plt.subplot(projection='polar')
+    ax.set_thetalim(0, 1)
+    assert tuple(np.radians((ax.get_thetamin(), ax.get_thetamax()))) == (0, 1)
+    ax.set_thetalim((2, 3))
+    assert tuple(np.radians((ax.get_thetamin(), ax.get_thetamax()))) == (2, 3)
