diff --git a/armi/utils/directoryChangers.py b/armi/utils/directoryChangers.py index 02518b0bc..4344a38c2 100644 --- a/armi/utils/directoryChangers.py +++ b/armi/utils/directoryChangers.py @@ -16,6 +16,7 @@ import random import shutil import string +import glob import armi from armi import runLog @@ -133,6 +134,19 @@ def _retrieveEntireFolder(self): def _transferFiles(initialPath, destinationPath, fileList): """ Transfer files into or out of the directory. + + Parameters + ---------- + initialPath : str + Path to the folder to find files in. + destinationPath: str + Path to the folder to move file to. + fileList : list of str or list of tuple + File names to move from initial to destination. If this is a + simple list of strings, the 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. .. warning:: On Windows the max number of characters in a path is 260. If you exceed this you will see FileNotFound errors here. @@ -142,17 +156,23 @@ def _transferFiles(initialPath, destinationPath, fileList): return if not os.path.exists(destinationPath): os.mkdir(destinationPath) - for ff in fileList: - if isinstance(ff, tuple): + for pattern in fileList: + if isinstance(pattern, tuple): # allow renames in transit - fromName, destName = ff + fromName, destName = pattern + copies = [(fromName, destName)] else: - fromName, destName = ff, ff - - fromPath = os.path.join(initialPath, fromName) - toPath = os.path.join(destinationPath, destName) - runLog.extra("Copying {} to {}".format(fromPath, toPath)) - shutil.copy(fromPath, toPath) + # expand globs if they're given + copies = [] + for ff in glob.glob(pattern): + # renaming not allowed with globs + copies.append((ff, ff)) + + for fromName, destName in copies: + fromPath = os.path.join(initialPath, fromName) + toPath = os.path.join(destinationPath, destName) + runLog.extra("Copying {} to {}".format(fromPath, toPath)) + shutil.copy(fromPath, toPath) class TemporaryDirectoryChanger(DirectoryChanger): diff --git a/armi/utils/tests/test_directoryChangers.py b/armi/utils/tests/test_directoryChangers.py index c4c1e498e..fee4ce73c 100644 --- a/armi/utils/tests/test_directoryChangers.py +++ b/armi/utils/tests/test_directoryChangers.py @@ -108,3 +108,36 @@ def test_temporary_cleans(self): tempName = dc.destination self.assertFalse(os.path.exists(tempName)) + + def test_file_retrieval(self): + """ + Make sure requested files and/or globs get copied back. + + * Checks basic copy feature + * Checks rename feature + * Checks glob expansion + """ + + def f(name): + """Utility to avoid test clashes during cleanups""" + return self._testMethodName + name + + with directoryChangers.TemporaryDirectoryChanger( + filesToRetrieve=[(f("file1.txt"), f("newfile1.txt"))] + ) as dc: + Path(f("file1.txt")).touch() + Path(f("file2.txt")).touch() + + self.assertTrue(os.path.exists(f("newfile1.txt"))) + os.remove(f("newfile1.txt")) + + with directoryChangers.TemporaryDirectoryChanger( + filesToRetrieve=[f("file*.txt")] + ) as dc: + Path(f("file1.txt")).touch() + Path(f("file2.txt")).touch() + + self.assertTrue(os.path.exists(f("file1.txt"))) + self.assertTrue(os.path.exists(f("file2.txt"))) + os.remove(f("file1.txt")) + os.remove(f("file2.txt"))