diff --git a/src/wormhole/test/test_util.py b/src/wormhole/test/test_util.py index fb58adc..ab6e1d6 100644 --- a/src/wormhole/test/test_util.py +++ b/src/wormhole/test/test_util.py @@ -38,3 +38,11 @@ class Utils(unittest.TestCase): d = util.bytes_to_dict(b) self.assertIsInstance(d, dict) self.assertEqual(d, {"a": "b", "c": 2}) + +class Space(unittest.TestCase): + def test_free_space(self): + free = util.estimate_free_space(".") + self.assert_(isinstance(free, (int, None)), free) + # some platforms (I think the VMs used by travis are in this + # category) return 0, and windows will return None, so don't assert + # anything more specific about the return value diff --git a/src/wormhole/util.py b/src/wormhole/util.py index 967243e..b5e39fb 100644 --- a/src/wormhole/util.py +++ b/src/wormhole/util.py @@ -1,5 +1,5 @@ # No unicode_literals -import json, unicodedata +import os, json, unicodedata from binascii import hexlify, unhexlify def to_bytes(u): @@ -24,3 +24,15 @@ def bytes_to_dict(b): d = json.loads(b.decode("utf-8")) assert isinstance(d, dict) return d + +def estimate_free_space(target): + # f_bfree is the blocks available to a root user. It might be more + # accurate to use f_bavail (blocks available to non-root user), but we + # don't know which user is running us, and a lot of installations don't + # bother with reserving extra space for root, so let's just stick to the + # basic (larger) estimate. + try: + s = os.statvfs(os.path.dirname(os.path.abspath(target))) + return s.f_frsize * s.f_bfree + except AttributeError: + return None