Enhance decorator to check returned tuple dataframes and input dataframes

README.md

@@ -1,6 +1,6 @@
 Bulwark's Documentation
 ========================================
-<a href="https://pypi.org/project/bulwark/"><img src="https://img.shields.io/pypi/dm/bulwark?style=for-the-badge" alt="downloads" /></a>  
+<a href="https://pypi.org/project/bulwark/"><img src="https://img.shields.io/pypi/dm/bulwark?style=for-the-badge" alt="downloads" /></a>
 
 <a href="https://pypi.org/project/bulwark/"><img src="https://img.shields.io/pypi/v/bulwark?style=for-the-badge" alt="latest release" /></a>
 <a href="https://pypi.org/project/bulwark/"><img src="https://img.shields.io/pypi/pyversions/bulwark?style=for-the-badge" alt="supported python versions" /></a>
@@ -63,7 +63,34 @@ on the functions you're already writing:
     @dc.HasNoNans()
     def compute(df):
         # complex operations to determine result
-        ...
+
+    return result_df
+```
+
+The checks can also be performed on an item in a returned tuple with a dataframe as
+one of the values:
+```python
+    import bulwark.decorators as dc
+
+    @dc.IsShape((-1, 10))
+    @dc.IsMonotonic(df=0,strict=True)
+    @dc.HasNoNans(df=0) # the number specifies the tuple of the returned list/tuple
+    @dc.HasNoNans(df=1)
+    def compute(df):
+        # complex operations to determine result
+
+    return (result_df1, result_df2)
+```
+The checks can also be applied to input parameters as well:
+```python
+    import bulwark.decorators as dc
+
+    @dc.IsShape(df='df_input',shape=(-1, 10))
+    @dc.IsMonotonic(df='df_input',strict=True)
+    @dc.HasNoNans(df='df_input') # the name specifies the name of the parameter
+    def compute(df_input):
+        # complex operations to determine result
+
     return result_df
 ```
 
@@ -83,7 +110,7 @@ Nope - just toggle the built-in "enabled" flag available for every decorator.
     @dc.IsShape((3, 2), enabled=False)
     def compute(df):
         # complex operations to determine result
-        ...
+
     return result_df
 ```
 
@@ -147,5 +174,5 @@ We work hard to make contributing as easy as possible,
 and previous open source experience is not required!
 Please see [contributing.md](docs/contributing.md) for how to get started.
 
-Thank you to all our past contributors, especially these folks:  
+Thank you to all our past contributors, especially these folks:
 [![](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/images/0)](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/links/0)[![](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/images/1)](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/links/1)[![](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/images/2)](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/links/2)[![](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/images/3)](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/links/3)[![](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/images/4)](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/links/4)[![](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/images/5)](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/links/5)[![](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/images/6)](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/links/6)[![](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/images/7)](https://sourcerer.io/fame/ZaxR/ZaxR/bulwark/links/7)

bulwark/decorators.py

@@ -1,30 +1,98 @@
 """Generates decorators for each check in `checks.py`."""
 import functools
 import sys
-from inspect import getfullargspec, getmembers, isfunction
+from inspect import (Parameter, getfullargspec, getmembers, isfunction,
+                     signature)
 
 import bulwark.checks as ck
 from bulwark.generic import snake_to_camel
 
 
 class BaseDecorator(object):
+    df_allowed_types = (int, str, type(None))
+
     def __init__(self, *args, **kwargs):
         self.enabled = kwargs.pop("enabled", True)  # setter to enforce bool would be a lot safer
         # self.warn = False ? No - put at func level for all funcs and pass through
+        self.params = getfullargspec(self.check_func).args[1:]
 
-        self.check_func_params = dict(
-            zip(getfullargspec(self.check_func).args[1:], args))
-        self.check_func_params.update(**kwargs)
+        self.__dict__.update(dict(zip(self.params, args)))
+        self.__dict__.update(**kwargs)
+
+        if type(self._df) not in self.df_allowed_types:
+            msg = ("'df' arg cannot by of type {}.\n"
+                   "Only allowed types are:\n"
+                   " str:  arg name of the input argument to the decorated function to check\n"
+                   " int:  the entry in a tuple returned by the decorated function to check\n"
+                   " None: check the single dataframe return by"
+                   " the decorated function".format(type(self._df))
+                   )
+            raise TypeError(msg)
+
+    @property
+    def _df(self):
+        return self.__dict__.get('df', None)
 
     def __call__(self, f):
+
+        if type(self._df) is str:
+            if not self._has_arg_name(self._df, f):
+                raise NameError("'{}' is not an arg to function '{}'".format(self._df, f.__name__))
+
         @functools.wraps(f)
         def decorated(*args, **kwargs):
-            df = f(*args, **kwargs)
-            if self.enabled:
-                self.check_func(df, **self.check_func_params)
-            return df
+
+            if not self.enabled:
+                return f(*args, **kwargs)
+
+            check_kwargs = {k: v for k, v in self.__dict__.items()
+                            if k not in ["check_func", "enabled", "params"]}
+
+            if self._df is None:
+                res = f(*args, **kwargs)
+                check_kwargs['df'] = res
+                self.check_func(**check_kwargs)
+                return res
+
+            if type(self._df) is int:
+                res = f(*args, **kwargs)
+                if type(res) is not tuple:
+                    raise TypeError('Your function needs to return'
+                                    'a tuple to use df=<int> in this decorator')
+                check_kwargs['df'] = res[self._df]
+                self.check_func(**check_kwargs)
+                return res
+
+            if type(self._df) is str:
+                df_default = self._get_arg_default(f, self._df)
+                df_value = self._get_arg_value(f, self._df, *args, **kwargs)
+                if isinstance(df_default, Parameter.empty) or (df_value is not None):
+                    check_kwargs['df'] = df_value
+                    self.check_func(**check_kwargs)
+                res = f(*args, **kwargs)
+                return res
         return decorated
 
+    @staticmethod
+    def _has_arg_name(arg_name, func):
+        sig = signature(func)
+        res = arg_name in sig.parameters
+        return res
+
+    @staticmethod
+    def _get_arg_default(func, arg_name):
+        func_sig = signature(func)
+        res = func_sig.parameters[arg_name].default
+        return res
+
+    @staticmethod
+    def _get_arg_value(func, arg_name, *args, **kwargs):
+        func_sig = signature(func)
+        bound_args = func_sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+        res = bound_args.arguments[arg_name]
+        return res
+
 
 def decorator_factory(decorator_name, func):
     """Takes in a function and outputs a class that can be used as a decorator."""

tests/test_checks.py

@@ -487,8 +487,8 @@ def test_multi_check():
     tm.assert_frame_equal(df, result)
 
     result = dc.MultiCheck(checks={ck.has_no_nans: {"columns": None},
-                                   ck.is_shape: {"shape": (3, 2)}},
-                           warn=False)(_noop)(df)
+                                   ck.is_shape: {"shape": (3, 2)}
+                                   }, warn=False)(_noop)(df)
     tm.assert_frame_equal(df, result)
 
     def total_sum_not_equal(df, amt):
@@ -543,3 +543,130 @@ def append_a_df(df, df2):
 
     with pytest.raises(AssertionError):
         dc.CustomCheck(f, 4)(_noop)(df)
+
+
+class TestFlexibleDecorator:
+
+    df_pass = pd.DataFrame([1])
+    df_fail = pd.DataFrame([0])
+
+    def test_return_single_df_none(self):
+
+        @dc.HasNoX(df=None, values=[0])
+        def _func(x):
+            return x
+
+        assert _func(self.df_pass) is self.df_pass
+        with pytest.raises(AssertionError):
+            _func(self.df_fail)
+
+        with pytest.raises(AttributeError):
+            _func(pd.Series([1]))
+
+    def test_return_single_df_default(self):
+
+        @dc.HasNoX(values=[0])
+        def _func(x):
+            return x
+
+        assert _func(self.df_pass) is self.df_pass
+        with pytest.raises(AssertionError):
+            _func(self.df_fail)
+
+        with pytest.raises(AttributeError):
+            _func(pd.Series([1]))
+
+    def test_required_dataframe_arg(self):
+
+        @dc.HasNoX(df='df0', values=[0])
+        def _func(df0):
+            return df0
+
+        assert _func(self.df_pass) is self.df_pass
+        assert _func(df0=self.df_pass) is self.df_pass
+
+        with pytest.raises(AssertionError):
+            _func(self.df_fail)
+        with pytest.raises(AssertionError):
+            _func(df0=self.df_fail)
+
+        with pytest.raises(TypeError):
+            _func()
+
+        assert _func(None) is None
+        assert _func(df0=None) is None
+
+        with pytest.raises(AttributeError):
+            _func(df0=pd.Series([1]))
+
+    def test_optional_dataframe_arg(self):
+
+        @dc.HasNoX(df='df_opt', values=[0])
+        def _func(x, df_opt=None):
+            return df_opt
+
+        assert _func(99) is None
+        assert _func(x=99) is None
+        assert _func(99, None) is None
+        assert _func(99, df_opt=None) is None
+        assert _func(df_opt=None, x=99) is None
+
+        assert _func(99, self.df_pass) is self.df_pass
+        assert _func(99, df_opt=self.df_pass) is self.df_pass
+        assert _func(df_opt=self.df_pass, x=99) is self.df_pass
+
+        with pytest.raises(AssertionError):
+            _func(99, self.df_fail)
+        with pytest.raises(AssertionError):
+            _func(99, df_opt=self.df_fail)
+        with pytest.raises(AssertionError):
+            _func(df_opt=self.df_fail, x=99)
+
+        with pytest.raises(AttributeError):
+            _func(99, df_opt='hello world')
+
+    def test_optional_dataframe_wrong_name(self):
+        with pytest.raises(NameError):
+            @dc.HasNoX(df='df0', values=[0])
+            def _func(x, df1=None):
+                return df1
+
+    def test_optional_dataframe_wrong_type(self):
+
+        @dc.HasNoX(df='df_opt', values=[0])
+        def _func(x, df_opt=1):
+            return df_opt
+
+        with pytest.raises(AttributeError):
+            _func(99)
+
+        assert _func(99, None) is None
+        assert _func(99, self.df_pass) is self.df_pass
+
+        with pytest.raises(AssertionError):
+            _func(99, self.df_fail)
+
+    def test_func_returns_tuple(self):
+        @dc.HasNoX(df=1, values=[0])
+        def _func(x):
+            return x
+
+        x = (99, self.df_pass)
+        assert _func(x) is x
+        with pytest.raises(AssertionError):
+            _func((99, self.df_fail))
+
+        with pytest.raises(IndexError):
+            _func((self.df_pass, ))
+
+        x = [99, self.df_pass]
+        with pytest.raises(TypeError):
+            _func(x)
+
+        x = pd.Series([None, self.df_pass], index=[1, 0])
+        with pytest.raises(TypeError):
+            _func(x)
+
+    def test_df_is_not_int(self):
+        with pytest.raises(TypeError):
+            dc.HasNoX(df=1.0, values=[0])