[Python-checkins] cpython: Issue #12442: add shutil.disk_usage()

[giampaolo.rodola]

http://hg.python.org/cpython/rev/2fc102ebaf73
changeset:   71110:2fc102ebaf73
user:        Giampaolo Rodola' <g.rodola at gmail.com>
date:        Fri Jul 01 13:55:36 2011 +0200
summary:
  Issue #12442: add shutil.disk_usage()

files:
  Doc/library/shutil.rst  |   8 ++++++++
  Doc/whatsnew/3.3.rst    |  11 ++++++++++-
  Lib/shutil.py           |  19 +++++++++++++++++++
  Lib/test/test_shutil.py |  10 ++++++++++
  Misc/NEWS               |   3 +++
  Modules/posixmodule.c   |  27 +++++++++++++++++++++++++++
  6 files changed, 77 insertions(+), 1 deletions(-)


diff --git a/Doc/library/shutil.rst b/Doc/library/shutil.rst
--- a/Doc/library/shutil.rst
+++ b/Doc/library/shutil.rst
@@ -164,6 +164,14 @@
    If the destination is on the current filesystem, then simply use rename.
    Otherwise, copy src (with :func:`copy2`) to the dst and then remove src.
 
+.. function:: disk_usage(path)
+
+   Return disk usage statistics about the given path as a namedtuple including
+   total, used and free space expressed in bytes.
+
+   .. versionadded:: 3.3
+
+   Availability: Unix, Windows.
 
 .. exception:: Error
 
diff --git a/Doc/whatsnew/3.3.rst b/Doc/whatsnew/3.3.rst
--- a/Doc/whatsnew/3.3.rst
+++ b/Doc/whatsnew/3.3.rst
@@ -200,7 +200,16 @@
 plaintex.  This can be useful to take advantage of firewalls that know how to
 handle NAT with non-secure FTP without opening fixed ports.
 
-(Patch submitted by Giampaolo Rodolà in :issue:`12139`.)
+(Contributed by Giampaolo Rodolà in :issue:`12139`)
+
+
+shutil
+------
+
+The :mod:`shutil` module has a new :func:`~shutil.disk_usage` providing total,
+used and free disk space statistics.
+
+(Contributed by Giampaolo Rodolà in :issue:`12442`)
 
 
 Optimizations
diff --git a/Lib/shutil.py b/Lib/shutil.py
--- a/Lib/shutil.py
+++ b/Lib/shutil.py
@@ -12,6 +12,7 @@
 import collections
 import errno
 import tarfile
+from collections import namedtuple
 
 try:
     import bz2
@@ -754,3 +755,21 @@
         func = _UNPACK_FORMATS[format][1]
         kwargs = dict(_UNPACK_FORMATS[format][2])
         func(filename, extract_dir, **kwargs)
+
+if hasattr(os, "statvfs") or os.name == 'nt':
+    _ntuple_diskusage = namedtuple('usage', 'total used free')
+
+    def disk_usage(path):
+        """Return disk usage statistics about the given path as a namedtuple
+        including total, used and free space expressed in bytes.
+        """
+        if hasattr(os, "statvfs"):
+            st = os.statvfs(path)
+            free = (st.f_bavail * st.f_frsize)
+            total = (st.f_blocks * st.f_frsize)
+            used = (st.f_blocks - st.f_bfree) * st.f_frsize
+        else:
+            import nt
+            total, free = nt._getdiskusage(path)
+            used = total - free
+        return _ntuple_diskusage(total, used, free)
diff --git a/Lib/test/test_shutil.py b/Lib/test/test_shutil.py
--- a/Lib/test/test_shutil.py
+++ b/Lib/test/test_shutil.py
@@ -728,6 +728,16 @@
         unregister_unpack_format('Boo2')
         self.assertEqual(get_unpack_formats(), formats)
 
+    @unittest.skipUnless(hasattr(shutil, 'disk_usage'),
+                         "disk_usage not available on this platform")
+    def test_disk_usage(self):
+        usage = shutil.disk_usage(os.getcwd())
+        self.assertTrue(usage.total > 0)
+        self.assertTrue(usage.used > 0)
+        self.assertTrue(usage.free >= 0)
+        self.assertTrue(usage.total >= usage.used)
+        self.assertTrue(usage.total > usage.free)
+
 
 class TestMove(unittest.TestCase):
 
diff --git a/Misc/NEWS b/Misc/NEWS
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -200,6 +200,9 @@
 Library
 -------
 
+- Issue #12442: new shutil.disk_usage function, providing total, used and free
+  disk space statistics.
+
 - Issue #12451: The XInclude default loader of xml.etree now decodes files from
   UTF-8 instead of the locale encoding if the encoding is not specified. It now
   also opens XML files for the parser in binary mode instead of the text mode
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
--- a/Modules/posixmodule.c
+++ b/Modules/posixmodule.c
@@ -7451,6 +7451,32 @@
 }
 #endif /* HAVE_STATVFS */
 
+#ifdef MS_WINDOWS
+PyDoc_STRVAR(win32__getdiskusage__doc__,
+"_getdiskusage(path) -> (total, free)\n\n\
+Return disk usage statistics about the given path as (total, free) tuple.");
+
+static PyObject *
+win32__getdiskusage(PyObject *self, PyObject *args)
+{
+    BOOL retval;
+    ULARGE_INTEGER _, total, free;
+    LPCTSTR path;
+
+    if (! PyArg_ParseTuple(args, "s", &path))
+        return NULL;
+
+    Py_BEGIN_ALLOW_THREADS
+    retval = GetDiskFreeSpaceEx(path, &_, &total, &free);
+    Py_END_ALLOW_THREADS
+    if (retval == 0)
+        return PyErr_SetFromWindowsErr(0);
+
+    return Py_BuildValue("(LL)", total.QuadPart, free.QuadPart);
+}
+#endif
+
+
 /* This is used for fpathconf(), pathconf(), confstr() and sysconf().
  * It maps strings representing configuration variable names to
  * integer values, allowing those functions to be called with the
@@ -9716,6 +9742,7 @@
     {"_getfinalpathname",       posix__getfinalpathname, METH_VARARGS, NULL},
     {"_getfileinformation",     posix__getfileinformation, METH_VARARGS, NULL},
     {"_isdir",                  posix__isdir, METH_VARARGS, posix__isdir__doc__},
+    {"_getdiskusage",           win32__getdiskusage, METH_VARARGS, win32__getdiskusage__doc__},
 #endif
 #ifdef HAVE_GETLOADAVG
     {"getloadavg",      posix_getloadavg, METH_NOARGS, posix_getloadavg__doc__},

-- 
Repository URL: http://hg.python.org/cpython