Move files from temporary directory changer (#2022)

Author: mgjarrett

armi/bookkeeping/db/database.py

@@ -78,6 +78,7 @@
 from armi.reactor.reactors import Core, Reactor
 from armi.settings.fwSettings.globalSettings import CONF_SORT_REACTOR
 from armi.utils import getNodesPerCycle
+from armi.utils import safeMove
 from armi.utils.textProcessors import resolveMarkupInclusions
 
 # CONSTANTS
@@ -312,7 +313,7 @@ def close(self, completedSuccessfully=False):
 
         if self._permission == "w":
             # move out of the FAST_PATH and into the working directory
-            newPath = shutil.move(self._fullPath, self._fileName)
+            newPath = safeMove(self._fullPath, self._fileName)
             self._fullPath = os.path.abspath(newPath)
 
     def splitDatabase(
@@ -351,7 +352,7 @@ def splitDatabase(
         backupDBPath = os.path.abspath(label.join(os.path.splitext(self._fileName)))
         runLog.info("Retaining full database history in {}".format(backupDBPath))
         if self._fullPath is not None:
-            shutil.move(self._fullPath, backupDBPath)
+            safeMove(self._fullPath, backupDBPath)
 
         self.h5db = h5py.File(self._fullPath, self._permission)
         dbOut = self.h5db

armi/utils/__init__.py

@@ -840,6 +840,50 @@ def safeCopy(src: str, dst: str) -> None:
     runLog.extra("Copied {} -> {}".format(src, dst))
 
 
+def safeMove(src: str, dst: str) -> None:
+    """Check that a file has been successfully moved before continuing."""
+    # Convert files to OS-independence
+    src = os.path.abspath(src)
+    dst = os.path.abspath(dst)
+    if os.path.isdir(dst):
+        dst = os.path.join(dst, os.path.basename(src))
+
+    srcSize = os.path.getsize(src)
+    if "win" in sys.platform:
+        # this covers Windows ("win32") and MacOS ("darwin")
+        shutil.move(src, dst)
+    elif "linux" in sys.platform:
+        cmd = f'mv "{src}" "{dst}"'
+        os.system(cmd)
+    else:
+        raise OSError(
+            "Cannot perform ``safeMove`` on files because ARMI only supports "
+            + "Linux, MacOS, and Windows."
+        )
+
+    waitTime = 0.01  # 10 ms
+    maxWaitTime = 6000  # 1 min
+    totalWaitTime = 0
+    while True:
+        try:
+            dstSize = os.path.getsize(dst)
+            if srcSize == dstSize:
+                break
+        except FileNotFoundError:
+            pass
+        time.sleep(waitTime)
+        totalWaitTime += waitTime
+        if totalWaitTime > maxWaitTime:
+            runLog.warning(
+                f"File move from {dst} to {src} has failed due to exceeding "
+                + f"a maximum wait time of {maxWaitTime/60} minutes."
+            )
+            return
+
+    runLog.extra("Moved {} -> {}".format(src, dst))
+    return dst
+
+
 # Allow us to check the copy operation is complete before continuing
 shutil_copy = shutil.copy
 shutil.copy = safeCopy

armi/utils/directoryChangers.py

@@ -22,6 +22,7 @@
 from armi import context
 from armi import runLog
 from armi.utils import pathTools
+from armi.utils import safeMove
 
 
 def _changeDirectory(destination):
@@ -128,19 +129,27 @@ def moveFiles(self):
         """Copy ``filesToMove`` into the destination directory on entry."""
         initialPath = self.initial
         destinationPath = self.destination
-        self._transferFiles(initialPath, destinationPath, self._filesToMove)
+        self._transferFiles(
+            initialPath, destinationPath, self._filesToMove, moveFiles=False
+        )
         if self.outputPath != self.initial:
             destinationPath = self.outputPath
-            self._transferFiles(initialPath, destinationPath, self._filesToMove)
+            self._transferFiles(
+                initialPath, destinationPath, self._filesToMove, moveFiles=False
+            )
 
     def retrieveFiles(self):
         """Copy ``filesToRetrieve`` back into the initial directory on exit."""
-        initialPath = self.destination
-        destinationPath = self.initial
-        self._transferFiles(initialPath, destinationPath, self._filesToRetrieve)
         if self.outputPath != self.initial:
-            destinationPath = self.outputPath
-            self._transferFiles(initialPath, destinationPath, self._filesToRetrieve)
+            self._transferFiles(
+                self.destination,
+                self.outputPath,
+                self._filesToRetrieve,
+                moveFiles=False,
+            )
+        self._transferFiles(
+            self.destination, self.initial, self._filesToRetrieve, moveFiles=True
+        )
 
     def _retrieveEntireFolder(self):
         """
@@ -173,7 +182,7 @@ def _createOutputDirectory(self):
             runLog.extra(f"Output folder already exists: {self.outputPath}")
 
     @staticmethod
-    def _transferFiles(initialPath, destinationPath, fileList):
+    def _transferFiles(initialPath, destinationPath, fileList, moveFiles=False):
         """
         Transfer files into or out of the directory.
 
@@ -192,6 +201,8 @@ def _transferFiles(initialPath, destinationPath, fileList):
             files will be transferred. Alternatively tuples of (initialName, finalName) are allowed
             if you want the file renamed during transit. In the non-tuple option, globs/wildcards
             are allowed.
+        moveFiles: bool, optional
+            Controls whether the files are "moved" (``mv``) or "copied" (``cp``)
 
         Warning
         -------
@@ -223,8 +234,12 @@ def _transferFiles(initialPath, destinationPath, fileList):
                     continue
 
                 toPath = os.path.join(destinationPath, destName)
-                runLog.extra("Copying {} to {}".format(fromPath, toPath))
-                shutil.copy(fromPath, toPath)
+                if moveFiles:
+                    runLog.extra("Moving {} to {}".format(fromPath, toPath))
+                    safeMove(fromPath, toPath)
+                else:
+                    runLog.extra("Copying {} to {}".format(fromPath, toPath))
+                    shutil.copy(fromPath, toPath)
 
 
 class TemporaryDirectoryChanger(DirectoryChanger):

armi/utils/tests/test_utils.py

@@ -41,6 +41,7 @@
     getStepLengths,
     hasBurnup,
     safeCopy,
+    safeMove,
 )
 
 
@@ -218,6 +219,53 @@ def test_safeCopy(self):
                 self.assertIn("Copied", mock.getStdout())
                 self.assertIn("file2", mock.getStdout())
                 self.assertIn("->", mock.getStdout())
+            self.assertTrue(os.path.exists(os.path.join("dir2", "file1.txt")))
+
+    def test_safeMove(self):
+        with directoryChangers.TemporaryDirectoryChanger():
+            os.mkdir("dir1")
+            os.mkdir("dir2")
+            file1 = "dir1/file1.txt"
+            with open(file1, "w") as f:
+                f.write("Hello")
+            file2 = "dir1\\file2.txt"
+            with open(file2, "w") as f:
+                f.write("Hello2")
+
+            with mockRunLogs.BufferLog() as mock:
+                # Test Linuxy file path
+                self.assertEqual("", mock.getStdout())
+                safeMove(file1, "dir2")
+                self.assertIn("Moved", mock.getStdout())
+                self.assertIn("file1", mock.getStdout())
+                self.assertIn("->", mock.getStdout())
+                # Clean up for next safeCopy
+                mock.emptyStdout()
+                # Test Windowsy file path
+                self.assertEqual("", mock.getStdout())
+                safeMove(file2, "dir2")
+                self.assertIn("Moved", mock.getStdout())
+                self.assertIn("file2", mock.getStdout())
+                self.assertIn("->", mock.getStdout())
+            self.assertTrue(os.path.exists(os.path.join("dir2", "file1.txt")))
+
+    def test_safeMoveDir(self):
+        with directoryChangers.TemporaryDirectoryChanger():
+            os.mkdir("dir1")
+            file1 = "dir1/file1.txt"
+            with open(file1, "w") as f:
+                f.write("Hello")
+            file2 = "dir1\\file2.txt"
+            with open(file2, "w") as f:
+                f.write("Hello2")
+
+            with mockRunLogs.BufferLog() as mock:
+                self.assertEqual("", mock.getStdout())
+                safeMove("dir1", "dir2")
+                self.assertIn("Moved", mock.getStdout())
+                self.assertIn("dir1", mock.getStdout())
+                self.assertIn("dir2", mock.getStdout())
+            self.assertTrue(os.path.exists(os.path.join("dir2", "file1.txt")))
 
 
 class CyclesSettingsTests(unittest.TestCase):

doc/release/0.5.rst

@@ -8,7 +8,7 @@ Release Date: TBD
 
 New Features
 ------------
-#. TBD
+#. Move instead of copy files from TemporaryDirectoryChanger. (`PR#2022 <https://github.com/terrapower/armi/pull/2022>`_)
 
 API Changes
 -----------