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:
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 ofmain()so I could possibly call them in other scripts. But left ones specific tomain()inside.Commenting. More? Less? This is just for me for now, but would you be able to see whats going on?
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()
Thanks for your great pointers. Below is an updated version as per your comments plus a couple other things that occurred to me while editing.
What I changed:
- Revamped if/else statements like you said to try and eliminate unneeded tests and redid all my Y/N/Whatever user prompts using your strategy.
- I did keep errordic global because of my obscure use of it in
unix_call. Is this even a good way to do this? - I ended up moving the filepath checking to its own function
file_path_check()(and usingos.sepwherever needed) because I found 2 other areas where I was trying to do something similar by checking for leading '/' or trailing '/'. Do you think that was a good move? While most of this functionality is done byos.path.joinof course, I couldn't getos.path.jointo add a leading '/' to the path. So I did this instead becauseopen()requires the leading '/' on the path. Was this a good move? - How are my if/elif/else designs for when something other then the expected Y/N input is given? i.e. "unrecognized input"? Couldn't I use a
whileloop to keep prompting the user for the correct input? - Finally, did I use the sys.exit() error numbers correctly? Nothing for success, 1 for normal exit, 2 for error exit?
As always, thanks for the help!
#!/usr/bin/python
##----------------
##Script Name: hd-diag
##Description of Script: Runs diagnostics on harddrives and writes results to
##file
##Date: 2013.07.22-09:39
##Version: 0.1.1
##Requirements: python2, smartmontools, posix environment
##TODO: add support for other incorrect input in YN questions
##----------------
import os
import sys
import subprocess
import getopt
#global variables
errordic = {}
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)
return p.stdout.read()
def write_str_to_file(string, filenamepath):
'''(str) -> string in file
Writes the contents of string to file as array at filenamepath.
Assumes filenamepath is full path.
'''
with open(filenamepath, 'w') as to_file:
for item in string:
to_file.write(str(item))
def file_path_check(path, filename):
'''(str, str) -> str
Concatenates filename and path, checks for completeness. "Current" is
accepted for path; current working directory is then return with filename.
Also, "Null" is accepted for filename; path alone is then returned.
In every case, file separators are checked and added where needed.
If unsure of path, current working direction is returned.
'''
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
if filename.lower() == 'null':
new_filename = ''
else:
new_filename = filename
return os.path.join(new_path, new_filename)
def main(argv):
#main variables
driveloc = ''
outputloc = ''
usage = "hd-diag.py -d <driveloc> -o <outputloc>"
global errordic
#get commandline arguments
#TODO: add option for verbosity
try:
opts, args = getopt.getopt(
argv, 'hd:o:', ['driveloc=', 'outputloc='])
except getopt.GetoptError:
print usage
sys.exit(2)
for opt, arg in opts:
if opt == '-h':
print usage
sys.exit()
elif opt in ('-d', '--driveloc'):
driveloc = arg
elif opt in ('-o', '--outputloc'):
outputloc = arg
#if no input, exit and try again
if not driveloc:
print "Need location of hard drive to test, try again."
print usage
sys.exit()
#inform user of selections
print "Selected drive is: ", driveloc
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
#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)
#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.lower()[0] == 'y':
#console_tstart = test_call("sudo smartctl -t long %s" % driveloc)
#print console_tstart
print "test skipped for debugging..."
elif testYN.lower()[0] == 'n':
print "Test cancelled. Exiting."
sys.exit(1)
else:
print "Unrecognized input. Exiting"
sys.exit(2)
#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.lower()[0] == 'c':
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.lower()[0] == 'e':
print "Test cancelled. Exiting."
sys.exit(1)
else:
print "Unrecognized input. Exiting"
sys.exit(2)
#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 = os.path.join(outputloc, filename)
write_confirm = raw_input(
"The data is ready to be written to:\n\t%s\n\n"
"would you like to continue? Note:File will be overwritten if it "
"already exists.\nContinue?[y/n]: " % finalwrite)
if write_confirm.lower()[0] == 'y':
#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()
elif write_confirm.lower()[0] == 'n':
print "File write canceled. Exiting"
sys.exit(1)
else:
print "Unrecognized input. Exiting"
sys.exit(2)
#run main function
if __name__ == '__main__':
main(sys.argv[1:])
