Better error handling and dir setting

status/diskusage.py

@@ -3,15 +3,17 @@
 # @Author: KevinMidboe
 # @Date:   2017-01-28 10:54:06
 # @Last Modified by:   KevinMidboe
-# @Last Modified time: 2017-01-31 22:29:43
+# @Last Modified time: 2017-02-18 11:39:00
 
 # f_avail = free blocks avail to non sudo
 # frsize = fundamental file sys byte size
 # f_blocks = total number of blocks in the filesystem
 
-from os import statvfs
+from os import statvfs, path
 from re import sub
 
+# TODO: Add multiple dirs.
+
 def sizeof(num, suffix='B'):
 	for unit in ['','K','M','G','T','P','E','Z']:
 		if abs(num) < 1024.0:
@@ -20,24 +22,18 @@ def sizeof(num, suffix='B'):
 	return "%.1f%s%s" % (num, 'Y', suffix)
 
 # diskUsage uses statvfs to calculate disk size and availability for given disk
-def diskUsage(optionalPath=None):
-	path = '/'
-	if optionalPath != None:
-		path += sub('\+', '/', optionalPath)
+def diskUsage(rootPath=None):
+	if rootPath == None: rootPath = '/'
+	
+	if not path.isdir(rootPath):
+		return {"Error": "Directory does not exist."}
 
-	try:
-		s = statvfs(path)
-		byteLeft = s.f_bavail * s.f_frsize
-		byteTotal = s.f_blocks * s.f_frsize
-		byteUsed = byteTotal-byteLeft
-
-		return { 'left':sizeof(byteLeft), 'used':sizeof(byteUsed), 'total':sizeof(byteTotal) }
-
-	except FileNotFoundError:
-		return None
-	except TypeError:
-		return None
+	s = statvfs(rootPath)
+	byteLeft = s.f_bavail * s.f_frsize
+	byteTotal = s.f_blocks * s.f_frsize
+	byteUsed = byteTotal-byteLeft
+	print(byteLeft, byteTotal, byteUsed)
+	return { 'left': sizeof(byteLeft), 'used': sizeof(byteUsed), 'total': sizeof(byteTotal) }
 
 if __name__=="__main__":
-	n = diskUsage('/media/hdd1')
-	print(n)
+	print(diskUsage())