<commit_before>import sys
import os
import logging
import time
from benchexec import systeminfo
from benchexec.model import Run
from benchexec.p4_run_setup import P4SetupHandler

from benchexec import tooladapter
from benchexec import util
import docker
from benchexec import benchexec
from benchexec import BenchExecException

STOPPED_BY_INTERRUPT = False

class P4Execution(object):
    """
        Class for executing p4 benchmarks
    """
    def __init__(self):
        self.nodes = None #Set by init
        self.switches = None #Set by init
        self.ptf_tester = None #Set by init

        self.client = None
        self.node_networks = []
        self.mgnt_network = None

    def init(self, config, benchmark):
        """
            This functions will set up the docker network to execute the test.
            As a result, it needs root permission for the setup part.
        """

        tool_locator = tooladapter.create_tool_locator(config)
        benchmark.executable = benchmark.tool.executable(tool_locator)
        benchmark.tool_version = benchmark.tool.version(benchmark.executable)

        switch_source_path, ptf_folder_path = self.read_folder_paths(benchmark)

        if not os.path.isdir(switch_source_path):
          logging.critical("Switch folder path not found: {0}".format(switch_source_path))
          raise BenchExecException("Switch folder path not found: {0}".format(switch_source_path))
        if not os.path.isdir(ptf_folder_path):
            logging.critical("Ptf test folder path not found: {0}".format(ptf_folder_path))
            raise( BenchExecException("Ptf test folder path not found: {0}".format(ptf_folder_path)))

        if not switch_source_path or not ptf_folder_path:
            raise BenchExecException("Switch or Ptf folder path not defined." +
                "Switch path: {0} Folder path: {1}".format(switch_source_path, ptf_folder_path))

        self.client = docker.from_env()

        #To be replaceded by input file
        NodeImageName = "basic_node"
        SwitchImageName = "switch_bmv2"
        SwitchTargetPath = "/app"
        nrOfNodes = 4

        try:
            #Setup networks
            self.setup_network(nrOfNodes)

            #Create node container
            self.nodes = []

            for device_nr in range(nrOfNodes):
                #Try get old node from previous run if it exits
                try:
                    self.nodes.append(self.client.containers.get("node{0}".format(device_nr+1)))
                    logging.debug("Old node container find with name: " + "node{0}".format(device_nr+1) + ". Using that")
                except docker.errors.APIError:
                    self.nodes.append(self.client.containers.create(NodeImageName,
                        command="python3 /usr/local/src/ptf/ptf_nn/ptf_nn_agent.py --device-socket {0}@tcp://172.19.0.{1}:10001 -i {2}-1@eth0".format(device_nr ,device_nr+3, device_nr),
                        detach=True,
                        name="node{0}".format(device_nr+1),
                        network="net{0}".format(device_nr+1)
                        ))

            #Switch containers
            self.switches = []

            mount_switch = docker.types.Mount(SwitchTargetPath, switch_source_path, type="bind")

            switch_command = "simple_switch "

            for device_nr in range(nrOfNodes):
                switch_command += "-i {0}@eth{0} ".format(device_nr)

            switch_command += "/app/P4/simple_switch.json"

            try:
                self.switches.append(self.client.containers.create("switch_bmv2",
                    command=switch_command,
                    detach=True,
                    network="net1",
                    name="switch1",
                    mounts = [mount_switch]
                    ))
            except docker.errors.APIError as e:
                logging.debug("Failed to create switch. Got error:" + str(e))
                self.close()
                raise BenchExecException("Switch contianer with name switch1 might already be running")

            #Connect all nodes   to the switch
            for device_nr in range(nrOfNodes)[1:]:
                self.client.networks.get("net{0}".format(device_nr+1)).connect(self.switches[0])

            #Ptf tester container
            mount_ptf_tester = docker.types.Mount("/app", ptf_folder_path, type="bind")
            try:
                self.ptf_tester = self.client.containers.create("ptf_tester",
                    detach=True,
                    name="ptfTester",
                    mounts=[mount_ptf_tester],
                    tty=True)
            except docker.errors.APIError:
                self.close()
                raise BenchExecException("Cannot create ptf tester container. Might already be running")

            #Setup mgnt network
            mgnt = self.client.networks.get("mgnt")
            mgnt.connect(self.ptf_tester)
            for node in self.nodes:
                mgnt.connect(node)
        except docker.errors.APIError as e:
            raise BenchExecException(str(e))

    def execute_benchmark(self, benchmark, output_handler):
        self.start_containers()

        test_dict = self.read_tests()

        # test_dict = {}
        # test_dict["Module1"] = ["Test1", "Test2", "Test3"]

        setup_handler = P4SetupHandler(benchmark, test_dict)
        setup_handler.update_runsets()

        for runSet in benchmark.run_sets:
            if STOPPED_BY_INTERRUPT:
                break

            if not runSet.should_be_executed():
                output_handler.output_for_skipping_run_set(runSet)

            elif not runSet.runs:
                output_handler.output_for_skipping_run_set(
                    runSet, "because it has no files"
                )

            output_handler.output_before_run_set(runSet)

            for run in runSet.runs:
                command = ("ptf --test-dir /app " + run.identifier + " " +
                "--device-socket 0-{0-64}@tcp://172.19.0.3:10001 " +
                "--device-socket 1-{0-64}@tcp://172.19.0.4:10001 " +
                "--device-socket 2-{0-64}@tcp://172.19.0.5:10001 " +
                "--device-socket 3-{0-64}@tcp://172.19.0.6:10001 " +
                "--platform nn")

                return_code, test_output = self._execute_benchmark(run, command)

                test_output = test_output.decode("utf-8")

                if os.path.exists("/home/sdn/benchexec/test.txt"):
                    os.remove("/home/sdn/benchexec/test.txt")
                f = open("/home/sdn/benchexec/test.txt", "w+")
                f.write(test_output)

                logging.debug("Logs: " + test_output)

                try:
                    with open(run.log_file, "w") as ouputFile:
                        for i in range(6):
                            ouputFile.write("Logging\n")

                        #for result in test_results:
                        ouputFile.write(test_output + "\n")
                except OSError:
                    print("Failed")

                values = {}
                values["exitcode"] = util.ProcessExitCode.from_raw(return_code)

                test = run.cmdline()
                run._cmdline = command.split(" ")

                for item in test:
                    f.write(item + "\n")

                f.write("\n---------- Run object ---------- \n")
                for item in dir(run):
                    f.write(item + "\n")

                f.write(run.identifier)
                f.close()

                run.set_result(values)

                output_handler.output_after_run(run)

            output_handler.output_after_benchmark(STOPPED_BY_INTERRUPT)

        #self.close()
        #self.clear_networks()

    def _execute_benchmark(self, run, command):

        return self.ptf_tester.exec_run(command, tty=True)

    def setup_network(self, nrOfNodes):
        """
            Creates the networks required to run the test. It will create 1 network for each node.
            This network represents the bridge between the node and the swtich. Further, it creates
            the management network(mgnt) used by the ptf_teter to inject packages.
        """
        for net_id in range(nrOfNodes+1)[1:]:
            try:
                network = self.client.networks.get("net{0}".format(net_id))
                self.node_networks.append(network)
            except docker.errors.NotFound:
                ipam_pool = docker.types.IPAMPool(
                    subnet="172.{0}.0.0/16".format(19+net_id),
                    gateway="172.{0}.0.1".format(19+net_id)
                )

                ipam_config = docker.types.IPAMConfig(
                    pool_configs=[ipam_pool]
                )

                self.client.networks.create(
                    "net{0}".format(net_id),
                    driver="bridge",
                    ipam=ipam_config
                )

        try:
            self.mgnt_network = self.client.networks.get("mgnt")
        except docker.errors.NotFound:
            ipam_pool = docker.types.IPAMPool(
                subnet="172.19.0.0/16",
                gateway="172.19.0.1"
                )

            ipam_config = docker.types.IPAMConfig(
                pool_configs=[ipam_pool]
            )

            self.mgnt_network = self.client.networks.create(
                "mgnt",
                driver="bridge",
                ipam=ipam_config
            )

    def read_tests(self):
        self.ptf_tester.start()

        _, test_info = self.ptf_tester.exec_run("ptf --test-dir /app --list")
        test_info = test_info.decode()

        test_dict = self.extract_info_from_test_info(test_info)

        return test_dict

    def extract_info_from_test_info(self,test_info):
        testing = test_info.split("\n")
        test_info = test_info.split("Test List:")[1]
        test_modules = test_info.split("Module ")
        nr_of_modules = len(test_modules) - 1
        test_modules[len(test_modules)-1] = test_modules[len(test_modules)-1].split("\n{0}".format(nr_of_modules))[0]

        test_dict = {}

        for i in range(nr_of_modules):
            test = test_modules[i+1].split("\n")
            module_name = test.pop(0).split(":")[0]

            test_names = []
            for test_string in test:
                if not str.isspace(test_string) and test_string:
                    test_names.append(test_string.split(":")[0].strip())

            test_dict[module_name] = test_names

        return test_dict

    def get_system_info(self):
        return systeminfo.SystemInfo()

    def read_folder_paths(self, benchmark):

        switch_folder = ""
        ptf_folder = ""
        option_index = 0

        while option_index < len(benchmark.options):
            try:
                if "switch" in benchmark.options[option_index].lower():
                    switch_folder = benchmark.options[option_index +1]
                elif "ptf" in benchmark.options[option_index].lower():
                    ptf_folder = benchmark.options[option_index +1]
            except:
                test = ""

            option_index += 2

        return switch_folder, ptf_folder

    def close(self):
        """
            Cleans up all the running containers. Should be called when test is done.
        """
        logging.debug("Closing containers")
        for container in self.nodes:
            container.remove(force=True)

        for container in self.switches:
            container.remove(force=True)
        if self.ptf_tester:
            self.ptf_tester.remove(force=True)

    def start_containers(self):
        self.ptf_tester.start()

        for container in self.nodes:
            container.start()

        for container in self.switches:
            container.start()

def main(argv=None):
    if argv is None:
        argv = sys.argv

    print("Hello World")

if __name__ == "__main__":
    main()<commit_msg>Cleaned up p4execution file<commit_after>import sys
import os
import logging
import time
from benchexec import systeminfo
from benchexec.model import Run
from benchexec.p4_run_setup import P4SetupHandler

from benchexec import tooladapter
from benchexec import util
import docker
from benchexec import benchexec
from benchexec import BenchExecException

STOPPED_BY_INTERRUPT = False

class P4Execution(object):
    """
        Class for executing p4 benchmarks
    """
    def __init__(self):
        self.nodes = None #Set by init
        self.switches = None #Set by init
        self.ptf_tester = None #Set by init

        self.client = None
        self.node_networks = []
        self.mgnt_network = None

    def init(self, config, benchmark):
        """
            This functions will set up the docker network to execute the test.
            As a result, it needs root permission for the setup part.
        """

        tool_locator = tooladapter.create_tool_locator(config)
        benchmark.executable = benchmark.tool.executable(tool_locator)
        benchmark.tool_version = benchmark.tool.version(benchmark.executable)

        switch_source_path, ptf_folder_path = self.read_folder_paths(benchmark)

        if not os.path.isdir(switch_source_path):
          logging.critical("Switch folder path not found: {0}".format(switch_source_path))
          raise BenchExecException("Switch folder path not found: {0}".format(switch_source_path))
        if not os.path.isdir(ptf_folder_path):
            logging.critical("Ptf test folder path not found: {0}".format(ptf_folder_path))
            raise( BenchExecException("Ptf test folder path not found: {0}".format(ptf_folder_path)))

        if not switch_source_path or not ptf_folder_path:
            raise BenchExecException("Switch or Ptf folder path not defined." +
                "Switch path: {0} Folder path: {1}".format(switch_source_path, ptf_folder_path))

        self.client = docker.from_env()

        #To be replaceded by input file
        NodeImageName = "basic_node"
        SwitchImageName = "switch_bmv2"
        SwitchTargetPath = "/app"
        nrOfNodes = 4

        try:
            #Setup networks
            self.setup_network(nrOfNodes)

            #Create node container
            self.nodes = []

            for device_nr in range(nrOfNodes):
                #Try get old node from previous run if it exits
                try:
                    self.nodes.append(self.client.containers.get("node{0}".format(device_nr+1)))
                    logging.debug("Old node container find with name: " + "node{0}".format(device_nr+1) + ". Using that")
                except docker.errors.APIError:
                    self.nodes.append(self.client.containers.create(NodeImageName,
                        command="python3 /usr/local/src/ptf/ptf_nn/ptf_nn_agent.py --device-socket {0}@tcp://172.19.0.{1}:10001 -i {2}-1@eth0".format(device_nr ,device_nr+3, device_nr),
                        detach=True,
                        name="node{0}".format(device_nr+1),
                        network="net{0}".format(device_nr+1)
                        ))

            #Switch containers
            self.switches = []

            mount_switch = docker.types.Mount(SwitchTargetPath, switch_source_path, type="bind")

            switch_command = "simple_switch "

            for device_nr in range(nrOfNodes):
                switch_command += "-i {0}@eth{0} ".format(device_nr)

            switch_command += "/app/P4/simple_switch.json"

            try:
                self.switches.append(self.client.containers.create("switch_bmv2",
                    command=switch_command,
                    detach=True,
                    network="net1",
                    name="switch1",
                    mounts = [mount_switch]
                    ))
            except docker.errors.APIError as e:
                logging.debug("Failed to create switch. Got error:" + str(e))
                self.close()
                raise BenchExecException("Switch contianer with name switch1 might already be running")

            #Connect all nodes   to the switch
            for device_nr in range(nrOfNodes)[1:]:
                self.client.networks.get("net{0}".format(device_nr+1)).connect(self.switches[0])

            #Ptf tester container
            mount_ptf_tester = docker.types.Mount("/app", ptf_folder_path, type="bind")
            try:
                self.ptf_tester = self.client.containers.create("ptf_tester",
                    detach=True,
                    name="ptfTester",
                    mounts=[mount_ptf_tester],
                    tty=True)
            except docker.errors.APIError:
                self.close()
                raise BenchExecException("Cannot create ptf tester container. Might already be running")

            #Setup mgnt network
            mgnt = self.client.networks.get("mgnt")
            mgnt.connect(self.ptf_tester)
            for node in self.nodes:
                mgnt.connect(node)
        except docker.errors.APIError as e:
            self.close()
            raise BenchExecException(str(e))

    def execute_benchmark(self, benchmark, output_handler):
        self.start_containers()

        test_dict = self.read_tests()

        setup_handler = P4SetupHandler(benchmark, test_dict)
        setup_handler.update_runsets()

        for runSet in benchmark.run_sets:
            if STOPPED_BY_INTERRUPT:
                break

            if not runSet.should_be_executed():
                output_handler.output_for_skipping_run_set(runSet)

            elif not runSet.runs:
                output_handler.output_for_skipping_run_set(
                    runSet, "because it has no files"
                )

            output_handler.output_before_run_set(runSet)

            for run in runSet.runs:
                command = ("ptf --test-dir /app " + run.identifier + " " +
                "--device-socket 0-{0-64}@tcp://172.19.0.3:10001 " +
                "--device-socket 1-{0-64}@tcp://172.19.0.4:10001 " +
                "--device-socket 2-{0-64}@tcp://172.19.0.5:10001 " +
                "--device-socket 3-{0-64}@tcp://172.19.0.6:10001 " +
                "--platform nn")

                return_code, test_output = self._execute_benchmark(run, command)

                test_output = test_output.decode("utf-8")

                if os.path.exists("/home/sdn/benchexec/test.txt"):
                    os.remove("/home/sdn/benchexec/test.txt")
                f = open("/home/sdn/benchexec/test.txt", "w+")
                f.write("Log file: " + run.log_file)
                logging.debug("Logs: " + test_output)

                try:
                    with open(run.log_file, "w") as ouputFile:
                        for i in range(6):
                            ouputFile.write("Logging\n")

                        #for result in test_results:
                        ouputFile.write(test_output + "\n")
                except OSError:
                    print("Failed")

                values = {}
                values["exitcode"] = util.ProcessExitCode.from_raw(return_code)

                test = run.cmdline()
                run._cmdline = command.split(" ")

                for item in test:
                    f.write(item + "\n")

                f.write("\n---------- Run object ---------- \n")
                for item in dir(run):
                    f.write(item + "\n")

                f.write(run.identifier)
                f.close()

                run.set_result(values)

                output_handler.output_after_run(run)

            output_handler.output_after_benchmark(STOPPED_BY_INTERRUPT)

        self.close()
        #self.clear_networks()

    def _execute_benchmark(self, run, command):

        return self.ptf_tester.exec_run(command, tty=True)

    def setup_network(self, nrOfNodes):
        """
            Creates the networks required to run the test. It will create 1 network for each node.
            This network represents the bridge between the node and the swtich. Further, it creates
            the management network(mgnt) used by the ptf_teter to inject packages.
        """
        for net_id in range(nrOfNodes+1)[1:]:
            try:
                network = self.client.networks.get("net{0}".format(net_id))
                self.node_networks.append(network)
            except docker.errors.NotFound:
                ipam_pool = docker.types.IPAMPool(
                    subnet="172.{0}.0.0/16".format(19+net_id),
                    gateway="172.{0}.0.1".format(19+net_id)
                )

                ipam_config = docker.types.IPAMConfig(
                    pool_configs=[ipam_pool]
                )

                self.client.networks.create(
                    "net{0}".format(net_id),
                    driver="bridge",
                    ipam=ipam_config
                )

        try:
            self.mgnt_network = self.client.networks.get("mgnt")
        except docker.errors.NotFound:
            ipam_pool = docker.types.IPAMPool(
                subnet="172.19.0.0/16",
                gateway="172.19.0.1"
                )

            ipam_config = docker.types.IPAMConfig(
                pool_configs=[ipam_pool]
            )

            self.mgnt_network = self.client.networks.create(
                "mgnt",
                driver="bridge",
                ipam=ipam_config
            )

    def read_tests(self):
        self.ptf_tester.start()

        _, test_info = self.ptf_tester.exec_run("ptf --test-dir /app --list")
        test_info = test_info.decode()

        test_dict = self.extract_info_from_test_info(test_info)

        return test_dict

    def extract_info_from_test_info(self,test_info):
        testing = test_info.split("\n")
        test_info = test_info.split("Test List:")[1]
        test_modules = test_info.split("Module ")
        nr_of_modules = len(test_modules) - 1
        test_modules[len(test_modules)-1] = test_modules[len(test_modules)-1].split("\n{0}".format(nr_of_modules))[0]

        test_dict = {}

        for i in range(nr_of_modules):
            test = test_modules[i+1].split("\n")
            module_name = test.pop(0).split(":")[0]

            test_names = []
            for test_string in test:
                if not str.isspace(test_string) and test_string:
                    test_names.append(test_string.split(":")[0].strip())

            test_dict[module_name] = test_names

        return test_dict

    def get_system_info(self):
        return systeminfo.SystemInfo()

    def read_folder_paths(self, benchmark):

        switch_folder = ""
        ptf_folder = ""
        option_index = 0

        while option_index < len(benchmark.options):
            try:
                if "switch" in benchmark.options[option_index].lower():
                    switch_folder = benchmark.options[option_index +1]
                elif "ptf" in benchmark.options[option_index].lower():
                    ptf_folder = benchmark.options[option_index +1]
            except:
                test = ""

            option_index += 2

        return switch_folder, ptf_folder

    def close(self):
        """
            Cleans up all the running containers. Should be called when test is done.
        """
        logging.debug("Closing containers")
        for container in self.nodes:
            container.remove(force=True)

        for container in self.switches:
            container.remove(force=True)
        if self.ptf_tester:
            self.ptf_tester.remove(force=True)

    def start_containers(self):
        self.ptf_tester.start()

        for container in self.nodes:
            container.start()

        for container in self.switches:
            container.start()

        #allow 2 sec of startuptime
        time.sleep(2)

def main(argv=None):
    if argv is None:
        argv = sys.argv

    print("Hello World")

if __name__ == "__main__":
    main()<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
:maintainer:    SaltStack
:maturity:      new
:platform:      all

Utilities supporting modules for Hashicorp Vault. Configuration instructions are
documented in the execution module docs.
'''

from __future__ import absolute_import, print_function, unicode_literals
import base64
import logging
import os
import requests

import salt.crypt
import salt.exceptions
import salt.utils.versions

log = logging.getLogger(__name__)
logging.getLogger("requests").setLevel(logging.WARNING)

# Load the __salt__ dunder if not already loaded (when called from utils-module)
__salt__ = None
def __virtual__():  # pylint: disable=expected-2-blank-lines-found-0
    try:
        global __salt__  # pylint: disable=global-statement
        if not __salt__:
            __salt__ = salt.loader.minion_mods(__opts__)
            return True
    except Exception as e:
        log.error("Could not load __salt__: %s", e)
        return False

def _get_token_and_url_from_master():
    '''
    Get a token with correct policies for the minion, and the url to the Vault
    service
    '''
    minion_id = __grains__['id']
    pki_dir = __opts__['pki_dir']

    # When rendering pillars, the module executes on the master, but the token
    # should be issued for the minion, so that the correct policies are applied
    if __opts__.get('__role', 'minion') == 'minion':
        private_key = '{0}/minion.pem'.format(pki_dir)
        log.debug('Running on minion, signing token request with key %s',
                  private_key)
        signature = base64.b64encode(salt.crypt.sign_message(
                                                             private_key,
                                                             minion_id
                                                            ))
        result = __salt__['publish.runner'](
                                            'vault.generate_token',
                                            arg=[minion_id, signature]
                                           )
    else:
        private_key = '{0}/master.pem'.format(pki_dir)
        log.debug('Running on master, signing token request for %s with key %s',
                  minion_id, private_key)
        signature = base64.b64encode(salt.crypt.sign_message(
                                                             private_key,
                                                             minion_id
                                                            ))
        result = __salt__['saltutil.runner'](
                                             'vault.generate_token',
                                             minion_id=minion_id,
                                             signature=signature,
                                             impersonated_by_master=True
                                            )

    if not result:
        log.error('Failed to get token from master! No result returned - '
                  'is the peer publish configuration correct?')
        raise salt.exceptions.CommandExecutionError(result)
    if not isinstance(result, dict):
        log.error('Failed to get token from master! '
                  'Response is not a dict: %s', result)
        raise salt.exceptions.CommandExecutionError(result)
    if 'error' in result:
        log.error('Failed to get token from master! '
                  'An error was returned: %s', result['error'])
        raise salt.exceptions.CommandExecutionError(result)
    return {
            'url': result['url'],
            'token': result['token'],
            'verify': result['verify'],
           }

def _get_vault_connection():
    '''
    Get the connection details for calling Vault, from local configuration if
    it exists, or from the master otherwise
    '''
    def _use_local_config():
        log.debug('Using Vault connection details from local config')
        try:
            if __opts__['vault']['auth']['method'] == 'approle':
                verify = __opts__['vault'].get('verify', None)
                if _selftoken_expired():
                    log.debug('Vault token expired. Recreating one')
                    # Requesting a short ttl token
                    url = '{0}/v1/auth/approle/login'.format(__opts__['vault']['url'])
                    payload = {'role_id': __opts__['vault']['auth']['role_id']}
                    if 'secret_id' in __opts__['vault']['auth']:
                        payload['secret_id'] = __opts__['vault']['auth']['secret_id']
                    response = requests.post(url, json=payload, verify=verify)
                    if response.status_code != 200:
                        errmsg = 'An error occured while getting a token from approle'
                        raise salt.exceptions.CommandExecutionError(errmsg)
                    __opts__['vault']['auth']['token'] = response.json()['auth']['client_token']
            return {
                'url': __opts__['vault']['url'],
                'token': __opts__['vault']['auth']['token'],
                'verify': __opts__['vault'].get('verify', None)
            }
        except KeyError as err:
            errmsg = 'Minion has "vault" config section, but could not find key "{0}" within'.format(err.message)
            raise salt.exceptions.CommandExecutionError(errmsg)

    if 'vault' in __opts__ and __opts__.get('__role', 'minion') == 'master':
        return _use_local_config()
    elif '_ssh_version' in __opts__:
        return _use_local_config()
    else:
        log.debug('Contacting master for Vault connection details')
        return _get_token_and_url_from_master()

def make_request(method, resource, profile=None, **args):
    '''
    Make a request to Vault
    '''
    if profile is not None and profile.keys().remove('driver') is not None:
        # Deprecated code path
        return make_request_with_profile(method, resource, profile, **args)

    connection = _get_vault_connection()
    token, vault_url = connection['token'], connection['url']
    if 'verify' not in args:
        args['verify'] = connection['verify']

    url = "{0}/{1}".format(vault_url, resource)
    headers = {'X-Vault-Token': token, 'Content-Type': 'application/json'}
    response = requests.request(method, url, headers=headers, **args)

    return response

def make_request_with_profile(method, resource, profile, **args):
    '''
    DEPRECATED! Make a request to Vault, with a profile including connection
    details.
    '''
    salt.utils.versions.warn_until(
        'Fluorine',
        'Specifying Vault connection data within a \'profile\' has been '
        'deprecated. Please see the documentation for details on the new '
        'configuration schema. Support for this function will be removed '
        'in Salt Fluorine.'
    )
    url = '{0}://{1}:{2}/v1/{3}'.format(
        profile.get('vault.scheme', 'https'),
        profile.get('vault.host'),
        profile.get('vault.port'),
        resource,
    )
    token = os.environ.get('VAULT_TOKEN', profile.get('vault.token'))
    if token is None:
        raise salt.exceptions.CommandExecutionError('A token was not configured')

    headers = {'X-Vault-Token': token, 'Content-Type': 'application/json'}
    response = requests.request(method, url, headers=headers, **args)

    return response

def _selftoken_expired():
    '''
    Validate the current token exists and is still valid
    '''
    try:
        verify = __opts__['vault'].get('verify', None)
        url = '{0}/v1/auth/token/lookup-self'.format(__opts__['vault']['url'])
        if 'token' not in __opts__['vault']['auth']:
            return True
        headers = {'X-Vault-Token': __opts__['vault']['auth']['token']}
        response = requests.get(url, headers=headers, verify=verify)
        if response.status_code != 200:
            return True
        return False
    except Exception as e:
        raise salt.exceptions.CommandExecutionError(
            'Error while looking up self token : {0}'.format(e)
        )
<commit_msg>use local config for vault when masterless<commit_after># -*- coding: utf-8 -*-
'''
:maintainer:    SaltStack
:maturity:      new
:platform:      all

Utilities supporting modules for Hashicorp Vault. Configuration instructions are
documented in the execution module docs.
'''

from __future__ import absolute_import, print_function, unicode_literals
import base64
import logging
import os
import requests

import salt.crypt
import salt.exceptions
import salt.utils.versions

log = logging.getLogger(__name__)
logging.getLogger("requests").setLevel(logging.WARNING)

# Load the __salt__ dunder if not already loaded (when called from utils-module)
__salt__ = None
def __virtual__():  # pylint: disable=expected-2-blank-lines-found-0
    try:
        global __salt__  # pylint: disable=global-statement
        if not __salt__:
            __salt__ = salt.loader.minion_mods(__opts__)
            return True
    except Exception as e:
        log.error("Could not load __salt__: %s", e)
        return False

def _get_token_and_url_from_master():
    '''
    Get a token with correct policies for the minion, and the url to the Vault
    service
    '''
    minion_id = __grains__['id']
    pki_dir = __opts__['pki_dir']

    # When rendering pillars, the module executes on the master, but the token
    # should be issued for the minion, so that the correct policies are applied
    if __opts__.get('__role', 'minion') == 'minion':
        private_key = '{0}/minion.pem'.format(pki_dir)
        log.debug('Running on minion, signing token request with key %s',
                  private_key)
        signature = base64.b64encode(salt.crypt.sign_message(
                                                             private_key,
                                                             minion_id
                                                            ))
        result = __salt__['publish.runner'](
                                            'vault.generate_token',
                                            arg=[minion_id, signature]
                                           )
    else:
        private_key = '{0}/master.pem'.format(pki_dir)
        log.debug('Running on master, signing token request for %s with key %s',
                  minion_id, private_key)
        signature = base64.b64encode(salt.crypt.sign_message(
                                                             private_key,
                                                             minion_id
                                                            ))
        result = __salt__['saltutil.runner'](
                                             'vault.generate_token',
                                             minion_id=minion_id,
                                             signature=signature,
                                             impersonated_by_master=True
                                            )

    if not result:
        log.error('Failed to get token from master! No result returned - '
                  'is the peer publish configuration correct?')
        raise salt.exceptions.CommandExecutionError(result)
    if not isinstance(result, dict):
        log.error('Failed to get token from master! '
                  'Response is not a dict: %s', result)
        raise salt.exceptions.CommandExecutionError(result)
    if 'error' in result:
        log.error('Failed to get token from master! '
                  'An error was returned: %s', result['error'])
        raise salt.exceptions.CommandExecutionError(result)
    return {
            'url': result['url'],
            'token': result['token'],
            'verify': result['verify'],
           }

def _get_vault_connection():
    '''
    Get the connection details for calling Vault, from local configuration if
    it exists, or from the master otherwise
    '''
    def _use_local_config():
        log.debug('Using Vault connection details from local config')
        try:
            if __opts__['vault']['auth']['method'] == 'approle':
                verify = __opts__['vault'].get('verify', None)
                if _selftoken_expired():
                    log.debug('Vault token expired. Recreating one')
                    # Requesting a short ttl token
                    url = '{0}/v1/auth/approle/login'.format(__opts__['vault']['url'])
                    payload = {'role_id': __opts__['vault']['auth']['role_id']}
                    if 'secret_id' in __opts__['vault']['auth']:
                        payload['secret_id'] = __opts__['vault']['auth']['secret_id']
                    response = requests.post(url, json=payload, verify=verify)
                    if response.status_code != 200:
                        errmsg = 'An error occured while getting a token from approle'
                        raise salt.exceptions.CommandExecutionError(errmsg)
                    __opts__['vault']['auth']['token'] = response.json()['auth']['client_token']
            return {
                'url': __opts__['vault']['url'],
                'token': __opts__['vault']['auth']['token'],
                'verify': __opts__['vault'].get('verify', None)
            }
        except KeyError as err:
            errmsg = 'Minion has "vault" config section, but could not find key "{0}" within'.format(err.message)
            raise salt.exceptions.CommandExecutionError(errmsg)

    if 'vault' in __opts__ and __opts__.get('__role', 'minion') == 'master':
        return _use_local_config()
    elif any((__opts__['local'], __opts__['file_client'] == 'local', __opts__['master_type'] == 'disable')):
        return _use_local_config()
    else:
        log.debug('Contacting master for Vault connection details')
        return _get_token_and_url_from_master()

def make_request(method, resource, profile=None, **args):
    '''
    Make a request to Vault
    '''
    if profile is not None and profile.keys().remove('driver') is not None:
        # Deprecated code path
        return make_request_with_profile(method, resource, profile, **args)

    connection = _get_vault_connection()
    token, vault_url = connection['token'], connection['url']
    if 'verify' not in args:
        args['verify'] = connection['verify']

    url = "{0}/{1}".format(vault_url, resource)
    headers = {'X-Vault-Token': token, 'Content-Type': 'application/json'}
    response = requests.request(method, url, headers=headers, **args)

    return response

def make_request_with_profile(method, resource, profile, **args):
    '''
    DEPRECATED! Make a request to Vault, with a profile including connection
    details.
    '''
    salt.utils.versions.warn_until(
        'Fluorine',
        'Specifying Vault connection data within a \'profile\' has been '
        'deprecated. Please see the documentation for details on the new '
        'configuration schema. Support for this function will be removed '
        'in Salt Fluorine.'
    )
    url = '{0}://{1}:{2}/v1/{3}'.format(
        profile.get('vault.scheme', 'https'),
        profile.get('vault.host'),
        profile.get('vault.port'),
        resource,
    )
    token = os.environ.get('VAULT_TOKEN', profile.get('vault.token'))
    if token is None:
        raise salt.exceptions.CommandExecutionError('A token was not configured')

    headers = {'X-Vault-Token': token, 'Content-Type': 'application/json'}
    response = requests.request(method, url, headers=headers, **args)

    return response

def _selftoken_expired():
    '''
    Validate the current token exists and is still valid
    '''
    try:
        verify = __opts__['vault'].get('verify', None)
        url = '{0}/v1/auth/token/lookup-self'.format(__opts__['vault']['url'])
        if 'token' not in __opts__['vault']['auth']:
            return True
        headers = {'X-Vault-Token': __opts__['vault']['auth']['token']}
        response = requests.get(url, headers=headers, verify=verify)
        if response.status_code != 200:
            return True
        return False
    except Exception as e:
        raise salt.exceptions.CommandExecutionError(
            'Error while looking up self token : {0}'.format(e)
        )
<|endoftext|>


<commit_before>"""
Notes:
    1) This is algorithm is primarily designed for a rectangular piece of
        paper lying flat on a flat surface, but may work in other situations
        assuming that the paper's corners are not obstructed.
    2) The camera's view of the paper must be unobstructed in first frame.

"""

# Basic Dependencies
from __future__ import division
from math import ceil, acos
from time import time

# External Dependencies
import numpy as np
from numpy.linalg import norm
import cv2

# Default User Parameters
VIDEO_FILE_LOCATION = 'sample.avi'
FPS_INTERVAL = 10  # updates FPS estimate after this many frames
MHI_DURATION = 10  # max frames remembered by motion history
FRAME_WIDTH, FRAME_HEIGHT = 640, 424
PAPER_RATIO = 11/8.5  # height/width of paper
ROT180 = False  # If paper is upside down, change this
REDUCE_DISPLAY_SIZE = False  # Use if output window is too big for your screen

# Internal Parameters
tol_corner_movement = 1
obst_tol = 10   # used to determine tolerance
closing_iterations = 10
show_thresholding = False  # Use to display thresholding

def rotate180(im):
    """Rotates an image by 180 degrees."""
    return cv2.flip(im, -1)

def persTransform(pts, H):
    """Transforms a list of points, `pts`,
    using the perspective transform `H`."""
    src = np.zeros((len(pts), 1, 2))
    src[:, 0] = pts
    dst = cv2.perspectiveTransform(src, H)
    return np.array(dst[:, 0, :], dtype='float32')

def affTransform(pts, A):
    """Transforms a list of points, `pts`,
    using the affine transform `A`."""
    src = np.zeros((len(pts), 1, 2))
    src[:, 0] = pts
    dst = cv2.transform(src, A)
    return np.array(dst[:, 0, :], dtype='float32')

def draw_polygon(im, vertices, vertex_colors=None, edge_colors=None,
                 alter_input_image=False, draw_edges=True, draw_vertices=True,
                 display=False, title='', pause=False):
    """returns image with polygon drawn on it."""
    _default_vertex_color = (255, 0, 0)
    _default_edge_color = (255, 0, 0)
    if not alter_input_image:
        im2 = im.copy()
    else:
        im2 = im
    if vertices is not None:
        N = len(vertices)
        vertices = [tuple(v) for v in vertices]
        if vertex_colors is None:
            vertex_colors = [_default_vertex_color] * N
        if edge_colors is None:
            edge_colors = [_default_edge_color] * N
        for i in range(N):
            startpt = vertices[(i - 1) % N]
            endpt = vertices[i]
            if draw_vertices:
                cv2.circle(im2, startpt, 3, vertex_colors[(i - 1) % N], -1)
            if draw_edges:
                cv2.line(im2, startpt, endpt, edge_colors[(i - 1) % N], 2)
    if display:
        cv2.imshow(title, im2)
        # Note: `0xFF == ord('q')`is apparently necessary for 64bit machines
        if pause and cv2.waitKey(0) & 0xFF == ord('q'):
            pass
    return im2

def run_main():

    # Initialize some variables
    frame = None
    old_homog = None
    old_inv_homog = None
    corner_history = []

    video_feed = cv2.VideoCapture(VIDEO_FILE_LOCATION)
    video_feed.set(cv2.cv.CV_CAP_PROP_FRAME_WIDTH, FRAME_WIDTH)
    video_feed.set(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT, FRAME_HEIGHT)

    frame_count = 0
    fps_time = time()
    while True:
        # initialize some stuff
        c_colors = [(0, 0, 255)] * 4

        # grab current frame from video_feed
        previous_frame = frame
        _, frame = video_feed.read()

        # Report FPS
        if not (frame_count % 10):
            fps = (time() - fps_time)/FPS_INTERVAL
            print('Frame:', frame_count, ' | FPS:', fps)
            fps_time = time()
        frame_count += 1

        # Convert to grayscale
        try:
            gray_img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        except:
            print("\nVideo feed ended.\n")
            break

        # get binary thresholding of image
        gray_smooth = cv2.GaussianBlur(gray_img, (15, 15), 0)
        _, bin_img = cv2.threshold(gray_smooth, 100, 255, cv2.THRESH_BINARY)

        # morphological closing
        kernel = np.ones((3, 3), np.uint8)
        bin_img = cv2.morphologyEx(bin_img, cv2.MORPH_CLOSE,
                                   kernel, iterations=closing_iterations)

        # Find corners.  To do this:
        # 1) Find the largest (area) contour in frame (after thresholding)
        # 2) get contours convex hull,
        # 3) reduce degree of convex hull with Douglas-Peucker algorithm,
        # 4) refine corners with subpixel corner finder

        # step 1
        contours, _ = cv2.findContours(bin_img, cv2.RETR_EXTERNAL,
                                                cv2.CHAIN_APPROX_SIMPLE)
        biggest_contour = max(contours, key=cv2.contourArea)

        # step 2
        hull = cv2.convexHull(biggest_contour)
        epsilon = 0.05 * cv2.arcLength(biggest_contour, True)

        # step 3
        hull = cv2.approxPolyDP(hull, epsilon, True)

        # step 4
        hull = np.float32(hull)
        method = cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT
        criteria = (method, 1000, 1e-4)
        cv2.cornerSubPix(gray_img, hull, (5, 5), (-1, -1), criteria)
        corners = [pt[0] for pt in hull]

        # Find top-right corner and use to label corners
        # Note: currently corners are in CW order
        # Note: ordering will be checked below against expected corners
        tr_index = np.argmin(c[0] + c[1] for c in corners)
        tl = corners[tr_index]
        bl = corners[(tr_index - 1) % 4]
        br = corners[(tr_index - 2) % 4]
        tr = corners[(tr_index - 3) % 4]

        # reformat and ensure that ordering is as expected below
        corners = np.float32([[c[0], c[1]] for c in [tl, bl, br, tr]])

        # IMPORTANT ASSUMPTIONS on paper tracking (used in code block below):
        # 1) if any one point is stationary from previous frame, then all
        #    are stationary with probability 1.
        # 2) if any corners are obstructed, assume the paper is still flat
        #    against the same plane as it was in the previous frame.
        #    I.e. the transformation from previous frame to this frame should
        #    be of the form of a translation and a rotation in said plane.
        # 3) see code comments for additional assumptions, haha, sorry

        def get_edge_lengths(topl, botl, botr, topr):
            """ Takes in list of four corners, returns four edge lengths
            in order top, right, bottom, left."""
            tbrl = [topr - topl, topr - botr, botr - botl, botl - topl]
            return [norm(edge) for edge in tbrl]

        if not corner_history:
            last_unob_corners = corners
        else:
            # determine expected corner locations and edge lengths
            expected_corners = last_unob_corners
            if last_unob_corners is None:
                expected_lengths = get_edge_lengths(*expected_corners)
            else:
                expected_lengths = get_edge_lengths(*last_unob_corners)

            # check ordering
            def cyclist(lst, k):
                if k:
                    return [lst[(i+k)%len(lst)] for i in xrange(len(lst))]
                return lst

            def _order_dist(offset):
                offset_corners = corners[cyclist(range(4), offset)]
                return norm(expected_corners - offset_corners), offset_corners
            if corner_history:
                corners = min(_order_dist(k) for k in range(4))[1]

            # Look for obstructions by looking for changes in edge lengths
            # Note: these lengths are not perspective invariant
            # TODO: checking by Hessian may be a better method
            new_lengths = get_edge_lengths(*corners)
            top_is_bad, rgt_is_bad, bot_is_bad, lft_is_bad = \
                [abs(l0 - l1) > obst_tol for l1, l0 in
                                zip(new_lengths, expected_lengths)]
            tl_ob = top_is_bad and lft_is_bad
            bl_ob = bot_is_bad and lft_is_bad
            br_ob = bot_is_bad and rgt_is_bad
            tr_ob = top_is_bad and rgt_is_bad

            is_obstr = [tl_ob, bl_ob, br_ob, tr_ob]
            ob_indices = [i for i, c in enumerate(is_obstr) if c]
            ob_corner_ct = sum(is_obstr)
            c_colors = [(0, 255, 0) if b else (0, 0, 255) for b in is_obstr]

            # Find difference of corners from expected location
            diffs = [norm(c - ec) for c, ec in zip(corners, expected_corners)]
            has_moved = [d > tol_corner_movement for d in diffs]

            # Check if paper has likely moved
            if sum(has_moved) < 4:
                # assume all is cool, just trust the corners found
                corners = last_unob_corners
                pass
            else:
                if sum(has_moved) == 1:
                    # only one corner has moved, just assume it's obstructed
                    # and replace it with the expected location
                    bad_corner_idx = np.argmax(diffs)
                    corners[bad_corner_idx] = expected_corners[bad_corner_idx]

                else:  # find paper's affine transformation in expected plane
                    print("frame={} | ob_corner_ct={}"
                          "".format(frame_count, ob_corner_ct))

                    if sum(is_obstr) in (1, 2, 3):
                        eco = zip(expected_corners, is_obstr)
                        exp_unob = np.float32([c for c, b in eco if not b])
                        exp_ob = np.float32([c for c, b in eco if b])
                        co = zip(corners, is_obstr)
                        new_unob = np.float32([c for c, b in co if not b])

                        exp_unob_pp = persTransform(exp_unob, old_homog)
                        exp_ob_pp = persTransform(exp_ob, old_homog)
                        new_unob_pp = persTransform(new_unob, old_homog)

                    # check for obstructions
                    if sum(is_obstr) == 0:  # yay! no obstructed corners!
                        pass

                    elif sum(is_obstr) == 1:
                        # Find the affine transformation in the paper's plane
                        # from expected locations of the three unobstructed
                        # corners to the found locations, then use this to
                        # estimate the obstructed corner's location
                        A = cv2.getAffineTransform(exp_unob_pp, new_unob_pp)
                        new_ob_pp = affTransform(exp_ob_pp, A)
                        new_ob = persTransform(new_ob_pp, old_inv_homog)
                        corners[np.ix_(ob_indices)] = new_ob

                    elif sum(is_obstr) == 2:
                        # Align the line between the good corners
                        # with the same line w.r.t the old corners
                        p1, q1 = new_unob_pp[0], new_unob_pp[1]
                        p0, q0 = exp_unob_pp[0], exp_unob_pp[1]
                        u0 = (q0 - p0) / norm(q0 - p0)
                        u1 = (q1 - p1) / norm(q1 - p1)
                        angle = acos(np.dot(u0, u1))  # unsigned
                        trans = p1 - p0

                        # Find rotation that moves u0 to u1
                        rotat = cv2.getRotationMatrix2D(tuple(p1), angle, 1)
                        rotat = rotat[:, :2]

                        # Expensive sign check for angle (could be improved)
                        if norm(np.dot(u0, rotat) - u1) > norm(np.dot(u1, rotat) - u0):
                            rotat = np.linalg.inv(rotat)

                        # transform the old coords of the hidden corners
                        # and map them back to the paper plane
                        exp_ob_pp += trans
                        new_ob_pp = affTransform(exp_ob_pp, rotat)
                        new_ob = persTransform(new_ob_pp, old_inv_homog)
                        corners[np.ix_(ob_indices)] = new_ob

                    elif sum(is_obstr) in (3, 4):
                        print("Uh oh, {} corners obstructed..."
                              "".format(ob_corner_ct))
                        corners = expected_corners
                    else:
                        raise Exception("This should never happen.")

        # Homography
        w = max(abs(br[0] - bl[0]),
                abs(tr[0] - tl[0]))  # width of paper in pixels
        h = PAPER_RATIO * w
        corners_pp = np.float32([[0, 0], [0, h], [w, h], [w, 0]])
        homog, mask = cv2.findHomography(corners, corners_pp)
        inv_homog, inv_mask = cv2.findHomography(corners_pp, corners)
        paper = cv2.warpPerspective(frame, homog, (int(ceil(w)), int(ceil(h))))
        if ROT180:
            paper = rotate180(paper)

        # Draw detected paper boundary on frame
        segmented_frame = draw_polygon(frame, corners, c_colors)

        # Resize paper to simplify display
        h = segmented_frame.shape[0]
        paper_w = int(round(h*paper.shape[1]/paper.shape[0]))
        resized_paper = cv2.resize(paper, (paper_w, h))

        # Display
        big_img = np.hstack((segmented_frame, resized_paper))
        if show_thresholding:
            bin_img = cv2.cvtColor(bin_img, cv2.COLOR_GRAY2BGR)
            big_img = np.hstack((big_img, bin_img))
        if REDUCE_DISPLAY_SIZE:
            reduced_size = tuple(np.array(big_img.shape[:2][::-1])//2)
            smaller_big_img = cv2.resize(big_img, reduced_size)
        cv2.imshow('', big_img)

        # Updates for next iteration
        corner_history.append(corners)
        old_homog = homog
        old_inv_homog = inv_homog

        # this is apparently necessary for 64bit machines
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    video_feed.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    try:
        run_main()
    except:
        cv2.destroyAllWindows()
        raise
<commit_msg>Fix FPS output<commit_after>"""
Notes:
    1) This is algorithm is primarily designed for a rectangular piece of
        paper lying flat on a flat surface, but may work in other situations
        assuming that the paper's corners are not obstructed.
    2) The camera's view of the paper must be unobstructed in first frame.

"""

# Basic Dependencies
from __future__ import division, print_function
from math import ceil, acos
from time import time

# External Dependencies
import numpy as np
from numpy.linalg import norm
import cv2

# Default User Parameters
VIDEO_FILE_LOCATION = 'sample.avi'
FPS_INTERVAL = 10  # updates FPS estimate after this many frames
MHI_DURATION = 10  # max frames remembered by motion history
FRAME_WIDTH, FRAME_HEIGHT = 640, 424
PAPER_RATIO = 11/8.5  # height/width of paper
ROT180 = False  # If paper is upside down, change this
REDUCE_DISPLAY_SIZE = False  # Use if output window is too big for your screen

# Internal Parameters
tol_corner_movement = 1
obst_tol = 10   # used to determine tolerance
closing_iterations = 10
show_thresholding = False  # Use to display thresholding

def rotate180(im):
    """Rotates an image by 180 degrees."""
    return cv2.flip(im, -1)

def persTransform(pts, H):
    """Transforms a list of points, `pts`,
    using the perspective transform `H`."""
    src = np.zeros((len(pts), 1, 2))
    src[:, 0] = pts
    dst = cv2.perspectiveTransform(src, H)
    return np.array(dst[:, 0, :], dtype='float32')

def affTransform(pts, A):
    """Transforms a list of points, `pts`,
    using the affine transform `A`."""
    src = np.zeros((len(pts), 1, 2))
    src[:, 0] = pts
    dst = cv2.transform(src, A)
    return np.array(dst[:, 0, :], dtype='float32')

def draw_polygon(im, vertices, vertex_colors=None, edge_colors=None,
                 alter_input_image=False, draw_edges=True, draw_vertices=True,
                 display=False, title='', pause=False):
    """returns image with polygon drawn on it."""
    _default_vertex_color = (255, 0, 0)
    _default_edge_color = (255, 0, 0)
    if not alter_input_image:
        im2 = im.copy()
    else:
        im2 = im
    if vertices is not None:
        N = len(vertices)
        vertices = [tuple(v) for v in vertices]
        if vertex_colors is None:
            vertex_colors = [_default_vertex_color] * N
        if edge_colors is None:
            edge_colors = [_default_edge_color] * N
        for i in range(N):
            startpt = vertices[(i - 1) % N]
            endpt = vertices[i]
            if draw_vertices:
                cv2.circle(im2, startpt, 3, vertex_colors[(i - 1) % N], -1)
            if draw_edges:
                cv2.line(im2, startpt, endpt, edge_colors[(i - 1) % N], 2)
    if display:
        cv2.imshow(title, im2)
        # Note: `0xFF == ord('q')`is apparently necessary for 64bit machines
        if pause and cv2.waitKey(0) & 0xFF == ord('q'):
            pass
    return im2

def run_main():

    # Initialize some variables
    frame = None
    old_homog = None
    old_inv_homog = None
    corner_history = []

    video_feed = cv2.VideoCapture(VIDEO_FILE_LOCATION)
    video_feed.set(cv2.cv.CV_CAP_PROP_FRAME_WIDTH, FRAME_WIDTH)
    video_feed.set(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT, FRAME_HEIGHT)

    frame_count = 0
    fps_time = time()
    while True:
        # initialize some stuff
        c_colors = [(0, 0, 255)] * 4

        # grab current frame from video_feed
        previous_frame = frame
        _, frame = video_feed.read()

        # Report FPS
        if not (frame_count % 10):
            fps = FPS_INTERVAL/(time() - fps_time)
            print('Frame:', frame_count, ' | FPS:', fps)
            fps_time = time()
        frame_count += 1

        # Convert to grayscale
        try:
            gray_img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        except:
            print("\nVideo feed ended.\n")
            break

        # get binary thresholding of image
        gray_smooth = cv2.GaussianBlur(gray_img, (15, 15), 0)
        _, bin_img = cv2.threshold(gray_smooth, 100, 255, cv2.THRESH_BINARY)

        # morphological closing
        kernel = np.ones((3, 3), np.uint8)
        bin_img = cv2.morphologyEx(bin_img, cv2.MORPH_CLOSE,
                                   kernel, iterations=closing_iterations)

        # Find corners.  To do this:
        # 1) Find the largest (area) contour in frame (after thresholding)
        # 2) get contours convex hull,
        # 3) reduce degree of convex hull with Douglas-Peucker algorithm,
        # 4) refine corners with subpixel corner finder

        # step 1
        contours, _ = cv2.findContours(bin_img, cv2.RETR_EXTERNAL,
                                                cv2.CHAIN_APPROX_SIMPLE)
        biggest_contour = max(contours, key=cv2.contourArea)

        # step 2
        hull = cv2.convexHull(biggest_contour)
        epsilon = 0.05 * cv2.arcLength(biggest_contour, True)

        # step 3
        hull = cv2.approxPolyDP(hull, epsilon, True)

        # step 4
        hull = np.float32(hull)
        method = cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT
        criteria = (method, 1000, 1e-4)
        cv2.cornerSubPix(gray_img, hull, (5, 5), (-1, -1), criteria)
        corners = [pt[0] for pt in hull]

        # Find top-right corner and use to label corners
        # Note: currently corners are in CW order
        # Note: ordering will be checked below against expected corners
        tr_index = np.argmin(c[0] + c[1] for c in corners)
        tl = corners[tr_index]
        bl = corners[(tr_index - 1) % 4]
        br = corners[(tr_index - 2) % 4]
        tr = corners[(tr_index - 3) % 4]

        # reformat and ensure that ordering is as expected below
        corners = np.float32([[c[0], c[1]] for c in [tl, bl, br, tr]])

        # IMPORTANT ASSUMPTIONS on paper tracking (used in code block below):
        # 1) if any one point is stationary from previous frame, then all
        #    are stationary with probability 1.
        # 2) if any corners are obstructed, assume the paper is still flat
        #    against the same plane as it was in the previous frame.
        #    I.e. the transformation from previous frame to this frame should
        #    be of the form of a translation and a rotation in said plane.
        # 3) see code comments for additional assumptions, haha, sorry

        def get_edge_lengths(topl, botl, botr, topr):
            """ Takes in list of four corners, returns four edge lengths
            in order top, right, bottom, left."""
            tbrl = [topr - topl, topr - botr, botr - botl, botl - topl]
            return [norm(edge) for edge in tbrl]

        if not corner_history:
            last_unob_corners = corners
        else:
            # determine expected corner locations and edge lengths
            expected_corners = last_unob_corners
            if last_unob_corners is None:
                expected_lengths = get_edge_lengths(*expected_corners)
            else:
                expected_lengths = get_edge_lengths(*last_unob_corners)

            # check ordering
            def cyclist(lst, k):
                if k:
                    return [lst[(i+k)%len(lst)] for i in xrange(len(lst))]
                return lst

            def _order_dist(offset):
                offset_corners = corners[cyclist(range(4), offset)]
                return norm(expected_corners - offset_corners), offset_corners
            if corner_history:
                corners = min(_order_dist(k) for k in range(4))[1]

            # Look for obstructions by looking for changes in edge lengths
            # Note: these lengths are not perspective invariant
            # TODO: checking by Hessian may be a better method
            new_lengths = get_edge_lengths(*corners)
            top_is_bad, rgt_is_bad, bot_is_bad, lft_is_bad = \
                [abs(l0 - l1) > obst_tol for l1, l0 in
                                zip(new_lengths, expected_lengths)]
            tl_ob = top_is_bad and lft_is_bad
            bl_ob = bot_is_bad and lft_is_bad
            br_ob = bot_is_bad and rgt_is_bad
            tr_ob = top_is_bad and rgt_is_bad

            is_obstr = [tl_ob, bl_ob, br_ob, tr_ob]
            ob_indices = [i for i, c in enumerate(is_obstr) if c]
            ob_corner_ct = sum(is_obstr)
            c_colors = [(0, 255, 0) if b else (0, 0, 255) for b in is_obstr]

            # Find difference of corners from expected location
            diffs = [norm(c - ec) for c, ec in zip(corners, expected_corners)]
            has_moved = [d > tol_corner_movement for d in diffs]

            # Check if paper has likely moved
            if sum(has_moved) < 4:
                # assume all is cool, just trust the corners found
                corners = last_unob_corners
                pass
            else:
                if sum(has_moved) == 1:
                    # only one corner has moved, just assume it's obstructed
                    # and replace it with the expected location
                    bad_corner_idx = np.argmax(diffs)
                    corners[bad_corner_idx] = expected_corners[bad_corner_idx]

                else:  # find paper's affine transformation in expected plane
                    print("frame={} | ob_corner_ct={}"
                          "".format(frame_count, ob_corner_ct))

                    if sum(is_obstr) in (1, 2, 3):
                        eco = zip(expected_corners, is_obstr)
                        exp_unob = np.float32([c for c, b in eco if not b])
                        exp_ob = np.float32([c for c, b in eco if b])
                        co = zip(corners, is_obstr)
                        new_unob = np.float32([c for c, b in co if not b])

                        exp_unob_pp = persTransform(exp_unob, old_homog)
                        exp_ob_pp = persTransform(exp_ob, old_homog)
                        new_unob_pp = persTransform(new_unob, old_homog)

                    # check for obstructions
                    if sum(is_obstr) == 0:  # yay! no obstructed corners!
                        pass

                    elif sum(is_obstr) == 1:
                        # Find the affine transformation in the paper's plane
                        # from expected locations of the three unobstructed
                        # corners to the found locations, then use this to
                        # estimate the obstructed corner's location
                        A = cv2.getAffineTransform(exp_unob_pp, new_unob_pp)
                        new_ob_pp = affTransform(exp_ob_pp, A)
                        new_ob = persTransform(new_ob_pp, old_inv_homog)
                        corners[np.ix_(ob_indices)] = new_ob

                    elif sum(is_obstr) == 2:
                        # Align the line between the good corners
                        # with the same line w.r.t the old corners
                        p1, q1 = new_unob_pp[0], new_unob_pp[1]
                        p0, q0 = exp_unob_pp[0], exp_unob_pp[1]
                        u0 = (q0 - p0) / norm(q0 - p0)
                        u1 = (q1 - p1) / norm(q1 - p1)
                        angle = acos(np.dot(u0, u1))  # unsigned
                        trans = p1 - p0

                        # Find rotation that moves u0 to u1
                        rotat = cv2.getRotationMatrix2D(tuple(p1), angle, 1)
                        rotat = rotat[:, :2]

                        # Expensive sign check for angle (could be improved)
                        if norm(np.dot(u0, rotat) - u1) > norm(np.dot(u1, rotat) - u0):
                            rotat = np.linalg.inv(rotat)

                        # transform the old coords of the hidden corners
                        # and map them back to the paper plane
                        exp_ob_pp += trans
                        new_ob_pp = affTransform(exp_ob_pp, rotat)
                        new_ob = persTransform(new_ob_pp, old_inv_homog)
                        corners[np.ix_(ob_indices)] = new_ob

                    elif sum(is_obstr) in (3, 4):
                        print("Uh oh, {} corners obstructed..."
                              "".format(ob_corner_ct))
                        corners = expected_corners
                    else:
                        raise Exception("This should never happen.")

        # Homography
        w = max(abs(br[0] - bl[0]),
                abs(tr[0] - tl[0]))  # width of paper in pixels
        h = PAPER_RATIO * w
        corners_pp = np.float32([[0, 0], [0, h], [w, h], [w, 0]])
        homog, mask = cv2.findHomography(corners, corners_pp)
        inv_homog, inv_mask = cv2.findHomography(corners_pp, corners)
        paper = cv2.warpPerspective(frame, homog, (int(ceil(w)), int(ceil(h))))
        if ROT180:
            paper = rotate180(paper)

        # Draw detected paper boundary on frame
        segmented_frame = draw_polygon(frame, corners, c_colors)

        # Resize paper to simplify display
        h = segmented_frame.shape[0]
        paper_w = int(round(h*paper.shape[1]/paper.shape[0]))
        resized_paper = cv2.resize(paper, (paper_w, h))

        # Display
        big_img = np.hstack((segmented_frame, resized_paper))
        if show_thresholding:
            bin_img = cv2.cvtColor(bin_img, cv2.COLOR_GRAY2BGR)
            big_img = np.hstack((big_img, bin_img))
        if REDUCE_DISPLAY_SIZE:
            reduced_size = tuple(np.array(big_img.shape[:2][::-1])//2)
            smaller_big_img = cv2.resize(big_img, reduced_size)
        cv2.imshow('', big_img)

        # Updates for next iteration
        corner_history.append(corners)
        old_homog = homog
        old_inv_homog = inv_homog

        # this is apparently necessary for 64bit machines
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    video_feed.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    try:
        run_main()
    except:
        cv2.destroyAllWindows()
        raise
<|endoftext|>


<commit_before><commit_msg>Update redpoint.py<commit_after><|endoftext|>


<commit_before>#!/usr/bin/env python
# coding=utf-8
from __future__ import (division, print_function, unicode_literals,
                        absolute_import)

import os
import datetime
import tempfile
import io

import pytest

from tinydb import TinyDB
from hashfs import HashFS

from sacred.dependencies import get_digest
from sacred.observers.tinydb_hashfs import TinyDbObserver, TinyDbOption
from sacred import optional as opt
from sacred.experiment import Experiment

T1 = datetime.datetime(1999, 5, 4, 3, 2, 1, 0)
T2 = datetime.datetime(1999, 5, 5, 5, 5, 5, 5)

@pytest.fixture()
def tinydb_obs(tmpdir):
    return TinyDbObserver.create(path=tmpdir.strpath)

@pytest.fixture()
def sample_run():
    exp = {'name': 'test_exp', 'sources': [], 'doc': '', 'base_dir': '/tmp'}
    host = {'hostname': 'test_host', 'cpu_count': 1, 'python_version': '3.4'}
    config = {'config': 'True', 'foo': 'bar', 'answer': 42}
    command = 'run'
    meta_info = {'comment': 'test run'}
    return {
        '_id': 'FEDCBA9876543210',
        'ex_info': exp,
        'command': command,
        'host_info': host,
        'start_time': T1,
        'config': config,
        'meta_info': meta_info,
    }

def test_tindb_observer_started_event_creates_run(tinydb_obs, sample_run):
    sample_run['_id'] = None
    _id = tinydb_obs.started_event(**sample_run)
    assert _id is not None
    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run == {
        '_id': _id,
        'experiment': sample_run['ex_info'],
        'format': tinydb_obs.VERSION,
        'command': sample_run['command'],
        'host': sample_run['host_info'],
        'start_time': sample_run['start_time'],
        'heartbeat': None,
        'info': {},
        'captured_out': '',
        'artifacts': [],
        'config': sample_run['config'],
        'meta': sample_run['meta_info'],
        'status': 'RUNNING',
        'resources': []
    }

def test_tinydb_observer_started_event_uses_given_id(tinydb_obs, sample_run):
    _id = tinydb_obs.started_event(**sample_run)
    assert _id == sample_run['_id']
    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['_id'] == sample_run['_id']

def test_tindb_observer_started_event_saves_given_sources(tinydb_obs,
                                                          sample_run):

    filename = 'setup.py'
    md5 = get_digest(filename)

    sample_run['ex_info']['sources'] = [[filename, md5]]
    _id = tinydb_obs.started_event(**sample_run)

    assert _id is not None
    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)

    # Check all but the experiment section
    db_run_copy = db_run.copy()
    del db_run_copy['experiment']
    assert db_run_copy == {
        '_id': _id,
        'format': tinydb_obs.VERSION,
        'command': sample_run['command'],
        'host': sample_run['host_info'],
        'start_time': sample_run['start_time'],
        'heartbeat': None,
        'info': {},
        'captured_out': '',
        'artifacts': [],
        'config': sample_run['config'],
        'meta': sample_run['meta_info'],
        'status': 'RUNNING',
        'resources': []
    }

    assert len(db_run['experiment']['sources']) == 1
    assert len(db_run['experiment']['sources'][0]) == 3
    assert db_run['experiment']['sources'][0][:2] == [filename, md5]
    assert isinstance(db_run['experiment']['sources'][0][2], io.BufferedReader)

    # Check that duplicate source files are still listed in ex_info
    tinydb_obs.db_run_id = None
    tinydb_obs.started_event(**sample_run)
    assert len(tinydb_obs.runs) == 2
    db_run2 = tinydb_obs.runs.get(eid=2)

    assert db_run['experiment']['sources'][0][:2] == db_run2['experiment']['sources'][0][:2]

def test_tindb_observer_started_event_generates_different_run_ids(tinydb_obs,
                                                                  sample_run):
    sample_run['_id'] = None
    _id = tinydb_obs.started_event(**sample_run)
    assert _id is not None

    # Check that duplicate source files are still listed in ex_info
    tinydb_obs.db_run_id = None
    sample_run['_id'] = None
    _id2 = tinydb_obs.started_event(**sample_run)

    assert len(tinydb_obs.runs) == 2
    # Check new random id is given for each run
    assert _id != _id2

def test_tinydb_observer_queued_event_is_not_implimented(tinydb_obs,
                                                         sample_run):

    sample_queued_run = sample_run.copy()
    del sample_queued_run['host_info']
    del sample_queued_run['start_time']
    sample_queued_run['queue_time'] = T1

    with pytest.raises(NotImplementedError):
        tinydb_obs.queued_event(**sample_queued_run)

def test_tinydb_observer_equality(tmpdir, tinydb_obs):

    db = TinyDB(os.path.join(tmpdir.strpath, 'metadata.json'))
    fs = HashFS(os.path.join(tmpdir.strpath, 'hashfs'), depth=3,
                width=2, algorithm='md5')
    m = TinyDbObserver(db, fs)

    assert tinydb_obs == m
    assert not tinydb_obs != m

    assert not tinydb_obs == 'foo'
    assert tinydb_obs != 'foo'

def test_tinydb_observer_heartbeat_event_updates_run(tinydb_obs, sample_run):
    tinydb_obs.started_event(**sample_run)

    info = {'my_info': [1, 2, 3], 'nr': 7}
    outp = 'some output'
    with tempfile.NamedTemporaryFile() as f:
        f.write(outp.encode())
        f.flush()
        tinydb_obs.heartbeat_event(info=info, cout_filename=f.name,
                                   beat_time=T2)

    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['heartbeat'] == T2
    assert db_run['info'] == info
    assert db_run['captured_out'] == outp

def test_tinydb_observer_completed_event_updates_run(tinydb_obs, sample_run):
    tinydb_obs.started_event(**sample_run)

    tinydb_obs.completed_event(stop_time=T2, result=42)

    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['stop_time'] == T2
    assert db_run['result'] == 42
    assert db_run['status'] == 'COMPLETED'

def test_tinydb_observer_interrupted_event_updates_run(tinydb_obs,
                                                       sample_run):
    tinydb_obs.started_event(**sample_run)

    tinydb_obs.interrupted_event(interrupt_time=T2, status='INTERRUPTED')

    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['stop_time'] == T2
    assert db_run['status'] == 'INTERRUPTED'

def test_tinydb_observer_failed_event_updates_run(tinydb_obs, sample_run):
    tinydb_obs.started_event(**sample_run)

    fail_trace = "lots of errors and\nso\non..."
    tinydb_obs.failed_event(fail_time=T2,
                            fail_trace=fail_trace)

    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['stop_time'] == T2
    assert db_run['status'] == 'FAILED'
    assert db_run['fail_trace'] == fail_trace

def test_tinydb_observer_artifact_event(tinydb_obs, sample_run):
    tinydb_obs.started_event(**sample_run)

    filename = "setup.py"
    name = 'mysetup'

    tinydb_obs.artifact_event(name, filename)

    assert tinydb_obs.fs.exists(filename)

    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['artifacts'][0][0] == name

    with open(filename, 'rb') as f:
        file_content = f.read()
    assert db_run['artifacts'][0][3].read() == file_content

def test_tinydb_observer_resource_event(tinydb_obs, sample_run):
    tinydb_obs.started_event(**sample_run)

    filename = "setup.py"
    md5 = get_digest(filename)

    tinydb_obs.resource_event(filename)

    assert tinydb_obs.fs.exists(filename)

    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['resources'][0][:2] == [filename, md5]

    with open(filename, 'rb') as f:
        file_content = f.read()
    assert db_run['resources'][0][2].read() == file_content

def test_tinydb_observer_resource_event_when_resource_present(tinydb_obs,
                                                              sample_run):
    tinydb_obs.started_event(**sample_run)

    filename = "setup.py"
    md5 = get_digest(filename)

    # Add file by other means
    tinydb_obs.fs.put(filename)

    tinydb_obs.resource_event(filename)

    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['resources'][0][:2] == [filename, md5]

@pytest.mark.skipif(not opt.has_numpy, reason='needs numpy')
def test_serialisation_of_numpy_ndarray(tmpdir):
    from sacred.observers.tinydb_hashfs import NdArraySerializer
    from tinydb_serialization import SerializationMiddleware
    import numpy as np

    # Setup Serialisation object for non list/dict objects
    serialization_store = SerializationMiddleware()
    serialization_store.register_serializer(NdArraySerializer(), 'TinyArray')

    db = TinyDB(os.path.join(tmpdir.strpath, 'metadata.json'),
                storage=serialization_store)

    eye_mat = np.eye(3)
    ones_array = np.ones(5)

    document = {
        'foo': 'bar',
        'some_array': eye_mat,
        'nested': {
            'ones': ones_array
        }
    }

    db.insert(document)
    returned_doc = db.all()[0]

    assert returned_doc['foo'] == 'bar'
    assert (returned_doc['some_array'] == eye_mat).all()
    assert (returned_doc['nested']['ones'] == ones_array).all()

@pytest.mark.skipif(not opt.has_pandas, reason='needs pandas')
def test_serialisation_of_pandas_dataframe(tmpdir):
    from sacred.observers.tinydb_hashfs import (DataFrameSerializer,
                                                SeriesSerializer)
    from tinydb_serialization import SerializationMiddleware

    import numpy as np
    import pandas as pd

    # Setup Serialisation object for non list/dict objects
    serialization_store = SerializationMiddleware()
    serialization_store.register_serializer(DataFrameSerializer(),
                                            'TinyDataFrame')
    serialization_store.register_serializer(SeriesSerializer(),
                                            'TinySeries')

    db = TinyDB(os.path.join(tmpdir.strpath, 'metadata.json'),
                storage=serialization_store)

    df = pd.DataFrame(np.eye(3), columns=list('ABC'))
    series = pd.Series(np.ones(5))

    document = {
        'foo': 'bar',
        'some_dataframe': df,
        'nested': {
            'ones': series
        }
    }

    db.insert(document)
    returned_doc = db.all()[0]

    assert returned_doc['foo'] == 'bar'
    assert (returned_doc['some_dataframe'] == df).all().all()
    assert (returned_doc['nested']['ones'] == series).all()

def test_parse_tinydb_arg():
    assert TinyDbOption.parse_tinydb_arg('foo') == 'foo'

def test_parse_tinydboption_apply(tmpdir):

    exp = Experiment()
    args = os.path.join(tmpdir.strpath)

    TinyDbOption.apply(args, exp)
    assert type(exp.observers[0]) == TinyDbObserver
<commit_msg>Fix typo in test names<commit_after>#!/usr/bin/env python
# coding=utf-8
from __future__ import (division, print_function, unicode_literals,
                        absolute_import)

import os
import datetime
import tempfile
import io

import pytest

from tinydb import TinyDB
from hashfs import HashFS

from sacred.dependencies import get_digest
from sacred.observers.tinydb_hashfs import TinyDbObserver, TinyDbOption
from sacred import optional as opt
from sacred.experiment import Experiment

T1 = datetime.datetime(1999, 5, 4, 3, 2, 1, 0)
T2 = datetime.datetime(1999, 5, 5, 5, 5, 5, 5)

@pytest.fixture()
def tinydb_obs(tmpdir):
    return TinyDbObserver.create(path=tmpdir.strpath)

@pytest.fixture()
def sample_run():
    exp = {'name': 'test_exp', 'sources': [], 'doc': '', 'base_dir': '/tmp'}
    host = {'hostname': 'test_host', 'cpu_count': 1, 'python_version': '3.4'}
    config = {'config': 'True', 'foo': 'bar', 'answer': 42}
    command = 'run'
    meta_info = {'comment': 'test run'}
    return {
        '_id': 'FEDCBA9876543210',
        'ex_info': exp,
        'command': command,
        'host_info': host,
        'start_time': T1,
        'config': config,
        'meta_info': meta_info,
    }

def test_tinydb_observer_started_event_creates_run(tinydb_obs, sample_run):
    sample_run['_id'] = None
    _id = tinydb_obs.started_event(**sample_run)
    assert _id is not None
    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run == {
        '_id': _id,
        'experiment': sample_run['ex_info'],
        'format': tinydb_obs.VERSION,
        'command': sample_run['command'],
        'host': sample_run['host_info'],
        'start_time': sample_run['start_time'],
        'heartbeat': None,
        'info': {},
        'captured_out': '',
        'artifacts': [],
        'config': sample_run['config'],
        'meta': sample_run['meta_info'],
        'status': 'RUNNING',
        'resources': []
    }

def test_tinydb_observer_started_event_uses_given_id(tinydb_obs, sample_run):
    _id = tinydb_obs.started_event(**sample_run)
    assert _id == sample_run['_id']
    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['_id'] == sample_run['_id']

def test_tinydb_observer_started_event_saves_given_sources(tinydb_obs,
                                                          sample_run):

    filename = 'setup.py'
    md5 = get_digest(filename)

    sample_run['ex_info']['sources'] = [[filename, md5]]
    _id = tinydb_obs.started_event(**sample_run)

    assert _id is not None
    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)

    # Check all but the experiment section
    db_run_copy = db_run.copy()
    del db_run_copy['experiment']
    assert db_run_copy == {
        '_id': _id,
        'format': tinydb_obs.VERSION,
        'command': sample_run['command'],
        'host': sample_run['host_info'],
        'start_time': sample_run['start_time'],
        'heartbeat': None,
        'info': {},
        'captured_out': '',
        'artifacts': [],
        'config': sample_run['config'],
        'meta': sample_run['meta_info'],
        'status': 'RUNNING',
        'resources': []
    }

    assert len(db_run['experiment']['sources']) == 1
    assert len(db_run['experiment']['sources'][0]) == 3
    assert db_run['experiment']['sources'][0][:2] == [filename, md5]
    assert isinstance(db_run['experiment']['sources'][0][2], io.BufferedReader)

    # Check that duplicate source files are still listed in ex_info
    tinydb_obs.db_run_id = None
    tinydb_obs.started_event(**sample_run)
    assert len(tinydb_obs.runs) == 2
    db_run2 = tinydb_obs.runs.get(eid=2)

    assert db_run['experiment']['sources'][0][:2] == db_run2['experiment']['sources'][0][:2]

def test_tinydb_observer_started_event_generates_different_run_ids(tinydb_obs,
                                                                  sample_run):
    sample_run['_id'] = None
    _id = tinydb_obs.started_event(**sample_run)
    assert _id is not None

    # Check that duplicate source files are still listed in ex_info
    tinydb_obs.db_run_id = None
    sample_run['_id'] = None
    _id2 = tinydb_obs.started_event(**sample_run)

    assert len(tinydb_obs.runs) == 2
    # Check new random id is given for each run
    assert _id != _id2

def test_tinydb_observer_queued_event_is_not_implimented(tinydb_obs,
                                                         sample_run):

    sample_queued_run = sample_run.copy()
    del sample_queued_run['host_info']
    del sample_queued_run['start_time']
    sample_queued_run['queue_time'] = T1

    with pytest.raises(NotImplementedError):
        tinydb_obs.queued_event(**sample_queued_run)

def test_tinydb_observer_equality(tmpdir, tinydb_obs):

    db = TinyDB(os.path.join(tmpdir.strpath, 'metadata.json'))
    fs = HashFS(os.path.join(tmpdir.strpath, 'hashfs'), depth=3,
                width=2, algorithm='md5')
    m = TinyDbObserver(db, fs)

    assert tinydb_obs == m
    assert not tinydb_obs != m

    assert not tinydb_obs == 'foo'
    assert tinydb_obs != 'foo'

def test_tinydb_observer_heartbeat_event_updates_run(tinydb_obs, sample_run):
    tinydb_obs.started_event(**sample_run)

    info = {'my_info': [1, 2, 3], 'nr': 7}
    outp = 'some output'
    with tempfile.NamedTemporaryFile() as f:
        f.write(outp.encode())
        f.flush()
        tinydb_obs.heartbeat_event(info=info, cout_filename=f.name,
                                   beat_time=T2)

    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['heartbeat'] == T2
    assert db_run['info'] == info
    assert db_run['captured_out'] == outp

def test_tinydb_observer_completed_event_updates_run(tinydb_obs, sample_run):
    tinydb_obs.started_event(**sample_run)

    tinydb_obs.completed_event(stop_time=T2, result=42)

    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['stop_time'] == T2
    assert db_run['result'] == 42
    assert db_run['status'] == 'COMPLETED'

def test_tinydb_observer_interrupted_event_updates_run(tinydb_obs,
                                                       sample_run):
    tinydb_obs.started_event(**sample_run)

    tinydb_obs.interrupted_event(interrupt_time=T2, status='INTERRUPTED')

    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['stop_time'] == T2
    assert db_run['status'] == 'INTERRUPTED'

def test_tinydb_observer_failed_event_updates_run(tinydb_obs, sample_run):
    tinydb_obs.started_event(**sample_run)

    fail_trace = "lots of errors and\nso\non..."
    tinydb_obs.failed_event(fail_time=T2,
                            fail_trace=fail_trace)

    assert len(tinydb_obs.runs) == 1
    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['stop_time'] == T2
    assert db_run['status'] == 'FAILED'
    assert db_run['fail_trace'] == fail_trace

def test_tinydb_observer_artifact_event(tinydb_obs, sample_run):
    tinydb_obs.started_event(**sample_run)

    filename = "setup.py"
    name = 'mysetup'

    tinydb_obs.artifact_event(name, filename)

    assert tinydb_obs.fs.exists(filename)

    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['artifacts'][0][0] == name

    with open(filename, 'rb') as f:
        file_content = f.read()
    assert db_run['artifacts'][0][3].read() == file_content

def test_tinydb_observer_resource_event(tinydb_obs, sample_run):
    tinydb_obs.started_event(**sample_run)

    filename = "setup.py"
    md5 = get_digest(filename)

    tinydb_obs.resource_event(filename)

    assert tinydb_obs.fs.exists(filename)

    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['resources'][0][:2] == [filename, md5]

    with open(filename, 'rb') as f:
        file_content = f.read()
    assert db_run['resources'][0][2].read() == file_content

def test_tinydb_observer_resource_event_when_resource_present(tinydb_obs,
                                                              sample_run):
    tinydb_obs.started_event(**sample_run)

    filename = "setup.py"
    md5 = get_digest(filename)

    # Add file by other means
    tinydb_obs.fs.put(filename)

    tinydb_obs.resource_event(filename)

    db_run = tinydb_obs.runs.get(eid=1)
    assert db_run['resources'][0][:2] == [filename, md5]

@pytest.mark.skipif(not opt.has_numpy, reason='needs numpy')
def test_serialisation_of_numpy_ndarray(tmpdir):
    from sacred.observers.tinydb_hashfs import NdArraySerializer
    from tinydb_serialization import SerializationMiddleware
    import numpy as np

    # Setup Serialisation object for non list/dict objects
    serialization_store = SerializationMiddleware()
    serialization_store.register_serializer(NdArraySerializer(), 'TinyArray')

    db = TinyDB(os.path.join(tmpdir.strpath, 'metadata.json'),
                storage=serialization_store)

    eye_mat = np.eye(3)
    ones_array = np.ones(5)

    document = {
        'foo': 'bar',
        'some_array': eye_mat,
        'nested': {
            'ones': ones_array
        }
    }

    db.insert(document)
    returned_doc = db.all()[0]

    assert returned_doc['foo'] == 'bar'
    assert (returned_doc['some_array'] == eye_mat).all()
    assert (returned_doc['nested']['ones'] == ones_array).all()

@pytest.mark.skipif(not opt.has_pandas, reason='needs pandas')
def test_serialisation_of_pandas_dataframe(tmpdir):
    from sacred.observers.tinydb_hashfs import (DataFrameSerializer,
                                                SeriesSerializer)
    from tinydb_serialization import SerializationMiddleware

    import numpy as np
    import pandas as pd

    # Setup Serialisation object for non list/dict objects
    serialization_store = SerializationMiddleware()
    serialization_store.register_serializer(DataFrameSerializer(),
                                            'TinyDataFrame')
    serialization_store.register_serializer(SeriesSerializer(),
                                            'TinySeries')

    db = TinyDB(os.path.join(tmpdir.strpath, 'metadata.json'),
                storage=serialization_store)

    df = pd.DataFrame(np.eye(3), columns=list('ABC'))
    series = pd.Series(np.ones(5))

    document = {
        'foo': 'bar',
        'some_dataframe': df,
        'nested': {
            'ones': series
        }
    }

    db.insert(document)
    returned_doc = db.all()[0]

    assert returned_doc['foo'] == 'bar'
    assert (returned_doc['some_dataframe'] == df).all().all()
    assert (returned_doc['nested']['ones'] == series).all()

def test_parse_tinydb_arg():
    assert TinyDbOption.parse_tinydb_arg('foo') == 'foo'

def test_parse_tinydboption_apply(tmpdir):

    exp = Experiment()
    args = os.path.join(tmpdir.strpath)

    TinyDbOption.apply(args, exp)
    assert type(exp.observers[0]) == TinyDbObserver
<|endoftext|>