From: Brian Warner <warner@allmydata.com>
Date: Tue, 3 Jul 2007 22:49:45 +0000 (-0700)
Subject: fileutil: add du() function
X-Git-Url: https://git.rkrishnan.org/?a=commitdiff_plain;h=9ddb92965161a067dabdffd838847305eb50591e;p=tahoe-lafs%2Ftahoe-lafs.git

fileutil: add du() function
---

diff --git a/src/allmydata/test/test_util.py b/src/allmydata/test/test_util.py
index 7951b1cb..0b7138cd 100644
--- 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 25ad9761..ac2391a9 100644
--- 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