BUG: fix GH #326, start unit testing plot methods

Author: wesm

pandas/core/frame.py

@@ -2738,7 +2738,7 @@ def clip_lower(self, threshold):
     # Plotting
 
     def plot(self, subplots=False, sharex=True, sharey=False, use_index=True,
-             figsize=None, grid=True, **kwds):  # pragma: no cover
+             figsize=None, grid=True, legend=True, ax=None, **kwds):
         """
         Make line plot of DataFrame's series with the index on the x-axis using
         matplotlib / pylab.
@@ -2764,12 +2764,15 @@ def plot(self, subplots=False, sharex=True, sharey=False, use_index=True,
         import matplotlib.pyplot as plt
 
         if subplots:
-            _, axes = plt.subplots(nrows=len(self.columns),
+            fig, axes = plt.subplots(nrows=len(self.columns),
                                    sharex=sharex, sharey=sharey,
                                    figsize=figsize)
         else:
-            fig = plt.figure(figsize=figsize)
-            ax = fig.add_subplot(111)
+            if ax is None:
+                fig = plt.figure(figsize=figsize)
+                ax = fig.add_subplot(111)
+            else:
+                fig = ax.get_figure()
 
         if use_index:
             x = self.index
@@ -2781,23 +2784,25 @@ def plot(self, subplots=False, sharex=True, sharey=False, use_index=True,
             y = self[col].values if not empty else np.zeros(x.shape)
             if subplots:
                 ax = axes[i]
-                ax.plot(x, y, 'k', label=col, **kwds)
+                ax.plot(x, y, 'k', label=str(col), **kwds)
                 ax.legend(loc='best')
             else:
-                ax.plot(x, y, label=col, **kwds)
+                ax.plot(x, y, label=str(col), **kwds)
 
             ax.grid(grid)
 
         # try to make things prettier
         try:
-            fig = plt.gcf()
             fig.autofmt_xdate()
         except Exception:
             pass
 
+        if legend and not subplots:
+            ax.legend(loc='best')
+
         plt.draw_if_interactive()
 
-    def hist(self, grid=True, **kwds):  # pragma: no cover
+    def hist(self, grid=True, **kwds):
         """
         Draw Histogram the DataFrame's series using matplotlib / pylab.
 

pandas/core/series.py

@@ -1636,7 +1636,7 @@ def isin(self, values):
 # Miscellaneous
 
     def plot(self, label=None, kind='line', use_index=True, rot=30, ax=None,
-             style='-', grid=True, **kwds):  # pragma: no cover
+             style='-', grid=True, **kwds):
         """
         Plot the input series with the index on the x-axis using matplotlib
 
@@ -1697,12 +1697,12 @@ def plot(self, label=None, kind='line', use_index=True, rot=30, ax=None,
         try:
             fig = plt.gcf()
             fig.autofmt_xdate()
-        except Exception:
+        except Exception:  # pragma: no cover
             pass
 
         plt.draw_if_interactive()
 
-    def hist(self, ax=None, grid=True, **kwds):  # pragma: no cover
+    def hist(self, ax=None, grid=True, **kwds):
         """
         Draw histogram of the input series using matplotlib
 

pandas/src/tseries.pyx

@@ -165,6 +165,18 @@ def isAllDates(ndarray[object, ndim=1] arr):
     return True
 
 def ismember(ndarray arr, set values):
+    '''
+    Checks whether
+
+    Parameters
+    ----------
+    arr : ndarray
+    values : set
+
+    Returns
+    -------
+    ismember : ndarray (boolean dtype)
+    '''
     cdef:
         Py_ssize_t i, n
         flatiter it

pandas/stats/tests/test_ols.py

@@ -47,7 +47,13 @@ class TestOLS(BaseTest):
     # TODO: Add tests for non pooled OLS.
 
     @classmethod
-    def setupClass(cls):
+    def setUpClass(cls):
+        try:
+            import matplotlib as mpl
+            mpl.use('Agg', warn=False)
+        except ImportError:
+            pass
+
         try:
             import scikits.statsmodels.api as _
         except ImportError:

pandas/tests/test_frame.py

@@ -3219,7 +3219,6 @@ def test_series_put_names(self):
             self.assertEqual(v.name, k)
 
 
-
 class TestDataFrameJoin(unittest.TestCase):
 
     def setUp(self):

pandas/tests/test_graphics.py

@@ -0,0 +1,84 @@
+import nose
+import os
+import unittest
+
+from pandas import Series, DataFrame
+import pandas.util.testing as tm
+
+import numpy as np
+
+class TestSeriesPlots(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        import sys
+        if 'IPython' in sys.modules:
+            raise nose.SkipTest
+
+        try:
+            import matplotlib as mpl
+            mpl.use('Agg', warn=False)
+        except ImportError:
+            raise nose.SkipTest
+
+    def setUp(self):
+        self.ts = tm.makeTimeSeries()
+        self.ts.name = 'ts'
+
+        self.series = tm.makeStringSeries()
+        self.series.name = 'series'
+
+    def test_plot(self):
+        _check_plot_works(self.ts.plot, label='foo')
+        _check_plot_works(self.ts.plot, use_index=False)
+        _check_plot_works(self.ts.plot, rot=0)
+        _check_plot_works(self.ts.plot, style='.')
+        _check_plot_works(self.ts[:10].plot, kind='bar')
+        _check_plot_works(self.series[:5].plot, kind='bar')
+
+    def test_hist(self):
+        _check_plot_works(self.ts.hist)
+        _check_plot_works(self.ts.hist, grid=False)
+
+
+class TestDataFramePlots(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        import sys
+        if 'IPython' in sys.modules:
+            raise nose.SkipTest
+
+        try:
+            import matplotlib as mpl
+            mpl.use('Agg', warn=False)
+        except ImportError:
+            raise nose.SkipTest
+
+    def test_plot(self):
+        pass
+
+    def test_plot_int_columns(self):
+        df = DataFrame(np.random.randn(100, 4)).cumsum()
+        _check_plot_works(df.plot, legend=True)
+
+
+PNG_PATH = 'tmp.png'
+
+def _check_plot_works(f, *args, **kwargs):
+    import matplotlib.pyplot as plt
+
+    fig = plt.gcf()
+    plt.clf()
+    ax = fig.add_subplot(211)
+    f(*args, **kwargs)
+
+    ax = fig.add_subplot(212)
+    kwargs['ax'] = ax
+    f(*args, **kwargs)
+    plt.savefig(PNG_PATH)
+    os.remove(PNG_PATH)
+
+if __name__ == '__main__':
+    nose.runmodule(argv=[__file__,'-vvs','-x','--pdb', '--pdb-failure'],
+                   exit=False)

pandas/tests/test_series.py

@@ -5,6 +5,8 @@
 import operator
 import unittest
 
+import nose
+
 from numpy import nan
 import numpy as np
 
@@ -1377,7 +1379,6 @@ def test_dropna_preserve_name(self):
         self.assertEquals(result.name, self.ts.name)
 
 if __name__ == '__main__':
-    import nose
     nose.runmodule(argv=[__file__,'-vvs','-x','--pdb', '--pdb-failure'],
                    exit=False)