Refactor check_fitted utility

Author: lixfz

ylearn/uplift/_model.py

@@ -1,21 +1,11 @@
 import numpy as np
 import pandas as pd
 
+from ylearn.utils import check_fitted as _check_fitted
 from ._metric import get_gain, get_qini, get_cumlift, auuc_score, qini_score
 from ._plot import plot_gain, plot_qini, plot_cumlift
 
 
-def _check_fitted(fn):
-    def _exec(obj, *args, **kwargs):
-        assert isinstance(obj, UpliftModel)
-        if obj.lift_ is None:
-            raise ValueError(f'fit {type(obj).__name__} before call {fn.__name__}() please.')
-
-        return fn(obj, *args, **kwargs)
-
-    return _exec
-
-
 class UpliftModel(object):
     def __init__(self):
         # fitted
@@ -57,6 +47,7 @@ def fit(self, df_lift, outcome='y', treatment='x', true_effect=None, treat=1, co
         )
         self.random_ = random
         self.lift_ = df_lift.copy()
+        self._is_fitted = True
 
         return self
 

ylearn/utils/__init__.py

@@ -13,3 +13,4 @@
 from ._common import to_df, context, is_notebook, view_pydot
 from ._common import to_list, join_list
 from ._common import to_snake_case, to_camel_case, drop_none
+from ._common import check_fitted, check_fitted_

ylearn/utils/_common.py

@@ -2,6 +2,7 @@
 import inspect
 import warnings
 from collections import OrderedDict
+from functools import partial
 
 import numpy as np
 import pandas as pd
@@ -262,3 +263,34 @@ def _get_array(cols):
     S = map(_get_array, S)
 
     return tuple(S)
+
+
+def check_fitted(fn, attr_name='_is_fitted', check=None, msg=None):
+    assert callable(fn)
+
+    sig = inspect.signature(fn)
+    assert 'self' in sig.parameters.keys()
+
+    def check_and_call(obj, *args, **kwargs):
+        if callable(check):
+            fitted = check(obj)
+        else:
+            fitted_tag = getattr(obj, attr_name, None)
+            if check is not None:
+                fitted = not (fitted_tag is check)
+            else:
+                fitted = fitted_tag
+
+        if not fitted:
+            if msg is not None:
+                raise ValueError(msg)
+            else:
+                raise ValueError(f'{type(obj).__name__} is not fitted.')
+
+        return fn(obj, *args, **kwargs)
+
+    return check_and_call
+
+
+def check_fitted_(attr_name='_is_fitted', check=None, msg=None):
+    return partial(check_fitted, attr_name=attr_name, check=check, msg=msg)