small fixes based on @sawyerbfuller, @bnavigator comments

murrayrm · murrayrm · commit 078e02856d15 · 2021-02-10T22:02:09.000-08:00
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -688,8 +688,7 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
         'nyquist', 'indent_direction', kwargs, _nyquist_defaults, pop=True)
 
     # If argument was a singleton, turn it into a list
-    isscalar = not hasattr(syslist, '__iter__')
-    if isscalar:
+    if not hasattr(syslist, '__iter__'):
         syslist = (syslist,)
 
     # Decide whether to go above Nyquist frequency
@@ -844,11 +843,12 @@ def nyquist_plot(syslist, omega=None, plot=True, omega_limits=None,
         ax.set_ylabel("Imaginary axis")
         ax.grid(color="lightgray")
 
+    # "Squeeze" the results
+    if len(syslist) == 1:
+        counts, contours = counts[0], contours[0]
+
     # Return counts and (optionally) the contour we used
-    if return_contour:
-        return (counts[0], contours[0]) if isscalar else (counts, contours)
-    else:
-        return counts[0] if isscalar else counts
+    return (counts, contours) if return_contour else counts
 
 
 # Internal function to add arrows to a curve
@@ -1101,7 +1101,7 @@ def _default_frequency_range(syslist, Hz=None, number_of_samples=None,
     freq_interesting = []
 
     # detect if single sys passed by checking if it is sequence-like
-    if not getattr(syslist, '__iter__', False):
+    if not hasattr(syslist, '__iter__'):
         syslist = (syslist,)
 
     for sys in syslist:
diff --git a/control/matlab/wrappers.py b/control/matlab/wrappers.py
@@ -59,7 +59,7 @@ def bode(*args, **kwargs):
     from ..freqplot import bode_plot
 
     # If first argument is a list, assume python-control calling format
-    if (getattr(args[0], '__iter__', False)):
+    if hasattr(args[0], '__iter__'):
         return bode_plot(*args, **kwargs)
 
     # Parse input arguments
@@ -97,7 +97,7 @@ def nyquist(*args, **kwargs):
     from ..freqplot import nyquist_plot
 
     # If first argument is a list, assume python-control calling format
-    if (getattr(args[0], '__iter__', False)):
+    if hasattr(args[0], '__iter__'):
         return nyquist_plot(*args, **kwargs)
 
     # Parse arguments
diff --git a/control/tests/nyquist_test.py b/control/tests/nyquist_test.py
@@ -17,6 +17,7 @@
 # In interactive mode, turn on ipython interactive graphics
 plt.ion()
 
+
 # Utility function for counting unstable poles of open loop (P in FBS)
 def _P(sys, indent='right'):
     if indent == 'right':
@@ -29,19 +30,14 @@ def _P(sys, indent='right'):
     else:
         raise TypeError("unknown indent value")
 
+
 # Utility function for counting unstable poles of closed loop (Z in FBS)
 def _Z(sys):
     return (sys.feedback().pole().real >= 0).sum()
 
-# Decorator to close figures when done with test (to avoid matplotlib warning)
-@pytest.fixture(scope="function")
-def figure_cleanup():
-    plt.close('all')
-    yield
-    plt.close('all')
 
 # Basic tests
-@pytest.mark.usefixtures("figure_cleanup")
+@pytest.mark.usefixtures("mplcleanup")
 def test_nyquist_basic():
     # Simple Nyquist plot
     sys = ct.rss(5, 1, 1)
@@ -116,7 +112,7 @@ def test_nyquist_basic():
 
 
 # Some FBS examples, for comparison
-@pytest.mark.usefixtures("figure_cleanup")
+@pytest.mark.usefixtures("mplcleanup")
 def test_nyquist_fbs_examples():
     s = ct.tf('s')
 
@@ -158,7 +154,7 @@ def test_nyquist_fbs_examples():
     1, 2, 3, 4,                 # specified number of arrows
     [0.1, 0.5, 0.9],            # specify arc lengths
 ])
-@pytest.mark.usefixtures("figure_cleanup")
+@pytest.mark.usefixtures("mplcleanup")
 def test_nyquist_arrows(arrows):
     sys = ct.tf([1.4], [1, 2, 1]) * ct.tf(*ct.pade(1, 4))
     plt.figure();
@@ -167,7 +163,7 @@ def test_nyquist_arrows(arrows):
     assert _Z(sys) == count + _P(sys)
 
 
-@pytest.mark.usefixtures("figure_cleanup")
+@pytest.mark.usefixtures("mplcleanup")
 def test_nyquist_encirclements():
     # Example 14.14: effect of friction in a cart-pendulum system
     s = ct.tf('s')
@@ -192,7 +188,7 @@ def test_nyquist_encirclements():
     assert _Z(sys) == count + _P(sys)
 
 
-@pytest.mark.usefixtures("figure_cleanup")
+@pytest.mark.usefixtures("mplcleanup")
 def test_nyquist_indent():
     # FBS Figure 10.10
     s = ct.tf('s')
