Performing SMART tests on many disks

Question

So I'm relatively new to programming in general, and I've finally strapped in and started learning Python. This is my first significant project using it and I'm trying to build good habits with layout of my code, readability, comments, when to make functions, when not to etc.

The problem I was solving was that I have about 20 hard drives that I wanted to run the same tests on using smartmontools on my ubuntu server system and make nice files detailing their status/health before I put them into a cluster system.

Some specific concerns of mine are:

1. Overall workflow. Did I implement use of my main() function correctly? I figured some functions that I could use outside of this script I left outside of main() so I could possibly call them in other scripts. But left ones specific to main() inside.

2. Commenting. More? Less? This is just for me for now, but would you be able to see whats going on?

3. I do lots of string manipulation and it looks sorta messy to me at a glance. Are there cleaner/easier ways to do this?

Any other critiques are welcome!

#!/usr/bin/python
##----------------
##Script Name: hd-diag
##Description of Script: Runs diagnostics on harddrives and writes results to
##file
##Date: 2013.07.09-11:48
##Version: 0.1.0
##Requirements: python2, smartmontools, posix environment
##TODO: add support for other incorrect input in YN questions
##----------------

import os
import sys
import subprocess
import getopt


def unix_call(cmd):
    '''(str) -> str

    Calls unix terminal process cmd and returns stdout upon process completion.
    '''

    output = []
    global errordic
    p = subprocess.Popen(
        cmd, shell=True,
        stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    for line in iter(p.stdout.readline, ''):
        output.append(line.strip('\n'))
    if p.stderr:
        for line in iter(p.stderr.readline, ''):
            errordic['unix_call'] = line.strip('\n')
    return '\n'.join(output).strip("'[]'")


def test_call(cmd):
    p = subprocess.Popen(
        cmd, shell=True, stdout=subprocess.PIPE)
    test_results_raw = p.stdout.read()
    return test_results_raw


def write_str_to_file(string, filenamepath):
    '''(str) -> string in file

    Writes the contents of str to file as array. Assumes current directory
    if just filename is given.

    Full path must be given.
    '''

    def do_write(string, destination):
        to_file = open(destination, 'w')
        for item in string:
            to_file.write(str(item))
        to_file.close()

    if '/' not in filenamepath:
        cur_dir = os.getcwd()
        dest = cur_dir + '/' + filenamepath
        do_write(string, dest)
    elif '/' in filenamepath:
        if filenamepath[0] != '/':
            dest = '/' + filenamepath
            do_write(string, dest)
        elif filenamepath[0] == '/':
            do_write(string, filenamepath)


#set global variables
errordic = {}
driveloc = ''
outputloc = ''


def main():
    global driveloc, outputloc, errordic

    def getargs(argv):
        global driveloc, outputloc, errordic
        try:
            opts, args = getopt.getopt(
                argv, 'hd:o:', ['driveloc=', 'outputloc='])
        except getopt.GetoptError:
            usage()
            sys.exit(2)
        for opt, arg in opts:
            if opt == '-h':
                usage()
                sys.exit()
            elif opt in ('-d', '--driveloc'):
                driveloc = arg
            elif opt in ('-o', '--outputloc'):
                outputloc = arg

    def usage():
        print "hd-diag.py -d <hdpath> -o <outputloc>"

    getargs(sys.argv[1:])

    print "Selected drive is: ", driveloc
    if outputloc != '':
        if outputloc[-1] != '/':
            outputloc = outputloc + '/'
            print "File output location is: ", outputloc
        elif outputloc[-1] == '/':
            print "File output location is: ", outputloc
    elif outputloc == '':
        print "No output location selected. Using current directory."
        outputloc = os.getcwd() + '/'
        print "File output location is: ", outputloc
    #TODO: detect here if drive is ata or sata and print

    #make sure harddrive is unmounted
    while unix_call("mount | grep '%s'" % driveloc):
        unix_call("sudo umount %s" % driveloc)
        if unix_call("mount | grep '%s'" % driveloc):
            print"Can't unmount %s" % driveloc
            print"Err:", errordic['unix_call']
            print"Exiting."
            sys.exit(1)
    else:
        print"%s is not mounted, continuing" % driveloc

    #check if drive supports SMART capability, enabled if need be
    if 'Available' in unix_call(
            "sudo smartctl -i %s | grep 'SMART support is:'" % driveloc):
        print "SMART support is available on %s..." % driveloc
        if 'Enabled' in unix_call(
                "sudo smartctl -i %s | grep 'SMART support is:'" % driveloc):
            print "...and enabled"
        else:
            print "...but disabled"
            while not 'Enabled' in unix_call("smartctl -i %s" % driveloc):
                enableYN = raw_input(
                    "Would you like to enable SMART support? [y/n]:")
                if enableYN in ('yYes'):
                    unix_call("sudo smartctl -s %s" % driveloc)
                elif enableYN in ('nNo'):
                    print "Script cannot continue without SMART capability "
                    "enabled"
                    print "Exiting."
                    sys.exit(1)
                else:
                    print 'Sorry, but "%s" is not a valid input.' % enableYN

    #run desired tests on harddrive
    #get estimated time for extended test
    est_ttime = unix_call(
        "sudo smartctl -c %s | grep -A 1 'Extended self-test'" % driveloc)
    est_ttime = est_ttime[est_ttime.rfind('('):].replace('( ', '(').rstrip('.')
    testYN = raw_input("The estimated time for the test is %s, would you "
                       "like to proceed?[y/n]: " % est_ttime)
    #run test
    if testYN in ('Yy'):
        console_tstart = test_call("sudo smartctl -t long %s" % driveloc)
        print console_tstart
    elif testYN in ('Nn'):
        print 'Test cancelled. Exiting.'
        sys.exit(1)

    #prompt for continue
    contCE = raw_input(
        "\nThe test is running. After the estimated time listed "
        "above has passed, press 'c' to coninue or type 'exit' to "
        "exit.[c/exit]: ")
    #extract test result data
    if contCE in 'cC':
        console_tinfo = test_call(
            "sudo smartctl -i %s" % driveloc).split('\n', 2)
        console_thealth = test_call(
            "sudo smartctl -H %s" % driveloc).split('\n', 2)
        console_tresults = test_call(
            "sudo smartctl -l selftest %s" % driveloc).split('\n', 2)
        console_terror = test_call(
            "sudo smartctl -l error %s" % driveloc).split('\n', 2)
        log_tinfo = str(console_tinfo[2]).lstrip('\n')
        log_thealth = str(console_thealth[2]).lstrip('\n')
        log_tresults = str(console_tresults[2]).lstrip('\n')
        log_terror = str(console_terror[2]).lstrip('\n')
        print log_tinfo + log_thealth + log_tresults + log_terror
    elif contCE in 'Exitexit':
        print 'Test cancelled. Exiting.'
        sys.exit(1)

    #write output to file
    #extract brand, model, serial no., and size for file naming
    deviceinfo = [i.strip('\n').lstrip() for i in test_call(
        "sudo smartctl -i %s" % driveloc).split(':')]
    devicebrand = deviceinfo[2].split(' ')[0]
    devicemodel = deviceinfo[3].split(' ')[0].split('\n')[0]
    deviceserial = deviceinfo[4].split(' ')[0].split('\n')[0]
    devicesize = deviceinfo[6].split(' ')[2].strip('[')
    filename = '-'.join([devicebrand, devicemodel, deviceserial, devicesize])

    #write data to file
    finalwrite = outputloc + filename
    write_confirm = raw_input(
        "The data is ready to be written to:\n\t%s\n\n"
        "would you like to continue?[y/n]: " % finalwrite)
    if write_confirm in 'Yy':
        #TODO test for already existing file
        write_str_to_file(
            log_tinfo + log_thealth + log_tresults + log_terror, finalwrite)
        print "File writen. Test complete. Exiting"
        sys.exit(1)
    elif write_confirm in 'Nn':
        print "File write canceled. Exiting"
        sys.exit(1)

#main
if __name__ == '__main__':
    main()

Answer 1

Having a look at your edited code :

a) In :

    output = []
    global errordic
    p = subprocess.Popen(
        cmd, shell=True,
        stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    for line in iter(p.stdout.readline, ''):
        output.append(line.strip('\n'))
    if p.stderr:
        for line in iter(p.stderr.readline, ''):
            errordic['unix_call'] = line.strip('\n')
    return '\n'.join(output).strip("'[]'")

1) the way you play with output seems way too complicated : if you want to do :

a=[]
for b in c
    a.append(f(b))

Python gives you an nice and efficient way to do so with List Comprehension : a=[f(b) for b in c].

In your case, you could write :

return '\n'.join([l.strip('\n') for l in iter(p.stdout.readline,'')]).strip("'[]'")

but I am not quite sure to understand why your remove newline characters only to put them back afterwards.

2) the way you populate errordic['unix_call'] is quite weird too as you keep replacing the value in such a way that at the end, you only have the latest line. Assuming this is wrong and that I understand what you want to do, list comprehension can help you once again :

errordic['unix_call'] = '\n'.join(l.strip('\n') for l in iter(p.stderr.readline, ''))

Now, this is to be verified but I'd expect my 2 pieces of code to be (roughly) the same as :

errordic['unix_call'] = ''.join(iter(p.stderr.readline, ''))
return ''.join(iter(p.stdout.readline,'')).strip("'[]'")

Also, you could return a pair of strings (stdout and stderr) instead of using global variables in a weird way

b) Then later, do you need unix_call and test_call even though they look quite similar ?

c) In

if path.lower() == 'current':
    new_path = os.getcwd()
else:
    if os.sep not in path:  # unrecognized filepath
        new_path = os.getcwd()
    else:
        if path[0] != os.sep:
            new_path = os.path.join(os.sep, path)
        else:
            new_path = path

You could make things slightly more straigthforward my using

if (path.lower() == 'current' or os.sep not in path):
    new_path = os.getcwd()
elif path[0] == os.sep:
    new_path = path
else:
    new_path = os.path.join(os.sep, path)

d) This :

if filename.lower() == 'null':
    new_filename = ''
else:
    new_filename = filename

could be a simple :

new_filename = '' if (filename.lower() == 'null') else filename

e) This :

if outputloc != '':
    outputloc = file_path_check(outputloc, 'null')
    print "File output location is: ", outputloc
else:
    print "No output location selected. Using current directory."
    outputloc = file_path_check('current', 'null')
    print "File output location is: ", outputloc

could be :

if outputloc == '':
    print "No output location selected. Using current directory."
    outputloc = 'current'
outputloc = file_path_check(outputloc, 'null')
print "File output location is: ", outputloc

f) This :

#make sure drive is unmounted
if not unix_call("mount | grep '%s'" % driveloc):
    print "%s is not mounted, continuing" % driveloc
else:
    while unix_call("mount | grep '%s'" % driveloc):
        unix_call("sudo umount %s" % driveloc)
        if unix_call("mount | grep '%s'" % driveloc):
            print"Can't unmount %s" % driveloc
            print"Err:", errordic['unix_call']
            print"Exiting."
            sys.exit(2)

could be written :

#make sure drive is unmounted
while unix_call("mount | grep '%s'" % driveloc):
    unix_call("sudo umount %s" % driveloc)
    if unix_call("mount | grep '%s'" % driveloc):
        print"Can't unmount %s" % driveloc
        print"Err:", errordic['unix_call']
        print"Exiting."
        sys.exit(2)
print "%s is not mounted, continuing" % driveloc

g) It would probably be worth defining a method taking a string as an argument, prompting it to the user and asking him for a y/n answer until a proper value is given and return a boolean.

Answer 2

Don't open files as this. Its a bad habit. It is easy to forget to close it.

to_file = open(destination, 'w')

Use it as this.

def do_write(string, destination):
    with open(destination, 'w') as to_file: 
        for item in string:
            to_file.write(item)

Also when you are passing a string then there is no need to use str again. Redundant.

You are making redundant elif checks. Obviously if '/' not in filenamepath: is False then you don't need to check whether it is in filepath. Use else instead.

Use os.path.join instead of string manipulation to join paths. It is os independent. Also checking for / isn't a good idea. You can use os.sep which again is os independent. But it would be a good idea to make sure that filename does not contain / by itself.

In the comments you are saying that full path must be given but in the function you are checking for that. Why? Decide what you want the function to do. It does both? Then you can use this

def write_str_to_file(string, filenamepath):
    '''Writes the contents of str to file as array.

    Assumes current directory if just filename is given
    otherwise full path must be given
    '''
    if os.sep not in filenamepath:
        dest = os.path.join(os.getcwd(), filenamepath)
    else:
        dest = filenamepath

    with open(dest, 'w') as to_file:
        for item in string:
            to_file.write(str(item))

Don't define global variables in the middle of nowhere. It makes it harder to find them. Put all globals at the top. Don't use global variables if you don't have to. Declare them in the main function and pass them around. You can return more than one values so problems with passing and returning all of them. Having global variables when you don't need them shows bad design. It might be necessary sometimes but it doesn't seem to be the in this case.

You don't need the additional variable test_results_raw in case of test_call. You can instead use

def test_call(cmd):
    p = subprocess.Popen(
        cmd, shell=True, stdout=subprocess.PIPE)
    return p.stdout.read()

Is the function usage actually necessary in the main function. You can define a constant at the beginning of the module and print that instead. Function calls are expensive and doing them for 1 print statement doesn't seem to be a good idea.

Again in the main function you are using if outputloc != '' and then elif outputloc == ''. It can be a simple else. Same goes for the nested if and elif. Checking for more conditions costs you in performance.

When checking for enableYN you can use enable.lower()[0] == 'y' instead of using enableYN in ('yYes'). In that case it checks for e and s also. People would be entering y, Y, yes or some other variation but with a y in the beginning. Same goes for checking for no.

I skipped a lot of the code because it is quite big.
Hope this helped.

---

About me being able to see what's going on by the comments. Not really. You missed what the main function is doing

About the strings being messy. Yes, they are. Try to change the things I have mentioned and I'll take a look again. Add your updated code like done in this question instead of modifying it.

Answer 3

Why are you using errordic? If you're trying to have a function return to variables, just do that.

def one_two():
    return 1, 2
one, two = one_two

It also seems like subprocess.communicate does what you want.

test_call seems to be a clone of subprocess.check_output.