Simple wrapper for OS utilities

diskinfo.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import os
import platform
import re
import subprocess
import sys
import tempfile

RE = re.compile(r'\s{2,}')


class DiskInfoWindows(object):

    def _run_utility(self, script):
        fh, path = tempfile.mkstemp(prefix='diskinfo_', text=True)
        with os.fdopen(fh, 'w') as f:
            f.write(script)
        output = subprocess.check_output(['diskpart', '/S', path])
        os.remove(path)
        self._parse(output.strip())

    def _parse(self, output, partitions=False):
        header = True
        for line in output.split('\n'):
            if header or 0 == len(line):
                if '-----' in line:
                    header = False
                continue
            label, status, size, free = RE.split(line.strip())
            print label, ':', size

    def hdd_list(self):
        self._run_utility('list disk\r\n')

    def hdd_info(self, index):
        self._run_utility('select disk %i\r\nlist partition\r\n' % index)


class DiskInfoLinux(object):
    def _run_utility(self, opts=[]):
        args = ['lsblk', '-rP', '-o', 'NAME,SIZE,TYPE'] + opts
        output = subprocess.check_output(args)
        return output.strip()

    def hdd_list(self):
        output = self._run_utility(['--nodeps'])
        for line in output.split('\n'):
            label, size, kind = line.split()
            if 'rom' == kind:
                continue
            print label, ':', size

    def hdd_info(self, index):
        output = self._run_utility()
        disks = []
        for line in output.split('\n'):
            label, size, kind = line.split()
            if 'disk' == kind:
                disks.append(label)
        for line in output.split('\n'):
            label, size, kind = line.split()
            if 'part' == kind:
                if disks[index] in label:
                    print label, ':', size


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Disks or their partition info.')
    parser.add_argument('index', metavar='I', type=int, default=-1, nargs='?',
                        help='show partitions of a disk by its index')
    args = parser.parse_args()


    osname = platform.system()
    supported = {'Windows': DiskInfoWindows,
                 'Linux': DiskInfoLinux}

    try:
        o = supported[osname]()
    except KeyError:
        print 'Platform %s is not supported yet.'
        sys.exit(-1)
    else:
        if -1 == args.index:
            o.hdd_list()
        else:
            o.hdd_info(args.index)

    sys.exit(0)