Clean up code

fancontrol.py

@@ -1,85 +1,116 @@
 #!/usr/bin/env python
+# -*- coding: utf8 -*-
+
+import os
 import sys
 import time
 import glob
+import stat
 
 try:
     from StringIO import StringIO
 except ImportError:
     from io import StringIO
-from subprocess import check_output, Popen, PIPE, STDOUT
+
+from subprocess import check_output, Popen, PIPE, STDOUT, CalledProcessError
+
+devices_to_check = ['/dev/sg*', '/dev/ses*']
+
+
+def usage():
+    print('python fancontrol.py 1-7')
+    sys.exit(-1)
+
+
+def get_requested_fan_speed():
+    if len(sys.argv) < 2:
+        return usage()
+    try:
+        speed = int(sys.argv[1])
+    except ValueError:
+        return usage()
+    if not 1 <= speed <= 7:
+        return usage()
+    return speed
 
 
 def print_speeds(device):
     for i in range(0, 6):
-        print("Fan {} speed: {}".format(i, check_output(
+        print('Fan {} speed: {}'.format(i, check_output(
             ['sg_ses', '--index=coo,{}'.format(i), '--get=1:2:11', device]).decode('utf-8').split('\n')[0]))
 
 
-if len(sys.argv) < 1:
+def find_sa120_devices():
-    print("python fancontrol.py 1-7")
+    devices = []
-    sys.exit(-1)
+    for device_glob in devices_to_check:
-
+        for device in glob.glob(device_glob):
-fan = int(sys.argv[1])
-
-if fan <= 0 or fan > 6:
-    raise Exception("Fan speed must be between 1 and 7")
-
-devices_to_check = ['/dev/sg*', '/dev/ses*']
-
-device = ""
-
-for device_glob in devices_to_check:
-    for dev_node in glob.glob(device_glob):
             try:
-            out = check_output(["sg_ses", dev_node], stderr=STDOUT)
+                output = check_output(['sg_ses', device], stderr=STDOUT)
-            if 'ThinkServerSA120' in out:
+                if b'ThinkServerSA120' in output:
-                device = dev_node
+                    print('Enclosure found on ' + device)
-                print("Enclosure found on " + device)
+                    devices.append(device)
+                else:
+                    print('Enclosure not found on ' + device)
+            except CalledProcessError:
+                print('Enclosure not found on ' + device)
+    return devices
 
+
+def set_fan_speeds(device, speed):
     print_speeds(device)
-                print("Reading current configuration...")
+    print('Reading current configuration...')
-                out = check_output(["sg_ses", "-p", "0x2", device, "--raw"]).decode('utf-8')
+    out = check_output(['sg_ses', '-p', '0x2', device, '--raw']).decode('utf-8')
+
     s = out.split()
 
     for i in range(0, 6):
-                    print("Setting fan {} to {}".format(i, fan))
+        print('Setting fan {} to {}'.format(i, speed))
         idx = 88 + 4 * i
-                    s[idx + 0] = "80"
+        s[idx + 0] = '80'
-                    s[idx + 1] = "00"
+        s[idx + 1] = '00'
-                    s[idx + 2] = "00"
+        s[idx + 2] = '00'
-                    s[idx + 3] = format(1 << 5 | fan & 7, "x")
+        s[idx + 3] = format(1 << 5 | speed & 7, 'x')
 
     output = StringIO()
     off = 0
     count = 0
-                line = ''
 
     while True:
         output.write(s[off])
         off = off + 1
         count = count + 1
         if count == 8:
-                        output.write("  ")
+            output.write('  ')
         elif count == 16:
-                        output.write("\n")
+            output.write('\n')
             count = 0
         else:
-                        output.write(" ")
+            output.write(' ')
         if off >= len(s):
             break
 
-                output.write("\n")
+    output.write('\n')
     p = Popen(['sg_ses', '-p', '0x2', device, '--control', '--data', '-'], stdout=PIPE, stdin=PIPE,
               stderr=PIPE)
-                print("Set fan speeds... Waiting to get fan speeds (ctrl+c to skip)")
+    print('Set fan speeds... Waiting to get fan speeds (ctrl+c to skip)')
     print(p.communicate(input=bytearray(output.getvalue(), 'utf-8'))[0].decode('utf-8'))
+    try:
         time.sleep(10)
         print_speeds(device)
-        except:
+    except KeyboardInterrupt:
-            print("Enclosure not found on " + dev_node)
+        pass
 
-if device == "":
+
-    print("Could not find enclosure")
+def main():
+    speed = get_requested_fan_speed()
+    devices = find_sa120_devices()
+    if not devices:
+        print('Could not find enclosure')
         sys.exit(1)
+    for device in devices:
+        set_fan_speeds(device, speed)
+    print('\nDone')
 
+
+if __name__ == '__main__':
+    main()