Merge branch '6090-docker-use-local-sso' closes #6138

[arvados.git] / sdk / python / arvados / util.py

diff --git a/sdk/python/arvados/util.py b/sdk/python/arvados/util.py
index 2532ee2cf1bbf6c1051ec053a4c42433990455ab..1316f2287f6e73bd1a2524ec331581bc5850f0b1 100644 (file)
--- a/sdk/python/arvados/util.py
+++ b/sdk/python/arvados/util.py
@@ -5,7 +5,9 @@ import re
 import subprocess
 import errno
 import sys
-from arvados.collection import *
+
+import arvados
+from arvados.collection import CollectionReader
 
 HEX_RE = re.compile(r'^[0-9a-fA-F]+$')
 
@@ -24,7 +26,7 @@ def clear_tmpdir(path=None):
     Ensure the given directory (or TASK_TMPDIR if none given)
     exists and is empty.
     """
-    if path == None:
+    if path is None:
         path = arvados.current_task().tmpdir
     if os.path.exists(path):
         p = subprocess.Popen(['rm', '-rf', path])
@@ -308,13 +310,31 @@ def stream_extract(stream, path, files=[], decompress=True):
     lockfile.close()
     return path
 
-def listdir_recursive(dirname, base=None):
+def listdir_recursive(dirname, base=None, max_depth=None):
+    """listdir_recursive(dirname, base, max_depth)
+
+    Return a list of file and directory names found under dirname.
+
+    If base is not None, prepend "{base}/" to each returned name.
+
+    If max_depth is None, descend into directories and return only the
+    names of files found in the directory tree.
+
+    If max_depth is a non-negative integer, stop descending into
+    directories at the given depth, and at that point return directory
+    names instead.
+
+    If max_depth==0 (and base is None) this is equivalent to
+    sorted(os.listdir(dirname)).
+    """
     allfiles = []
     for ent in sorted(os.listdir(dirname)):
         ent_path = os.path.join(dirname, ent)
         ent_base = os.path.join(base, ent) if base else ent
-        if os.path.isdir(ent_path):
-            allfiles += listdir_recursive(ent_path, ent_base)
+        if os.path.isdir(ent_path) and max_depth != 0:
+            allfiles += listdir_recursive(
+                ent_path, base=ent_base,
+                max_depth=(max_depth-1 if max_depth else None))
         else:
             allfiles += [ent_base]
     return allfiles
@@ -331,8 +351,8 @@ def is_hex(s, *length_args):
     """
     num_length_args = len(length_args)
     if num_length_args > 2:
-        raise ArgumentError("is_hex accepts up to 3 arguments ({} given)".
-                            format(1 + num_length_args))
+        raise errors.ArgumentError("is_hex accepts up to 3 arguments ({} given)"
+                                   .format(1 + num_length_args))
     elif num_length_args == 2:
         good_len = (length_args[0] <= len(s) <= length_args[1])
     elif num_length_args == 1: