support clean up

rgf/sklearn.py

@@ -128,6 +128,22 @@ def _cleanup():
                 os.remove(fn)
 
 
+def _cleanup_partial(uuid):
+    if uuid not in _UUIDS:
+        return
+    model_glob = os.path.join(_TEMP_PATH, uuid + "*")
+    for fn in glob(model_glob):
+        os.remove(fn)
+    _UUIDS.remove(uuid)
+
+
+def _get_temp_path():
+    """
+    For test
+    """
+    return _TEMP_PATH
+
+
 def _sigmoid(x):
     return 1.0 / (1.0 + np.exp(-x))
 
@@ -683,6 +699,17 @@ def predict(self, X):
         y = np.argmax(y, axis=1)
         return np.asarray(list(self._classes_map.values()))[np.searchsorted(list(self._classes_map.keys()), y)]
 
+    def cleanup(self):
+        """
+        Clean tempfile used by this model.
+        """
+        if self._estimators is not None:
+            for est in self._estimators:
+                _cleanup_partial(est._file_prefix)
+
+        # No more able to predict without refitting.
+        self._fitted = False
+
 
 class _RGFBinaryClassifier(BaseEstimator, ClassifierMixin):
     """
@@ -1208,3 +1235,9 @@ def __setstate__(self, state):
             with open(self._latest_model_loc, 'wb') as fw:
                 fw.write(self.__dict__["model"])
             del self.__dict__["model"]
+
+    def cleanup(self):
+        """
+        Clean tempfile used by this model.
+        """
+        _cleanup_partial(self._file_prefix)

rgf/test/test.py

@@ -1,3 +1,5 @@
+import glob
+import os
 import pickle
 import unittest
 
@@ -11,7 +13,7 @@
 from sklearn.utils.estimator_checks import check_estimator
 from sklearn.utils.validation import check_random_state
 
-from rgf.sklearn import RGFClassifier, RGFRegressor, _cleanup
+from rgf.sklearn import RGFClassifier, RGFRegressor, _cleanup, _get_temp_path
 
 
 class TestRGFClassfier(unittest.TestCase):
@@ -241,6 +243,15 @@ def test_joblib_pickle(self):
 
         np.testing.assert_allclose(y_pred1, y_pred2)
 
+    def test_cleanup(self):
+        clf = RGFClassifier()
+        clf.fit(self.X_train, self.y_train)
+        clf.cleanup()
+
+        for est in clf.estimators_:
+            glob_file = os.path.join(_get_temp_path(), est._file_prefix + "*")
+            self.assertFalse(glob.glob(glob_file))
+
 
 class TestRGFRegressor(unittest.TestCase):
     def setUp(self):
@@ -408,6 +419,14 @@ def test_joblib_pickle(self):
 
         np.testing.assert_allclose(y_pred1, y_pred2)
 
+    def test_cleanup(self):
+        reg = RGFRegressor()
+        reg.fit(self.X_train, self.y_train)
+        reg.cleanup()
+
+        glob_file = os.path.join(_get_temp_path(), reg._file_prefix + "*")
+        self.assertFalse(glob.glob(glob_file))
+
 
 if __name__ == '__main__':
     unittest.main()