Repository: unionai-oss/pandera

Issue:
handle `self` argument when using check_input in class method when str obj_getter is supplied
Check is the wrapped function is a class method:

https://stackoverflow.com/questions/5963729/check-if-a-function-is-a-method-of-some-object

if so, ignore the first argument.


Discussion:

Resolving patch:
diff --git a/pandera/pandera.py b/pandera/pandera.py
--- a/pandera/pandera.py
+++ b/pandera/pandera.py
@@ -684,9 +684,8 @@ def __init__(
             then assumes one of the valid pandas string values:
             http://pandas.pydata.org/pandas-docs/stable/basics.html#dtypes
         :type pandas_dtype: str|PandasDtype
-        :param checks: If element_wise is True, then callable signature should
-            be:
-            x -> x where x is a scalar element in the column. Otherwise,
+        :param checks: if element_wise is True, then callable signature should
+            be: x -> bool where x is a scalar element in the column. Otherwise,
             x is assumed to be a pandas.Series object.
         :type checks: callable
         :param nullable: Whether or not column can contain null values.
@@ -741,6 +740,18 @@ def __repr__(self):
         return "<Schema Column: '%s' type=%s>" % (self._name, dtype)
 
 
+def _get_fn_argnames(fn):
+    if sys.version_info.major >= 3:
+        arg_spec_args = inspect.getfullargspec(fn).args
+    else:
+        arg_spec_args = inspect.getargspec(fn).args
+
+    if inspect.ismethod(fn) and arg_spec_args[0] == "self":
+        # don't include "self" argument
+        arg_spec_args = arg_spec_args[1:]
+    return arg_spec_args
+
+
 def check_input(schema, obj_getter=None):
     """Validate function argument when function is called.
 
@@ -770,10 +781,7 @@ def _wrapper(fn, instance, args, kwargs):
             if obj_getter in kwargs:
                 kwargs[obj_getter] = schema.validate(kwargs[obj_getter])
             else:
-                if sys.version_info.major >= 3:
-                    arg_spec_args = inspect.getfullargspec(fn).args
-                else:
-                    arg_spec_args = inspect.getargspec(fn).args
+                arg_spec_args = _get_fn_argnames(fn)
                 args_dict = OrderedDict(
                     zip(arg_spec_args, args))
                 args_dict[obj_getter] = schema.validate(args_dict[obj_getter])


Test patch:
diff --git a/tests/test_pandera.py b/tests/test_pandera.py
--- a/tests/test_pandera.py
+++ b/tests/test_pandera.py
@@ -272,6 +272,50 @@ def func_input_transform2(_, df):
     assert "column2" in result2[1]
 
 
+def test_check_input_method_decorators():
+    in_schema = DataFrameSchema({"column1": Column(String)})
+    out_schema = DataFrameSchema({"column2": Column(Int)})
+    dataframe = pd.DataFrame({"column1": ["a", "b", "c"]})
+
+    def _transform_helper(df):
+        return df.assign(column2=[1, 2, 3])
+
+    class TransformerClass(object):
+
+        @check_input(in_schema)
+        @check_output(out_schema)
+        def transform_first_arg(self, df):
+            return _transform_helper(df)
+
+        @check_input(in_schema, 0)
+        @check_output(out_schema)
+        def transform_first_arg_with_list_getter(self, df):
+            return _transform_helper(df)
+
+        @check_input(in_schema, 1)
+        @check_output(out_schema)
+        def transform_secord_arg_with_list_getter(self, x, df):
+            return _transform_helper(df)
+
+        @check_input(in_schema, "df")
+        @check_output(out_schema)
+        def transform_secord_arg_with_dict_getter(self, x, df):
+            return _transform_helper(df)
+
+    def _assert_expectation(result_df):
+        assert isinstance(result_df, pd.DataFrame)
+        assert "column2" in result_df.columns
+
+    transformer = TransformerClass()
+    _assert_expectation(transformer.transform_first_arg(dataframe))
+    _assert_expectation(
+        transformer.transform_first_arg_with_list_getter(dataframe))
+    _assert_expectation(
+        transformer.transform_secord_arg_with_list_getter(None, dataframe))
+    _assert_expectation(
+        transformer.transform_secord_arg_with_dict_getter(None, dataframe))
+
+
 def test_string_dtypes():
     # TODO: add tests for all datatypes
     schema = DataFrameSchema(