Add glob option to directory changer file retrieval. (#68)

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):

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"))