From 57729c3615aed9f688f30f04f7caa5c3ba42374e Mon Sep 17 00:00:00 2001
From: Scott Zawalski <scottz@google.com>
Date: Wed, 8 Sep 2010 15:59:23 -0700
Subject: [PATCH] Add ListFiles a function to recursively list files in a
 directory. Add unittest for ListFiles to cover common use cases.

BUG=
TEST=

Review URL: http://codereview.chromium.org/3325017
---
 chromite/lib/cros_build_lib.py          | 25 +++++++++
 chromite/lib/cros_build_lib_unittest.py | 69 +++++++++++++++++++++++++
 2 files changed, 94 insertions(+)
 create mode 100755 chromite/lib/cros_build_lib_unittest.py

diff --git a/chromite/lib/cros_build_lib.py b/chromite/lib/cros_build_lib.py
index a0cd73c6ed..825e43cc39 100644
--- a/chromite/lib/cros_build_lib.py
+++ b/chromite/lib/cros_build_lib.py
@@ -4,6 +4,7 @@
 
 """Common python commands used by various build scripts."""
 
+import os
 import subprocess
 import sys
 
@@ -125,3 +126,27 @@ def Info(message):
   """
   print >> sys.stderr, (
       Color(_STDOUT_IS_TTY).Color(Color.BLUE, '\nINFO: ' + message))
+
+
+def ListFiles(base_dir):
+  """Recurively list files in a directory.
+
+  Keyword arguments:
+    base_dir: directory to start recursively listing in.
+
+  Returns:
+    A list of files relative to the base_dir path or
+    An empty list of there are no files in the directories.
+  """
+  directories = [base_dir]
+  files_list = []
+  while directories:
+    directory = directories.pop()
+    for name in os.listdir(directory):
+      fullpath = os.path.join(directory, name)
+      if os.path.isfile(fullpath):
+        files_list.append(fullpath)
+      elif os.path.isdir(fullpath):
+        directories.append(fullpath)
+
+  return files_list
diff --git a/chromite/lib/cros_build_lib_unittest.py b/chromite/lib/cros_build_lib_unittest.py
new file mode 100755
index 0000000000..ef6ee9e27a
--- /dev/null
+++ b/chromite/lib/cros_build_lib_unittest.py
@@ -0,0 +1,69 @@
+#!/usr/bin/python
+
+import errno
+import os 
+import unittest 
+import shutil
+import tempfile 
+
+import cros_build_lib
+
+class TestListFiles(unittest.TestCase):
+
+  def setUp(self):
+    self.root_dir = tempfile.mkdtemp(prefix='listfiles_unittest')
+
+  def tearDown(self):
+    shutil.rmtree(self.root_dir)
+
+  def _createNestedDir(self, dir_structure):
+    for entry in dir_structure:
+      full_path = os.path.join(os.path.join(self.root_dir, entry))
+      # ensure dirs are created
+      try:
+        os.makedirs(os.path.dirname(full_path))
+        if full_path.endswith('/'):
+          # we only want to create directories
+          return
+      except OSError, err:
+        if err.errno == errno.EEXIST:
+          # we don't care if the dir already exists
+          pass
+        else:
+          raise
+      # create dummy files
+      tmp = open(full_path, 'w')
+      tmp.close()
+
+  def testTraverse(self):
+    """
+    Test that we are traversing the directory properly
+    """
+    dir_structure = ['one/two/test.txt', 'one/blah.py',
+                     'three/extra.conf']
+    self._createNestedDir(dir_structure)
+
+    files = cros_build_lib.ListFiles(self.root_dir)
+    for file in files:
+      file = file.replace(self.root_dir, '').lstrip('/')
+      if file not in dir_structure:
+        self.fail('%s was not found in %s' % (file, dir_structure))
+
+  def testEmptyFilePath(self):
+    """
+    Test that we return nothing when directories are empty
+    """
+    dir_structure = ['one/', 'two/', 'one/a/']
+    self._createNestedDir(dir_structure)
+    files = cros_build_lib.ListFiles(self.root_dir)
+    self.assertEqual(files, [])
+
+  def testNoSuchDir(self):
+    try:
+      cros_build_lib.ListFiles('/me/no/existe')
+    except OSError, err:
+      self.assertEqual(err.errno, errno.ENOENT)
+    
+
+if __name__ == '__main__':
+  unittest.main()