<commit_before>#!/usr/bin/env python3

# This is a python script that emulates a terminal session and runs
# commands from a supplied markdown file..

import difflib
import optparse
import os
import pexpect
import random
import time
import shlex
import sys
import json
import colorama
colorama.init(strip=None)

class Demo(object):
    def __init__(self, script_dir = "demo_scripts", filename = "script.md", env = None, is_simulation = True, is_automated = False, is_testing = False):
        self.filename = filename
        self.script_dir = script_dir
        self.env = env
        self.is_simulation = is_simulation
        self.is_automated = is_automated
        self.is_testing = is_testing

    def run(self):
        # Reads a script.md file in the indicated directoy and runs the
        # commands contained within. If simulation == True then human
        # entry will be simulated (looks like typing and waits for
        # keyboard input before proceeding to the next command). This is
        # useful if you want to run a fully automated demo.
        #
        # The script.md file will be parsed as follows:
        #
        # ``` marks the start or end of a code block
        #
        # Each line in a code block will be treated as a separate command.
        #
        # All other lines will be ignored

        in_code_block = False
        in_results_section = False

        if not self.script_dir.endswith('/'):
            self.script_dir = script_dir + "/"

        file = self.script_dir + self.filename

        lines = list(open(file))
        expected_results = ""
        actual_results = ""
        passed_tests = 0
        failed_tests = 0
        is_first_line = True
        executed_code_in_this_section = False

        for line in lines:
            if in_results_section and in_code_block and not line.startswith("```"):
                expected_results += line

            if line.startswith("Results:"):
                # Entering results section
                in_results_section = True
                pos = line.lower().find("expected similarity: ")
                if pos >= 0:
                    pos = pos + len("expected similarity: ")
                    s = line[pos:]
                    expected_similarity = float(s)
                else:
                    expected_similarity = 0.66
            elif line.startswith("```") and not in_code_block:
                # Entering a code block, if in_results_section = True then it's a results block
                in_code_block = True
            elif line.startswith("```") and in_code_block and in_results_section:
                # Finishing results section
                if in_results_section and is_testing:
                    if test_results(expected_results, actual_results, expected_similarity):
                        passed_tests += 1
                    else:
                        failed_tests += 1
                        expected_results = ""
                        actual_results = ""
                        in_results_section = False
                        in_code_block = False
            elif line.startswith("```") and in_code_block and not in_results_section:
                # Finishing executable code block
                in_code_block = False
            elif in_code_block and not in_results_section:
                # Executable line
                print("$ ", end="", flush=True)
                check_for_interactive_command(self.script_dir, self.is_automated)
                actual_results = simulate_command(line, self.script_dir, self.env, self.is_simulation, self.is_automated)
                executed_code_in_this_section = True
            elif line.startswith("#") and not in_code_block and not in_results_section and not self.is_automated:
                # Heading in descriptive text, indicating a new section
                if is_first_line:
                    run_command("clear", script_dir, env)
                elif executed_code_in_this_section:
                    executed_code_in_this_section = False
                    print("$ ", end="", flush=True)
                    check_for_interactive_command(script_dir, self.is_automated)
                    simulate_command("clear", script_dir, env, self.is_simulation, self.is_automated)
                    if not is_simulation:
                        print("$ ", end="", flush=True)
                        # Since this is a heading we are not really simulating a command, it appears as a comment
                        simulate_command(line, script_dir, env, is_simulation, is_automated)
            elif not self.is_simulation and not in_results_section:
                # Descriptinve text
                print(colorama.Fore.CYAN, end = "")
                print(line, end="", flush=True)
                print(colorama.Style.RESET_ALL, end = "")

            is_first_line = False

        if self.is_testing:
            print("\n\n=============================\n\n")
            print("Test Run Complete.")
            if failed_tests > 0:
                print(colorama.Fore.RED + colorama.Style.BRIGHT)
                print("Failed Tests: " + str(failed_tests))
                print(colorama.Style.RESET_ALL)
            else:
                print(colorama.Fore.GREEN + colorama.Style.BRIGHT)
                print("No failed tests")
                print(colorama.Style.RESET_ALL)
                print("Passed Tests: " + str(passed_tests))
            if failed_tests > 0:
                print("\n\n")
                print("View failure reports in context in the above output.")
                print("\n\n=============================\n\n")

def type_command(command, script_dir, simulation):
    # Displays the command on the screen
    # If simulation == True then it will look like someone is typing the command
    print(colorama.Fore.WHITE + colorama.Style.BRIGHT, end="")
    for char in command:
        if (char != '\n'):
            print(char, end="", flush=True)
        if simulation:
            delay = random.uniform(0.01, 0.04)
            time.sleep(delay)
    print(colorama.Style.RESET_ALL, end="")

def simulate_command(command, script_dir, env = None, simulation = True, is_automatic=False):
    # Types the command on the screen, executes it and outputs the
    # results if simulation == True then system will make the "typing"
    # look real and will wait for keyboard entry before proceeding to
    # the next command
    type_command(command, script_dir, simulation)
    check_for_interactive_command(script_dir, is_automatic)
    print()
    output = run_command(command, script_dir, env)

    return output

def get_simdem_environment(directory):
    # Populates each shell environment with a set of environment vars
    # loaded via env.json file stored either in the project root
    # directory

    env = {}

    if not directory.endswith('/'):
        directory = directory + "/"

    filename = directory + "../env.json"
    if os.path.isfile(directory + "../env.json"):
        with open(filename) as env_file:
            app_env = json.load(env_file)
            env.update(app_env)

    filename = directory + "env.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            script_env = json.load(env_file)
            env.update(script_env)

    filename = directory + "../env.local.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            local_env = json.load(env_file)
            env.update(local_env)

    filename = directory + "env.local.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            local_env = json.load(env_file)
            env.update(local_env)

    return env

def run_command(command, script_dir, env=None):
    if command.startswith("sudo "):
        is_docker_command = 'if [ -f /.dockerenv ]; then echo "True"; else echo "False"; fi'
        shell = pexpect.spawnu('/bin/bash', ['-c', is_docker_command], env=env, cwd=script_dir, timeout=None)
        shell.expect(pexpect.EOF)
        is_docker = shell.before.strip() == "True"
        if is_docker:
            command = command[5:]

    print(colorama.Fore.GREEN+colorama.Style.BRIGHT)
    shell = pexpect.spawnu('/bin/bash', ['-c', command], env=env, cwd=script_dir, timeout=None)
    shell.logfile = sys.stdout
    shell.expect(pexpect.EOF)
    print(colorama.Style.RESET_ALL)
    return shell.before

def check_for_interactive_command(script_dir, is_automated=False):
    # Wait for a key to be pressed. Most keys result in the script
    # progressing, but a few have special meaning. See the
    # documentation or code for a description of the special keys.
    if not is_automated:
        key = get_instruction_key()

        if key == 'b' or key == 'B':
            command = input()
            run_command(command, script_dir)
            print("$ ", end="", flush=True)
            check_for_interactive_command(script_dir, is_automated)

def get_instruction_key():
    """Waits for a single keypress on stdin.

    This is a silly function to call if you need to do it a lot because it has
    to store stdin's current setup, setup stdin for reading single keystrokes
    then read the single keystroke then revert stdin back after reading the
    keystroke.

    Returns the character of the key that was pressed (zero on
    KeyboardInterrupt which can happen when a signal gets handled)

    This method is licensed under cc by-sa 3.0
    Thanks to mheyman http://stackoverflow.com/questions/983354/how-do-i-make-python-to-wait-for-a-pressed-key\
    """
    import termios, fcntl, sys, os
    fd = sys.stdin.fileno()
    # save old state
    flags_save = fcntl.fcntl(fd, fcntl.F_GETFL)
    attrs_save = termios.tcgetattr(fd)
    # make raw - the way to do this comes from the termios(3) man page.
    attrs = list(attrs_save) # copy the stored version to update
    # iflag
    attrs[0] &= ~(termios.IGNBRK | termios.BRKINT | termios.PARMRK
                  | termios.ISTRIP | termios.INLCR | termios. IGNCR
                  | termios.ICRNL | termios.IXON )
    # oflag
    attrs[1] &= ~termios.OPOST
    # cflag
    attrs[2] &= ~(termios.CSIZE | termios. PARENB)
    attrs[2] |= termios.CS8
    # lflag
    attrs[3] &= ~(termios.ECHONL | termios.ECHO | termios.ICANON
                  | termios.ISIG | termios.IEXTEN)
    termios.tcsetattr(fd, termios.TCSANOW, attrs)
    # turn off non-blocking
    fcntl.fcntl(fd, fcntl.F_SETFL, flags_save & ~os.O_NONBLOCK)
    # read a single keystroke
    try:
        ret = sys.stdin.read(1) # returns a single character
    except KeyboardInterrupt:
        ret = 0
    finally:
        # restore old state
        termios.tcsetattr(fd, termios.TCSAFLUSH, attrs_save)
        fcntl.fcntl(fd, fcntl.F_SETFL, flags_save)
    return ret

def test_results(expected_results, actual_results, expected_similarity = 0.66):
    differ = difflib.Differ()
    comparison = differ.compare(actual_results, expected_results)
    diff = differ.compare(actual_results, expected_results)
    seq = difflib.SequenceMatcher(lambda x: x in " \t\n\r", actual_results, expected_results)

    is_pass = seq.ratio() >= expected_similarity

    if not is_pass:
        print("\n\n=============================\n\n")
        print(colorama.Fore.RED + colorama.Style.BRIGHT)
        print("FAILED")
        print(colorama.Style.RESET_ALL)
        print("Similarity ratio:    " + str(seq.ratio()))
        print("Expected Similarity: " + str(expected_similarity))
        print("\n\n=============================\n\n")
        print("Expected results:")
        print(colorama.Fore.GREEN + colorama.Style.BRIGHT)
        print(expected_results)
        print(colorama.Style.RESET_ALL)
        print("Actual results:")
        print(colorama.Fore.RED + colorama.Style.BRIGHT)
        print(actual_results)
        print(colorama.Style.RESET_ALL)
        print("\n\n=============================\n\n")
        print(colorama.Style.RESET_ALL)
    return is_pass

def get_bash_script(script_dir, env=None, is_simulation = True, is_automated=False, is_testing=False):
    # Reads a script.md file in the indicated directoy and builds an
    # executable bash script from the commands contained within.

    in_code_block = False
    in_results_section = False

    if not script_dir.endswith('/'):
        script_dir = script_dir + "/"
    filename = script_dir + "script.md"

    script = ""
    env = get_simdem_environment(script_dir)
    for key, value in env.items():
        script += key + "='" + value + "'\n"

    lines = list(open(filename))
    for line in lines:
        if line.startswith("Results:"):
            # Entering results section
            in_results_section = True
        elif line.startswith("```") and not in_code_block:
            # Entering a code block, if in_results_section = True then it's a results block
            in_code_block = True
        elif line.startswith("```") and in_code_block and in_results_section:
            # Finishing results section
            in_results_section = False
            in_code_block = False
        elif line.startswith("```") and in_code_block and not in_results_section:
            # Finishing executable code block
            in_code_block = False
        elif in_code_block and not in_results_section:
            # Executable line
            script += line
        elif line.startswith("#") and not in_code_block and not in_results_section and not is_automated:
            # Heading in descriptive text
            script +="\n"
    return script

def main():
    """SimDem CLI interpreter"""

    p = optparse.OptionParser("%prog [run|test|script] <options> DEMO_NAME", version="%prog 0.2.3")
    p.add_option('--style', '-s', default="tutorial",
                 help="The style of simulation you want to run. 'tutorial' (the default) will print out all text and pause for user input before running commands. 'simulate' will not print out the text but will still pause for input.")
    p.add_option('--path', '-p', default="demo_scripts/",
                 help="The Path to the demo scripts directory.")
    p.add_option('--auto', '-a', default="False",
                 help="If set to anything other than False the application will not wait for user keypresses between commands.")
    p.add_option('--test', '-t', default="False",
                 help="If set to anything other than False the output of the command will be compared to the expected results in the sript. Any failures will be reported")

    options, arguments = p.parse_args()

    if len(arguments) == 0:
        arguments.append("run")

    if not options.path.endswith("/"):
        options.path += "/"

    if options.auto == "False":
        is_automatic = False
    else:
        is_automatic = True

    if options.test == "False":
        is_test = False
    else:
        is_test = True

    if options.style == "simulate":
        simulate = True
    elif options.style == 'tutorial':
        simulate = False
    else:
        print("Unkown style (--style, -s): " + options.style)
        exit(1)

    if len(arguments) == 2:
        script_dir = options.path + arguments[1]
    else:
        script_dir = options.path

    env = os.environ.copy()
    env.update(get_simdem_environment(script_dir))
    cmd = arguments[0]

    filename = "script.md"
    if cmd == "run":
        demo = Demo(script_dir, filename, env, simulate, is_automatic, is_test);
        demo.run()
    elif cmd == "test":
        is_automatic = True and options.auto
        is_test = True and options.test
        demo = Demo(script_dir, filename, env, simulate, is_automatic, is_test);
        demo.run()
    elif cmd == "script":
        print(get_bash_script(script_dir, env))
    else:
        print("Unknown command: " + cmd)
        print("Run with --help for guidance.")
main()
<commit_msg>WIP: no longer passing variables to simulate command function<commit_after>#!/usr/bin/env python3

# This is a python script that emulates a terminal session and runs
# commands from a supplied markdown file..

import difflib
import optparse
import os
import pexpect
import random
import time
import shlex
import sys
import json
import colorama
colorama.init(strip=None)

class Demo(object):
    def __init__(self, script_dir = "demo_scripts", filename = "script.md", env = None, is_simulation = True, is_automated = False, is_testing = False):
        self.filename = filename
        self.script_dir = script_dir
        self.env = env
        self.is_simulation = is_simulation
        self.is_automated = is_automated
        self.is_testing = is_testing
        self.current_comment = ""

    def run(self):
        # Reads a script.md file in the indicated directoy and runs the
        # commands contained within. If simulation == True then human
        # entry will be simulated (looks like typing and waits for
        # keyboard input before proceeding to the next command). This is
        # useful if you want to run a fully automated demo.
        #
        # The script.md file will be parsed as follows:
        #
        # ``` marks the start or end of a code block
        #
        # Each line in a code block will be treated as a separate command.
        #
        # All other lines will be ignored

        in_code_block = False
        in_results_section = False

        if not self.script_dir.endswith('/'):
            self.script_dir = script_dir + "/"

        file = self.script_dir + self.filename

        lines = list(open(file))
        expected_results = ""
        actual_results = ""
        passed_tests = 0
        failed_tests = 0
        is_first_line = True
        executed_code_in_this_section = False

        for line in lines:
            if in_results_section and in_code_block and not line.startswith("```"):
                expected_results += line

            if line.startswith("Results:"):
                # Entering results section
                in_results_section = True
                pos = line.lower().find("expected similarity: ")
                if pos >= 0:
                    pos = pos + len("expected similarity: ")
                    s = line[pos:]
                    expected_similarity = float(s)
                else:
                    expected_similarity = 0.66
            elif line.startswith("```") and not in_code_block:
                # Entering a code block, if in_results_section = True then it's a results block
                in_code_block = True
            elif line.startswith("```") and in_code_block and in_results_section:
                # Finishing results section
                if in_results_section and is_testing:
                    if test_results(expected_results, actual_results, expected_similarity):
                        passed_tests += 1
                    else:
                        failed_tests += 1
                        expected_results = ""
                        actual_results = ""
                        in_results_section = False
                        in_code_block = False
            elif line.startswith("```") and in_code_block and not in_results_section:
                # Finishing executable code block
                in_code_block = False
            elif in_code_block and not in_results_section:
                # Executable line
                print("$ ", end="", flush=True)
                check_for_interactive_command(self.script_dir, self.is_automated)
                self.current_command = line
                actual_results = simulate_command(self)
                executed_code_in_this_section = True
            elif line.startswith("#") and not in_code_block and not in_results_section and not self.is_automated:
                # Heading in descriptive text, indicating a new section
                if is_first_line:
                    run_command("clear", script_dir, env)
                elif executed_code_in_this_section:
                    executed_code_in_this_section = False
                    print("$ ", end="", flush=True)
                    check_for_interactive_command(script_dir, self.is_automated)
                    self.current_command = "clear"
                    simulate_command(self)
                    if not is_simulation:
                        print("$ ", end="", flush=True)
                        # Since this is a heading we are not really simulating a command, it appears as a comment
                        simulate_command(line, script_dir, env, is_simulation, is_automated)
            elif not self.is_simulation and not in_results_section:
                # Descriptinve text
                print(colorama.Fore.CYAN, end = "")
                print(line, end="", flush=True)
                print(colorama.Style.RESET_ALL, end = "")

            is_first_line = False

        if self.is_testing:
            print("\n\n=============================\n\n")
            print("Test Run Complete.")
            if failed_tests > 0:
                print(colorama.Fore.RED + colorama.Style.BRIGHT)
                print("Failed Tests: " + str(failed_tests))
                print(colorama.Style.RESET_ALL)
            else:
                print(colorama.Fore.GREEN + colorama.Style.BRIGHT)
                print("No failed tests")
                print(colorama.Style.RESET_ALL)
                print("Passed Tests: " + str(passed_tests))
            if failed_tests > 0:
                print("\n\n")
                print("View failure reports in context in the above output.")
                print("\n\n=============================\n\n")

def type_command(command, script_dir, simulation):
    # Displays the command on the screen
    # If simulation == True then it will look like someone is typing the command
    print(colorama.Fore.WHITE + colorama.Style.BRIGHT, end="")
    for char in command:
        if (char != '\n'):
            print(char, end="", flush=True)
        if simulation:
            delay = random.uniform(0.01, 0.04)
            time.sleep(delay)
    print(colorama.Style.RESET_ALL, end="")

def simulate_command(demo):
    # Types the command on the screen, executes it and outputs the
    # results if simulation == True then system will make the "typing"
    # look real and will wait for keyboard entry before proceeding to
    # the next command
    type_command(demo.current_command, demo.script_dir, demo.is_simulation)
    check_for_interactive_command(demo.script_dir, demo.is_automated)
    print()
    output = run_command(demo.current_command, demo.script_dir, demo.env)

    return output

def get_simdem_environment(directory):
    # Populates each shell environment with a set of environment vars
    # loaded via env.json file stored either in the project root
    # directory

    env = {}

    if not directory.endswith('/'):
        directory = directory + "/"

    filename = directory + "../env.json"
    if os.path.isfile(directory + "../env.json"):
        with open(filename) as env_file:
            app_env = json.load(env_file)
            env.update(app_env)

    filename = directory + "env.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            script_env = json.load(env_file)
            env.update(script_env)

    filename = directory + "../env.local.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            local_env = json.load(env_file)
            env.update(local_env)

    filename = directory + "env.local.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            local_env = json.load(env_file)
            env.update(local_env)

    return env

def run_command(command, script_dir, env=None):
    if command.startswith("sudo "):
        is_docker_command = 'if [ -f /.dockerenv ]; then echo "True"; else echo "False"; fi'
        shell = pexpect.spawnu('/bin/bash', ['-c', is_docker_command], env=env, cwd=script_dir, timeout=None)
        shell.expect(pexpect.EOF)
        is_docker = shell.before.strip() == "True"
        if is_docker:
            command = command[5:]

    print(colorama.Fore.GREEN+colorama.Style.BRIGHT)
    shell = pexpect.spawnu('/bin/bash', ['-c', command], env=env, cwd=script_dir, timeout=None)
    shell.logfile = sys.stdout
    shell.expect(pexpect.EOF)
    print(colorama.Style.RESET_ALL)
    return shell.before

def check_for_interactive_command(script_dir, is_automated=False):
    # Wait for a key to be pressed. Most keys result in the script
    # progressing, but a few have special meaning. See the
    # documentation or code for a description of the special keys.
    if not is_automated:
        key = get_instruction_key()

        if key == 'b' or key == 'B':
            command = input()
            run_command(command, script_dir)
            print("$ ", end="", flush=True)
            check_for_interactive_command(script_dir, is_automated)

def get_instruction_key():
    """Waits for a single keypress on stdin.

    This is a silly function to call if you need to do it a lot because it has
    to store stdin's current setup, setup stdin for reading single keystrokes
    then read the single keystroke then revert stdin back after reading the
    keystroke.

    Returns the character of the key that was pressed (zero on
    KeyboardInterrupt which can happen when a signal gets handled)

    This method is licensed under cc by-sa 3.0
    Thanks to mheyman http://stackoverflow.com/questions/983354/how-do-i-make-python-to-wait-for-a-pressed-key\
    """
    import termios, fcntl, sys, os
    fd = sys.stdin.fileno()
    # save old state
    flags_save = fcntl.fcntl(fd, fcntl.F_GETFL)
    attrs_save = termios.tcgetattr(fd)
    # make raw - the way to do this comes from the termios(3) man page.
    attrs = list(attrs_save) # copy the stored version to update
    # iflag
    attrs[0] &= ~(termios.IGNBRK | termios.BRKINT | termios.PARMRK
                  | termios.ISTRIP | termios.INLCR | termios. IGNCR
                  | termios.ICRNL | termios.IXON )
    # oflag
    attrs[1] &= ~termios.OPOST
    # cflag
    attrs[2] &= ~(termios.CSIZE | termios. PARENB)
    attrs[2] |= termios.CS8
    # lflag
    attrs[3] &= ~(termios.ECHONL | termios.ECHO | termios.ICANON
                  | termios.ISIG | termios.IEXTEN)
    termios.tcsetattr(fd, termios.TCSANOW, attrs)
    # turn off non-blocking
    fcntl.fcntl(fd, fcntl.F_SETFL, flags_save & ~os.O_NONBLOCK)
    # read a single keystroke
    try:
        ret = sys.stdin.read(1) # returns a single character
    except KeyboardInterrupt:
        ret = 0
    finally:
        # restore old state
        termios.tcsetattr(fd, termios.TCSAFLUSH, attrs_save)
        fcntl.fcntl(fd, fcntl.F_SETFL, flags_save)
    return ret

def test_results(expected_results, actual_results, expected_similarity = 0.66):
    differ = difflib.Differ()
    comparison = differ.compare(actual_results, expected_results)
    diff = differ.compare(actual_results, expected_results)
    seq = difflib.SequenceMatcher(lambda x: x in " \t\n\r", actual_results, expected_results)

    is_pass = seq.ratio() >= expected_similarity

    if not is_pass:
        print("\n\n=============================\n\n")
        print(colorama.Fore.RED + colorama.Style.BRIGHT)
        print("FAILED")
        print(colorama.Style.RESET_ALL)
        print("Similarity ratio:    " + str(seq.ratio()))
        print("Expected Similarity: " + str(expected_similarity))
        print("\n\n=============================\n\n")
        print("Expected results:")
        print(colorama.Fore.GREEN + colorama.Style.BRIGHT)
        print(expected_results)
        print(colorama.Style.RESET_ALL)
        print("Actual results:")
        print(colorama.Fore.RED + colorama.Style.BRIGHT)
        print(actual_results)
        print(colorama.Style.RESET_ALL)
        print("\n\n=============================\n\n")
        print(colorama.Style.RESET_ALL)
    return is_pass

def get_bash_script(script_dir, env=None, is_simulation = True, is_automated=False, is_testing=False):
    # Reads a script.md file in the indicated directoy and builds an
    # executable bash script from the commands contained within.

    in_code_block = False
    in_results_section = False

    if not script_dir.endswith('/'):
        script_dir = script_dir + "/"
    filename = script_dir + "script.md"

    script = ""
    env = get_simdem_environment(script_dir)
    for key, value in env.items():
        script += key + "='" + value + "'\n"

    lines = list(open(filename))
    for line in lines:
        if line.startswith("Results:"):
            # Entering results section
            in_results_section = True
        elif line.startswith("```") and not in_code_block:
            # Entering a code block, if in_results_section = True then it's a results block
            in_code_block = True
        elif line.startswith("```") and in_code_block and in_results_section:
            # Finishing results section
            in_results_section = False
            in_code_block = False
        elif line.startswith("```") and in_code_block and not in_results_section:
            # Finishing executable code block
            in_code_block = False
        elif in_code_block and not in_results_section:
            # Executable line
            script += line
        elif line.startswith("#") and not in_code_block and not in_results_section and not is_automated:
            # Heading in descriptive text
            script +="\n"
    return script

def main():
    """SimDem CLI interpreter"""

    p = optparse.OptionParser("%prog [run|test|script] <options> DEMO_NAME", version="%prog 0.2.3")
    p.add_option('--style', '-s', default="tutorial",
                 help="The style of simulation you want to run. 'tutorial' (the default) will print out all text and pause for user input before running commands. 'simulate' will not print out the text but will still pause for input.")
    p.add_option('--path', '-p', default="demo_scripts/",
                 help="The Path to the demo scripts directory.")
    p.add_option('--auto', '-a', default="False",
                 help="If set to anything other than False the application will not wait for user keypresses between commands.")
    p.add_option('--test', '-t', default="False",
                 help="If set to anything other than False the output of the command will be compared to the expected results in the sript. Any failures will be reported")

    options, arguments = p.parse_args()

    if len(arguments) == 0:
        arguments.append("run")

    if not options.path.endswith("/"):
        options.path += "/"

    if options.auto == "False":
        is_automatic = False
    else:
        is_automatic = True

    if options.test == "False":
        is_test = False
    else:
        is_test = True

    if options.style == "simulate":
        simulate = True
    elif options.style == 'tutorial':
        simulate = False
    else:
        print("Unkown style (--style, -s): " + options.style)
        exit(1)

    if len(arguments) == 2:
        script_dir = options.path + arguments[1]
    else:
        script_dir = options.path

    env = os.environ.copy()
    env.update(get_simdem_environment(script_dir))
    cmd = arguments[0]

    filename = "script.md"
    if cmd == "run":
        demo = Demo(script_dir, filename, env, simulate, is_automatic, is_test);
        demo.run()
    elif cmd == "test":
        is_automatic = True and options.auto
        is_test = True and options.test
        demo = Demo(script_dir, filename, env, simulate, is_automatic, is_test);
        demo.run()
    elif cmd == "script":
        print(get_bash_script(script_dir, env))
    else:
        print("Unknown command: " + cmd)
        print("Run with --help for guidance.")
main()
<|endoftext|>


<commit_before>#!/usr/bin/env python3

# This is a python script that emulates a terminal session and runs
# commands from a supplied markdown file..

import difflib
import optparse
import os
import pexpect
import random
import time
import shlex
import sys
import json

def type_command(command, script_dir, simulation):
    # Displays the command on the screen
    # If simulation == True then it will look like someone is typing the command
    for char in command:
        if (char != '\n'):
            print(char, end="", flush=True)
        if simulation:
            delay = random.uniform(0.02, 0.08)
            time.sleep(delay)

def simulate_command(command, script_dir, env = None, simulation = True, is_automatic=False):
    # Types the command on the screen, executes it and outputs the
    # results if simulation == True then system will make the "typing"
    # look real and will wait for keyboard entry before proceeding to
    # the next command
    type_command(command, script_dir, simulation)
    check_for_interactive_command(script_dir, is_automatic)
    print()
    output = run_command(command, script_dir, env)

    return output

def environment_setup(directory):
    # Populates each shell environment with a set of environment vars
    # loaded via env.json file stored either in the project root
    # directory
    env = os.environ.copy()

    if not directory.endswith('/'):
        directory = directory + "/"

    filename = directory + "../env.json"
    if os.path.isfile(directory + "../env.json"):
        with open(filename) as env_file:
            app_env = json.load(env_file)
            env.update(app_env)

    filename = directory + "env.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            script_env = json.load(env_file)
            env.update(script_env)

    filename = directory + "../env.local.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            local_env = json.load(env_file)
            env.update(local_env)

    filename = directory + "env.local.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            local_env = json.load(env_file)
            env.update(local_env)

    return env

def run_command(command, script_dir, env=None):
    if command.startswith("sudo "):
        is_docker_command = 'if [ -f /.dockerenv ]; then echo "True"; else echo "False"; fi'
        shell = pexpect.spawnu('/bin/bash', ['-c', is_docker_command], env=env, cwd=script_dir, timeout=None)
        shell.expect(pexpect.EOF)
        is_docker = shell.before.strip() == "True"
        if is_docker:
            command = command[5:]

    shell = pexpect.spawnu('/bin/bash', ['-c', command], env=env, cwd=script_dir, timeout=None)
    shell.logfile = sys.stdout
    shell.expect(pexpect.EOF)
    return shell.before

def check_for_interactive_command(script_dir, is_automated=False):
    # Wait for a key to be pressed. Most keys result in the script
    # progressing, but a few have special meaning. See the
    # documentation or code for a description of the special keys.
    if not is_automated:
        key = get_instruction_key()

        if key == 'b' or key == 'B':
            command = input()
            run_command(command, script_dir)
            print("$ ", end="", flush=True)
            check_for_interactive_command(script_dir, is_automated)

def get_instruction_key():
    """Waits for a single keypress on stdin.

    This is a silly function to call if you need to do it a lot because it has
    to store stdin's current setup, setup stdin for reading single keystrokes
    then read the single keystroke then revert stdin back after reading the
    keystroke.

    Returns the character of the key that was pressed (zero on
    KeyboardInterrupt which can happen when a signal gets handled)

    This method is licensed under cc by-sa 3.0
    Thanks to mheyman http://stackoverflow.com/questions/983354/how-do-i-make-python-to-wait-for-a-pressed-key\
    """
    import termios, fcntl, sys, os
    fd = sys.stdin.fileno()
    # save old state
    flags_save = fcntl.fcntl(fd, fcntl.F_GETFL)
    attrs_save = termios.tcgetattr(fd)
    # make raw - the way to do this comes from the termios(3) man page.
    attrs = list(attrs_save) # copy the stored version to update
    # iflag
    attrs[0] &= ~(termios.IGNBRK | termios.BRKINT | termios.PARMRK
                  | termios.ISTRIP | termios.INLCR | termios. IGNCR
                  | termios.ICRNL | termios.IXON )
    # oflag
    attrs[1] &= ~termios.OPOST
    # cflag
    attrs[2] &= ~(termios.CSIZE | termios. PARENB)
    attrs[2] |= termios.CS8
    # lflag
    attrs[3] &= ~(termios.ECHONL | termios.ECHO | termios.ICANON
                  | termios.ISIG | termios.IEXTEN)
    termios.tcsetattr(fd, termios.TCSANOW, attrs)
    # turn off non-blocking
    fcntl.fcntl(fd, fcntl.F_SETFL, flags_save & ~os.O_NONBLOCK)
    # read a single keystroke
    try:
        ret = sys.stdin.read(1) # returns a single character
    except KeyboardInterrupt:
        ret = 0
    finally:
        # restore old state
        termios.tcsetattr(fd, termios.TCSAFLUSH, attrs_save)
        fcntl.fcntl(fd, fcntl.F_SETFL, flags_save)
    return ret

def test_results(expected_results, actual_results, expected_similarity = 0.66):
    differ = difflib.Differ()
    comparison = differ.compare(actual_results, expected_results)
    diff = differ.compare(actual_results, expected_results)
    seq = difflib.SequenceMatcher(lambda x: x in " \t\n\r", actual_results, expected_results)

    is_pass = seq.ratio() >= expected_similarity

    if not is_pass:
        print("\n\n=============================\n\n")
        print("FAILED")
        print("Similarity ratio:    " + str(seq.ratio()))
        print("Expected Similarity: " + str(expected_similarity))
        print("\n\n=============================\n\n")
        print("Expected results:")
        print(expected_results)
        print("Actual results:")
        print(actual_results)
        print("\n\n=============================\n\n")

    return is_pass

def run_script(script_dir, env=None, is_simulation = True, is_automated=False, is_testing=False):
    # Reads a script.md file in the indicated directoy and runs the
    # commands contained within. If simulation == True then human
    # entry will be simulated (looks like typing and waits for
    # keyboard input before proceeding to the next command). This is
    # useful if you want to run a fully automated demo.
    #
    # The script.md file will be parsed as follows:
    #
    # ``` marks the start or end of a code block
    #
    # Each line in a code block will be treated as a separate command.
    #
    # All other lines will be ignored

    in_code_block = False
    in_results_section = False

    if not script_dir.endswith('/'):
        script_dir = script_dir + "/"
    filename = script_dir + "script.md"

    lines = list(open(filename))
    expected_results = ""
    actual_results = ""
    passed_tests = 0
    failed_tests = 0
    is_first_line = True

    for line in lines:
        if in_results_section and in_code_block and not line.startswith("```"):
            expected_results += line

        if line.startswith("Results:"):
            # Entering results section
            in_results_section = True
            pos = line.lower().find("expected similarity: ")
            if pos >= 0:
                pos = pos + len("expected similarity: ")
                s = line[pos:]
                expected_similarity = float(s)
            else:
                expected_similarity = 0.66
        elif line.startswith("```") and not in_code_block:
            # Entering a code block, if in_results_section = True then it's a results block
            in_code_block = True
        elif line.startswith("```") and in_code_block and in_results_section:
            # Finishing results section
            if in_results_section and is_testing:
                if test_results(expected_results, actual_results, expected_similarity):
                    passed_tests += 1
                else:
                    failed_tests += 1
            expected_results = ""
            actual_results = ""
            in_results_section = False
            in_code_block = False
        elif line.startswith("```") and in_code_block and not in_results_section:
            # Finishing executable code block
            in_code_block = False
        elif in_code_block and not in_results_section:
            # Executable line
            print("$ ", end="", flush=True)
            check_for_interactive_command(script_dir, is_automated)
            actual_results = simulate_command(line, script_dir, env, is_simulation, is_automated)
        elif line.startswith("#") and not in_code_block and not in_results_section and not is_automated:
            # Heading in descriptive text
            if is_first_line:
                run_command("clear", script_dir, env)
            else:
                print("$ ", end="", flush=True)
                check_for_interactive_command(script_dir, is_automated)
                simulate_command("clear", script_dir, env, is_simulation, is_automated)
            print("$ ", end="", flush=True)
            simulate_command(line, script_dir, env, is_simulation, is_automated)
        elif not is_simulation and not in_results_section:
            # Descriptinve text
            print(line, end="", flush=True)

        is_first_line = False

    if is_testing:
        print("\n\n=============================\n\n")
        print("Test Run Complete.")
        print("Failed Tests: " + str(failed_tests))
        print("Passed Tests: " + str(passed_tests))
        if failed_tests > 0:
            print("\n\n")
            print("View failure reports in context in the above output.")
        print("\n\n=============================\n\n")

def get_bash_script(script_dir, env=None, is_simulation = True, is_automated=False, is_testing=False):
    # Reads a script.md file in the indicated directoy and builds an
    # executable bash script from the commands contained within.

    in_code_block = False
    in_results_section = False

    if not script_dir.endswith('/'):
        script_dir = script_dir + "/"
    filename = script_dir + "script.md"

    lines = list(open(filename))
    script = ""

    for line in lines:
        if line.startswith("Results:"):
            # Entering results section
            in_results_section = True
        elif line.startswith("```") and not in_code_block:
            # Entering a code block, if in_results_section = True then it's a results block
            in_code_block = True
        elif line.startswith("```") and in_code_block and in_results_section:
            # Finishing results section
            in_results_section = False
            in_code_block = False
        elif line.startswith("```") and in_code_block and not in_results_section:
            # Finishing executable code block
            in_code_block = False
        elif in_code_block and not in_results_section:
            # Executable line
            script += line
        elif line.startswith("#") and not in_code_block and not in_results_section and not is_automated:
            # Heading in descriptive text
            script +="\n"
    return script

def main():
    """SimDem CLI interpreter"""

    p = optparse.OptionParser("%prog [run|test|script] <options> DEMO_NAME", version="%prog 0.2.3")
    p.add_option('--style', '-s', default="tutorial",
                 help="The style of simulation you want to run. 'tutorial' (the default) will print out all text and pause for user input before running commands. 'simulate' will not print out the text but will still pause for input.")
    p.add_option('--path', '-p', default="demo_scripts/",
                 help="The Path to the demo scripts directory.")
    p.add_option('--auto', '-a', default="False",
                 help="If set to anything other than False the application will not wait for user keypresses between commands.")
    p.add_option('--test', '-t', default="False",
                 help="If set to anything other than False the output of the command will be compared to the expected results in the sript. Any failures will be reported")

    options, arguments = p.parse_args()

    if len(arguments) == 0:
        arguments.append("run")

    if not options.path.endswith("/"):
        options.path += "/"

    if options.auto == "False":
        is_automatic = False
    else:
        is_automatic = True

    if options.test == "False":
        is_test = False
    else:
        is_test = True

    if options.style == "simulate":
        simulate = True
    elif options.style == 'tutorial':
        simulate = False
    else:
        print("Unkown style (--style, -s): " + options.style)
        exit(1)

    if len(arguments) == 2:
        script_dir = options.path + arguments[1]
    else:
        script_dir = options.path

    env = environment_setup(script_dir)
    cmd = arguments[0]

    if cmd == "run":
        run_script(script_dir, env, simulate, is_automatic, is_test)
    elif cmd == "test":
        is_automatic = True and options.auto
        is_test = True and options.test
        run_script(script_dir, env, simulate, is_automatic, is_test)
    elif cmd == "script":
        print(get_bash_script(script_dir, env))
    else:
        print("Unknown command: " + cmd)
        print("Run with --help for guidance.")
main()
<commit_msg>Add environment variables setup to the script output. Rename environment_setup directory<commit_after>#!/usr/bin/env python3

# This is a python script that emulates a terminal session and runs
# commands from a supplied markdown file..

import difflib
import optparse
import os
import pexpect
import random
import time
import shlex
import sys
import json

def type_command(command, script_dir, simulation):
    # Displays the command on the screen
    # If simulation == True then it will look like someone is typing the command
    for char in command:
        if (char != '\n'):
            print(char, end="", flush=True)
        if simulation:
            delay = random.uniform(0.02, 0.08)
            time.sleep(delay)

def simulate_command(command, script_dir, env = None, simulation = True, is_automatic=False):
    # Types the command on the screen, executes it and outputs the
    # results if simulation == True then system will make the "typing"
    # look real and will wait for keyboard entry before proceeding to
    # the next command
    type_command(command, script_dir, simulation)
    check_for_interactive_command(script_dir, is_automatic)
    print()
    output = run_command(command, script_dir, env)

    return output

def get_simdem_environment(directory):
    # Populates each shell environment with a set of environment vars
    # loaded via env.json file stored either in the project root
    # directory

    env = {}

    if not directory.endswith('/'):
        directory = directory + "/"

    filename = directory + "../env.json"
    if os.path.isfile(directory + "../env.json"):
        with open(filename) as env_file:
            app_env = json.load(env_file)
            env.update(app_env)

    filename = directory + "env.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            script_env = json.load(env_file)
            env.update(script_env)

    filename = directory + "../env.local.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            local_env = json.load(env_file)
            env.update(local_env)

    filename = directory + "env.local.json"
    if os.path.isfile(filename):
        with open(filename) as env_file:
            local_env = json.load(env_file)
            env.update(local_env)

    return env

def run_command(command, script_dir, env=None):
    if command.startswith("sudo "):
        is_docker_command = 'if [ -f /.dockerenv ]; then echo "True"; else echo "False"; fi'
        shell = pexpect.spawnu('/bin/bash', ['-c', is_docker_command], env=env, cwd=script_dir, timeout=None)
        shell.expect(pexpect.EOF)
        is_docker = shell.before.strip() == "True"
        if is_docker:
            command = command[5:]

    shell = pexpect.spawnu('/bin/bash', ['-c', command], env=env, cwd=script_dir, timeout=None)
    shell.logfile = sys.stdout
    shell.expect(pexpect.EOF)
    return shell.before

def check_for_interactive_command(script_dir, is_automated=False):
    # Wait for a key to be pressed. Most keys result in the script
    # progressing, but a few have special meaning. See the
    # documentation or code for a description of the special keys.
    if not is_automated:
        key = get_instruction_key()

        if key == 'b' or key == 'B':
            command = input()
            run_command(command, script_dir)
            print("$ ", end="", flush=True)
            check_for_interactive_command(script_dir, is_automated)

def get_instruction_key():
    """Waits for a single keypress on stdin.

    This is a silly function to call if you need to do it a lot because it has
    to store stdin's current setup, setup stdin for reading single keystrokes
    then read the single keystroke then revert stdin back after reading the
    keystroke.

    Returns the character of the key that was pressed (zero on
    KeyboardInterrupt which can happen when a signal gets handled)

    This method is licensed under cc by-sa 3.0
    Thanks to mheyman http://stackoverflow.com/questions/983354/how-do-i-make-python-to-wait-for-a-pressed-key\
    """
    import termios, fcntl, sys, os
    fd = sys.stdin.fileno()
    # save old state
    flags_save = fcntl.fcntl(fd, fcntl.F_GETFL)
    attrs_save = termios.tcgetattr(fd)
    # make raw - the way to do this comes from the termios(3) man page.
    attrs = list(attrs_save) # copy the stored version to update
    # iflag
    attrs[0] &= ~(termios.IGNBRK | termios.BRKINT | termios.PARMRK
                  | termios.ISTRIP | termios.INLCR | termios. IGNCR
                  | termios.ICRNL | termios.IXON )
    # oflag
    attrs[1] &= ~termios.OPOST
    # cflag
    attrs[2] &= ~(termios.CSIZE | termios. PARENB)
    attrs[2] |= termios.CS8
    # lflag
    attrs[3] &= ~(termios.ECHONL | termios.ECHO | termios.ICANON
                  | termios.ISIG | termios.IEXTEN)
    termios.tcsetattr(fd, termios.TCSANOW, attrs)
    # turn off non-blocking
    fcntl.fcntl(fd, fcntl.F_SETFL, flags_save & ~os.O_NONBLOCK)
    # read a single keystroke
    try:
        ret = sys.stdin.read(1) # returns a single character
    except KeyboardInterrupt:
        ret = 0
    finally:
        # restore old state
        termios.tcsetattr(fd, termios.TCSAFLUSH, attrs_save)
        fcntl.fcntl(fd, fcntl.F_SETFL, flags_save)
    return ret

def test_results(expected_results, actual_results, expected_similarity = 0.66):
    differ = difflib.Differ()
    comparison = differ.compare(actual_results, expected_results)
    diff = differ.compare(actual_results, expected_results)
    seq = difflib.SequenceMatcher(lambda x: x in " \t\n\r", actual_results, expected_results)

    is_pass = seq.ratio() >= expected_similarity

    if not is_pass:
        print("\n\n=============================\n\n")
        print("FAILED")
        print("Similarity ratio:    " + str(seq.ratio()))
        print("Expected Similarity: " + str(expected_similarity))
        print("\n\n=============================\n\n")
        print("Expected results:")
        print(expected_results)
        print("Actual results:")
        print(actual_results)
        print("\n\n=============================\n\n")

    return is_pass

def run_script(script_dir, env=None, is_simulation = True, is_automated=False, is_testing=False):
    # Reads a script.md file in the indicated directoy and runs the
    # commands contained within. If simulation == True then human
    # entry will be simulated (looks like typing and waits for
    # keyboard input before proceeding to the next command). This is
    # useful if you want to run a fully automated demo.
    #
    # The script.md file will be parsed as follows:
    #
    # ``` marks the start or end of a code block
    #
    # Each line in a code block will be treated as a separate command.
    #
    # All other lines will be ignored

    in_code_block = False
    in_results_section = False

    if not script_dir.endswith('/'):
        script_dir = script_dir + "/"
    filename = script_dir + "script.md"

    lines = list(open(filename))
    expected_results = ""
    actual_results = ""
    passed_tests = 0
    failed_tests = 0
    is_first_line = True

    for line in lines:
        if in_results_section and in_code_block and not line.startswith("```"):
            expected_results += line

        if line.startswith("Results:"):
            # Entering results section
            in_results_section = True
            pos = line.lower().find("expected similarity: ")
            if pos >= 0:
                pos = pos + len("expected similarity: ")
                s = line[pos:]
                expected_similarity = float(s)
            else:
                expected_similarity = 0.66
        elif line.startswith("```") and not in_code_block:
            # Entering a code block, if in_results_section = True then it's a results block
            in_code_block = True
        elif line.startswith("```") and in_code_block and in_results_section:
            # Finishing results section
            if in_results_section and is_testing:
                if test_results(expected_results, actual_results, expected_similarity):
                    passed_tests += 1
                else:
                    failed_tests += 1
            expected_results = ""
            actual_results = ""
            in_results_section = False
            in_code_block = False
        elif line.startswith("```") and in_code_block and not in_results_section:
            # Finishing executable code block
            in_code_block = False
        elif in_code_block and not in_results_section:
            # Executable line
            print("$ ", end="", flush=True)
            check_for_interactive_command(script_dir, is_automated)
            actual_results = simulate_command(line, script_dir, env, is_simulation, is_automated)
        elif line.startswith("#") and not in_code_block and not in_results_section and not is_automated:
            # Heading in descriptive text
            if is_first_line:
                run_command("clear", script_dir, env)
            else:
                print("$ ", end="", flush=True)
                check_for_interactive_command(script_dir, is_automated)
                simulate_command("clear", script_dir, env, is_simulation, is_automated)
            print("$ ", end="", flush=True)
            simulate_command(line, script_dir, env, is_simulation, is_automated)
        elif not is_simulation and not in_results_section:
            # Descriptinve text
            print(line, end="", flush=True)

        is_first_line = False

    if is_testing:
        print("\n\n=============================\n\n")
        print("Test Run Complete.")
        print("Failed Tests: " + str(failed_tests))
        print("Passed Tests: " + str(passed_tests))
        if failed_tests > 0:
            print("\n\n")
            print("View failure reports in context in the above output.")
        print("\n\n=============================\n\n")

def get_bash_script(script_dir, env=None, is_simulation = True, is_automated=False, is_testing=False):
    # Reads a script.md file in the indicated directoy and builds an
    # executable bash script from the commands contained within.

    in_code_block = False
    in_results_section = False

    if not script_dir.endswith('/'):
        script_dir = script_dir + "/"
    filename = script_dir + "script.md"

    script = ""
    env = get_simdem_environment(script_dir)
    for key, value in env.items():
        script += key + "='" + value + "'\n"

    lines = list(open(filename))
    for line in lines:
        if line.startswith("Results:"):
            # Entering results section
            in_results_section = True
        elif line.startswith("```") and not in_code_block:
            # Entering a code block, if in_results_section = True then it's a results block
            in_code_block = True
        elif line.startswith("```") and in_code_block and in_results_section:
            # Finishing results section
            in_results_section = False
            in_code_block = False
        elif line.startswith("```") and in_code_block and not in_results_section:
            # Finishing executable code block
            in_code_block = False
        elif in_code_block and not in_results_section:
            # Executable line
            script += line
        elif line.startswith("#") and not in_code_block and not in_results_section and not is_automated:
            # Heading in descriptive text
            script +="\n"
    return script

def main():
    """SimDem CLI interpreter"""

    p = optparse.OptionParser("%prog [run|test|script] <options> DEMO_NAME", version="%prog 0.2.3")
    p.add_option('--style', '-s', default="tutorial",
                 help="The style of simulation you want to run. 'tutorial' (the default) will print out all text and pause for user input before running commands. 'simulate' will not print out the text but will still pause for input.")
    p.add_option('--path', '-p', default="demo_scripts/",
                 help="The Path to the demo scripts directory.")
    p.add_option('--auto', '-a', default="False",
                 help="If set to anything other than False the application will not wait for user keypresses between commands.")
    p.add_option('--test', '-t', default="False",
                 help="If set to anything other than False the output of the command will be compared to the expected results in the sript. Any failures will be reported")

    options, arguments = p.parse_args()

    if len(arguments) == 0:
        arguments.append("run")

    if not options.path.endswith("/"):
        options.path += "/"

    if options.auto == "False":
        is_automatic = False
    else:
        is_automatic = True

    if options.test == "False":
        is_test = False
    else:
        is_test = True

    if options.style == "simulate":
        simulate = True
    elif options.style == 'tutorial':
        simulate = False
    else:
        print("Unkown style (--style, -s): " + options.style)
        exit(1)

    if len(arguments) == 2:
        script_dir = options.path + arguments[1]
    else:
        script_dir = options.path

    env = os.environ.copy()
    env.update(get_simdem_environment(script_dir))
    cmd = arguments[0]

    if cmd == "run":
        run_script(script_dir, env, simulate, is_automatic, is_test)
    elif cmd == "test":
        is_automatic = True and options.auto
        is_test = True and options.test
        run_script(script_dir, env, simulate, is_automatic, is_test)
    elif cmd == "script":
        print(get_bash_script(script_dir, env))
    else:
        print("Unknown command: " + cmd)
        print("Run with --help for guidance.")
main()
<|endoftext|>


<commit_before>from carrot.connection import DjangoBrokerConnection

from celery.events import EventReceiver
from celery.monitoring.state import monitor_state
from celery.monitoring.web import WebServerThread

class MonitorListener(object):

    def __init__(self, state):
        self.connection = DjangoBrokerConnection()
        self.receiver = EventReceiver(self.connection, handlers={
            "worker-heartbeat": state.receive_heartbeat,
            "worker-online": state.receive_worker_event,
            "worker-offline": state.receive_worker_event,
            "task-received": state.receive_task_received,
            "task-accepted": state.receive_task_event,
            "task-succeeded": state.receive_task_event,
            "task-failed": state.receive_task_event,
            "task-retried": state.receive_task_event
        })

    def start(self):
        self.receiver.consume()

class MonitorService(object):

    def __init__(self, logger, is_detached=False, http_port=8989):
        self.logger = logger
        self.is_detached = is_detached
        self.http_port = http_port

    def start(self):
        listener = MonitorListener(monitor_state)
        webthread = WebServerThread(port=self.http_port)
        webthread.start()

        listener.start()
<commit_msg>Cosmetic changes<commit_after>from carrot.connection import DjangoBrokerConnection

from celery.events import EventReceiver
from celery.monitoring.state import monitor_state
from celery.monitoring.web import WebServerThread

class MonitorListener(object):
    """Capture events sent by messages and store them in memory."""

    def __init__(self, state):
        self.connection = DjangoBrokerConnection()
        self.receiver = EventReceiver(self.connection, handlers={
            "task-received": state.receive_task_received,
            "task-accepted": state.receive_task_event,
            "task-succeeded": state.receive_task_event,
            "task-retried": state.receive_task_event
            "task-failed": state.receive_task_event,
            "worker-online": state.receive_worker_event,
            "worker-offline": state.receive_worker_event,
            "worker-heartbeat": state.receive_heartbeat,
        })

    def start(self):
        self.receiver.capture()

class MonitorService(object):
    """celerymon"""

    def __init__(self, logger, is_detached=False, http_port=8989):
        self.logger = logger
        self.is_detached = is_detached
        self.http_port = http_port

    def start(self):
        MonitorListener(monitor_state).start()
        WebServerThread(port=self.http_port).start()
<|endoftext|>


<commit_before>from utils.subdetect_utils import dilate_to_cytoring, dilate_to_cytoring_buffer
from utils.concave_seg import levelset_geo_separete
from utils.concave_seg import run_concave_cut
from utils.filters import MultiSnakesCombined
from utils.concave_seg import levelset_lap
from utils.filters import label, adaptive_thresh
from scipy.ndimage.filters import minimum_filter
import numpy as np
from scipy.ndimage import morphology
from skimage.morphology import remove_small_objects
from utils.labels_handling import convert_labels
from utils.subdetect_utils import label_high_pass, label_nearest, repair_sal
from utils.global_holder import holder
from segment_operation import constant_thres
np.random.seed(0)

def ring_dilation(labels, MARGIN=0, RINGWIDTH=4):
    """Create a ring around label.
    :param RINGWIDTH (int): Width of rings
    :param MARGIN (int): A region of rings is ignored if they are within MARGIN pixels away from label.

    Examples:
        >>> arr = np.zeros((5, 5));arr[2, 2] = 10
        >>> ring_dilation(None, arr, None, MARGIN=1, RINGWIDTH=2)
        array([[ 0, 10, 10, 10,  0],
               [10,  0,  0,  0, 10],
               [10,  0,  0,  0, 10],
               [10,  0,  0,  0, 10],
               [ 0, 10, 10, 10,  0]], dtype=uint16)
    """
    return dilate_to_cytoring(labels, RINGWIDTH, MARGIN)

def ring_dilation_buffer(labels, MARGIN=0, RINGWIDTH=4, BUFFER=2):
    return dilate_to_cytoring_buffer(labels, RINGWIDTH, MARGIN, BUFFER)

def ring_dilation_above_thres(labels, img,  MARGIN=2, RINGWIDTH=4,
                              EXTRA_RINGWIDTH=15, THRES=50):
    sub_label = dilate_to_cytoring(labels, RINGWIDTH, MARGIN)
    extra_sub_label = dilate_to_cytoring(labels, EXTRA_RINGWIDTH, RINGWIDTH)
    extra_sub_label[img < THRES] = 0
    return sub_label + extra_sub_label

def ring_dilation_above_offset_buffer(labels, img, MARGIN=0, RINGWIDTH=2, BUFFER=2,
                                      OFFSET=100, FILSIZE=50):
    """Dilate from label to make a ring.
    Calculate the local minimum as a background, and if image is less brighter
    than background + offset, remove the region from the ring.
    """
    sub_label = dilate_to_cytoring_buffer(labels, RINGWIDTH, MARGIN, BUFFER)
    minimg = minimum_filter(img, size=FILSIZE)
    sub_label[img < (minimg + OFFSET)] = 0
    return sub_label

def ring_dilation_above_adaptive(labels, img, MARGIN=0, RINGWIDTH=4, BUFFER=2, RATIO=1.05, FILSIZE=10):
    sub_labels = dilate_to_cytoring_buffer(labels, RINGWIDTH, MARGIN, BUFFER)
    bw = adaptive_thresh(img, R=RATIO, FILTERINGSIZE=FILSIZE)
    sub_labels[-bw] = 0
    return sub_labels

def geodesic_levelset(labels, img, NITER=10, PROP=1):
    """propagated outwards until it sticks to the shape boundaries in an image.
    Generally used for making/repairing too small objects bigger.
    Larger NITER will give more propagation and roundy object.
    """
    return levelset_geo_separete(img, labels, niter=NITER, prop=PROP)

def concave_cut(labels, img, SMALL_RAD=7, LARGE_RAD=14, EDGELEN=6, THRES=180):
    """
    Attempt a cut for objects larger than np.pi * large_rad ** 2.
    For each pixel, the angle of two vectors pointing to boundary pixels distant by EDGELEN
    is calculated to find strong concavity more than THRES angle.
    The cut line is chosen by minimizing intensity below a line between two pixels.
    To reduce calculation costs, only pixels on a watershed line are assessed.

        SMALL_RAD (int): minimum radius of nuclei
        LARGE_RAD (int): maximum radius of nuclei
        EDGELEN (int): length of triangle edges to calculate angle
        THRES (int): threshold for concavity angles
    """
    labels = run_concave_cut(img, labels, small_rad=SMALL_RAD, large_rad=LARGE_RAD,
                             EDGELEN=EDGELEN, THRES=THRES)
    return labels

def watershed_cut(labels, img, MIN_SIGMA=2, MAX_SIGMA=10, THRES=1000):
    from utils.filters import lap_local_max, sitk_watershed_intensity
    sigma_list = range(int(MIN_SIGMA), int(MAX_SIGMA))
    local_maxima = lap_local_max(img, sigma_list, THRES)
    return sitk_watershed_intensity(labels, local_maxima)

def propagate_multisnakes(labels, img, NITER=3, SMOOTHING=1, lambda1=1, lambda2=1):
    """
    Higher lambda2 relative to lambda1 gives more outward propagation.
    """
    ms = MultiSnakesCombined(img, labels, smoothing=SMOOTHING, lambda1=lambda1, lambda2=lambda2)
    labels = ms.multi_step(niter=NITER)
    return labels

def laplacian_levelset(labels, img, NITER=100, CURVE=3, PROP=-1):
    return label(levelset_lap(img, labels, NITER, CURVE, PROP))

def voronoi_cut(labels):
    from utils.subdetect_utils import voronoi_expand
    return voronoi_expand(labels)

def detect_puncta_voronoi(labels, img, level=7, PERC=50, FILSIZE=1):
    from utils.subdetect_utils import voronoi_expand
    from utils.fish_detect import detect_puncta
    vor = voronoi_expand(labels)
    puncta = detect_puncta(img, level=level, PERC=PERC, FILSIZE=FILSIZE)
    vor[puncta == 0] = 0
    return vor

def morphological(labels, func='grey_opening', size=3, iterations=1):
    morph_operation = getattr(morphology, func)
    for i in range(iterations):
        labels = morph_operation(labels, size=(size, size))
    return labels

def watershed_divide(labels, regmax=10, min_size=100):
    """
    divide objects in labels with watershed segmentation.
        regmax:
        min_size: objects smaller than this size will not be divided.
    """
    from utils.subdetect_utils import watershed_labels

    large_labels = remove_small_objects(labels, min_size, connectivity=4)
    labels[large_labels > 0] = 0
    ws_large = watershed_labels(large_labels, regmax)
    ws_large += labels.max()
    ws_large[ws_large == labels.max()] = 0
    return labels + ws_large

def cytoplasm_levelset(labels, img, niter=20, dt=-0.5, thres=0.5):
    """ Segment cytoplasm from supplied nuclear labels and probability map
        Expand using level sets method from nuclei to membrane.
        Uses an implementation of Level set method. See:
        https://wiseodd.github.io/techblog/2016/11/05/levelset-method/

    Args:
        labels (numpy.ndarray): nuclear mask labels
        img (numpy.ndarray): probability map
        niter (int): step size to expand mask, number of iterations to run the levelset algorithm
        dt (float): negative values for porgation, positive values for shrinking
        thres (float): threshold of probability value to extend the nuclear mask to

    Returns:
        cytolabels (numpy.ndarray): cytoplasm mask labels

    """
    from skimage.morphology import closing, disk, remove_small_holes
    from utils.dlevel_set import dlevel_set
    phi = labels.copy()
    phi[labels == 0] = 1
    phi[labels > 0] = -1

    outlines = img.copy()
    outlines = -outlines
    outlines = outlines - outlines.min()
    outlines = outlines/outlines.max()

    mask = outlines < thres
    phi = dlevel_set(phi, outlines, niter=niter, dt=dt, mask=mask)

    cytolabels = label(remove_small_holes(label(phi < 0)))
    cytolabels = closing(cytolabels, disk(3))

    temp = cytolabels.copy()
    temp[labels == 0] = 0
    cytolabels = convert_labels(temp, labels, cytolabels)
    return cytolabels

def segment_bacteria(nuc, img, slen=3, SIGMA=0.5,THRES=20, CLOSE=20, THRESCHANGE=1000, MINAREA=5, dist=25):
    """ Segment bacteria using high pass filter and assign to closest nucleus

    Args:
        nuc (numpy.ndarray): nuclear mask labels
        img (numpy.ndarray): image in bacterial channel
        slen (int): Size of Gaussian kernel
        SIGMA (float): Standard deviation for Gaussian kernel
        THRES (int): Threshold pixel intensity fo real signal
        CLOSE (int): Radius for disk used to return morphological closing of the image (dilation followed by erosion to remove dark spots and connect bright cracks)
        THRESCHANGE (int): argument unnecessary?
        MINAREA (int): minimum area in pixels for a bacterium
    dist (int): acceptable distance bac can be from mask
    Returns:
        labels (numpy.ndarray[np.uint16]): bacterial mask labels

    """
    labels = label_high_pass(img, slen=slen, SIGMA=SIGMA, THRES=THRES, CLOSE=3)
    if labels.any():
        labels, comb, nuc_prop, nuc_loc = label_nearest(img, labels, nuc,dist)
    from skimage.morphology import remove_small_objects
    labels = remove_small_objects(labels, MINAREA)
    return labels.astype(np.uint16)

def segment_bacteria_cst_filter(nuc, img, slen=3, SIGMA=0.5,THRES=20, CLOSE=20, THRESCHANGE=1000, MINAREA=5, dist=25):
    """ Segment bacteria using constant threshold and assign to closest nucleus

    Args:
        nuc (numpy.ndarray): nuclear mask labels
        img (numpy.ndarray): image in bacterial channel
        slen (int): Size of Gaussian kernel
        SIGMA (float): Standard deviation for Gaussian kernel
        THRES (int): Threshold pixel intensity fo real signal
        CLOSE (int): Radius for disk used to return morphological closing of the image (dilation followed by erosion to remove dark spots and connect bright cracks)
        THRESCHANGE (int): argument unnecessary?
        MINAREA (int): minimum area in pixels for a bacterium
    distance (int): acceptable distance bacteria can be away from mask
    Returns:
        labels (numpy.ndarray[np.uint16]): bacterial mask labels

    """
    labels = constant_thres(img, THRES=THRES)
    if labels.any():
        labels, comb, nuc_prop, nuc_loc = label_nearest(img, labels, nuc,dist)
    from skimage.morphology import remove_small_objects
    labels = remove_small_objects(labels, MINAREA)
    return labels.astype(np.uint16)

def segment_bacteria_return_cyto_no_bac(nuc, img, slen=3, SIGMA=0.5,THRES=20, CLOSE=20, THRESCHANGE=1000, MINAREA=5, dist=25):
    """ Segment bacteria and assign to closest nucleus and return the mask back without regions containing bacteria

    Args:
        nuc (numpy.ndarray): nuclear mask labels
        img (numpy.ndarray): image in bacterial channel
        slen (int): Size of Gaussian kernel
        SIGMA (float): Standard deviation for Gaussian kernel
        THRES (int): Threshold pixel intensity fo real signal
        CLOSE (int): Radius for disk used to return morphological closing of the image (dilation followed by erosion to remove dark spots and connect bright cracks)
        THRESCHANGE (int): argument unnecessary?
        MINAREA (int): minimum area in pixels for a bacterium

    Returns:
        labels (numpy.ndarray[np.uint16]): bacterial mask labels

    """
    labels = label_high_pass(img, slen=slen, SIGMA=SIGMA, THRES=THRES, CLOSE=3)
    if labels.any():
        labels, comb, nuc_prop, nuc_loc = label_nearest(img, labels, nuc,dist)
    from skimage.morphology import remove_small_objects
    labels = remove_small_objects(labels, MINAREA)
    labels = comb - labels
    return labels.astype(np.uint16)

def segment_bacteria_no_near(img, slen=3, SIGMA=0.5,THRES=20, CLOSE=20, THRESCHANGE=1000, MINAREA=5, dist=25):
    """ Segment all bacteria in image, don't assign to closest nucleus

    Args:
        nuc (numpy.ndarray): nuclear mask labels
        img (numpy.ndarray): image in bacterial channel
        slen (int): Size of Gaussian kernel
        SIGMA (float): Standard deviation for Gaussian kernel
        THRES (int): Threshold pixel intensity fo real signal
        CLOSE (int): Radius for disk used to return morphological closing of the image (dilation followed by erosion to remove dark spots and connect bright cracks)
        THRESCHANGE (int): argument unnecessary?
        MINAREA (int): minimum area in pixels for a bacterium

    Returns:
        labels (numpy.ndarray[np.uint16]): bacterial mask labels

    """
    labels = label_high_pass(img, slen=slen, SIGMA=SIGMA, THRES=THRES, CLOSE=3)
    return labels.astype(np.uint16)

def phage_image_only_under_bac(bac,img):
    """ Modify image so that only regions under bacteria mask have values and rest of pixels are set to 0
    bac (numpy.ndarray): bacterial mask labels
    img (numpy.ndarray): img in any channel
    """
    indices = bac==0
    img[indices]=0 # set regions where there's no bacteria to 0
    return img.astype(float)

def remove_bac_from_img(img,label):
    """ Modify an image so that regions containing bacteria have value 0 and the rest of the pixels have their original value
    img (numpy.ndarray): img in any channel
    label (numpy.ndarray): bacterial mask labels
    """
    new_image = np.zeros(img.shape,np.float32)
    coords = label > 0
    img[coords] = 0
    new_image = img
    return new_image

def segment_bacteria_repair(nuc, img, slen=3, SIGMA=0.5,THRES=20, CLOSE=20, THRESCHANGE=1000, MINAREA=5, dist=25):
    """ Segment bacteria and assign to closest nucleus ## This is still under construction

    Args:
        nuc (numpy.ndarray): nuclear mask labels
        img (numpy.ndarray): image in bacterial channel
        slen (int): Size of Gaussian kernel
        SIGMA (float): Standard deviation for Gaussian kernel
        THRES (int): Threshold pixel intensity fo real signal
        CLOSE (int): Radius for disk used to return morphological closing of the image (dilation followed by erosion to remove dark spots and connect bright cracks)
        THRESCHANGE (int): argument unnecessary?
        MINAREA (int): minimum area in pixels for a bacterium

    Returns:
        labels (numpy.ndarray[np.uint16]): bacterial mask labels

    """
    labels = label_high_pass(img, slen=slen, SIGMA=SIGMA, THRES=THRES, CLOSE=3)
    if labels.any():
        labels, comb, nuc_prop, nuc_loc = label_nearest(img, labels, nuc,dist)
        if not hasattr(holder,'plabel'): #not sure what this is doing
            holder.plabel = labels
            holder.pcomb = comb
            holder.pimg = img
            return labels
            #.astype(np.unit16)
        labels = repair_sal(img, holder.pimg,comb, holder.pcomb,labels,nuc_prop,nuc_loc,THRESCHANGE)
    from skimage.morphology import remove_small_objects
    labels = remove_small_objects(labels, MINAREA)
    return labels
    #labels.astype(np.uint16)
<commit_msg>phage_image_only_under_bac and remove_bac_from_img are now filter_with_labels. segment_bacteria_cst_filter/segment_bacteria_no_near are integrated with segment_bacteria. ring_dilation_buffer is integrated with ring_dilation<commit_after>from utils.subdetect_utils import dilate_to_cytoring, dilate_to_cytoring_buffer
from utils.concave_seg import levelset_geo_separete
from utils.concave_seg import run_concave_cut
from utils.filters import MultiSnakesCombined
from utils.concave_seg import levelset_lap
from utils.filters import label, adaptive_thresh
from scipy.ndimage.filters import minimum_filter
import numpy as np
from scipy.ndimage import morphology
from skimage.morphology import remove_small_objects
from utils.labels_handling import convert_labels
from utils.subdetect_utils import label_high_pass, label_nearest, repair_sal
from utils.global_holder import holder
from segment_operation import constant_thres
np.random.seed(0)

def ring_dilation(labels, MARGIN=0, RINGWIDTH=4, BUFFER=0):
    """Create a ring around label.
    :param RINGWIDTH (int): Width of rings
    :param MARGIN (int): A region of rings is ignored if they are within MARGIN pixels away from label.

    Examples:
        >>> arr = np.zeros((5, 5));arr[2, 2] = 10
        >>> ring_dilation(None, arr, None, MARGIN=1, RINGWIDTH=2)
        array([[ 0, 10, 10, 10,  0],
               [10,  0,  0,  0, 10],
               [10,  0,  0,  0, 10],
               [10,  0,  0,  0, 10],
               [ 0, 10, 10, 10,  0]], dtype=uint16)
    """
    if BUFFER == 0:
        return dilate_to_cytoring(labels, RINGWIDTH, MARGIN)
    else:
        return dilate_to_cytoring_buffer(labels, RINGWIDTH, MARGIN, BUFFER)

def ring_dilation_above_thres(labels, img,  MARGIN=2, RINGWIDTH=4,
                              EXTRA_RINGWIDTH=15, THRES=50):
    sub_label = dilate_to_cytoring(labels, RINGWIDTH, MARGIN)
    extra_sub_label = dilate_to_cytoring(labels, EXTRA_RINGWIDTH, RINGWIDTH)
    extra_sub_label[img < THRES] = 0
    return sub_label + extra_sub_label

def ring_dilation_above_offset_buffer(labels, img, MARGIN=0, RINGWIDTH=2, BUFFER=2,
                                      OFFSET=100, FILSIZE=50):
    """Dilate from label to make a ring.
    Calculate the local minimum as a background, and if image is less brighter
    than background + offset, remove the region from the ring.
    """
    sub_label = dilate_to_cytoring_buffer(labels, RINGWIDTH, MARGIN, BUFFER)
    minimg = minimum_filter(img, size=FILSIZE)
    sub_label[img < (minimg + OFFSET)] = 0
    return sub_label

def ring_dilation_above_adaptive(labels, img, MARGIN=0, RINGWIDTH=4, BUFFER=2, RATIO=1.05, FILSIZE=10):
    sub_labels = dilate_to_cytoring_buffer(labels, RINGWIDTH, MARGIN, BUFFER)
    bw = adaptive_thresh(img, R=RATIO, FILTERINGSIZE=FILSIZE)
    sub_labels[-bw] = 0
    return sub_labels

def geodesic_levelset(labels, img, NITER=10, PROP=1):
    """propagated outwards until it sticks to the shape boundaries in an image.
    Generally used for making/repairing too small objects bigger.
    Larger NITER will give more propagation and roundy object.
    """
    return levelset_geo_separete(img, labels, niter=NITER, prop=PROP)

def concave_cut(labels, img, SMALL_RAD=7, LARGE_RAD=14, EDGELEN=6, THRES=180):
    """
    Attempt a cut for objects larger than np.pi * large_rad ** 2.
    For each pixel, the angle of two vectors pointing to boundary pixels distant by EDGELEN
    is calculated to find strong concavity more than THRES angle.
    The cut line is chosen by minimizing intensity below a line between two pixels.
    To reduce calculation costs, only pixels on a watershed line are assessed.

        SMALL_RAD (int): minimum radius of nuclei
        LARGE_RAD (int): maximum radius of nuclei
        EDGELEN (int): length of triangle edges to calculate angle
        THRES (int): threshold for concavity angles
    """
    labels = run_concave_cut(img, labels, small_rad=SMALL_RAD, large_rad=LARGE_RAD,
                             EDGELEN=EDGELEN, THRES=THRES)
    return labels

def watershed_cut(labels, img, MIN_SIGMA=2, MAX_SIGMA=10, THRES=1000):
    from utils.filters import lap_local_max, sitk_watershed_intensity
    sigma_list = range(int(MIN_SIGMA), int(MAX_SIGMA))
    local_maxima = lap_local_max(img, sigma_list, THRES)
    return sitk_watershed_intensity(labels, local_maxima)

def propagate_multisnakes(labels, img, NITER=3, SMOOTHING=1, lambda1=1, lambda2=1):
    """
    Higher lambda2 relative to lambda1 gives more outward propagation.
    """
    ms = MultiSnakesCombined(img, labels, smoothing=SMOOTHING, lambda1=lambda1, lambda2=lambda2)
    labels = ms.multi_step(niter=NITER)
    return labels

def laplacian_levelset(labels, img, NITER=100, CURVE=3, PROP=-1):
    return label(levelset_lap(img, labels, NITER, CURVE, PROP))

def voronoi_cut(labels):
    from utils.subdetect_utils import voronoi_expand
    return voronoi_expand(labels)

def detect_puncta_voronoi(labels, img, level=7, PERC=50, FILSIZE=1):
    from utils.subdetect_utils import voronoi_expand
    from utils.fish_detect import detect_puncta
    vor = voronoi_expand(labels)
    puncta = detect_puncta(img, level=level, PERC=PERC, FILSIZE=FILSIZE)
    vor[puncta == 0] = 0
    return vor

def morphological(labels, func='grey_opening', size=3, iterations=1):
    morph_operation = getattr(morphology, func)
    for i in range(iterations):
        labels = morph_operation(labels, size=(size, size))
    return labels

def watershed_divide(labels, regmax=10, min_size=100):
    """
    divide objects in labels with watershed segmentation.
        regmax:
        min_size: objects smaller than this size will not be divided.
    """
    from utils.subdetect_utils import watershed_labels

    large_labels = remove_small_objects(labels, min_size, connectivity=4)
    labels[large_labels > 0] = 0
    ws_large = watershed_labels(large_labels, regmax)
    ws_large += labels.max()
    ws_large[ws_large == labels.max()] = 0
    return labels + ws_large

def cytoplasm_levelset(labels, img, niter=20, dt=-0.5, thres=0.5):
    """ Segment cytoplasm from supplied nuclear labels and probability map
        Expand using level sets method from nuclei to membrane.
        Uses an implementation of Level set method. See:
        https://wiseodd.github.io/techblog/2016/11/05/levelset-method/

    Args:
        labels (numpy.ndarray): nuclear mask labels
        img (numpy.ndarray): probability map
        niter (int): step size to expand mask, number of iterations to run the levelset algorithm
        dt (float): negative values for porgation, positive values for shrinking
        thres (float): threshold of probability value to extend the nuclear mask to

    Returns:
        cytolabels (numpy.ndarray): cytoplasm mask labels

    """
    from skimage.morphology import closing, disk, remove_small_holes
    from utils.dlevel_set import dlevel_set
    phi = labels.copy()
    phi[labels == 0] = 1
    phi[labels > 0] = -1

    outlines = img.copy()
    outlines = -outlines
    outlines = outlines - outlines.min()
    outlines = outlines/outlines.max()

    mask = outlines < thres
    phi = dlevel_set(phi, outlines, niter=niter, dt=dt, mask=mask)

    cytolabels = label(remove_small_holes(label(phi < 0)))
    cytolabels = closing(cytolabels, disk(3))

    temp = cytolabels.copy()
    temp[labels == 0] = 0
    cytolabels = convert_labels(temp, labels, cytolabels)
    return cytolabels

def segment_bacteria(nuc, img, slen=3, SIGMA=0.5, THRES=20, CLOSE=20, MINAREA=5,
                     dist=25, SEGMENT='highpass', ASSIGN=True):
    """ Segment bacteria using high pass filter or constant thresholding and assign to closest nucleus

    Args:
        nuc (numpy.ndarray): nuclear mask labels
        img (numpy.ndarray): image in bacterial channel
        slen (int): Size of Gaussian kernel
        SIGMA (float): Standard deviation for Gaussian kernel
        THRES (int): Threshold pixel intensity fo real signal
        CLOSE (int): Radius for disk used to return morphological closing of the image
                     (dilation followed by erosion to remove dark spots and connect bright cracks)
        THRESCHANGE (int): argument unnecessary?
        MINAREA (int): minimum area in pixels for a bacterium
        dist (int): acceptable distance bac can be from mask
        SEGMENT: choose one of 'highpass' or 'constant'.
        ASSIGN: If False, it does not assign to closest nucleus.
    Returns:
        labels (numpy.ndarray[np.uint16]): bacterial mask labels

    """
    if SEGMENT == 'highpass':
        labels = label_high_pass(img, slen=slen, SIGMA=SIGMA, THRES=THRES, CLOSE=3)
    elif SEGMENT == 'constant':
        labels = constant_thres(img, THRES=THRES)
    if not ASSIGN:
        return labels.astype(np.uint16)
    if labels.any():
        labels, comb, nuc_prop, nuc_loc = label_nearest(img, labels, nuc, dist)
    from skimage.morphology import remove_small_objects
    labels = remove_small_objects(labels, MINAREA)
    return labels.astype(np.uint16)

def segment_bacteria_return_cyto_no_bac(nuc, img, slen=3, SIGMA=0.5,THRES=20, CLOSE=20, THRESCHANGE=1000, MINAREA=5, dist=25):
    """ Segment bacteria and assign to closest nucleus and return the mask back without regions containing bacteria

    Args:
        nuc (numpy.ndarray): nuclear mask labels
        img (numpy.ndarray): image in bacterial channel
        slen (int): Size of Gaussian kernel
        SIGMA (float): Standard deviation for Gaussian kernel
        THRES (int): Threshold pixel intensity fo real signal
        CLOSE (int): Radius for disk used to return morphological closing of the image (dilation followed by erosion to remove dark spots and connect bright cracks)
        THRESCHANGE (int): argument unnecessary?
        MINAREA (int): minimum area in pixels for a bacterium

    Returns:
        labels (numpy.ndarray[np.uint16]): bacterial mask labels

    """
    labels = label_high_pass(img, slen=slen, SIGMA=SIGMA, THRES=THRES, CLOSE=3)
    if labels.any():
        labels, comb, nuc_prop, nuc_loc = label_nearest(img, labels, nuc,dist)
    from skimage.morphology import remove_small_objects
    labels = remove_small_objects(labels, MINAREA)
    labels = comb - labels
    return labels.astype(np.uint16)

def filter_with_labels(labels, img, BG=True):
    """
    Modify an image so that background regions in labels have value 0
    and the rest of the pixels have their original value.
    Setting BG to False will set foreground regions to be 0.
    """
    if BG:
        img[labels == 0] = 0
    if not BG:
        img[labels > 0] = 0
    return img

def segment_bacteria_repair(nuc, img, slen=3, SIGMA=0.5,THRES=20, CLOSE=20, THRESCHANGE=1000, MINAREA=5, dist=25):
    """ Segment bacteria and assign to closest nucleus ## This is still under construction

    Args:
        nuc (numpy.ndarray): nuclear mask labels
        img (numpy.ndarray): image in bacterial channel
        slen (int): Size of Gaussian kernel
        SIGMA (float): Standard deviation for Gaussian kernel
        THRES (int): Threshold pixel intensity fo real signal
        CLOSE (int): Radius for disk used to return morphological closing of the image (dilation followed by erosion to remove dark spots and connect bright cracks)
        THRESCHANGE (int): argument unnecessary?
        MINAREA (int): minimum area in pixels for a bacterium

    Returns:
        labels (numpy.ndarray[np.uint16]): bacterial mask labels

    """
    labels = label_high_pass(img, slen=slen, SIGMA=SIGMA, THRES=THRES, CLOSE=3)
    if labels.any():
        labels, comb, nuc_prop, nuc_loc = label_nearest(img, labels, nuc,dist)
        if not hasattr(holder,'plabel'): #not sure what this is doing
            holder.plabel = labels
            holder.pcomb = comb
            holder.pimg = img
            return labels
            #.astype(np.unit16)
        labels = repair_sal(img, holder.pimg,comb, holder.pcomb,labels,nuc_prop,nuc_loc,THRESCHANGE)
    from skimage.morphology import remove_small_objects
    labels = remove_small_objects(labels, MINAREA)
    return labels
    #labels.astype(np.uint16)
<|endoftext|>