Split atomic_save in folder and directory variant.

annif/backend/omikuji.py

@@ -95,10 +95,9 @@ def _create_model(self, params, jobs):
 
         self._model = omikuji.Model.train_on_data(
             train_path, hyper_param, jobs or None)
-        annif.util.atomic_save(
+        annif.util.atomic_save_folder(
             self._model,
-            model_path,
-            None)
+            model_path)
 
     def _train(self, corpus, params, jobs=0):
         if corpus != 'cached':

annif/backend/xtransformer.py

@@ -10,7 +10,8 @@
 from annif.suggestion import ListSuggestionResult, SubjectSuggestion
 from . import mixins
 from . import backend
-from annif.util import boolean, apply_param_parse_config, atomic_save
+from annif.util import boolean, apply_param_parse_config, atomic_save_folder, \
+    atomic_save
 
 from pecos.xmc.xtransformer.model import XTransformer
 from pecos.xmc.xtransformer.module import MLProblemWithText
@@ -203,7 +204,7 @@ def _create_model(self, params, jobs):
             steps_scale=None,
             label_feat=None,
             )
-        atomic_save(self._model, model_path, None)
+        atomic_save_folder(self._model, model_path)
 
     def _train(self, corpus, params, jobs=0):
         if corpus == 'cached':

annif/util.py

@@ -14,27 +14,47 @@ def atomic_save(obj, dirname, filename, method=None):
     """Save the given object (which must have a .save() method, unless the
     method parameter is given) into the given directory with the given
     filename, using a temporary file and renaming the temporary file to the
-    final name. To save a directory explicitly set filename=None."""
-
-    if filename:
-        prefix, suffix = os.path.splitext(filename)
-        tempfd, tempfilename = tempfile.mkstemp(
-            prefix=prefix, suffix=suffix, dir=dirname)
-        os.close(tempfd)
-        target_pth = os.path.join(dirname, filename)
-    else:
-        tldir = os.path.dirname(dirname.rstrip('/'))
-        os.makedirs(dirname, exist_ok=tldir)
-        tempdir = tempfile.TemporaryDirectory(dir=tldir)
-        tempfilename = tempdir.name
-        target_pth = dirname
+    final name.
+    The .save() mehod or the function provided in the method argument
+    will be called with the path to the temporary file."""
+
+    prefix, suffix = os.path.splitext(filename)
+    tempfd, tempfilename = tempfile.mkstemp(
+        prefix=prefix, suffix=suffix, dir=dirname)
+    os.close(tempfd)
     logger.debug('saving %s to temporary file %s', str(obj)[:90], tempfilename)
     if method is not None:
         method(obj, tempfilename)
     else:
         obj.save(tempfilename)
     for fn in glob.glob(tempfilename + '*'):
-        newname = fn.replace(tempfilename, target_pth)
+        newname = fn.replace(tempfilename, os.path.join(dirname, filename))
+        logger.debug('renaming temporary file %s to %s', fn, newname)
+        os.rename(fn, newname)
+
+
+def atomic_save_folder(obj, dirname, method=None):
+    """Save the given object (which must have a .save() method, unless the
+    method parameter is given) into the given directory,
+    using a temporary directory and renaming the temporary directory to the
+    final name.
+    The .save() mehod or the function provided in the method argument
+    will be called with the path to the temporary directory."""
+
+    tldir = os.path.dirname(dirname.rstrip('/'))
+    os.makedirs(dirname, exist_ok=tldir)
+    tempdir = tempfile.TemporaryDirectory(dir=tldir)
+    temp_dir_name = tempdir.name
+    target_pth = dirname
+    logger.debug(
+        'saving %s to temporary file %s', str(obj)[:90],
+        temp_dir_name)
+    if method is not None:
+        method(obj, temp_dir_name)
+    else:
+        obj.save(temp_dir_name)
+    for fn in glob.glob(temp_dir_name + '*'):
+        newname = fn.replace(temp_dir_name, target_pth)
         logger.debug('renaming temporary file %s to %s', fn, newname)
         if os.path.isdir(newname):
             rmtree(newname)

tests/test_util.py

@@ -83,7 +83,10 @@ def save_folder(obj, pth):
         _save(None, osp.join(pth, fname_0))
         _save(None, osp.join(pth, fname_1))
     folder_path = tmpdir.join(folder_name)
-    annif.util.atomic_save(None, folder_path.strpath, None, method=save_folder)
+    annif.util.atomic_save_folder(
+        None,
+        folder_path.strpath,
+        method=save_folder)
     assert folder_path.exists()
     for f_name in [fname_0, fname_1]:
         f_pth = folder_path.join(f_name)