<commit_before>import argparse
import os
import re
os.system('cls' if os.name=='nt' else 'clear')

#=============FUNCTION DEFINITIONS===============#
def hline():
    """This function simply prints a line of 60 "-" chars, to add divisions to screen output"""
    print "-" * 60

def parse_vars_file(file):
    """This function takes the file with all the variable and their values, listed for each host, and
        puts it into a data structure that can later be used to write out each config script.  It also
        tracks the names of hosts found in the file and returns those for quick review prior to exporting
        the configurations.

        The data structure is a dictionary with each hostname as the key, and the value being another
        dictionary (sub-dictionary?) with each find/replace key and its associated value)"""

    #Reset file position to beginnig (just in case)
    file.seek(0)
    host_list = []
    dict_vars = {}
    #This variable tracks whether the line being parsed is part of a host "block" of variables or not.
    hostblock = 0

    for line in file:
        #Ignore comment lines
        if line.startswith("#"):
                pass

        #Use a marker so we know if we are inside of a particular host's block of variables.  When we see a
        #HOSTNAME var, it is set to 1, at <--END--> it is set back to 0.
        elif "<--END-->" in line:
            hostblock = 0

        #This string "::" is used to separate the key and value in the variable file.  If any line doesn't
        #have this (with exception of the END line, shown above), ignore it and print an error.
        elif "::" not in line:
            print "Did NOT find Seperator '%r'. Line ignored" % line

        #If we find a HOSTNAME variable, but the function thinks we are still in a host's block of variables,
        #then throw an error and exit.
        elif "<HOSTNAME>" in line and hostblock == 1:
            hline()
            print "ERROR: <HOSTNAME> found inside another host's variable block."
            print "ERROR LINE: '%r'" % line
            print "Exiting program. Please correct line and try again."
            exit(0)

        #HOSTNAME variable after an <--END--> line (hostblock=0) triggers the start of a new host block, and
        #creates an entry in the dictionary for the host.
        elif "<HOSTNAME>" in line and hostblock == 0:
            #Split the variable (<HOSTNAME>) from the value
            label, hostname = line.split('::')
            hostname = hostname.rstrip()
            for imported_host in host_list:
                if hostname == imported_host:
                    print "ERROR: DUPLICATE HOSTNAME IN VARIABLES FILE. EXITING..."
                    exit(0)
            hostblock = 1
            #Create an entry in the dictionary for the host.  The value will be another dictionary.
            host_list.append(hostname)
            dict_vars[hostname] = {}
            #The first value of the sub-dictionary will be the <HOSTNAME> variable and actual hostname for
            #this device
            dict_vars[hostname][label] = hostname

        #If the other statements don't match, and we are inside a host's block of variables, then assume
        #this line is a variable and add it to the sub-dictionary for the particular host.
        elif hostblock == 1:
            #This should match all find/replace lines under a particular hostname.  This code just adds
            #the pairs to the sub-dictionary for the associated host.
            key, value = line.split('::')
            value = value.rstrip()
            dict_vars[hostname][key] = value

        else:
            #If this section matches a line, something was wrong with the file.  It prints an output of
            #the line and under which host this occurred to help track down the error (for example, it
            #could be HOSTNAME line that had an error, so all lines for that host will match this)
            print "INVALID LINE in %s: '%r'" % (hostname, line)

    #Return the list of hosts, the host count, and the data structure to the main program.
    return host_list, dict_vars

def write_configs(vars, template):
    """This template writes each host's configuration file.  It should be passed both the template file
        to use, as well as the data structure generated by the parse_vars_file() function."""

    if args.verbose:
        hline()
    for host in vars:
        #reset the template read position to start for each host
        template.seek(0)
        #create file name based on hostname of the device
        filename = host + ".txt"
        hostinfo = vars[host]
        #Check that the "configs" directory exists.  If not, create it.
        if not os.path.exists("configs"):
            os.makedirs("configs")
        dstfile = open("./configs/" + filename, 'w')
        if args.verbose:
            print "Starting write of file %s." % filename
            hline()
        for line in template:
            # For each line of the template, do a search for each find/replace "key".  If it is found
            # replace it with the actual value.  Each line is process for every find/replace key in case
            # the line has more than one.  i.e.  ip address <INSIDE_IP> <INSIDE_MASK>
            ok_to_write = True
            for key, value in hostinfo.iteritems():
                if key in line:
                    if value == "":
                        #Do not write the line if there is no value for the variable that matched the line
                        ok_to_write = False
                        if args.verbose:
                            print "Found empty value for variable %s. Line '%s' removed." % (key, line)
                    else:
                        line = line.replace(key,value)
                        if args.verbose:
                            print "Found an instance of %s and replaced it with '%s'" % (key, value)
            else:
                #After the line has been checked/modified for all applicable keys (for loop completed),
                #and the value for the variable wasn't empty, write the line to the output file.
                if ok_to_write:
                    dstfile.write(line)
        if args.verbose:
            print "Configuration %s has been completed" % filename
            hline()
    else:
        #Once the entire loop has finished, all lines in the template have been modified and written.
        #Close the file.
        print "Configuration files exported.\n"
        dstfile.close()

def print_vars_per_host(varlist):
    """This should be passed the data structure with all the hostname find/replace variables in it.
        This function will crawl the data structure to list all find/replace variables and values
        grouped by the host they are meant for."""

    hline()
    #Loop through the primary dictionary (one entry per host)
    for host in host_list:  #References global host_list
        print "Variables for device %s:" % host
        hline()
        #Loop through the secondary dictionary, in alpha order, and list key/value pairs
        keylist = sorted(varlist[host])
        for key in keylist:
            print "KEY: %s , VALUE: %s" % (key, varlist[host][key])
        else:
            hline()
            #Pause after each host's output.  Allow a method to quit should the user see some
            #incorrect data
            pause = raw_input("Press ENTER to Continue (q to QUIT)")
            if pause.lower() == "q":
                print "EXITING DUE TO USER REQUEST"
                exit(0)
            hline()
    else:
        hline()
        print "DONE"
        hline()

def find_unique_vars(file):
    """This function searches the file for all unique find/replace variables, and returns a list
        of unique variables"""

    #A list to collect every variable delimited by < >.  All instances (duplicates) will be added.
    all_vars = []

    file.seek(0)
    #Compile a regex that searches for non-whitespace chars between a < and >
    reg=re.compile(r"<\S*?>")
    for line in file:
        #Use findall because some lines may have more than 1 variable on it. Output is a list.
        linematch = reg.findall(line)
        for item in linematch:
            #ignore this special string that only marks the end of a hosts variable block
            if item == "<--END-->":
                pass
            else:
                #If it is anything else, append it to the list of variables found
                all_vars.append(item)
    else:
        #Sets can't have duplicate items, so by converting the list to a set and back, we will
        #be left with a list with only one copy of each unique variable.
        find_unique_vars = list(set(all_vars))
    return find_unique_vars

def var_list_compare(filename1, filename2, vars_1, vars_2):
    """This function will take the list of unique variables found in two files and report if they match.
        If they do not match, then this function will output the difference.  This function will print
        the results and returns a boolean expression on whether the two var lists are identical"""
    if args.verbose:
        print "Comparing lists of variables from %s and %s" % (filename1, filename2)
    #Using the "&" operation between 2 sets gives the common items.  If the common item set isn't the
    #same length as both the var_1 and var_2 sets, then they aren't exactly the same.
    common_set = set(vars_1) & set(vars_2)
    if len(common_set) == len(set(vars_1)) and len(common_set) == len(set(vars_2)):
        if args.verbose:
            print "All variables match between %s and %s\n" % (filename1, filename2)
        return True #because both lists were identical
    else:
        print "WARNING:"
        #Figure out which vars are left when you subtract one set from the other.  Whatever is left only exists
        #in that one file.  Those variables will not be replaced in the output script.
        if len(set(vars_1) - set(vars_2)) != 0:
            print "Variables that only exist in %s are: %s\n" % (filename1, list(set(vars_1) - set(vars_2)))
        if len(set(vars_2) - set(vars_1)) != 0:
            print "Variables that only exist in %s are: %s\n" % (filename2, list(set(vars_2) - set(vars_1)))
        return False #because one list contained more (or different) vars than the other

def per_host_var_check(template_list, vardata):
    """This function will check the variables found in the template against each host's imported list of
         variables.  If a host is missing variables, the user will be alerted about which host has missing
         variables."""

    error = 0
    if args.verbose:
        print "Running per-host variable checks..."
        hline()
    for host in vardata:
        host_vars = []
        for key in vardata[host]:
            host_vars.append(key)
        diff_set = set(template_list) - set(host_vars)
        if len(diff_set) != 0:
            error = 1
            print "Host %s is missing the following vars found in the template:\n%s" % (host, ", ".join(diff_set))
            hline()
        elif len(diff_set) == 0 and args.verbose:
            print "Host %s has all the required variables in its block" % host
            hline()
        else:
            pass
    if error == 1:
        print
        print "The host(s) mentioned above are missing variables that are found in the template file.  Please add these variables and try again.\n"
        print "If the host does not require the missing variable, add an entry for the variable but leave the value blank i.e. '<MISSINGVAR>:: '\n"
        print "This will tell the script to remove any lines in that host's configuration file that contains <MISSINGVAR>.\n"
        exit(0)
    else:
        if args.verbose:
            print "Variable check completed successfully\n"

def q_review_vars():
    """This is a decision tree on whether the user wants to review the variables before
        allowing the program to export all the config files.  It will loop indefinitely until
        a valid respose is received."""

    if args.verbose:
        #If verbose flag was selected, automatically review instead of asking, then move onto config write.
        print_vars_per_host(importdata)
        q_start_export()
        exit(0)
    while True:
        answ2 = raw_input("Would you like to review the imported vars before generating configs? ")
        if answ2.lower() == "yes" or answ2.lower() == "y":
            print_vars_per_host(importdata)
            q_start_export()
            break
        elif answ2.lower() == "no" or answ2.lower() == "n":
            q_start_export()
            break
        else:
            print "I did not understand that response"

def q_start_export():
    """This is a decision tree on whether the user wants to export the configuration files."""
    while True:
        answ3 = raw_input("Would you like to export the config files? ")
        if answ3.lower() == "yes" or answ3.lower() == "y":
            write_configs(importdata, template)
            template.close()
            break
        elif answ3.lower() == "no" or answ3.lower() == "n":
            print "Configuration files have NOT been exported.  Exiting..."
            exit(0)
        else:
            print "I did not understand that answer"

#=================BEGIN MAIN PROGRAM=====================#

#This section manages the (-h) help output, with the help of ArgParse, and makes sure that all required arguments are passed to the script when run.
parser = argparse.ArgumentParser(description="This script will generate a configuration file for multiple devices.  It needs a 'template' configuration file that uses variable names in every location where each device needs a different value.  This script also needs a 'variables' file that lists the variables and associated replacement value for the device.  This variables file can contain sections for multiple configuration files.")
group = parser.add_mutually_exclusive_group()
group.add_argument("-v", "--verbose", action="store_true", help="Will provide a more verbose output and automatically trigger a review of all parsed variables")
group.add_argument("-q", "--quiet", action="store_true", help="Will provide minimal output and will automatically skip all review questions")
parser.add_argument("template", help="Name of the file that is the configuration template")
parser.add_argument("variables", help="Name of the file that contains find/replace variables for each device")
#Assign input arguments to the "args" variable for reference later in the script.
args = parser.parse_args()

#print "\n" * 10
print "Start of ConfigMerge output:"
hline()

var_file = open(args.variables, 'r')
template = open(args.template, 'r')

#Call function to parse variable files and assign outputs to variables
host_list, importdata = parse_vars_file(var_file)

#Gather unique variables from each file, so they can be compared.
var_file_u = find_unique_vars(var_file)
template_u = find_unique_vars(template)
identical = var_list_compare(args.variables, args.template, var_file_u, template_u)

#If there were errors in the comparison, allow the user to decide if they want to continue or not.
if not identical:
    while True:
        answ = raw_input("Would you like to continue with the merge anyway? ")
        if answ.lower() == "yes" or answ.lower() == "y":
            break
        elif answ.lower() == "no" or answ.lower() == "n":
            print "Please modify your input files and try again.  Thanks!"
            var_file.close()
            template.close()
            exit(0)
        else:
            print "I did not understand that response."

#Close variables file -- no longer needed
var_file.close()

#Now make sure each host has all variables found in the template configuration
per_host_var_check(template_u, importdata)

if args.quiet:
    #If quiet flag was set, skip the review and get straight to the config writing
    write_configs(importdata, template)
else:
    #Otherwise, give brief review of the hosts found
    print "Found settings for %d config files.  They are: " % len(host_list)
    x = 1
    for host in host_list:
        print "%d. %s" % (x, host)
        x = x + 1
    print ""
    #First Decision tree.  Allows user to decide if brief review looks accurate.
    #Other decision tree functions are called depending on the response.
    while True:
        answ = raw_input("Is this correct? (Type 'yes' to continue) ")
        if answ.lower() == "yes" or answ.lower() == "y":
            q_review_vars()
            break
        elif answ.lower() == "no" or answ.lower() == "n":
            print "Please modify your vars file and try again.  Thanks!"
            template.close()
            exit(0)
        else:
            print "I did not understand that response."

<commit_msg>Minor cleanup<commit_after>import argparse
import os
import re
os.system('cls' if os.name=='nt' else 'clear')

#=============FUNCTION DEFINITIONS===============#
def hline():
    """This function simply prints a line of 60 "-" chars, to add divisions to screen output"""
    print "-" * 60

def parse_vars_file(file):
    """This function takes the file with all the variable and their values, listed for each host, and
        puts it into a data structure that can later be used to write out each config script.  It also
        tracks the names of hosts found in the file and returns those for quick review prior to exporting
        the configurations.

        The data structure is a dictionary with each hostname as the key, and the value being another
        dictionary (sub-dictionary?) with each find/replace key and its associated value)"""

    #Reset file position to beginnig (just in case)
    file.seek(0)
    host_list = []
    dict_vars = {}
    #This variable tracks whether the line being parsed is part of a host "block" of variables or not.
    hostblock = 0

    for line in file:
        #Ignore comment lines
        if line.startswith("#"):
                pass

        #Use a marker so we know if we are inside of a particular host's block of variables.  When we see a
        #HOSTNAME var, it is set to 1, at <--END--> it is set back to 0.
        elif "<--END-->" in line:
            hostblock = 0

        #This string "::" is used to separate the key and value in the variable file.  If any line doesn't
        #have this (with exception of the END line, shown above), ignore it and print an error.
        elif "::" not in line:
            print "Did NOT find Seperator '%r'. Line ignored" % line

        #If we find a HOSTNAME variable, but the function thinks we are still in a host's block of variables,
        #then throw an error and exit.
        elif "<HOSTNAME>" in line and hostblock == 1:
            hline()
            print "ERROR: <HOSTNAME> found inside another host's variable block."
            print "ERROR LINE: '%r'" % line
            print "Exiting program. Please correct line and try again."
            exit(0)

        #HOSTNAME variable after an <--END--> line (hostblock=0) triggers the start of a new host block, and
        #creates an entry in the dictionary for the host.
        elif "<HOSTNAME>" in line and hostblock == 0:
            #Split the variable (<HOSTNAME>) from the value
            label, hostname = line.split('::')
            hostname = hostname.rstrip()
            for imported_host in host_list:
                if hostname == imported_host:
                    print "ERROR: DUPLICATE HOSTNAME IN VARIABLES FILE. EXITING..."
                    exit(0)
            hostblock = 1
            #Create an entry in the dictionary for the host.  The value will be another dictionary.
            host_list.append(hostname)
            dict_vars[hostname] = {}
            #The first value of the sub-dictionary will be the <HOSTNAME> variable and actual hostname for
            #this device
            dict_vars[hostname][label] = hostname

        #If the other statements don't match, and we are inside a host's block of variables, then assume
        #this line is a variable and add it to the sub-dictionary for the particular host.
        elif hostblock == 1:
            #This should match all find/replace lines under a particular hostname.  This code just adds
            #the pairs to the sub-dictionary for the associated host.
            key, value = line.split('::')
            value = value.rstrip()
            dict_vars[hostname][key] = value

        else:
            #If this section matches a line, something was wrong with the file.  It prints an output of
            #the line and under which host this occurred to help track down the error (for example, it
            #could be HOSTNAME line that had an error, so all lines for that host will match this)
            print "INVALID LINE in %s: '%r'" % (hostname, line)

    #Return the list of hosts, the host count, and the data structure to the main program.
    return host_list, dict_vars

def write_configs(vars, template):
    """This template writes each host's configuration file.  It should be passed both the template file
        to use, as well as the data structure generated by the parse_vars_file() function."""

    if args.verbose:
        hline()
    for host in vars:
        #reset the template read position to start for each host
        template.seek(0)
        #create file name based on hostname of the device
        filename = host + ".txt"
        hostinfo = vars[host]
        #Check that the "configs" directory exists.  If not, create it.
        if not os.path.exists("configs"):
            os.makedirs("configs")
        dstfile = open("./configs/" + filename, 'w')
        if args.verbose:
            print "Starting write of file %s." % filename
            hline()
        for line in template:
            # For each line of the template, do a search for each find/replace "key".  If it is found
            # replace it with the actual value.  Each line is process for every find/replace key in case
            # the line has more than one.  i.e.  ip address <INSIDE_IP> <INSIDE_MASK>
            ok_to_write = True
            for key, value in hostinfo.iteritems():
                if key in line:
                    if value == "":
                        #Do not write the line if there is no value for the variable that matched the line
                        ok_to_write = False
                        if args.verbose:
                            print "Found empty value for variable %s. Line '%s' removed." % (key, line)
                    else:
                        line = line.replace(key,value)
                        if args.verbose:
                            print "Found an instance of %s and replaced it with '%s'" % (key, value)
            else:
                #After the line has been checked/modified for all applicable keys (for loop completed),
                #and the value for the variable wasn't empty, write the line to the output file.
                if ok_to_write:
                    dstfile.write(line)
        if args.verbose:
            print "Configuration %s has been completed" % filename
            hline()
    else:
        #Once the entire loop has finished, all lines in the template have been modified and written.
        #Close the file.
        print "Configuration files exported.\n"
        dstfile.close()

def print_vars_per_host(varlist):
    """This should be passed the data structure with all the hostname find/replace variables in it.
        This function will crawl the data structure to list all find/replace variables and values
        grouped by the host they are meant for."""

    hline()
    #Loop through the primary dictionary (one entry per host)
    for host in host_list:  #References global host_list
        print "Variables for device %s:" % host
        hline()
        #Loop through the secondary dictionary, in alpha order, and list key/value pairs
        keylist = sorted(varlist[host])
        for key in keylist:
            print "KEY: %s , VALUE: %s" % (key, varlist[host][key])
        else:
            hline()
            #Pause after each host's output.  Allow a method to quit should the user see some
            #incorrect data
            pause = raw_input("Press ENTER to Continue (q to QUIT)")
            if pause.lower() == "q":
                print "EXITING DUE TO USER REQUEST"
                exit(0)
            hline()
    else:
        hline()
        print "DONE"
        hline()

def find_unique_vars(file):
    """This function searches the file for all unique find/replace variables, and returns a list
        of unique variables"""

    #A list to collect every variable delimited by < >.  All instances (duplicates) will be added.
    all_vars = []

    file.seek(0)
    #Compile a regex that searches for non-whitespace chars between a < and >
    reg=re.compile(r"<\S*?>")
    for line in file:
        #Use findall because some lines may have more than 1 variable on it. Output is a list.
        linematch = reg.findall(line)
        for item in linematch:
            #ignore this special string that only marks the end of a hosts variable block
            if item == "<--END-->":
                pass
            else:
                #If it is anything else, append it to the list of variables found
                all_vars.append(item)
    else:
        #Sets can't have duplicate items, so by converting the list to a set and back, we will
        #be left with a list with only one copy of each unique variable.
        find_unique_vars = list(set(all_vars))
    return find_unique_vars

def var_list_compare(filename1, filename2, vars_1, vars_2):
    """This function will take the list of unique variables found in two files and report if they match.
        If they do not match, then this function will output the difference.  This function will print
        the results and returns a boolean expression on whether the two var lists are identical"""
    if args.verbose:
        print "Comparing lists of variables from %s and %s" % (filename1, filename2)
    #Using the "&" operation between 2 sets gives the common items.  If the common item set isn't the
    #same length as both the var_1 and var_2 sets, then they aren't exactly the same.
    common_set = set(vars_1) & set(vars_2)
    if len(common_set) == len(set(vars_1)) and len(common_set) == len(set(vars_2)):
        if args.verbose:
            print "All variables match between %s and %s\n" % (filename1, filename2)
        return True #because both lists were identical
    else:
        print "WARNING:"
        #Figure out which vars are left when you subtract one set from the other.  Whatever is left only exists
        #in that one file.  Those variables will not be replaced in the output script.
        if len(set(vars_1) - set(vars_2)) != 0:
            print "Variables that only exist in %s are: %s\n" % (filename1, list(set(vars_1) - set(vars_2)))
        if len(set(vars_2) - set(vars_1)) != 0:
            print "Variables that only exist in %s are: %s\n" % (filename2, list(set(vars_2) - set(vars_1)))
        return False #because one list contained more (or different) vars than the other

def per_host_var_check(template_list, vardata):
    """This function will check the variables found in the template against each host's imported list of
         variables.  If a host is missing variables, the user will be alerted about which host has missing
         variables."""

    error = 0
    if args.verbose:
        print "Running per-host variable checks..."
        hline()
    for host in vardata:
        host_vars = []
        for key in vardata[host]:
            host_vars.append(key)
        diff_set = set(template_list) - set(host_vars)
        if len(diff_set) != 0:
            error = 1
            print "Host %s is missing the following vars found in the template:\n%s" % (host, ", ".join(diff_set))
            hline()
        elif len(diff_set) == 0 and args.verbose:
            print "Host %s has all the required variables in its block" % host
            hline()
        else:
            pass
    if error == 1:
        print
        print "The host(s) mentioned above are missing variables that are found in the template file.  Please add these variables and try again.\n"
        print "If the host does not require the missing variable, add an entry for the variable but leave the value blank i.e. '<MISSINGVAR>:: '\n"
        print "This will tell the script to remove any lines in that host's configuration file that contains <MISSINGVAR>.\n"
        exit(0)
    else:
        if args.verbose:
            print "Variable check completed successfully\n"

def q_review_vars():
    """This is a decision tree on whether the user wants to review the variables before
        allowing the program to export all the config files.  It will loop indefinitely until
        a valid respose is received."""

    if args.verbose:
        #If verbose flag was selected, automatically review instead of asking, then move onto config write.
        print_vars_per_host(importdata)
        q_start_export()
        exit(0)
    while True:
        answ2 = raw_input("Would you like to review the imported vars before generating configs? ")
        if answ2.lower() == "yes" or answ2.lower() == "y":
            print_vars_per_host(importdata)
            q_start_export()
            break
        elif answ2.lower() == "no" or answ2.lower() == "n":
            q_start_export()
            break
        else:
            print "I did not understand that response"

def q_start_export():
    """This is a decision tree on whether the user wants to export the configuration files."""
    while True:
        answ3 = raw_input("Would you like to export the config files? ")
        if answ3.lower() == "yes" or answ3.lower() == "y":
            write_configs(importdata, template)
            template.close()
            break
        elif answ3.lower() == "no" or answ3.lower() == "n":
            print "Configuration files have NOT been exported.  Exiting..."
            exit(0)
        else:
            print "I did not understand that answer"

#=================BEGIN MAIN PROGRAM=====================#

#This section manages the (-h) help output, with the help of ArgParse, and makes sure that all required arguments are passed to the script when run.
parser = argparse.ArgumentParser(description="This script will generate a configuration file for multiple devices.  It needs a 'template' configuration file that uses variable names in every location where each device needs a different value.  This script also needs a 'variables' file that lists the variables and associated replacement value for the device.  This variables file can contain sections for multiple configuration files.")
group = parser.add_mutually_exclusive_group()
group.add_argument("-v", "--verbose", action="store_true", help="Will provide a more verbose output and automatically trigger a review of all parsed variables")
group.add_argument("-q", "--quiet", action="store_true", help="Will provide minimal output and will automatically skip all review questions")
parser.add_argument("template", help="Name of the file that is the configuration template")
parser.add_argument("variables", help="Name of the file that contains find/replace variables for each device")
#Assign input arguments to the "args" variable for reference later in the script.
args = parser.parse_args()

print "Start of ConfigMerge output:"
hline()

var_file = open(args.variables, 'r')
template = open(args.template, 'r')

#Call function to parse variable files and assign outputs to variables
host_list, importdata = parse_vars_file(var_file)

#Gather unique variables from each file, so they can be compared.
var_file_u = find_unique_vars(var_file)
template_u = find_unique_vars(template)
identical = var_list_compare(args.variables, args.template, var_file_u, template_u)

#If there were errors in the comparison, allow the user to decide if they want to continue or not.
if not identical:
    while True:
        answ = raw_input("Would you like to continue with the merge anyway? ")
        if answ.lower() == "yes" or answ.lower() == "y":
            break
        elif answ.lower() == "no" or answ.lower() == "n":
            print "Please modify your input files and try again.  Thanks!"
            var_file.close()
            template.close()
            exit(0)
        else:
            print "I did not understand that response."

#Close variables file -- no longer needed
var_file.close()

#Now make sure each host has all variables found in the template configuration
per_host_var_check(template_u, importdata)

if args.quiet:
    #If quiet flag was set, skip the review and get straight to the config writing
    write_configs(importdata, template)
else:
    #Otherwise, give brief review of the hosts found
    print "Found settings for %d config files.  They are: " % len(host_list)
    x = 1
    for host in host_list:
        print "%d. %s" % (x, host)
        x = x + 1
    print ""
    #First Decision tree.  Allows user to decide if brief review looks accurate.
    #Other decision tree functions are called depending on the response.
    while True:
        answ = raw_input("Is this correct? (Type 'yes' to continue) ")
        if answ.lower() == "yes" or answ.lower() == "y":
            q_review_vars()
            break
        elif answ.lower() == "no" or answ.lower() == "n":
            print "Please modify your vars file and try again.  Thanks!"
            template.close()
            exit(0)
        else:
            print "I did not understand that response."

<|endoftext|>


<commit_before>import json
import logging
import sys
import urllib.parse
import urllib.request
from datetime import datetime
from zipfile import ZipFile
from objectstore.objectstore import ObjectStore
import requests
from fme.sql_utils import run_sql_script
from setup import FME_SERVER_API, FME_SERVER, INSTANCE_ID, SCRIPT_ROOT

from fme import fme_server
from fme.fme_utils import (
    run_transformation_job, fme_api_auth, upload_repository, upload,
    wait_for_job_to_complete)

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

def start_transformation_gebieden():
    log.info("Starting transformation")
    return run_transformation_job(
        'BGT-DB',
        'inlezen_gebieden_uit_Shape_en_WFS.fmw',
        {
            "subsection": "REST_SERVICE",
            "FMEDirectives": {},
            "NMDirectives": {"successTopics": [], "failureTopics": []},
            "TMDirectives": {"tag": "linux", "description": "DB BGT kaartbladen"},
            "publishedParameters": [{"name": "DestDataset_POSTGIS_4", "value": "bgt"},
                                    {"name": "bron_BGT_kaartbladen",
                                     "value": ["$(FME_SHAREDRESOURCE_DATA)/Import_kaartbladen/BGT_kaartbladen.shp"]}]})

def start_transformation_db():
    log.info("Starting transformation")
    return run_transformation_job(
        'BGT-DB',
        'inlezen_DB_BGT_uit_citygml.fmw',
        {
            "subsection": "REST_SERVICE",
            "FMEDirectives": {},
            "NMDirectives": {"successTopics": [], "failureTopics": []},
            "TMDirectives": {"tag": "linux", "description": "DB BGT uit city gml"
                             },
            "publishedParameters": [{"name": "DestDataset_POSTGIS_4", "value": "bgt"},
                                    {"name": "CITYGML_IN_ADE_XSD_DOC_CITYGML",
                                     "value": ["$(FME_SHAREDRESOURCE_DATA)/Import_XSD/imgeo.xsd"]},
                                    {"name": "SourceDataset_CITYGML",
                                     "value": ["$(FME_SHAREDRESOURCE_DATA)/Import_GML/bgt_buurt.gml"]}, ]})

def start_transformation_shapes():
    log.info("Starting transformation")

    # update data in `Export shapes` and  `Export_Shapes_Totaalgebied` directories
    return run_transformation_job(
        'BGT-SHAPES',
        'aanmaak_esrishape_uit_DB_BGT.fmw',
        {
            "subsection": "REST_SERVICE",
            "FMEDirectives": {},
            "NMDirectives": {"successTopics": [], "failureTopics": []},
            "TMDirectives": {"tag": "linux", "description": "Aanmaak Shapes uit DB"},
            "publishedParameters": [
                {"name": "SourceDataset_POSTGIS", "value": "bgt"},
                {"name": "SourceDataset_POSTGIS_3", "value": "bgt"},
                {"name": "DestDataset_ESRISHAPE2", "value": "$(FME_SHAREDRESOURCE_DATA)/Export_Shapes.zip"},
                {"name": "DestDataset_ESRISHAPE3",
                 "value": "$(FME_SHAREDRESOURCE_DATA)/Export_Shapes_Totaalgebied.zip"}]})

def start_test_transformation():
    return run_transformation_job(
        'BGT-SHAPES',
        'aanmaak_esrishape_test_zip.fmw',
        {
            "subsection": "REST_SERVICE",
            "FMEDirectives": {},
            "NMDirectives": {"successTopics": [], "failureTopics": []},
            "TMDirectives": {"tag": "linux", "description": "Aanmaak Shapes uit DB"},
            "publishedParameters": [
                {"name": "SourceDataset_POSTGIS", "value": "bgt"},
                {"name": "SourceDataset_POSTGIS_3", "value": "bgt"},
                {"name": "DestDataset_ESRISHAPE2", "value": "$(FME_SHAREDRESOURCE_DATA)/TestExport.zip"},
                {"name": "DestDataset_ESRISHAPE3", "value":
                    "$(FME_SHAREDRESOURCE_DATA)/TestExportTotaalgebied.zip"}]})

def upload_resulting_shapes_to_objectstore():
    """
    Uploads the resulting shapes to BGT objectstore
    :return:
    """
    store = ObjectStore('BGT')
    log.info("Upload resulting shapes to BGT objectstore")
    files = ['Export_Shapes.zip', 'Export_Shapes_Totaalgebied.zip']

    for path in files:
        log.info("Download {} for storing in objectstore".format(path))
        download_url = '{FME_SERVER}{url_connect}/FME_SHAREDRESOURCE_DATA/filesys/{DIR}?detail=low' \
                       'createDirectories=false&detail=low&overwrite=false'.format(
            DIR=path, FME_SERVER=FME_SERVER, url_connect="/fmerest/v2/resources/connections")
        res = requests.get(download_url, headers=fme_api_auth())
        res.raise_for_status()
        if res.status_code == 200:
            res_name = path.split('/')[-1].split('.')
            timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
            store.put_to_objectstore('shapes/{}-{}.{}'.format(res_name[0], timestamp, res_name[-1]),
                                     res.content,
                                     res.headers['Content-Type'])
            log.info("Uploaded {} to objectstore BGT/shapes".format(res_name))
    log.info("Uploaded resulting shapes to BGT objectstore")

def unzip_pdok_file():
    """
    Unzip the extract_bgt.zip file into the `/tmp/data` directory
    :return:
    """
    log.info("Start unzipping contents")
    with ZipFile('extract_bgt.zip', 'r') as myzip:
        myzip.extractall('/tmp/data/')
    log.info("Unzip complete")
    return 0

def download_bgt():
    """
    Downloads the PDOK files to extract_bgt.zip and calls unzip_pdok_file()
    :return: 0
    """
    target = "extract_bgt.zip"
    pdok_extract = "https://www.pdok.nl/download/service/extract.zip"
    tiles = {
        "layers": [
            {"aggregateLevel": 0,
             "codes": [
                 38121, 38115, 38120, 38077, 38076, 38078, 38079, 38421, 38423, 38429, 38431, 38474, 38475, 38478,
                 38479, 38490, 38491, 38494, 38516, 38518, 38690, 38525, 38696, 38688, 38666, 38667, 38670, 38671,
                 38682, 38680, 38681, 38675, 38673, 38331, 38329, 38323, 38321, 38299, 38298, 38287, 38285, 38284,
                 38281, 38275, 38274, 38103, 38109, 38280, 38108, 38105, 38104, 38093, 38092, 38094, 38116, 38674,
                 38672, 38330, 38328, 38317, 38476, 38477, 38488, 38489, 38492, 38493, 38495, 38664, 38665, 38668,
                 38517, 38466, 38467, 38118, 38119, 38117, 38095, 38106, 38107, 38110, 38111, 38282, 38283, 38316,
                 38319, 38472, 38464, 38465, 38122, 38470, 38471, 38468, 38469, 38126, 38127, 38124, 38125, 38482,
                 38480, 38138, 38136, 38483, 38481, 38139, 38140, 38142, 38484, 38486, 38487, 38485, 38143, 38141,
                 38134, 38132, 38133, 38135, 38306, 38304, 38312, 38314, 38130, 38128, 38129, 38131, 38137, 38313,
                 38307, 38308, 38310, 38315, 38318, 38656, 38658, 38659, 38657, 38662, 38660, 38661, 38663, 38311,
                 38309, 38320, 38322]}]}
    tiles_as_json = json.dumps(tiles)
    datum = datetime.now().strftime("%d-%m-%Y")

    source = pdok_extract + "?" + urllib.parse.urlencode({
        "extractset": "citygml",
        "excludedtypes": "plaatsbepalingspunt",
        "history": "false",
        "enddate": datum,
        "tiles": tiles_as_json,
    })

    log.info("Starting download from %s to %s", source, target)

    def progress(count, block_size, total_size):
        approximate = ''
        if total_size == -1:
            # evil PDOK not telling us the content size :-(
            # educated guess is ~350Mb
            total_size = 350 * 1024 * 1024
            approximate = '~'

        percentage = float(count * block_size * 100.0 / total_size)
        log.info("%s%2.2f%%", approximate, percentage)

    urllib.request.urlretrieve(source, target, reporthook=progress)
    unzip_pdok_file()
    log.info("Download complete")
    return 0

if __name__ == '__main__':
    logging.basicConfig(level='DEBUG')
    logging.getLogger('requests').setLevel('WARNING')

    server_manager = fme_server.Server(FME_SERVER, INSTANCE_ID, FME_SERVER_API)

    log.info("Starting script, current server status is %s", server_manager.get_status())

    download_bgt()
    try:
        server_manager.start()
        run_sql_script("{app}/020_create_schema.sql".format(app=SCRIPT_ROOT))

        # upload the GML files and FMW scripts
        upload('/tmp/data', 'resources/connections', 'Import_GML', '*.*', recreate_dir=True)
        upload_repository(
            '{app}/030_inlezen_BGT/fme'.format(app=SCRIPT_ROOT), 'repositories', 'BGT-DB', '*.*', register_fmejob=True)

        # DB Connection in FMECLoud needs to be set manually - NOT POSSIBLE IN CURRENT API VERSION
        upload('{app}/030_inlezen_BGT/xsd'.format(app=SCRIPT_ROOT), 'resources/connections', 'Import_XSD', 'imgeo.xsd')
        upload('{app}/030_inlezen_BGT/bron_shapes'.format(app=SCRIPT_ROOT), 'resources/connections',
               'Import_kaartbladen',
               '*.*')
        try:
            wait_for_job_to_complete(start_transformation_gebieden())
            wait_for_job_to_complete(start_transformation_db())
        except Exception:
            sys.exit(1)

        run_sql_script("{app}/060_aanmaak_tabellen_BGT.sql".format(app=SCRIPT_ROOT))

        # aanmaak db-views shapes_bgt
        run_sql_script("{app}/090_aanmaak_DB_views_BGT.sql".format(app=SCRIPT_ROOT))
        run_sql_script("{app}/090_aanmaak_DB_views_IMGEO.sql".format(app=SCRIPT_ROOT))

        # upload shapes fmw scripts naar reposiory
        upload_repository(
            '{app}/100_aanmaak_producten_BGT'.format(app=SCRIPT_ROOT), 'BGT-SHAPES', '*.*', register_fmejob=True)
        # run the `aanmaak_esrishape_uit_DB_BGT` scrop
        try:
            wait_for_job_to_complete(start_transformation_shapes())
        except Exception:
            sys.exit(1)

        upload_resulting_shapes_to_objectstore()
        # TODO: 1) download db and do telling OR do it on remote database
        # TODO: Telling: 040

        # import controle db
        # subprocess.call("{app}/070_import_gml_controledb.sh".format(app=SCRIPT_ROOT), shell=True)

        # import csv / mapping db
        # run_sql_script("{app}/075_aanmaak_mapping.sql".format(app=SCRIPT_ROOT))

        # TODO: Frequentie verdeling: 080
    except:
        log.exception("Could not process server jobs")
    finally:
        server_manager.stop()
        sys.exit(0)
<commit_msg>keep history of GML downloads in objectstore<commit_after>import json
import logging
import sys
import subprocess
import urllib.parse
import urllib.request
from datetime import datetime
from zipfile import ZipFile
from objectstore.objectstore import ObjectStore
import requests
from fme.sql_utils import run_sql_script
from setup import FME_SERVER_API, FME_SERVER, INSTANCE_ID, SCRIPT_ROOT

from fme import fme_server
from fme.fme_utils import (
    run_transformation_job, fme_api_auth, upload_repository, upload,
    wait_for_job_to_complete)

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

def start_transformation_gebieden():
    log.info("Starting transformation")
    return run_transformation_job(
        'BGT-DB',
        'inlezen_gebieden_uit_Shape_en_WFS.fmw',
        {
            "subsection": "REST_SERVICE",
            "FMEDirectives": {},
            "NMDirectives": {"successTopics": [], "failureTopics": []},
            "TMDirectives": {"tag": "linux", "description": "DB BGT kaartbladen"},
            "publishedParameters": [{"name": "DestDataset_POSTGIS_4", "value": "bgt"},
                                    {"name": "bron_BGT_kaartbladen",
                                     "value": ["$(FME_SHAREDRESOURCE_DATA)/Import_kaartbladen/BGT_kaartbladen.shp"]}]})

def start_transformation_db():
    log.info("Starting transformation")
    return run_transformation_job(
        'BGT-DB',
        'inlezen_DB_BGT_uit_citygml.fmw',
        {
            "subsection": "REST_SERVICE",
            "FMEDirectives": {},
            "NMDirectives": {"successTopics": [], "failureTopics": []},
            "TMDirectives": {"tag": "linux", "description": "DB BGT uit city gml"
                             },
            "publishedParameters": [{"name": "DestDataset_POSTGIS_4", "value": "bgt"},
                                    {"name": "CITYGML_IN_ADE_XSD_DOC_CITYGML",
                                     "value": ["$(FME_SHAREDRESOURCE_DATA)/Import_XSD/imgeo.xsd"]},
                                    {"name": "SourceDataset_CITYGML",
                                     "value": ["$(FME_SHAREDRESOURCE_DATA)/Import_GML/bgt_buurt.gml"]}, ]})

def start_transformation_shapes():
    log.info("Starting transformation")

    # update data in `Export shapes` and  `Export_Shapes_Totaalgebied` directories
    return run_transformation_job(
        'BGT-SHAPES',
        'aanmaak_esrishape_uit_DB_BGT.fmw',
        {
            "subsection": "REST_SERVICE",
            "FMEDirectives": {},
            "NMDirectives": {"successTopics": [], "failureTopics": []},
            "TMDirectives": {"tag": "linux", "description": "Aanmaak Shapes uit DB"},
            "publishedParameters": [
                {"name": "SourceDataset_POSTGIS", "value": "bgt"},
                {"name": "SourceDataset_POSTGIS_3", "value": "bgt"},
                {"name": "DestDataset_ESRISHAPE2", "value": "$(FME_SHAREDRESOURCE_DATA)/Export_Shapes.zip"},
                {"name": "DestDataset_ESRISHAPE3",
                 "value": "$(FME_SHAREDRESOURCE_DATA)/Export_Shapes_Totaalgebied.zip"}]})

def start_test_transformation():
    return run_transformation_job(
        'BGT-SHAPES',
        'aanmaak_esrishape_test_zip.fmw',
        {
            "subsection": "REST_SERVICE",
            "FMEDirectives": {},
            "NMDirectives": {"successTopics": [], "failureTopics": []},
            "TMDirectives": {"tag": "linux", "description": "Aanmaak Shapes uit DB"},
            "publishedParameters": [
                {"name": "SourceDataset_POSTGIS", "value": "bgt"},
                {"name": "SourceDataset_POSTGIS_3", "value": "bgt"},
                {"name": "DestDataset_ESRISHAPE2", "value": "$(FME_SHAREDRESOURCE_DATA)/TestExport.zip"},
                {"name": "DestDataset_ESRISHAPE3", "value":
                    "$(FME_SHAREDRESOURCE_DATA)/TestExportTotaalgebied.zip"}]})

def upload_bgt_source_zip():
    store = ObjectStore('BGT')
    log.info("Upload BGT source zip")
    content = open('extract_bgt.zip', 'rb').read()

    timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
    filename = 'bron-gml/extract_bgt-{}.zip'.format(timestamp)
    store.put_to_objectstore(filename, content, 'application/octet-stream')
    log.info("Uploaded {} to objectstore BGT/bron-gml".format(filename))

def upload_resulting_shapes_to_objectstore():
    """
    Uploads the resulting shapes to BGT objectstore
    :return:
    """
    store = ObjectStore('BGT')
    log.info("Upload resulting shapes to BGT objectstore")
    files = ['Export_Shapes.zip', 'Export_Shapes_Totaalgebied.zip']

    for path in files:
        log.info("Download {} for storing in objectstore".format(path))
        download_url = '{FME_SERVER}{url_connect}/FME_SHAREDRESOURCE_DATA/filesys/{DIR}?detail=low' \
                       'createDirectories=false&detail=low&overwrite=false'.format(
            DIR=path, FME_SERVER=FME_SERVER, url_connect="/fmerest/v2/resources/connections")
        res = requests.get(download_url, headers=fme_api_auth())
        res.raise_for_status()
        if res.status_code == 200:
            res_name = path.split('/')[-1].split('.')
            timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
            store.put_to_objectstore('shapes/{}-{}.{}'.format(res_name[0], timestamp, res_name[-1]),
                                     res.content,
                                     res.headers['Content-Type'])
            log.info("Uploaded {} to objectstore BGT/shapes".format(res_name))
    log.info("Uploaded resulting shapes to BGT objectstore")

def unzip_pdok_file():
    """
    Unzip the extract_bgt.zip file into the `/tmp/data` directory
    :return:
    """
    log.info("Start unzipping contents")
    with ZipFile('extract_bgt.zip', 'r') as myzip:
        myzip.extractall('/tmp/data/')
    log.info("Unzip complete")
    return 0

def download_bgt():
    """
    Downloads the PDOK files to extract_bgt.zip and calls unzip_pdok_file()
    :return: 0
    """
    target = "extract_bgt.zip"
    pdok_extract = "https://www.pdok.nl/download/service/extract.zip"
    tiles = {
        "layers": [
            {"aggregateLevel": 0,
             "codes": [
                 38121, 38115, 38120, 38077, 38076, 38078, 38079, 38421, 38423, 38429, 38431, 38474, 38475, 38478,
                 38479, 38490, 38491, 38494, 38516, 38518, 38690, 38525, 38696, 38688, 38666, 38667, 38670, 38671,
                 38682, 38680, 38681, 38675, 38673, 38331, 38329, 38323, 38321, 38299, 38298, 38287, 38285, 38284,
                 38281, 38275, 38274, 38103, 38109, 38280, 38108, 38105, 38104, 38093, 38092, 38094, 38116, 38674,
                 38672, 38330, 38328, 38317, 38476, 38477, 38488, 38489, 38492, 38493, 38495, 38664, 38665, 38668,
                 38517, 38466, 38467, 38118, 38119, 38117, 38095, 38106, 38107, 38110, 38111, 38282, 38283, 38316,
                 38319, 38472, 38464, 38465, 38122, 38470, 38471, 38468, 38469, 38126, 38127, 38124, 38125, 38482,
                 38480, 38138, 38136, 38483, 38481, 38139, 38140, 38142, 38484, 38486, 38487, 38485, 38143, 38141,
                 38134, 38132, 38133, 38135, 38306, 38304, 38312, 38314, 38130, 38128, 38129, 38131, 38137, 38313,
                 38307, 38308, 38310, 38315, 38318, 38656, 38658, 38659, 38657, 38662, 38660, 38661, 38663, 38311,
                 38309, 38320, 38322]}]}
    tiles_as_json = json.dumps(tiles)
    datum = datetime.now().strftime("%d-%m-%Y")

    source = pdok_extract + "?" + urllib.parse.urlencode({
        "extractset": "citygml",
        "excludedtypes": "plaatsbepalingspunt",
        "history": "false",
        "enddate": datum,
        "tiles": tiles_as_json,
    })

    log.info("Starting download from %s to %s", source, target)

    def progress(count, block_size, total_size):
        approximate = ''
        if total_size == -1:
            # evil PDOK not telling us the content size :-(
            # educated guess is ~350Mb
            total_size = 350 * 1024 * 1024
            approximate = '~'

        percentage = float(count * block_size * 100.0 / total_size)
        log.info("%s%2.2f%%", approximate, percentage)

    urllib.request.urlretrieve(source, target, reporthook=progress)
    unzip_pdok_file()
    log.info("Download complete")
    return 0

if __name__ == '__main__':
    logging.basicConfig(level='DEBUG')
    logging.getLogger('requests').setLevel('WARNING')

    server_manager = fme_server.Server(FME_SERVER, INSTANCE_ID, FME_SERVER_API)

    log.info("Starting script, current server status is %s", server_manager.get_status())

    download_bgt()
    try:
        server_manager.start()
        run_sql_script("{app}/020_create_schema.sql".format(app=SCRIPT_ROOT))

        # upload the GML files and FMW scripts
        upload('/tmp/data', 'resources/connections', 'Import_GML', '*.*', recreate_dir=True)
        upload_repository(
            '{app}/030_inlezen_BGT/fme'.format(app=SCRIPT_ROOT), 'repositories', 'BGT-DB', '*.*', register_fmejob=True)

        # DB Connection in FMECLoud needs to be set manually - NOT POSSIBLE IN CURRENT API VERSION
        upload('{app}/030_inlezen_BGT/xsd'.format(app=SCRIPT_ROOT), 'resources/connections', 'Import_XSD', 'imgeo.xsd')
        upload('{app}/030_inlezen_BGT/bron_shapes'.format(app=SCRIPT_ROOT), 'resources/connections',
               'Import_kaartbladen',
               '*.*')
        try:
            wait_for_job_to_complete(start_transformation_gebieden())
            wait_for_job_to_complete(start_transformation_db())
        except Exception:
            sys.exit(1)

        run_sql_script("{app}/060_aanmaak_tabellen_BGT.sql".format(app=SCRIPT_ROOT))

        # aanmaak db-views shapes_bgt
        run_sql_script("{app}/090_aanmaak_DB_views_BGT.sql".format(app=SCRIPT_ROOT))
        run_sql_script("{app}/090_aanmaak_DB_views_IMGEO.sql".format(app=SCRIPT_ROOT))

        # upload shapes fmw scripts naar reposiory
        upload_repository(
            '{app}/100_aanmaak_producten_BGT'.format(app=SCRIPT_ROOT), 'BGT-SHAPES', '*.*', register_fmejob=True)
        # run the `aanmaak_esrishape_uit_DB_BGT` scrop
        try:
            wait_for_job_to_complete(start_transformation_shapes())
        except Exception:
            sys.exit(1)

        # upload the resulting shapes an the source GML zip to objectstore
        upload_resulting_shapes_to_objectstore()
        upload_bgt_source_zip()

        # TODO: 1) download db and do telling OR do it on remote database
        # TODO: Telling: 040
        # import controle db
        # subprocess.call("{app}/070_import_gml_controledb.sh".format(app=SCRIPT_ROOT), shell=True)

        # import csv / mapping db
        # run_sql_script("{app}/075_aanmaak_mapping.sql".format(app=SCRIPT_ROOT))

        # TODO: Frequentie verdeling: 080
    except:
        log.exception("Could not process server jobs")
    finally:
        server_manager.stop()
        sys.exit(0)
<|endoftext|>


<commit_before><commit_msg>A better test for file existence<commit_after><|endoftext|>


<commit_before># Analyze Color of Object

import os
import cv2
import numpy as np
from . import print_image
from . import plot_image
from . import fatal_error
from . import plot_colorbar

def _pseudocolored_image(device, histogram, bins, img, mask, background, channel, filename, resolution,
                         analysis_images, debug):
    """Pseudocolor image.

    Inputs:
    histogram       = a normalized histogram of color values from one color channel
    bins            = number of color bins the channel is divided into
    img             = input image
    mask            = binary mask image
    background      = what background image?: channel image (img) or white
    channel         = color channel name
    filename        = input image filename
    resolution      = output image resolution
    analysis_images = list of analysis image filenames
    debug           = print or plot. Print = save to file, Plot = print to screen.

    Returns:
    analysis_images = list of analysis image filenames

    :param histogram: list
    :param bins: int
    :param img: numpy array
    :param mask: numpy array
    :param background: str
    :param channel: str
    :param filename: str
    :param resolution: int
    :param analysis_images: list
    :return analysis_images: list
    """
    mask_inv = cv2.bitwise_not(mask)

    cplant = cv2.applyColorMap(histogram, colormap=2)
    cplant1 = cv2.bitwise_and(cplant, cplant, mask=mask)

    output_imgs = {"pseudo_on_img": {"background": "img", "img": None},
                   "pseudo_on_white": {"background": "white", "img": None}}

    if background == 'img' or background == 'both':
        # mask the background and color the plant with color scheme 'jet'
        img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        img_back = cv2.bitwise_and(img_gray, img_gray, mask=mask_inv)
        img_back3 = np.dstack((img_back, img_back, img_back))

        output_imgs["pseudo_on_img"]["img"] = cv2.add(cplant1, img_back3)

    if background == 'white' or background == 'both':
        # Get the image size
        if np.shape(img)[2] == 3:
            ix, iy, iz = np.shape(img)
        else:
            ix, iy = np.shape(img)
        size = ix, iy
        back = np.zeros(size, dtype=np.uint8)
        w_back = back + 255
        w_back3 = np.dstack((w_back, w_back, w_back))
        img_back3 = cv2.bitwise_and(w_back3, w_back3, mask=mask_inv)
        output_imgs["pseudo_on_white"]["img"] = cv2.add(cplant1, img_back3)

    if filename:
        for key in output_imgs:
            if output_imgs[key]["img"] is not None:
                fig_name_pseudo = str(filename[0:-4]) + '_' + str(channel) + '_pseudo_on_' + \
                                  output_imgs[key]["background"] + '.jpg'
                path = os.path.dirname(filename)
                print_image(output_imgs[key]["img"], fig_name_pseudo)
                analysis_images.append(['IMAGE', 'pseudo', fig_name_pseudo])
    else:
        path = "."

    if debug is not None:
        if debug == 'print':
            for key in output_imgs:
                if output_imgs[key]["img"] is not None:
                    print_image(output_imgs[key]["img"], (str(device) + "_" + output_imgs[key]["background"] +
                                                          '_pseudocolor.jpg'))
            fig_name = 'VIS_pseudocolor_colorbar_' + str(channel) + '_channel.svg'
            if not os.path.isfile(os.path.join(path, fig_name)):
                plot_colorbar(path, fig_name, bins)
        elif debug == 'plot':
            for key in output_imgs:
                if output_imgs[key]["img"] is not None:
                    plot_image(output_imgs[key]["img"])

    return analysis_images

def analyze_color(img, imgname, mask, bins, device, debug=None, hist_plot_type=None, pseudo_channel='v',
                  pseudo_bkg='img', resolution=300, filename=False):
    """Analyze the color properties of an image object

    Inputs:
    img              = image
    imgname          = name of input image
    mask             = mask made from selected contours
    device           = device number. Used to count steps in the pipeline
    debug            = None, print, or plot. Print = save to file, Plot = print to screen.
    hist_plot_type   = 'None', 'all', 'rgb','lab' or 'hsv'
    color_slice_type = 'None', 'rgb', 'hsv' or 'lab'
    pseudo_channel   = 'None', 'l', 'm' (green-magenta), 'y' (blue-yellow), h','s', or 'v', creates pseduocolored image
                       based on the specified channel
    pseudo_bkg       = 'img' => channel image, 'white' => white background image, 'both' => both img and white options
    filename         = False or image name. If defined print image

    Returns:
    device           = device number
    hist_header      = color histogram data table headers
    hist_data        = color histogram data table values
    analysis_images  = list of output images

    :param img: numpy array
    :param imgname: str
    :param mask: numpy array
    :param bins: int
    :param device: int
    :param debug: str
    :param hist_plot_type: str
    :param pseudo_channel: str
    :param pseudo_bkg: str
    :param resolution: int
    :param filename: str
    :return device: int
    :return hist_header: list
    :return hist_data: list
    :return analysis_images: list
    """
    device += 1

    masked = cv2.bitwise_and(img, img, mask=mask)
    b, g, r = cv2.split(masked)
    lab = cv2.cvtColor(masked, cv2.COLOR_BGR2LAB)
    l, m, y = cv2.split(lab)
    hsv = cv2.cvtColor(masked, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)

    # Color channel dictionary
    norm_channels = {"b": b / (256 / bins),
                     "g": g / (256 / bins),
                     "r": r / (256 / bins),
                     "l": l / (256 / bins),
                     "m": m / (256 / bins),
                     "y": y / (256 / bins),
                     "h": h / (256 / bins),
                     "s": s / (256 / bins),
                     "v": v / (256 / bins)
                     }
    # Histogram plot types
    hist_types = {"all": ("b", "g", "r", "l", "m", "y", "h", "s", "v"),
                  "rgb": ("b", "g", "r"),
                  "lab": ("l", "m", "y"),
                  "hsv": ("h", "s", "v")}

    # If the user-input pseudo_channel is not None and is not found in the list of accepted channels, exit
    if pseudo_channel is not None and pseudo_channel not in norm_channels:
        fatal_error("Pseudocolor channel was " + str(pseudo_channel) +
                    ', but can only be one of the following: None, "l", "m", "y", "h", "s" or "v"!')
    # If the user-input pseudocolored image background is not in the accepted input list, exit
    if pseudo_bkg not in ["white", "img", "both"]:
        fatal_error("The pseudocolored image background was " + str(pseudo_bkg) +
                    ', but can only be one of the following: "white", "img", or "both"!')
    # If the user-input histogram color-channel plot type is not in the list of accepted channels, exit
    if hist_plot_type is not None and hist_plot_type not in hist_types:
        fatal_error("The histogram plot type was " + str(hist_plot_type) +
                    ', but can only be one of the following: None, "all", "rgb", "lab", or "hsv"!')

    histograms = {
        "b": {"label": "blue", "graph_color": "blue",
              "hist": cv2.calcHist([norm_channels["b"]], [0], mask, [bins], [0, (bins - 1)])},
        "g": {"label": "green", "graph_color": "forestgreen",
              "hist": cv2.calcHist([norm_channels["g"]], [0], mask, [bins], [0, (bins - 1)])},
        "r": {"label": "red", "graph_color": "red",
              "hist": cv2.calcHist([norm_channels["r"]], [0], mask, [bins], [0, (bins - 1)])},
        "l": {"label": "lightness", "graph_color": "dimgray",
              "hist": cv2.calcHist([norm_channels["l"]], [0], mask, [bins], [0, (bins - 1)])},
        "m": {"label": "green-magenta", "graph_color": "magenta",
              "hist": cv2.calcHist([norm_channels["m"]], [0], mask, [bins], [0, (bins - 1)])},
        "y": {"label": "blue-yellow", "graph_color": "yellow",
              "hist": cv2.calcHist([norm_channels["y"]], [0], mask, [bins], [0, (bins - 1)])},
        "h": {"label": "hue", "graph_color": "blueviolet",
              "hist": cv2.calcHist([norm_channels["h"]], [0], mask, [bins], [0, (bins - 1)])},
        "s": {"label": "saturation", "graph_color": "cyan",
              "hist": cv2.calcHist([norm_channels["s"]], [0], mask, [bins], [0, (bins - 1)])},
        "v": {"label": "value", "graph_color": "orange",
              "hist": cv2.calcHist([norm_channels["v"]], [0], mask, [bins], [0, (bins - 1)])}
    }

    hist_data_b = [l[0] for l in histograms["b"]["hist"]]
    hist_data_g = [l[0] for l in histograms["g"]["hist"]]
    hist_data_r = [l[0] for l in histograms["r"]["hist"]]
    hist_data_l = [l[0] for l in histograms["l"]["hist"]]
    hist_data_m = [l[0] for l in histograms["m"]["hist"]]
    hist_data_y = [l[0] for l in histograms["y"]["hist"]]
    hist_data_h = [l[0] for l in histograms["h"]["hist"]]
    hist_data_s = [l[0] for l in histograms["s"]["hist"]]
    hist_data_v = [l[0] for l in histograms["v"]["hist"]]

    binval = np.arange(0, bins)
    bin_values = [l for l in binval]

    # Store Color Histogram Data
    hist_header = [
        'HEADER_HISTOGRAM',
        'bin-number',
        'bin-values',
        'blue',
        'green',
        'red',
        'lightness',
        'green-magenta',
        'blue-yellow',
        'hue',
        'saturation',
        'value'
    ]

    hist_data = [
        'HISTOGRAM_DATA',
        bins,
        bin_values,
        hist_data_b,
        hist_data_g,
        hist_data_r,
        hist_data_l,
        hist_data_m,
        hist_data_y,
        hist_data_h,
        hist_data_s,
        hist_data_v
    ]

    analysis_images = []

    if pseudo_channel is not None:
        analysis_images = _pseudocolored_image(device, norm_channels[pseudo_channel], bins, img, mask, pseudo_bkg,
                                               pseudo_channel, filename, resolution, analysis_images, debug)

    if hist_plot_type is not None and filename:
        import matplotlib
        matplotlib.use('Agg')
        from matplotlib import pyplot as plt

        # Create Histogram Plot
        for channel in hist_types[hist_plot_type]:
            plt.plot(histograms[channel]["hist"], color=histograms[channel]["graph_color"],
                     label=histograms[channel]["label"])
            plt.xlim([0, bins - 1])
            plt.legend()

        # Print plot
        fig_name = (str(filename[0:-4]) + '_' + str(hist_plot_type) + '_hist.svg')
        plt.savefig(fig_name)
        analysis_images.append(['IMAGE', 'hist', fig_name])
        if debug == 'print':
            fig_name = (str(device) + '_' + str(hist_plot_type) + '_hist.svg')
            plt.savefig(fig_name)
        plt.clf()

    return device, hist_header, hist_data, analysis_images
<commit_msg>Fixed NumPy Python2/3 type casting issue<commit_after># Analyze Color of Object

import os
import cv2
import numpy as np
from . import print_image
from . import plot_image
from . import fatal_error
from . import plot_colorbar

def _pseudocolored_image(device, histogram, bins, img, mask, background, channel, filename, resolution,
                         analysis_images, debug):
    """Pseudocolor image.

    Inputs:
    histogram       = a normalized histogram of color values from one color channel
    bins            = number of color bins the channel is divided into
    img             = input image
    mask            = binary mask image
    background      = what background image?: channel image (img) or white
    channel         = color channel name
    filename        = input image filename
    resolution      = output image resolution
    analysis_images = list of analysis image filenames
    debug           = print or plot. Print = save to file, Plot = print to screen.

    Returns:
    analysis_images = list of analysis image filenames

    :param histogram: list
    :param bins: int
    :param img: numpy array
    :param mask: numpy array
    :param background: str
    :param channel: str
    :param filename: str
    :param resolution: int
    :param analysis_images: list
    :return analysis_images: list
    """
    mask_inv = cv2.bitwise_not(mask)

    cplant = cv2.applyColorMap(histogram, colormap=2)
    cplant1 = cv2.bitwise_and(cplant, cplant, mask=mask)

    output_imgs = {"pseudo_on_img": {"background": "img", "img": None},
                   "pseudo_on_white": {"background": "white", "img": None}}

    if background == 'img' or background == 'both':
        # mask the background and color the plant with color scheme 'jet'
        img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        img_back = cv2.bitwise_and(img_gray, img_gray, mask=mask_inv)
        img_back3 = np.dstack((img_back, img_back, img_back))

        output_imgs["pseudo_on_img"]["img"] = cv2.add(cplant1, img_back3)

    if background == 'white' or background == 'both':
        # Get the image size
        if np.shape(img)[2] == 3:
            ix, iy, iz = np.shape(img)
        else:
            ix, iy = np.shape(img)
        size = ix, iy
        back = np.zeros(size, dtype=np.uint8)
        w_back = back + 255
        w_back3 = np.dstack((w_back, w_back, w_back))
        img_back3 = cv2.bitwise_and(w_back3, w_back3, mask=mask_inv)
        output_imgs["pseudo_on_white"]["img"] = cv2.add(cplant1, img_back3)

    if filename:
        for key in output_imgs:
            if output_imgs[key]["img"] is not None:
                fig_name_pseudo = str(filename[0:-4]) + '_' + str(channel) + '_pseudo_on_' + \
                                  output_imgs[key]["background"] + '.jpg'
                path = os.path.dirname(filename)
                print_image(output_imgs[key]["img"], fig_name_pseudo)
                analysis_images.append(['IMAGE', 'pseudo', fig_name_pseudo])
    else:
        path = "."

    if debug is not None:
        if debug == 'print':
            for key in output_imgs:
                if output_imgs[key]["img"] is not None:
                    print_image(output_imgs[key]["img"], (str(device) + "_" + output_imgs[key]["background"] +
                                                          '_pseudocolor.jpg'))
            fig_name = 'VIS_pseudocolor_colorbar_' + str(channel) + '_channel.svg'
            if not os.path.isfile(os.path.join(path, fig_name)):
                plot_colorbar(path, fig_name, bins)
        elif debug == 'plot':
            for key in output_imgs:
                if output_imgs[key]["img"] is not None:
                    plot_image(output_imgs[key]["img"])

    return analysis_images

def analyze_color(img, imgname, mask, bins, device, debug=None, hist_plot_type=None, pseudo_channel='v',
                  pseudo_bkg='img', resolution=300, filename=False):
    """Analyze the color properties of an image object

    Inputs:
    img              = image
    imgname          = name of input image
    mask             = mask made from selected contours
    device           = device number. Used to count steps in the pipeline
    debug            = None, print, or plot. Print = save to file, Plot = print to screen.
    hist_plot_type   = 'None', 'all', 'rgb','lab' or 'hsv'
    color_slice_type = 'None', 'rgb', 'hsv' or 'lab'
    pseudo_channel   = 'None', 'l', 'm' (green-magenta), 'y' (blue-yellow), h','s', or 'v', creates pseduocolored image
                       based on the specified channel
    pseudo_bkg       = 'img' => channel image, 'white' => white background image, 'both' => both img and white options
    filename         = False or image name. If defined print image

    Returns:
    device           = device number
    hist_header      = color histogram data table headers
    hist_data        = color histogram data table values
    analysis_images  = list of output images

    :param img: numpy array
    :param imgname: str
    :param mask: numpy array
    :param bins: int
    :param device: int
    :param debug: str
    :param hist_plot_type: str
    :param pseudo_channel: str
    :param pseudo_bkg: str
    :param resolution: int
    :param filename: str
    :return device: int
    :return hist_header: list
    :return hist_data: list
    :return analysis_images: list
    """
    device += 1

    masked = cv2.bitwise_and(img, img, mask=mask)
    b, g, r = cv2.split(masked)
    lab = cv2.cvtColor(masked, cv2.COLOR_BGR2LAB)
    l, m, y = cv2.split(lab)
    hsv = cv2.cvtColor(masked, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)

    # Color channel dictionary
    norm_channels = {"b": np.divide(b, (256 / bins)).astype(np.uint8),
                     "g": np.divide(g, (256 / bins)).astype(np.uint8),
                     "r": np.divide(r, (256 / bins)).astype(np.uint8),
                     "l": np.divide(l, (256 / bins)).astype(np.uint8),
                     "m": np.divide(m, (256 / bins)).astype(np.uint8),
                     "y": np.divide(y, (256 / bins)).astype(np.uint8),
                     "h": np.divide(h, (256 / bins)).astype(np.uint8),
                     "s": np.divide(s, (256 / bins)).astype(np.uint8),
                     "v": np.divide(v, (256 / bins)).astype(np.uint8)
                     }

    # Histogram plot types
    hist_types = {"all": ("b", "g", "r", "l", "m", "y", "h", "s", "v"),
                  "rgb": ("b", "g", "r"),
                  "lab": ("l", "m", "y"),
                  "hsv": ("h", "s", "v")}

    # If the user-input pseudo_channel is not None and is not found in the list of accepted channels, exit
    if pseudo_channel is not None and pseudo_channel not in norm_channels:
        fatal_error("Pseudocolor channel was " + str(pseudo_channel) +
                    ', but can only be one of the following: None, "l", "m", "y", "h", "s" or "v"!')
    # If the user-input pseudocolored image background is not in the accepted input list, exit
    if pseudo_bkg not in ["white", "img", "both"]:
        fatal_error("The pseudocolored image background was " + str(pseudo_bkg) +
                    ', but can only be one of the following: "white", "img", or "both"!')
    # If the user-input histogram color-channel plot type is not in the list of accepted channels, exit
    if hist_plot_type is not None and hist_plot_type not in hist_types:
        fatal_error("The histogram plot type was " + str(hist_plot_type) +
                    ', but can only be one of the following: None, "all", "rgb", "lab", or "hsv"!')
    histograms = {
        "b": {"label": "blue", "graph_color": "blue",
              "hist": cv2.calcHist([norm_channels["b"]], [0], mask, [bins], [0, (bins - 1)])},
        "g": {"label": "green", "graph_color": "forestgreen",
              "hist": cv2.calcHist([norm_channels["g"]], [0], mask, [bins], [0, (bins - 1)])},
        "r": {"label": "red", "graph_color": "red",
              "hist": cv2.calcHist([norm_channels["r"]], [0], mask, [bins], [0, (bins - 1)])},
        "l": {"label": "lightness", "graph_color": "dimgray",
              "hist": cv2.calcHist([norm_channels["l"]], [0], mask, [bins], [0, (bins - 1)])},
        "m": {"label": "green-magenta", "graph_color": "magenta",
              "hist": cv2.calcHist([norm_channels["m"]], [0], mask, [bins], [0, (bins - 1)])},
        "y": {"label": "blue-yellow", "graph_color": "yellow",
              "hist": cv2.calcHist([norm_channels["y"]], [0], mask, [bins], [0, (bins - 1)])},
        "h": {"label": "hue", "graph_color": "blueviolet",
              "hist": cv2.calcHist([norm_channels["h"]], [0], mask, [bins], [0, (bins - 1)])},
        "s": {"label": "saturation", "graph_color": "cyan",
              "hist": cv2.calcHist([norm_channels["s"]], [0], mask, [bins], [0, (bins - 1)])},
        "v": {"label": "value", "graph_color": "orange",
              "hist": cv2.calcHist([norm_channels["v"]], [0], mask, [bins], [0, (bins - 1)])}
    }

    hist_data_b = [l[0] for l in histograms["b"]["hist"]]
    hist_data_g = [l[0] for l in histograms["g"]["hist"]]
    hist_data_r = [l[0] for l in histograms["r"]["hist"]]
    hist_data_l = [l[0] for l in histograms["l"]["hist"]]
    hist_data_m = [l[0] for l in histograms["m"]["hist"]]
    hist_data_y = [l[0] for l in histograms["y"]["hist"]]
    hist_data_h = [l[0] for l in histograms["h"]["hist"]]
    hist_data_s = [l[0] for l in histograms["s"]["hist"]]
    hist_data_v = [l[0] for l in histograms["v"]["hist"]]

    binval = np.arange(0, bins)
    bin_values = [l for l in binval]

    # Store Color Histogram Data
    hist_header = [
        'HEADER_HISTOGRAM',
        'bin-number',
        'bin-values',
        'blue',
        'green',
        'red',
        'lightness',
        'green-magenta',
        'blue-yellow',
        'hue',
        'saturation',
        'value'
    ]

    hist_data = [
        'HISTOGRAM_DATA',
        bins,
        bin_values,
        hist_data_b,
        hist_data_g,
        hist_data_r,
        hist_data_l,
        hist_data_m,
        hist_data_y,
        hist_data_h,
        hist_data_s,
        hist_data_v
    ]

    analysis_images = []

    if pseudo_channel is not None:
        analysis_images = _pseudocolored_image(device, norm_channels[pseudo_channel], bins, img, mask, pseudo_bkg,
                                               pseudo_channel, filename, resolution, analysis_images, debug)

    if hist_plot_type is not None and filename:
        import matplotlib
        matplotlib.use('Agg')
        from matplotlib import pyplot as plt

        # Create Histogram Plot
        for channel in hist_types[hist_plot_type]:
            plt.plot(histograms[channel]["hist"], color=histograms[channel]["graph_color"],
                     label=histograms[channel]["label"])
            plt.xlim([0, bins - 1])
            plt.legend()

        # Print plot
        fig_name = (str(filename[0:-4]) + '_' + str(hist_plot_type) + '_hist.svg')
        plt.savefig(fig_name)
        analysis_images.append(['IMAGE', 'hist', fig_name])
        if debug == 'print':
            fig_name = (str(device) + '_' + str(hist_plot_type) + '_hist.svg')
            plt.savefig(fig_name)
        plt.clf()

    return device, hist_header, hist_data, analysis_images
<|endoftext|>