Skip to content

Instantly share code, notes, and snippets.

@notbanker notbanker/chained.py
Last active Apr 3, 2019

Embed
What would you like to do?
Context manager to temporarily pandas set chained assignment warning to None,'warn' or 'raise, then revert
import pandas as pd
class ChainedAssignent:
""" Context manager to temporarily set pandas chained assignment warning. Usage:
with ChainedAssignment():
blah
with ChainedAssignment('error'):
run my code and figure out which line causes the error!
"""
def __init__(self, chained = None):
acceptable = [ None, 'warn','raise']
assert chained in acceptable, "chained must be in " + str(acceptable)
self.swcw = chained
def __enter__( self ):
self.saved_swcw = pd.options.mode.chained_assignment
pd.options.mode.chained_assignment = self.swcw
return self
def __exit__(self, *args):
pd.options.mode.chained_assignment = self.saved_swcw
@shaybensasson

This comment has been minimized.

Copy link

commented Apr 3, 2019

Great!
Here is my small contribution (pytest unittests)

import pytest
import pandas as pd
import numpy as np
from pandas.core.common import SettingWithCopyWarning, SettingWithCopyError

from pandas_chained_assignment_warn_handler import PandasChainedAssignmentWarnHandler


def reproduce_warn():
    df = pd.DataFrame(data={'col1': np.arange(11), 'col2': [1] * 11, 'label': [1] * 11})
    df.loc[0, 'col2'] = np.nan

    X, y = df[['col1', 'col2']], df[['label']]
    X.loc[0, 'col2'] = 1.

    assert X.loc[0, 'col2'] == 1.
    assert np.isnan(df.loc[0, 'col2'])  # still nan

def test_reproduce_warn_raises_warn():
    with pytest.warns(SettingWithCopyWarning):
        reproduce_warn()

def test_context_handler_blocks_warn():
    with pytest.warns(None) as record:
        with PandasChainedAssignmentWarnHandler(chained=None):
            reproduce_warn()
        assert len(record) == 0

def test_context_handler_display_warn():
    with pytest.warns(SettingWithCopyWarning):
        with PandasChainedAssignmentWarnHandler(chained='warn'):
            reproduce_warn()

def test_context_handler_display_error():
    with pytest.raises(SettingWithCopyError):
        with PandasChainedAssignmentWarnHandler(chained='raise'):
            reproduce_warn()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.