Average Man's Segment Test

#! /usr/bin/python

#This code tests my temperature monitor display using 7-segment display modules with 2x MCP23017 chips
#(see http://AverageManVsRaspberryPi.com for the full blog)
#Each LED within each 7-segment display should light up in turn, confirming correct sequence (see blog for sequence)
#
#MCP23017 ports can be addressed as followed:
#
#port 0 = 1
#port 1 = 2
#port 2 = 4
#port 3 = 8
#port 4 = 16
#port 5 = 32
#port 6 = 64
#port 7 = 128
#
#MCP23017 has 2 banks/ports, A and B. Each bank has ports 0-7 as above.
#

#Imports
import smbus
import sys
import getopt
import time
import os

#Set up SMBus
bus = smbus.SMBus(0)  #for revision 2 board...use 1

#Set the I2C address
address1 = 0x20 (Address pins all to GND)
address2 = 0x21 (First 2 address pins to GND, last address pin to 5v)

# Set both banks to outputs
bus.write_byte_data(address1,0x00,0x00)
bus.write_byte_data(address1,0x01,0x00)
bus.write_byte_data(address2,0x00,0x00)
bus.write_byte_data(address2,0x01,0x00)

#Bank mapping
bank1 = 1 #Port A on 0x20
bank2 = 2 #Port B on 0x20
bank3 = 3 #Port A on 0x21

#More Bank Mapping
def set_led(data,bank):
  if bank == 1:
   bus.write_byte_data(address1,0x12,data) #Bank 1 = MCP port A
  if bank == 2:
   bus.write_byte_data(address1,0x13,data) #Bank 2 = MCP port B
  if bank == 3:
   bus.write_byte_data(address2,0x12,data) #Bank 3 = MCP port A (second chip)

#Clear all Segments

set_led(0,bank1)
set_led(0,bank2)
set_led(0,bank3)
time.sleep(2)

def main():

  count = 1
  while 1:
    set_led(count,bank1)
    time.sleep(2)
    set_led(count,bank2)
    time.sleep(2)
    set_led(count,bank3)
    time.sleep(2)

    #Clear all Segments
    set_led(0,bank1)
    set_led(0,bank2)
    set_led(0,bank3)
    time.sleep(2)

    #Add count to count, to get the next port number
    count = count + count

    #Check for last segment
    if count == 256:
      return

if __name__ == "__main__":
   main()