/pandas/tests/test_graphics.py
Python | 125 lines | 98 code | 27 blank | 0 comment | 9 complexity | 2e023a81985f4be24c00df9746e48cd8 MD5 | raw file
Possible License(s): BSD-3-Clause
- import nose
- import os
- import string
- 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):
- df = tm.makeTimeDataFrame()
- _check_plot_works(df.plot, grid=False)
- _check_plot_works(df.plot, subplots=True)
- _check_plot_works(df.plot, subplots=True, use_index=False)
- def test_plot_bar(self):
- df = DataFrame(np.random.randn(6, 4),
- index=list(string.ascii_letters[:6]),
- columns=['one', 'two', 'three', 'four'])
- _check_plot_works(df.plot, kind='bar')
- _check_plot_works(df.plot, kind='bar', legend=False)
- _check_plot_works(df.plot, kind='bar', subplots=True)
- df = DataFrame(np.random.randn(10, 15),
- index=list(string.ascii_letters[:10]),
- columns=range(15))
- _check_plot_works(df.plot, kind='bar')
- def test_boxplot(self):
- df = DataFrame(np.random.randn(6, 4),
- index=list(string.ascii_letters[:6]),
- columns=['one', 'two', 'three', 'four'])
- df['indic'] = ['foo', 'bar'] * 3
- df['indic2'] = ['foo', 'bar', 'foo'] * 2
- _check_plot_works(df.boxplot)
- _check_plot_works(df.boxplot, column='one', by='indic')
- _check_plot_works(df.boxplot, column='one', by=['indic', 'indic2'])
- _check_plot_works(df.boxplot, by='indic')
- _check_plot_works(df.boxplot, by=['indic', 'indic2'])
- def test_hist(self):
- df = DataFrame(np.random.randn(100, 4))
- _check_plot_works(df.hist)
- _check_plot_works(df.hist, grid=False)
- 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)
- ret = f(*args, **kwargs)
- assert(ret is not None) # do something more intelligent
- ax = fig.add_subplot(212)
- try:
- kwargs['ax'] = ax
- ret = f(*args, **kwargs)
- assert(ret is not None) # do something more intelligent
- except Exception:
- pass
- plt.savefig(PNG_PATH)
- os.remove(PNG_PATH)
- if __name__ == '__main__':
- nose.runmodule(argv=[__file__,'-vvs','-x','--pdb', '--pdb-failure'],
- exit=False)