Repository: terrapower/armi
Pull request #2022 (2024-11-26 00:52:10)

Move files from temporary directory changer
## What is the change?

Use `shutil.move` / `mv` instead of `shutil.copy` / `cp` to move files from temporary directory back to the original working directory on `__exit__`.

## Why is the change being made?

Copying files from one location to another keeps the original file intact, but the original files in a temporary directory do not need to be kept around after the `__exit__` from the `TemporaryDirectoryChanger`. These files are immediately deleted after the copy operation is completed.

Using `shutil.move()` is significantly faster than `shutil.copy()` when the source and destination locations are on the same drive, since the data doesn't have to physically be copied to a new place in memory.

<!-- MANDATORY: Explain why the change is necessary -->
<!-- Optional: Link to any related GitHub Issues -->


---

## Checklist

<!--
    You (the pull requester) should put an `x` in the boxes below you have completed.

    (If a checkbox doesn't apply to your PR, check it anyway.)
-->

- [x] This PR has only [one purpose or idea](https://terrapower.github.io/armi/developer/tooling.html#one-idea-one-pr).
- [x] [Tests](https://terrapower.github.io/armi/developer/tooling.html#test-it) have been added/updated to verify any new/changed code.

<!-- Check the code quality -->

- [x] The code style follows [good practices](https://terrapower.github.io/armi/developer/standards_and_practices.html).
- [x] The commit message(s) follow [good practices](https://terrapower.github.io/armi/developer/tooling.html).

<!-- Check the project-level cruft -->

- [x] The [release notes](https://terrapower.github.io/armi/developer/tooling.html#add-release-notes) have been updated if necessary.
- [x] The [documentation](https://terrapower.github.io/armi/developer/tooling.html#document-it) is still up-to-date in the `doc` folder.
- [x] The dependencies are still up-to-date in `pyproject.toml`.

mgjarrett: I didn't touch `safeCopy` at all, but I added some `assert` statements to verify that the `safeCopy` call is actually copying the files. The unit test for `safeCopy` passes locally for me, but here it is now failing on those assertions, which is strange.

It's also weird that the existing unit test of `safeCopy` doesn't actually assert that the files are copied successfully. Makes me wonder if there is a known issue here.
opotowsky: > > The secret overwriting of shutil.copy is something I want to remove from ARMI
> 
> I note we already have a ticket for this: #1690
> 
> But I think Michael's PR doesn't touch that, because he didn't alter `safeCopy()`, but added a new `saveMove()`. Am I wrong?

Yep! the first iteration in this PR had a similar setup, (where shutil.move was overwritten with `safeMove`) but that was undone in favor of not overwriting `shutil.move`. 

Patch:
diff --git a/armi/bookkeeping/db/database.py b/armi/bookkeeping/db/database.py
index 884198ae..9a067581 100644
--- a/armi/bookkeeping/db/database.py
+++ b/armi/bookkeeping/db/database.py
@@ -78,6 +78,7 @@ from armi.reactor.reactorParameters import makeParametersReadOnly
 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 @@ class Database:
 
         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 @@ class Database:
         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
diff --git a/armi/bookkeeping/report/reportingUtils.py b/armi/bookkeeping/report/reportingUtils.py
index 815918c3..9ac13287 100644
--- a/armi/bookkeeping/report/reportingUtils.py
+++ b/armi/bookkeeping/report/reportingUtils.py
@@ -377,28 +377,28 @@ def _getSystemInfoLinux():
 
 
 def getSystemInfo():
-    """Get system information, assuming the system is Windows or Linux.
+    """Get system information, assuming the system is Linux, MacOS, and Windows.
 
     Notes
     -----
-    The format of the system information will be different on Windows vs Linux.
+    The format of the system information will be different on Linux, MacOS, and Windows.
 
     Returns
     -------
     str
         Basic system information: OS name, OS version, basic processor information
     """
-    # Get basic system information (on Windows and Linux)
-    if "win" in sys.platform:
+    # Get basic system information (on Linux, MacOS, and Windows)
+    if "darwin" in sys.platform:
+        return _getSystemInfoMac()
+    elif "win" in sys.platform:
         return _getSystemInfoWindows()
     elif "linux" in sys.platform:
         return _getSystemInfoLinux()
-    elif "darwin" in sys.platform:
-        return _getSystemInfoMac()
     else:
         runLog.warning(
             f"Cannot get system information for {sys.platform} because ARMI only "
-            + "supports Linux, Windows, and MacOS."
+            + "supports Linux, MacOS, and Windows."
         )
         return ""
 
diff --git a/armi/context.py b/armi/context.py
index 3dfd7072..259b87d2 100644
--- a/armi/context.py
+++ b/armi/context.py
@@ -82,7 +82,8 @@ START_TIME = time.ctime()
 # Set batch mode if not a TTY, which means you're on a cluster writing to a stdout file. In this
 # mode you cannot respond to prompts. (This does not work reliably for both Windows and Linux so an
 # os-specific solution is applied.)
-isatty = sys.stdout.isatty() if "win" in sys.platform else sys.stdin.isatty()
+IS_WINDOWS = ("win" in sys.platform) and ("darwin" not in sys.platform)
+isatty = sys.stdout.isatty() if IS_WINDOWS else sys.stdin.isatty()
 CURRENT_MODE = Mode.INTERACTIVE if isatty else Mode.BATCH
 Mode.setMode(CURRENT_MODE)
 
diff --git a/armi/utils/__init__.py b/armi/utils/__init__.py
index ae9c6844..6e211813 100644
--- a/armi/utils/__init__.py
+++ b/armi/utils/__init__.py
@@ -806,8 +806,10 @@ def safeCopy(src: str, dst: str) -> None:
     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.copyfile(src, dst)
         shutil.copymode(src, dst)
     elif "linux" in sys.platform:
@@ -816,8 +818,9 @@ def safeCopy(src: str, dst: str) -> None:
     else:
         raise OSError(
             "Cannot perform ``safeCopy`` on files because ARMI only supports "
-            + "Linux and Windows."
+            + "Linux, MacOs, and Windows."
         )
+
     waitTime = 0.01  # 10 ms
     maxWaitTime = 300  # 5 min
     totalWaitTime = 0
@@ -832,11 +835,55 @@ def safeCopy(src: str, dst: str) -> None:
                 f"File copy from {dst} to {src} has failed due to exceeding "
                 + f"a maximum wait time of {maxWaitTime/60} minutes."
             )
-            break
+            Return
 
     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
diff --git a/armi/utils/directoryChangers.py b/armi/utils/directoryChangers.py
index 4f705943..773452d5 100644
--- a/armi/utils/directoryChangers.py
+++ b/armi/utils/directoryChangers.py
@@ -22,6 +22,7 @@ import string
 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 @@ class DirectoryChanger:
         """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 @@ class DirectoryChanger:
             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 @@ class DirectoryChanger:
             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 @@ class DirectoryChanger:
                     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):
diff --git a/doc/release/0.5.rst b/doc/release/0.5.rst
index aaf57f59..91a786ee 100644
--- a/doc/release/0.5.rst
+++ b/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
 -----------


Test patch:
diff --git a/armi/utils/tests/test_utils.py b/armi/utils/tests/test_utils.py
index f0e0b650..86736275 100644
--- a/armi/utils/tests/test_utils.py
+++ b/armi/utils/tests/test_utils.py
@@ -41,6 +41,7 @@ from armi.utils import (
     getStepLengths,
     hasBurnup,
     safeCopy,
+    safeMove,
 )
 
 
@@ -218,6 +219,53 @@ class TestGeneralUtils(unittest.TestCase):
                 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):