TST: provide a clear argument to assert_produces_warning to guarantee that we are clearing

     specific warnings before trying to catch them (in case they have already happened)
     requires the user to specificy the exact class where they are coming

Author: jreback

pandas/tests/test_indexing.py

@@ -1,4 +1,5 @@
 # pylint: disable-msg=W0612,E1101
+import sys
 import nose
 import itertools
 import warnings
@@ -19,6 +20,7 @@
                                  assert_frame_equal, assert_panel_equal,
                                  assert_attr_equal)
 from pandas import concat
+from pandas.io.common import PerformanceWarning
 
 import pandas.util.testing as tm
 from pandas import date_range
@@ -1489,8 +1491,6 @@ def test_loc_multiindex(self):
         assert_series_equal(result, expected)
 
     def test_multiindex_perf_warn(self):
-        import sys
-        from pandas.io.common import PerformanceWarning
 
         if sys.version_info < (2, 7):
             raise nose.SkipTest('python version < 2.7')
@@ -1499,14 +1499,18 @@ def test_multiindex_perf_warn(self):
                         'joe':['x', 'x', 'z', 'y'],
                         'jolie':np.random.rand(4)}).set_index(['jim', 'joe'])
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        with tm.assert_produces_warning(PerformanceWarning, clear=[pd.core.index]):
             _ = df.loc[(1, 'z')]
 
         df = df.iloc[[2,1,3,0]]
         with tm.assert_produces_warning(PerformanceWarning):
             _ = df.loc[(0,)]
 
     def test_multiindex_get_loc(self):  # GH7724, GH2646
+
+        # ignore the warning here
+        warnings.simplefilter('ignore', PerformanceWarning)
+
         # test indexing into a multi-index before & past the lexsort depth
         from numpy.random import randint, choice, randn
         cols = ['jim', 'joe', 'jolie', 'joline', 'jolia']
@@ -1568,6 +1572,9 @@ def loop(mi, df, keys):
                 assert not mi.index.lexsort_depth < i
                 loop(mi, df, keys)
 
+        # restore
+        warnings.simplefilter('always', PerformanceWarning)
+
     def test_series_getitem_multiindex(self):
 
         # GH 6018
@@ -1621,6 +1628,8 @@ def test_ix_general(self):
                 'year': {0: 2012, 1: 2011, 2: 2012, 3: 2012, 4: 2012}}
         df = DataFrame(data).set_index(keys=['col', 'year'])
         key = 4.0, 2012
+
+        # emits a PerformanceWarning, ok
         tm.assert_frame_equal(df.ix[key], df.iloc[2:])
 
         # this is ok

pandas/util/testing.py

@@ -24,7 +24,7 @@
 from numpy.testing import assert_array_equal
 
 import pandas as pd
-from pandas.core.common import _is_sequence, array_equivalent
+from pandas.core.common import _is_sequence, array_equivalent, is_list_like
 import pandas.core.index as index
 import pandas.core.series as series
 import pandas.core.frame as frame
@@ -1619,7 +1619,7 @@ def handle_success(self, exc_type, exc_value, traceback):
 
 
 @contextmanager
-def assert_produces_warning(expected_warning=Warning, filter_level="always"):
+def assert_produces_warning(expected_warning=Warning, filter_level="always", clear=None):
     """
     Context manager for running code that expects to raise (or not raise)
     warnings.  Checks that code raises the expected warning and only the
@@ -1646,6 +1646,19 @@ def assert_produces_warning(expected_warning=Warning, filter_level="always"):
     ..warn:: This is *not* thread-safe.
     """
     with warnings.catch_warnings(record=True) as w:
+
+        if clear is not None:
+            # make sure that we are clearning these warnings
+            # if they have happened before
+            # to guarantee that we will catch them
+            if not is_list_like(clear):
+                clear = [ clear ]
+            for m in clear:
+                try:
+                    m.__warningregistry__.clear()
+                except:
+                    pass
+
         saw_warning = False
         warnings.simplefilter(filter_level)
         yield w