diff --git a/bulwark/checks.py b/bulwark/checks.py index 6c83484..53988aa 100644 --- a/bulwark/checks.py +++ b/bulwark/checks.py @@ -7,6 +7,7 @@ - return the original, unaltered pd.DataFrame """ +import collections import warnings import numpy as np @@ -477,12 +478,16 @@ def is_same_as(df, df_to_compare, **kwargs): return df -def multi_check(df, checks, warn=False): +def multi_check(df, checks, *more_checks, warn=False): """Asserts that all checks pass. - Args: df (pd.DataFrame): Any pd.DataFrame. - checks (dict): Mapping of check functions to parameters for those check functions. + checks: A single function specification, as below. + This exists for backwards compatibility, and to ensure there's at least one check. + more_checks: Any number of check function specifications. + The simplest specification is a check function, which assumes no parameters. + Or it may be a Mapping of check functions to parameters for those functions. + In the event of duplicates, the last specification will win. warn (bool): Indicates whether an error should be raised or only a warning notification should be displayed. Default is to error. @@ -491,8 +496,16 @@ def multi_check(df, checks, warn=False): Original `df`. """ + + checks_map = {} + for check in [checks, *more_checks]: + if isinstance(check, collections.Mapping): + checks_map.update(check) + else: + checks_map[check] = {} + error_msgs = [] - for func, params in checks.items(): + for func, params in checks_map.items(): try: func(df, **params) except AssertionError as e: diff --git a/bulwark/decorators.py b/bulwark/decorators.py index 4718a05..66f8708 100644 --- a/bulwark/decorators.py +++ b/bulwark/decorators.py @@ -45,6 +45,33 @@ class decorator_name(BaseDecorator): setattr(this_module, decorator_name, decorator_factory(decorator_name, func)) +class MultiCheck: + """ + Notes: + - This code is purposefully located below the auto-generation of decorators, + so this overwrites the auto-generated MultiCheck. + - `MultiCheck`'s __init__ diverges from `BaseDecorator`, due to the use of *args. + + TODO: Work this into BaseDecorator? + + """ + + def __init__(self, checks, *more_checks, warn=False, enabled=True): + self.checks = checks + self.more_checks = more_checks + self.warn = warn + self.enabled = enabled + + def __call__(self, f): + @functools.wraps(f) + def decorated(*args, **kwargs): + df = f(*args, **kwargs) + if self.enabled: + ck.multi_check(df, self.checks, *self.more_checks, warn=self.warn) + return df + return decorated + + class CustomCheck: """ Notes: diff --git a/tests/test_checks.py b/tests/test_checks.py index c1707b6..b0c1690 100644 --- a/tests/test_checks.py +++ b/tests/test_checks.py @@ -515,6 +515,52 @@ def total_sum_not_equal(df, amt): # with pytest.raises(AssertionError): +def test_multi_check_varargs(): + df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) + result = ck.multi_check(df, ck.has_no_nans, warn=False) + tm.assert_frame_equal(df, result) + + result = ck.multi_check(df, + ck.has_no_nans, + {ck.is_shape: {"shape": (3, 2)}}, + warn=False) + tm.assert_frame_equal(df, result) + + result = dc.MultiCheck(ck.has_no_nans, + {ck.is_shape: {"shape": (3, 2)}}, + warn=False)(_noop)(df) + tm.assert_frame_equal(df, result) + + def total_sum_not_equal(df, amt): + if amt == 51: + raise AssertionError("Test") + return True + + result = ck.multi_check(df, + ck.has_no_nans, + {total_sum_not_equal: {"amt": 52}}, + warn=False) + tm.assert_frame_equal(df, result) + + result = dc.MultiCheck(ck.has_no_nans, + {total_sum_not_equal: {"amt": 52}}, + warn=False)(_noop)(df) + tm.assert_frame_equal(df, result) + + result = dc.MultiCheck(ck.has_no_nans, + {total_sum_not_equal: {"amt": 51}}, + {total_sum_not_equal: {"amt": 52}}, # last wins + warn=False)(_noop)(df) + tm.assert_frame_equal(df, result) + + with pytest.raises(AssertionError): + result = dc.MultiCheck(ck.has_no_nans, + {total_sum_not_equal: {"amt": 51}}, + warn=False)(_noop)(df) + + # with pytest.raises(AssertionError): + + def test_custom_check(): def f(df, length): if len(df) <= length: