fileutil: add du() function

diff --git a/src/allmydata/test/test_util.py b/src/allmydata/test/test_util.py
index 7951b1cb04279b00346e012fcb38a2cecb1c0d56..0b7138cde4f6989700ecfdf886a2994015f070bb 100644 (file)
--- a/src/allmydata/test/test_util.py
+++ b/src/allmydata/test/test_util.py
@@ -276,10 +276,10 @@ class FileUtil(unittest.TestCase):
         fn = os.path.join(basedir, path)
         fileutil.make_dirs(fn, mode)
 
-    def touch(self, basedir, path, mode=None):
+    def touch(self, basedir, path, mode=None, data="touch\n"):
         fn = os.path.join(basedir, path)
         f = open(fn, "w")
-        f.write("touch\n")
+        f.write(data)
         f.close()
         if mode is not None:
             os.chmod(fn, mode)
@@ -358,3 +358,18 @@ class FileUtil(unittest.TestCase):
         fileutil.rename(fn, fn2)
         self.failIf(os.path.exists(fn))
         self.failUnless(os.path.exists(fn2))
+
+    def test_du(self):
+        basedir = "util/FileUtil/test_du"
+        fileutil.make_dirs(basedir)
+        d = os.path.join(basedir, "space-consuming")
+        self.mkdir(d, "a/b")
+        self.touch(d, "a/b/1.txt", data="a"*10)
+        self.touch(d, "a/b/2.txt", data="b"*11)
+        self.mkdir(d, "a/c")
+        self.touch(d, "a/c/1.txt", data="c"*12)
+        self.touch(d, "a/c/2.txt", data="d"*13)
+
+        used = fileutil.du(basedir)
+        self.failUnlessEqual(10+11+12+13, used)
+

diff --git a/src/allmydata/util/fileutil.py b/src/allmydata/util/fileutil.py
index 25ad9761f8636a56759bd0aec6e2332082abfc65..ac2391a98b69873510fd0683ddd4422e37f64473 100644 (file)
--- a/src/allmydata/util/fileutil.py
+++ b/src/allmydata/util/fileutil.py
@@ -184,3 +184,14 @@ def open_or_create(fname, binarymode=True):
         return open(fname, binarymode and "r+b" or "r+")
     except EnvironmentError:
         return open(fname, binarymode and "w+b" or "w+")
+
+
+def du(basedir):
+    size = 0
+
+    for root, dirs, files in os.walk(basedir):
+        for f in files:
+            fn = os.path.join(root, f)
+            size += os.path.getsize(fn)
+
+    return size