imglob made into a standalone function

fsl/scripts/imglob.py

@@ -26,24 +26,33 @@ exts   = fslimage.ALLOWED_EXTENSIONS
 groups = fslimage.FILE_GROUPS
 
 
-def main(argv=None):
-    """The ``imglob`` utility. Given a list of file names, identifies and prints
-    the unique NIFTI/ANALYZE image files.
+def imglob(paths, output=None):
+    """Given a list of file names, identifies and returns the unique
+    NIFTI/ANALYZE image files that exist.
+
+    :arg paths:  Sequence of paths/prefixes to glob.
+
+    :arg output: One of ``'prefix'`` (the default), ``'all'``, or
+                 ``'primary'``:
+
+                  - ``'prefix'``:  Returns the files without extensions.
+                  - ``'all'``:     Returns all files that match (e.g. both
+                                   ``.img`` and ``.hdr`` files will be
+                                   returned).
+                  - ``'primary'``: Returns only the primary file of each
+                                   matching file group, e.g. only the
+                                   ``.hdr`` file would be returned from
+                                   an ``.img``/``.hdr`` pair.
+
+    :returns: A sequence of resolved path names, in the form specified
+              by the ``output`` parameter.
     """
 
-    if argv is None:
-        argv = sys.argv[1:]
-
-    if len(argv) < 1:
-        print(usage)
-        return 1
-
-    if   argv[0] == '-extension':  output = 'primary'
-    elif argv[0] == '-extensions': output = 'all'
-    else:                          output = 'prefix'
+    if output is None:
+        output = 'prefix'
 
-    if output == 'prefix': paths = argv
-    else:                  paths = argv[1:]
+    if output not in ('prefix', 'all', 'primary'):
+        raise ValueError('Unsupported output format: {}'.format(output))
 
     imgfiles = []
 
@@ -51,6 +60,7 @@ def main(argv=None):
     # hdr and img and otherwise) that match
     for path in paths:
         try:
+            path = fslimage.removeExt(path)
             imgfiles.extend(fslimage.addExt(path, unambiguous=False))
         except fslpath.PathError:
             continue
@@ -66,9 +76,32 @@ def main(argv=None):
                                             allowedExts=exts,
                                             fileGroups=groups)
 
-    imgfiles = sorted(set(imgfiles))
+    return list(sorted(set(imgfiles)))
+
+
+def main(argv=None):
+    """The ``imglob`` application. Given a list of file names, identifies and
+    prints the unique NIFTI/ANALYZE image files.
+    """
+
+    if argv is None:
+        argv = sys.argv[1:]
+
+    if len(argv) < 1:
+        print(usage)
+        return 1
+
+    if   argv[0] == '-extension':  output = 'primary'
+    elif argv[0] == '-extensions': output = 'all'
+    else:                          output = 'prefix'
+
+    if output == 'prefix': paths = argv
+    else:                  paths = argv[1:]
+
+    imgfiles = imglob(paths, output)
 
-    print(' '.join(imgfiles))
+    if len(imgfiles) > 0:
+        print(' '.join(imgfiles))
 
     return 0