Merge pull request MIT-LCP#308 from biggates/bugfix/plot/create_figure

briangow · web-flow · commit 76eab3592c7b · 2021-06-21T15:00:15.000-04:00
bugfix: unable to plot single channel record
diff --git a/tests/test_plot.py b/tests/test_plot.py
@@ -0,0 +1,31 @@
+import wfdb
+from wfdb.plot import plot
+import unittest
+
+class TestPlot(unittest.TestCase):
+
+    def test_get_plot_dims(self):
+        sampfrom = 0
+        sampto = 3000
+        record = wfdb.rdrecord('sample-data/100', physical=True, sampfrom=sampfrom, sampto=sampto)
+        ann = wfdb.rdann('sample-data/100', 'atr', sampfrom=sampfrom, sampto=sampto)
+        sig_len, n_sig, n_annot, n_subplots = plot.get_plot_dims(signal=record.p_signal, ann_samp=[ann.sample])
+
+        assert sig_len == sampto - sampfrom
+        assert n_sig == record.n_sig
+        assert n_annot == 1
+        assert n_subplots == record.n_sig
+
+    def test_create_figure_single_subplots(self):
+        n_subplots = 1
+        fig, axes = plot.create_figure(n_subplots, sharex=True, sharey=True, figsize=None)
+        assert fig is not None
+        assert axes is not None
+        assert len(axes) == n_subplots
+
+    def test_create_figure_multiple_subplots(self):
+        n_subplots = 5
+        fig, axes = plot.create_figure(n_subplots, sharex=True, sharey=True, figsize=None)
+        assert fig is not None
+        assert axes is not None
+        assert len(axes) == n_subplots
diff --git a/wfdb/plot/plot.py b/wfdb/plot/plot.py
@@ -230,6 +230,8 @@ def create_figure(n_subplots, sharex, sharey, figsize):
     fig, axes = plt.subplots(
         nrows=n_subplots, ncols=1, sharex=sharex, sharey=sharey, figsize=figsize
     )
+    if n_subplots == 1:
+        axes = [axes]
     return fig, axes
 
 

Original file line number	Diff line number	Diff line change
`@@ -230,6 +230,8 @@ def create_figure(n_subplots, sharex, sharey, figsize):`
`230`	`230`	`fig, axes = plt.subplots(`
`231`	`231`	`nrows=n_subplots, ncols=1, sharex=sharex, sharey=sharey, figsize=figsize`
`232`	`232`	`)`
	`233`	`+ if n_subplots == 1:`
	`234`	`+ axes = [axes]`
`233`	`235`	`return fig, axes`
`234`	`236`
`235`	`237`