From 0f6b64b102b8df716dd6e4249cd18583711cac3e Mon Sep 17 00:00:00 2001
From: CRImier <crimier@yandex.ru>
Date: Sat, 19 Nov 2016 02:48:53 +0200
Subject: [PATCH] Added size calculation function, redefined all functions

---
 pyrtitions/__init__.py   |  6 ++++-
 pyrtitions/pyrtitions.py | 49 +++++++++++++++++++++++++++++++---------
 2 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/pyrtitions/__init__.py b/pyrtitions/__init__.py
index 2de73b5..3dd7ed0 100644
--- a/pyrtitions/__init__.py
+++ b/pyrtitions/__init__.py
@@ -1 +1,5 @@
-from pyrtitions import get_partitions, get_partitions_and_mounts, generate_mountpoint, label_filter, pprint_partitions
+from pyrtitions import get_partitions, get_uuids_and_labels, generate_mountpoint, get_device_sizes, get_size_from_block_count, get_mounts, label_filter, pprint_partitions
+
+
+if __name__ == "__main__":             
+    pprint_partitions(get_partitions())
diff --git a/pyrtitions/pyrtitions.py b/pyrtitions/pyrtitions.py
index d216cfa..3685f79 100755
--- a/pyrtitions/pyrtitions.py
+++ b/pyrtitions/pyrtitions.py
@@ -2,7 +2,7 @@
 import shlex
 import string
 
-def get_partitions():
+def get_uuids_and_labels():
     partitions = []
     labels = {}
     dbu_dir = "/dev/disk/by-uuid/"
@@ -27,15 +27,31 @@ def get_partitions():
         #[{"uuid":"5OUU1DMUCHUNIQUEW0W", "path":"/dev/sda1"}, {"label":"label1", "uuid":"MANYLETTER5SUCH1ONGWOW", "path":"/dev/sdc1"}]
     return partitions
 
-def get_partitions_and_mounts():
-    partitions = get_partitions()
-    mounted_partitions = {}
+def get_size_from_block_count(block_count_str, size_step=1000, sizes = ["K", "M", "G", "T"], format_spec=":2.2f"):
+    block_count = float(block_count_str)
+    size_counter = 0
+    while block_count >= float(size_step):
+        block_count /= size_step
+        size_counter += 1
+    return ("{"+format_spec+"}{}").format(block_count, sizes[size_counter])
 
-    #Good source of information about mounted partitions is /etc/mtab
-    mtab_path = "/etc/mtab" 
-    f = open(mtab_path, "r")
-    lines = f.readlines()
-    f.close()
+def get_device_sizes():
+    with open("/proc/partitions", "r") as f:
+        lines = f.readlines()
+    info_dict = {}
+    lines = lines[1:] #First line is the header
+    for line in lines:
+        line = line.strip()
+        if line: 
+            major, minor, block_count_str, name = line.split()
+            hr_size = get_size_from_block_count(block_count_str)
+            info_dict[name] = [hr_size, block_count_str, int(major), int(minor)]
+    return info_dict
+            
+def get_mounts(mounts_file="/etc/mtab"):
+    mounted_partitions = {}
+    with open(mounts_file, "r") as f:
+        lines = f.readlines()
     for line in lines:
         line = line.strip()
         if line: 
@@ -52,6 +68,12 @@ def get_partitions_and_mounts():
             if device_path.startswith("/dev"):
                 device_path = os.path.realpath(device_path) #Can be a symlink?
                 mounted_partitions[device_path] = [mountpoint, part_type, mount_opts]
+    return mounted_partitions
+
+def get_partitions():
+    partitions = get_uuids_and_labels()
+    mounted_partitions = get_mounts()
+    partition_sizes = get_device_sizes()
     for entry in partitions:
          if entry["path"] in mounted_partitions:
              mpoint, ptype, opts = mounted_partitions[entry["path"]]
@@ -62,8 +84,14 @@ def get_partitions_and_mounts():
          else:
              entry["mounted"] = False
              entry["mountpoint"] = None
+         basename = os.path.basename(entry["path"])
+         if basename in partition_sizes:
+             entry["size"] = partition_sizes[basename][0]
+         else:
+             entry["size"] = None
     return partitions
 
+
 def generate_mountpoint(part_info, base_dir="/media"):
     """Generates a valid mountpoint path, for example, for automatic mount purposes"""
     #We could use either label (easier and prettier)
@@ -102,5 +130,4 @@ def pprint_partitions(partitions):
             print("\t{}:{}".format(key, value))
 
 if __name__ == "__main__":
-    partitions = get_partitions_and_mounts()
-    pprint_partitions(partitions)
+    pprint_partitions(get_partitions())