<commit_before>#!/usr/bin/env python3
"""
This file is part of Mbed TLS (https://tls.mbed.org)

Copyright (c) 2018, Arm Limited, All Rights Reserved

Purpose

This script checks the current state of the source code for minor issues,
including incorrect file permissions, presence of tabs, non-Unix line endings,
trailing whitespace, presence of UTF-8 BOM, and TODO comments.
Note: requires python 3, must be run from Mbed TLS root.
"""

import os
import argparse
import logging
import codecs
import sys

class IssueTracker(object):
    """Base class for issue tracking. Issues should inherit from this and
    overwrite either issue_with_line if they check the file line by line, or
    overwrite check_file_for_issue if they check the file as a whole."""

    def __init__(self):
        self.heading = ""
        self.files_exemptions = []
        self.files_with_issues = {}

    def should_check_file(self, filepath):
        for files_exemption in self.files_exemptions:
            if filepath.endswith(files_exemption):
                return False
        return True

    def issue_with_line(self, line):
        raise NotImplementedError

    def check_file_for_issue(self, filepath):
        with open(filepath, "rb") as f:
            for i, line in enumerate(iter(f.readline, b"")):
                self.check_file_line(filepath, line, i + 1)

    def check_file_line(self, filepath, line, line_number):
        if self.issue_with_line(line):
            if filepath not in self.files_with_issues.keys():
                self.files_with_issues[filepath] = []
            self.files_with_issues[filepath].append(line_number)

    def output_file_issues(self, logger):
        if self.files_with_issues.values():
            logger.info(self.heading)
            for filename, lines in sorted(self.files_with_issues.items()):
                if lines:
                    logger.info("{}: {}".format(
                        filename, ", ".join(str(x) for x in lines)
                    ))
                else:
                    logger.info(filename)
            logger.info("")

class PermissionIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "Incorrect permissions:"

    def check_file_for_issue(self, filepath):
        if not (os.access(filepath, os.X_OK) ==
                filepath.endswith((".sh", ".pl", ".py"))):
            self.files_with_issues[filepath] = None

class EndOfFileNewlineIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "Missing newline at end of file:"

    def check_file_for_issue(self, filepath):
        with open(filepath, "rb") as f:
            if not f.read().endswith(b"\n"):
                self.files_with_issues[filepath] = None

class Utf8BomIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "UTF-8 BOM present:"

    def check_file_for_issue(self, filepath):
        with open(filepath, "rb") as f:
            if f.read().startswith(codecs.BOM_UTF8):
                self.files_with_issues[filepath] = None

class LineEndingIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "Non Unix line endings:"

    def issue_with_line(self, line):
        return b"\r" in line

class TrailingWhitespaceIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "Trailing whitespace:"
        self.files_exemptions = [".md"]

    def issue_with_line(self, line):
        return line.rstrip(b"\r\n") != line.rstrip()

class TabIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "Tabs present:"
        self.files_exemptions = [
            "Makefile", "generate_visualc_files.pl"
        ]

    def issue_with_line(self, line):
        return b"\t" in line

class TodoIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "TODO present:"
        self.files_exemptions = [
            __file__, "benchmark.c", "pull_request_template.md"
        ]

    def issue_with_line(self, line):
        return b"todo" in line.lower()

class IntegrityChecker(object):

    def __init__(self, log_file):
        self.check_repo_path()
        self.logger = None
        self.setup_logger(log_file)
        self.files_to_check = (
            ".c", ".h", ".sh", ".pl", ".py", ".md", ".function", ".data",
            "Makefile", "CMakeLists.txt", "ChangeLog"
        )
        self.excluded_directories = ['.git', 'mbed-os']
        self.excluded_paths = list(map(os.path.normpath, [
            'cov-int',
            'examples',
            'yotta/module'
        ]))
        self.issues_to_check = [
            PermissionIssueTracker(),
            EndOfFileNewlineIssueTracker(),
            Utf8BomIssueTracker(),
            LineEndingIssueTracker(),
            TrailingWhitespaceIssueTracker(),
            TabIssueTracker(),
            TodoIssueTracker(),
        ]

    def check_repo_path(self):
        if not all(os.path.isdir(d) for d in ["include", "library", "tests"]):
            raise Exception("Must be run from Mbed TLS root")

    def setup_logger(self, log_file, level=logging.INFO):
        self.logger = logging.getLogger()
        self.logger.setLevel(level)
        if log_file:
            handler = logging.FileHandler(log_file)
            self.logger.addHandler(handler)
        else:
            console = logging.StreamHandler()
            self.logger.addHandler(console)

    def prune_branch(self, root, d):
        if d in self.excluded_directories:
            return True
        if os.path.normpath(os.path.join(root, d)) in self.excluded_paths:
            return True
        return False

    def check_files(self):
        for root, dirs, files in os.walk("."):
            dirs[:] = sorted(d for d in dirs if not self.prune_branch(root, d))
            for filename in sorted(files):
                filepath = os.path.join(root, filename)
                if not filepath.endswith(self.files_to_check):
                    continue
                for issue_to_check in self.issues_to_check:
                    if issue_to_check.should_check_file(filepath):
                        issue_to_check.check_file_for_issue(filepath)

    def output_issues(self):
        integrity_return_code = 0
        for issue_to_check in self.issues_to_check:
            if issue_to_check.files_with_issues:
                integrity_return_code = 1
            issue_to_check.output_file_issues(self.logger)
        return integrity_return_code

def run_main():
    parser = argparse.ArgumentParser(
        description=(
            "This script checks the current state of the source code for "
            "minor issues, including incorrect file permissions, "
            "presence of tabs, non-Unix line endings, trailing whitespace, "
            "presence of UTF-8 BOM, and TODO comments. "
            "Note: requires python 3, must be run from Mbed TLS root."
        )
    )
    parser.add_argument(
        "-l", "--log_file", type=str, help="path to optional output log",
    )
    check_args = parser.parse_args()
    integrity_check = IntegrityChecker(check_args.log_file)
    integrity_check.check_files()
    return_code = integrity_check.output_issues()
    sys.exit(return_code)

if __name__ == "__main__":
    run_main()
<commit_msg>Remove yotta support from check-files.py<commit_after>#!/usr/bin/env python3
"""
This file is part of Mbed TLS (https://tls.mbed.org)

Copyright (c) 2018, Arm Limited, All Rights Reserved

Purpose

This script checks the current state of the source code for minor issues,
including incorrect file permissions, presence of tabs, non-Unix line endings,
trailing whitespace, presence of UTF-8 BOM, and TODO comments.
Note: requires python 3, must be run from Mbed TLS root.
"""

import os
import argparse
import logging
import codecs
import sys

class IssueTracker(object):
    """Base class for issue tracking. Issues should inherit from this and
    overwrite either issue_with_line if they check the file line by line, or
    overwrite check_file_for_issue if they check the file as a whole."""

    def __init__(self):
        self.heading = ""
        self.files_exemptions = []
        self.files_with_issues = {}

    def should_check_file(self, filepath):
        for files_exemption in self.files_exemptions:
            if filepath.endswith(files_exemption):
                return False
        return True

    def issue_with_line(self, line):
        raise NotImplementedError

    def check_file_for_issue(self, filepath):
        with open(filepath, "rb") as f:
            for i, line in enumerate(iter(f.readline, b"")):
                self.check_file_line(filepath, line, i + 1)

    def check_file_line(self, filepath, line, line_number):
        if self.issue_with_line(line):
            if filepath not in self.files_with_issues.keys():
                self.files_with_issues[filepath] = []
            self.files_with_issues[filepath].append(line_number)

    def output_file_issues(self, logger):
        if self.files_with_issues.values():
            logger.info(self.heading)
            for filename, lines in sorted(self.files_with_issues.items()):
                if lines:
                    logger.info("{}: {}".format(
                        filename, ", ".join(str(x) for x in lines)
                    ))
                else:
                    logger.info(filename)
            logger.info("")

class PermissionIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "Incorrect permissions:"

    def check_file_for_issue(self, filepath):
        if not (os.access(filepath, os.X_OK) ==
                filepath.endswith((".sh", ".pl", ".py"))):
            self.files_with_issues[filepath] = None

class EndOfFileNewlineIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "Missing newline at end of file:"

    def check_file_for_issue(self, filepath):
        with open(filepath, "rb") as f:
            if not f.read().endswith(b"\n"):
                self.files_with_issues[filepath] = None

class Utf8BomIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "UTF-8 BOM present:"

    def check_file_for_issue(self, filepath):
        with open(filepath, "rb") as f:
            if f.read().startswith(codecs.BOM_UTF8):
                self.files_with_issues[filepath] = None

class LineEndingIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "Non Unix line endings:"

    def issue_with_line(self, line):
        return b"\r" in line

class TrailingWhitespaceIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "Trailing whitespace:"
        self.files_exemptions = [".md"]

    def issue_with_line(self, line):
        return line.rstrip(b"\r\n") != line.rstrip()

class TabIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "Tabs present:"
        self.files_exemptions = [
            "Makefile", "generate_visualc_files.pl"
        ]

    def issue_with_line(self, line):
        return b"\t" in line

class TodoIssueTracker(IssueTracker):

    def __init__(self):
        super().__init__()
        self.heading = "TODO present:"
        self.files_exemptions = [
            __file__, "benchmark.c", "pull_request_template.md"
        ]

    def issue_with_line(self, line):
        return b"todo" in line.lower()

class IntegrityChecker(object):

    def __init__(self, log_file):
        self.check_repo_path()
        self.logger = None
        self.setup_logger(log_file)
        self.files_to_check = (
            ".c", ".h", ".sh", ".pl", ".py", ".md", ".function", ".data",
            "Makefile", "CMakeLists.txt", "ChangeLog"
        )
        self.excluded_directories = ['.git', 'mbed-os']
        self.excluded_paths = list(map(os.path.normpath, [
            'cov-int',
            'examples',
        ]))
        self.issues_to_check = [
            PermissionIssueTracker(),
            EndOfFileNewlineIssueTracker(),
            Utf8BomIssueTracker(),
            LineEndingIssueTracker(),
            TrailingWhitespaceIssueTracker(),
            TabIssueTracker(),
            TodoIssueTracker(),
        ]

    def check_repo_path(self):
        if not all(os.path.isdir(d) for d in ["include", "library", "tests"]):
            raise Exception("Must be run from Mbed TLS root")

    def setup_logger(self, log_file, level=logging.INFO):
        self.logger = logging.getLogger()
        self.logger.setLevel(level)
        if log_file:
            handler = logging.FileHandler(log_file)
            self.logger.addHandler(handler)
        else:
            console = logging.StreamHandler()
            self.logger.addHandler(console)

    def prune_branch(self, root, d):
        if d in self.excluded_directories:
            return True
        if os.path.normpath(os.path.join(root, d)) in self.excluded_paths:
            return True
        return False

    def check_files(self):
        for root, dirs, files in os.walk("."):
            dirs[:] = sorted(d for d in dirs if not self.prune_branch(root, d))
            for filename in sorted(files):
                filepath = os.path.join(root, filename)
                if not filepath.endswith(self.files_to_check):
                    continue
                for issue_to_check in self.issues_to_check:
                    if issue_to_check.should_check_file(filepath):
                        issue_to_check.check_file_for_issue(filepath)

    def output_issues(self):
        integrity_return_code = 0
        for issue_to_check in self.issues_to_check:
            if issue_to_check.files_with_issues:
                integrity_return_code = 1
            issue_to_check.output_file_issues(self.logger)
        return integrity_return_code

def run_main():
    parser = argparse.ArgumentParser(
        description=(
            "This script checks the current state of the source code for "
            "minor issues, including incorrect file permissions, "
            "presence of tabs, non-Unix line endings, trailing whitespace, "
            "presence of UTF-8 BOM, and TODO comments. "
            "Note: requires python 3, must be run from Mbed TLS root."
        )
    )
    parser.add_argument(
        "-l", "--log_file", type=str, help="path to optional output log",
    )
    check_args = parser.parse_args()
    integrity_check = IntegrityChecker(check_args.log_file)
    integrity_check.check_files()
    return_code = integrity_check.output_issues()
    sys.exit(return_code)

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


<commit_before># Copyright (c) 2017 Ultimaker B.V.
# Cura is released under the terms of the AGPLv3 or higher.

from UM.Logger import Logger

from UM.Settings.DefinitionContainer import DefinitionContainer
from UM.Settings.InstanceContainer import InstanceContainer
from UM.Settings.ContainerRegistry import ContainerRegistry

from .GlobalStack import GlobalStack
from .ExtruderStack import ExtruderStack
from typing import Optional

##  Contains helper functions to create new machines.
class CuraStackBuilder:
    ##  Create a new instance of a machine.
    #
    #   \param name The name of the new machine.
    #   \param definition_id The ID of the machine definition to use.
    #
    #   \return The new global stack or None if an error occurred.
    @classmethod
    def createMachine(cls, name: str, definition_id: str) -> Optional[GlobalStack]:
        registry = ContainerRegistry.getInstance()
        definitions = registry.findDefinitionContainers(id = definition_id)
        if not definitions:
            Logger.log("w", "Definition {definition} was not found!", definition = definition_id)
            return None

        machine_definition = definitions[0]
        name = registry.createUniqueName("machine", "", name, machine_definition.name)
        # Make sure the new name does not collide with any definition or (quality) profile
        # createUniqueName() only looks at other stacks, but not at definitions or quality profiles
        # Note that we don't go for uniqueName() immediately because that function matches with ignore_case set to true
        if registry.findContainers(id = name):
            name = registry.uniqueName(name)

        new_global_stack = cls.createGlobalStack(
            new_stack_id = name,
            definition = machine_definition,
            quality = "default",
            material = "default",
            variant = "default",
        )

        for extruder_definition in registry.findDefinitionContainers(machine = machine_definition.id):
            position = extruder_definition.getMetaDataEntry("position", None)
            if not position:
                Logger.log("w", "Extruder definition %s specifies no position metadata entry.", extruder_definition.id)

            new_extruder_id = registry.uniqueName(extruder_definition.id)
            new_extruder = cls.createExtruderStack(
                new_extruder_id,
                definition = extruder_definition,
                machine_definition = machine_definition,
                quality = "default",
                material = "default",
                variant = "default",
                next_stack = new_global_stack
            )

        return new_global_stack

    ##  Create a new Extruder stack
    #
    #   \param new_stack_id The ID of the new stack.
    #   \param definition The definition to base the new stack on.
    #   \param machine_definition The machine definition to use for the user container.
    #   \param kwargs You can add keyword arguments to specify IDs of containers to use for a specific type, for example "variant": "0.4mm"
    #
    #   \return A new Global stack instance with the specified parameters.
    @classmethod
    def createExtruderStack(cls, new_stack_id: str, definition: DefinitionContainer, machine_definition: DefinitionContainer, **kwargs) -> ExtruderStack:
        stack = ExtruderStack(new_stack_id)
        stack.setName(definition.getName())
        stack.setDefinition(definition)
        stack.addMetaDataEntry("position", definition.getMetaDataEntry("position"))

        user_container = InstanceContainer(new_stack_id + "_user")
        user_container.addMetaDataEntry("type", "user")
        user_container.addMetaDataEntry("extruder", new_stack_id)
        from cura.CuraApplication import CuraApplication
        user_container.addMetaDataEntry("setting_version", CuraApplication.SettingVersion)
        user_container.setDefinition(machine_definition)

        stack.setUserChanges(user_container)

        if "next_stack" in kwargs:
            stack.setNextStack(kwargs["next_stack"])

        # Important! The order here matters, because that allows the stack to
        # assume the material and variant have already been set.
        if "definition_changes" in kwargs:
            stack.setDefinitionChangesById(kwargs["definition_changes"])

        if "variant" in kwargs:
            stack.setVariantById(kwargs["variant"])

        if "material" in kwargs:
            stack.setMaterialById(kwargs["material"])

        if "quality" in kwargs:
            stack.setQualityById(kwargs["quality"])

        if "quality_changes" in kwargs:
            stack.setQualityChangesById(kwargs["quality_changes"])

        # Only add the created containers to the registry after we have set all the other
        # properties. This makes the create operation more transactional, since any problems
        # setting properties will not result in incomplete containers being added.
        registry = ContainerRegistry.getInstance()
        registry.addContainer(stack)
        registry.addContainer(user_container)

        return stack

    ##  Create a new Global stack
    #
    #   \param new_stack_id The ID of the new stack.
    #   \param definition The definition to base the new stack on.
    #   \param kwargs You can add keyword arguments to specify IDs of containers to use for a specific type, for example "variant": "0.4mm"
    #
    #   \return A new Global stack instance with the specified parameters.
    @classmethod
    def createGlobalStack(cls, new_stack_id: str, definition: DefinitionContainer, **kwargs) -> GlobalStack:
        stack = GlobalStack(new_stack_id)
        stack.setDefinition(definition)

        user_container = InstanceContainer(new_stack_id + "_user")
        user_container.addMetaDataEntry("type", "user")
        user_container.addMetaDataEntry("machine", new_stack_id)
        from cura.CuraApplication import CuraApplication
        user_container.addMetaDataEntry("setting_version", CuraApplication.SettingVersion)
        user_container.setDefinition(definition)

        stack.setUserChanges(user_container)

        # Important! The order here matters, because that allows the stack to
        # assume the material and variant have already been set.
        if "definition_changes" in kwargs:
            stack.setDefinitionChangesById(kwargs["definition_changes"])

        if "variant" in kwargs:
            stack.setVariantById(kwargs["variant"])

        if "material" in kwargs:
            stack.setMaterialById(kwargs["material"])

        if "quality" in kwargs:
            stack.setQualityById(kwargs["quality"])

        if "quality_changes" in kwargs:
            stack.setQualityChangesById(kwargs["quality_changes"])

        registry = ContainerRegistry.getInstance()
        registry.addContainer(stack)
        registry.addContainer(user_container)

        return stack
<commit_msg>Always create a definition changes container for a newly created stack<commit_after># Copyright (c) 2017 Ultimaker B.V.
# Cura is released under the terms of the AGPLv3 or higher.

from UM.Logger import Logger

from UM.Settings.DefinitionContainer import DefinitionContainer
from UM.Settings.InstanceContainer import InstanceContainer
from UM.Settings.ContainerRegistry import ContainerRegistry

from .GlobalStack import GlobalStack
from .ExtruderStack import ExtruderStack
from typing import Optional

##  Contains helper functions to create new machines.
class CuraStackBuilder:
    ##  Create a new instance of a machine.
    #
    #   \param name The name of the new machine.
    #   \param definition_id The ID of the machine definition to use.
    #
    #   \return The new global stack or None if an error occurred.
    @classmethod
    def createMachine(cls, name: str, definition_id: str) -> Optional[GlobalStack]:
        registry = ContainerRegistry.getInstance()
        definitions = registry.findDefinitionContainers(id = definition_id)
        if not definitions:
            Logger.log("w", "Definition {definition} was not found!", definition = definition_id)
            return None

        machine_definition = definitions[0]
        name = registry.createUniqueName("machine", "", name, machine_definition.name)
        # Make sure the new name does not collide with any definition or (quality) profile
        # createUniqueName() only looks at other stacks, but not at definitions or quality profiles
        # Note that we don't go for uniqueName() immediately because that function matches with ignore_case set to true
        if registry.findContainers(id = name):
            name = registry.uniqueName(name)

        new_global_stack = cls.createGlobalStack(
            new_stack_id = name,
            definition = machine_definition,
            quality = "default",
            material = "default",
            variant = "default",
        )

        for extruder_definition in registry.findDefinitionContainers(machine = machine_definition.id):
            position = extruder_definition.getMetaDataEntry("position", None)
            if not position:
                Logger.log("w", "Extruder definition %s specifies no position metadata entry.", extruder_definition.id)

            new_extruder_id = registry.uniqueName(extruder_definition.id)
            new_extruder = cls.createExtruderStack(
                new_extruder_id,
                definition = extruder_definition,
                machine_definition = machine_definition,
                quality = "default",
                material = "default",
                variant = "default",
                next_stack = new_global_stack
            )

        return new_global_stack

    ##  Create a new Extruder stack
    #
    #   \param new_stack_id The ID of the new stack.
    #   \param definition The definition to base the new stack on.
    #   \param machine_definition The machine definition to use for the user container.
    #   \param kwargs You can add keyword arguments to specify IDs of containers to use for a specific type, for example "variant": "0.4mm"
    #
    #   \return A new Global stack instance with the specified parameters.
    @classmethod
    def createExtruderStack(cls, new_stack_id: str, definition: DefinitionContainer, machine_definition: DefinitionContainer, **kwargs) -> ExtruderStack:
        stack = ExtruderStack(new_stack_id)
        stack.setName(definition.getName())
        stack.setDefinition(definition)
        stack.addMetaDataEntry("position", definition.getMetaDataEntry("position"))

        user_container = InstanceContainer(new_stack_id + "_user")
        user_container.addMetaDataEntry("type", "user")
        user_container.addMetaDataEntry("extruder", new_stack_id)
        from cura.CuraApplication import CuraApplication
        user_container.addMetaDataEntry("setting_version", CuraApplication.SettingVersion)
        user_container.setDefinition(machine_definition)

        stack.setUserChanges(user_container)

        if "next_stack" in kwargs:
            stack.setNextStack(kwargs["next_stack"])

        # Important! The order here matters, because that allows the stack to
        # assume the material and variant have already been set.
        if "definition_changes" in kwargs:
            stack.setDefinitionChangesById(kwargs["definition_changes"])
        else:
            stack.setDefinitionChanges(cls._createDefinitionChangesContainer(stack, new_stack_id + "_settings"))

        if "variant" in kwargs:
            stack.setVariantById(kwargs["variant"])

        if "material" in kwargs:
            stack.setMaterialById(kwargs["material"])

        if "quality" in kwargs:
            stack.setQualityById(kwargs["quality"])

        if "quality_changes" in kwargs:
            stack.setQualityChangesById(kwargs["quality_changes"])

        # Only add the created containers to the registry after we have set all the other
        # properties. This makes the create operation more transactional, since any problems
        # setting properties will not result in incomplete containers being added.
        registry = ContainerRegistry.getInstance()
        registry.addContainer(stack)
        registry.addContainer(user_container)

        return stack

    ##  Create a new Global stack
    #
    #   \param new_stack_id The ID of the new stack.
    #   \param definition The definition to base the new stack on.
    #   \param kwargs You can add keyword arguments to specify IDs of containers to use for a specific type, for example "variant": "0.4mm"
    #
    #   \return A new Global stack instance with the specified parameters.
    @classmethod
    def createGlobalStack(cls, new_stack_id: str, definition: DefinitionContainer, **kwargs) -> GlobalStack:
        stack = GlobalStack(new_stack_id)
        stack.setDefinition(definition)

        user_container = InstanceContainer(new_stack_id + "_user")
        user_container.addMetaDataEntry("type", "user")
        user_container.addMetaDataEntry("machine", new_stack_id)
        from cura.CuraApplication import CuraApplication
        user_container.addMetaDataEntry("setting_version", CuraApplication.SettingVersion)
        user_container.setDefinition(definition)

        stack.setUserChanges(user_container)

        # Important! The order here matters, because that allows the stack to
        # assume the material and variant have already been set.
        if "definition_changes" in kwargs:
            stack.setDefinitionChangesById(kwargs["definition_changes"])
        else:
            stack.setDefinitionChanges(cls._createDefinitionChangesContainer(stack, new_stack_id + "_settings"))

        if "variant" in kwargs:
            stack.setVariantById(kwargs["variant"])

        if "material" in kwargs:
            stack.setMaterialById(kwargs["material"])

        if "quality" in kwargs:
            stack.setQualityById(kwargs["quality"])

        if "quality_changes" in kwargs:
            stack.setQualityChangesById(kwargs["quality_changes"])

        registry = ContainerRegistry.getInstance()
        registry.addContainer(stack)
        registry.addContainer(user_container)

        return stack

    @classmethod
    def _createDefinitionChangesContainer(cls, container_stack, container_name, container_index = None):
        from cura.CuraApplication import CuraApplication
        definition_changes_container = InstanceContainer(container_name)
        definition = container_stack.getBottom()
        definition_changes_container.setDefinition(definition)
        definition_changes_container.addMetaDataEntry("type", "definition_changes")
        definition_changes_container.addMetaDataEntry("setting_version", CuraApplication.SettingVersion)

        ContainerRegistry.getInstance().addContainer(definition_changes_container)
        container_stack.definitionChanges = definition_changes_container

        return definition_changes_container
<|endoftext|>


<commit_before>__version__ = '4.0.1.dev1'

if __name__ == '__main__':
    print(__version__)
<commit_msg>bump version to 4.1.0<commit_after>__version__ = '4.1.0'

if __name__ == '__main__':
    print(__version__)
<|endoftext|>


<commit_before># -*- coding: UTF-8 -*-
from __future__ import unicode_literals

from django.core.management.base import BaseCommand

from tx_highered.models import Institution

class Command(BaseCommand):
    help = (
        'Create a report card of our data quality. Run it through '
        'http://setosa.io/csv-fingerprint/ for a quick visualization.'
    )

    def handle(self, *args, **options):

        def latest_year(institution, relation):
            """Latest year institution has data, returns year or ''."""
            try:
                return getattr(institution, relation).all().order_by('-year')[0].year
            except IndexError:
                return ''

        header = ['name']
        print ','.join(header)
        for institution in Institution.objects.published():
            row = [institution.slug]
            header.append('Admissions IPEDS')
            row.append(latest_year(institution, 'admissions'))
            header.append('Admissions THECB')
            if institution.is_private:
                row.append('NA')
            else:
                row.append(latest_year(institution, 'publicadmissions'))

            header.append('Test Scores')
            row.append(latest_year(institution, 'testscores'))

            header.append('Enrollment IPEDS')
            row.append(latest_year(institution, 'enrollment'))
            header.append('Enrollment THECB')
            if institution.is_private:
                row.append('NA')
            else:
                row.append(latest_year(institution, 'publicenrollment'))

            header.append('Price Trends')
            row.append(latest_year(institution, 'pricetrends'))

            header.append('Graduation IPEDS')
            row.append(latest_year(institution, 'graduationrates'))
            header.append('Graduation THECB')
            if institution.is_private:
                row.append('NA')
            else:
                row.append(latest_year(institution, 'publicgraduationrates'))

            print ','.join(map(unicode, row))

        # print header
        # print rows
<commit_msg>split enrollment into more practical FTE/demographics columns<commit_after># -*- coding: UTF-8 -*-
from __future__ import unicode_literals

from django.core.management.base import BaseCommand
from django.db.models import Model, Q

from tx_highered.models import Institution

class Command(BaseCommand):
    help = (
        'Create a report card of our data quality. Run it through '
        'http://setosa.io/csv-fingerprint/ for a quick visualization.'
    )

    def handle(self, *args, **options):

        def latest_year(model_or_queryset, relation=None):
            """Latest year institution has data, returns year or ''."""
            if relation is not None and isinstance(model_or_queryset, Model):
                queryset = getattr(model_or_queryset, relation).all()
            else:
                queryset = model_or_queryset
            try:
                return queryset.order_by('-year')[0].year
            except IndexError:
                return ''

        header = ['name']
        print ','.join(header)
        for institution in Institution.objects.published():
            row = [institution.slug]
            header.append('Admissions IPEDS')
            row.append(latest_year(institution, 'admissions'))
            header.append('Admissions THECB')
            if institution.is_private:
                row.append('NA')
            else:
                row.append(latest_year(institution, 'publicadmissions'))

            header.append('Test Scores')
            row.append(latest_year(institution, 'testscores'))

            # Enrollment
            #
            # We only care about fulltime equivalent, which is only in IPEDS
            # for now.
            header.append('Full-Time Equivalent')
            row.append(latest_year(institution.enrollment.filter(
                fulltime_equivalent__isnull=False)))

            # Demographics
            #
            # Assume that every school at least has white or black students
            header.append('Demographics IPEDS')
            row.append(latest_year(institution.enrollment.filter(
                Q(total_percent_white__isnull=False) |
                Q(total_percent_black__isnull=False)
            )))
            header.append('Demographics THECB')
            if institution.is_private:
                row.append('NA')
            else:
                row.append(latest_year(institution.publicenrollment.filter(
                    Q(white_percent__isnull=False) |
                    Q(african_american_percent=False)
                )))

            header.append('Price Trends')
            row.append(latest_year(institution, 'pricetrends'))

            header.append('Graduation IPEDS')
            row.append(latest_year(institution, 'graduationrates'))
            header.append('Graduation THECB')
            if institution.is_private:
                row.append('NA')
            else:
                row.append(latest_year(institution, 'publicgraduationrates'))

            print ','.join(map(unicode, row))

        # print header
        # print rows
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
    :codeauthor: :email:`Alexander Schwartz <alexander.schwartz@gmx.net>`
'''

# Import python libs
from __future__ import absolute_import
import os
import tempfile

# Import Salt Testing libs
from salttesting import TestCase, skipIf
from salttesting.helpers import ensure_in_syspath
from salttesting.mock import (
    NO_MOCK,
    NO_MOCK_REASON,
    MagicMock,
    patch
)

ensure_in_syspath('../../')

# Import Salt Libs
from salt.states import archive as archive
from salt.ext.six.moves import zip  # pylint: disable=import-error,redefined-builtin

# Globals
archive.__salt__ = {}
archive.__opts__ = {"cachedir": "/tmp", "test": False}
archive.__env__ = 'test'

@skipIf(NO_MOCK, NO_MOCK_REASON)
class ArchiveTestCase(TestCase):

    def setUp(self):
        super(ArchiveTestCase, self).setUp()

    def tearDown(self):
        super(ArchiveTestCase, self).tearDown()

    def test_extracted_tar(self):
        '''
        archive.extracted tar options
        '''

        source = 'file.tar.gz'
        tmp_dir = os.path.join(tempfile.gettempdir(), 'test_archive', '')
        test_tar_opts = [
            '--no-anchored foo',
            'v -p --opt',
            '-v -p',
            '--long-opt -z',
            'z -v -weird-long-opt arg',
        ]
        ret_tar_opts = [
            ['tar', 'x', '--no-anchored', 'foo', '-f'],
            ['tar', 'xv', '-p', '--opt', '-f'],
            ['tar', 'x', '-v', '-p', '-f'],
            ['tar', 'x', '--long-opt', '-z', '-f'],
            ['tar', 'xz', '-v', '-weird-long-opt', 'arg', '-f'],
        ]

        mock_true = MagicMock(return_value=True)
        mock_false = MagicMock(return_value=False)
        ret = {'stdout': ['saltines', 'cheese'], 'stderr': 'biscuits', 'retcode': '31337', 'pid': '1337'}
        mock_run = MagicMock(return_value=ret)
        mock_source_list = MagicMock(return_value=source)

        with patch('os.path.exists', mock_true):
            with patch.dict(archive.__opts__, {'test': False,
                                               'cachedir': tmp_dir}):
                with patch.dict(archive.__salt__, {'file.directory_exists': mock_false,
                                                   'file.file_exists': mock_false,
                                                   'file.makedirs': mock_true,
                                                   'cmd.run_all': mock_run,
                                                   'file.source_list': mock_source_list}):
                    filename = os.path.join(
                        tmp_dir,
                        'files/test/_tmp_test_archive_.tar'
                    )
                    for test_opts, ret_opts in zip(test_tar_opts, ret_tar_opts):
                        ret = archive.extracted(tmp_dir,
                                                source,
                                                'tar',
                                                tar_options=test_opts)
                        ret_opts.append(filename)
                        mock_run.assert_called_with(ret_opts, cwd=tmp_dir, python_shell=False)

    def test_tar_gnutar(self):
        '''
        Tests the call of extraction with gnutar
        '''
        gnutar = MagicMock(return_value='tar (GNU tar)')
        source = 'GNU tar'
        missing = MagicMock(return_value=False)
        nop = MagicMock(return_value=True)
        run_all = MagicMock(return_value={'retcode': 0, 'stdout': 'stdout', 'stderr': 'stderr'})
        mock_source_list = MagicMock(return_value=source)

        with patch.dict(archive.__salt__, {'cmd.run': gnutar,
                                           'file.directory_exists': missing,
                                           'file.file_exists': missing,
                                           'state.single': nop,
                                           'file.makedirs': nop,
                                           'cmd.run_all': run_all,
                                           'file.source_list': mock_source_list}):
            ret = archive.extracted('/tmp/out', '/tmp/foo.tar.gz', 'tar', tar_options='xvzf', keep=True)
            self.assertEqual(ret['changes']['extracted_files'], 'stdout')

    def test_tar_bsdtar(self):
        '''
        Tests the call of extraction with bsdtar
        '''
        bsdtar = MagicMock(return_value='tar (bsdtar)')
        source = 'bsdtar'
        missing = MagicMock(return_value=False)
        nop = MagicMock(return_value=True)
        run_all = MagicMock(return_value={'retcode': 0, 'stdout': 'stdout', 'stderr': 'stderr'})
        mock_source_list = MagicMock(return_value=source)

        with patch.dict(archive.__salt__, {'cmd.run': bsdtar,
                                           'file.directory_exists': missing,
                                           'file.file_exists': missing,
                                           'state.single': nop,
                                           'file.makedirs': nop,
                                           'cmd.run_all': run_all,
                                           'file.source_list': mock_source_list}):
            ret = archive.extracted('/tmp/out', '/tmp/foo.tar.gz', 'tar', tar_options='xvzf', keep=True)
            self.assertEqual(ret['changes']['extracted_files'], 'stderr')

if __name__ == '__main__':
    from integration import run_tests
    run_tests(ArchiveTestCase, needs_daemon=False)
<commit_msg>Fixup unit test for archive module<commit_after># -*- coding: utf-8 -*-
'''
    :codeauthor: :email:`Alexander Schwartz <alexander.schwartz@gmx.net>`
'''

# Import python libs
from __future__ import absolute_import
import os
import tempfile

# Import Salt Testing libs
from salttesting import TestCase, skipIf
from salttesting.helpers import ensure_in_syspath
from salttesting.mock import (
    NO_MOCK,
    NO_MOCK_REASON,
    MagicMock,
    patch
)

ensure_in_syspath('../../')

# Import Salt Libs
from salt.states import archive as archive
from salt.ext.six.moves import zip  # pylint: disable=import-error,redefined-builtin

# Globals
archive.__salt__ = {}
archive.__grains__ = {'os': 'FooOS!'}
archive.__opts__ = {"cachedir": "/tmp", "test": False}
archive.__env__ = 'test'

@skipIf(NO_MOCK, NO_MOCK_REASON)
class ArchiveTestCase(TestCase):

    def setUp(self):
        super(ArchiveTestCase, self).setUp()

    def tearDown(self):
        super(ArchiveTestCase, self).tearDown()

    def test_extracted_tar(self):
        '''
        archive.extracted tar options
        '''

        source = 'file.tar.gz'
        tmp_dir = os.path.join(tempfile.gettempdir(), 'test_archive', '')
        test_tar_opts = [
            '--no-anchored foo',
            'v -p --opt',
            '-v -p',
            '--long-opt -z',
            'z -v -weird-long-opt arg',
        ]
        ret_tar_opts = [
            ['tar', 'x', '--no-anchored', 'foo', '-f'],
            ['tar', 'xv', '-p', '--opt', '-f'],
            ['tar', 'x', '-v', '-p', '-f'],
            ['tar', 'x', '--long-opt', '-z', '-f'],
            ['tar', 'xz', '-v', '-weird-long-opt', 'arg', '-f'],
        ]

        mock_true = MagicMock(return_value=True)
        mock_false = MagicMock(return_value=False)
        ret = {'stdout': ['saltines', 'cheese'], 'stderr': 'biscuits', 'retcode': '31337', 'pid': '1337'}
        mock_run = MagicMock(return_value=ret)
        mock_source_list = MagicMock(return_value=source)

        with patch('os.path.exists', mock_true):
            with patch.dict(archive.__opts__, {'test': False,
                                               'cachedir': tmp_dir}):
                with patch.dict(archive.__salt__, {'file.directory_exists': mock_false,
                                                   'file.file_exists': mock_false,
                                                   'file.makedirs': mock_true,
                                                   'cmd.run_all': mock_run,
                                                   'file.source_list': mock_source_list}):
                    filename = os.path.join(
                        tmp_dir,
                        'files/test/_tmp_test_archive_.tar'
                    )
                    for test_opts, ret_opts in zip(test_tar_opts, ret_tar_opts):
                        ret = archive.extracted(tmp_dir,
                                                source,
                                                'tar',
                                                tar_options=test_opts)
                        ret_opts.append(filename)
                        mock_run.assert_called_with(ret_opts, cwd=tmp_dir, python_shell=False)

    def test_tar_gnutar(self):
        '''
        Tests the call of extraction with gnutar
        '''
        gnutar = MagicMock(return_value='tar (GNU tar)')
        source = 'GNU tar'
        missing = MagicMock(return_value=False)
        nop = MagicMock(return_value=True)
        run_all = MagicMock(return_value={'retcode': 0, 'stdout': 'stdout', 'stderr': 'stderr'})
        mock_source_list = MagicMock(return_value=source)

        with patch.dict(archive.__salt__, {'cmd.run': gnutar,
                                           'file.directory_exists': missing,
                                           'file.file_exists': missing,
                                           'state.single': nop,
                                           'file.makedirs': nop,
                                           'cmd.run_all': run_all,
                                           'file.source_list': mock_source_list}):
            ret = archive.extracted('/tmp/out', '/tmp/foo.tar.gz', 'tar', tar_options='xvzf', keep=True)
            self.assertEqual(ret['changes']['extracted_files'], 'stdout')

    def test_tar_bsdtar(self):
        '''
        Tests the call of extraction with bsdtar
        '''
        bsdtar = MagicMock(return_value='tar (bsdtar)')
        source = 'bsdtar'
        missing = MagicMock(return_value=False)
        nop = MagicMock(return_value=True)
        run_all = MagicMock(return_value={'retcode': 0, 'stdout': 'stdout', 'stderr': 'stderr'})
        mock_source_list = MagicMock(return_value=source)

        with patch.dict(archive.__salt__, {'cmd.run': bsdtar,
                                           'file.directory_exists': missing,
                                           'file.file_exists': missing,
                                           'state.single': nop,
                                           'file.makedirs': nop,
                                           'cmd.run_all': run_all,
                                           'file.source_list': mock_source_list}):
            ret = archive.extracted('/tmp/out', '/tmp/foo.tar.gz', 'tar', tar_options='xvzf', keep=True)
            self.assertEqual(ret['changes']['extracted_files'], 'stderr')

if __name__ == '__main__':
    from integration import run_tests
    run_tests(ArchiveTestCase, needs_daemon=False)
<|endoftext|>


<commit_before>#! /usr/local/bin/python
"""
    See LICENSE file for copyright and license details.
"""
from datetime import datetime

from database.databaseaccess import DatabaseAccess
from modules.core_module import CoreModule
from modules.statement import Statement
from modules.constant import *
from modules.function import *
from generic.modules.function import *

class Finance(CoreModule):
    """
        Finance class.
    """

    def create_statements(self, input_fields):
        """
            Creates the records needed for Table.FINANCE
            and returns them as a Statement object.
        """
        #TODO: move all create statements tot their own module?
        #This code here is already way to big to my liking.
        try:
            dba = DatabaseAccess(self.config)
            date_created = current_date()
            date_modified = current_date()
            statement_finance = Statement(Table.FINANCE)
            records = 0
            currency_exchange_id = dba.first_currency_exchange_id_from_latest()
            rate_id = dba.first_rate_id_from_latest()
            for fields in input_fields:
                #TODO: what to do with the subcategory_id? and cat_id?
                account_id = dba.account_id_from_account(fields['i_account'])

                #NOTE: in the database, the first values in the tables of the
                #below id's, are empty/dummy values, used for when we are not
                #dealing with stocks.
                market_id = 1
                stock_name_id = 1
                rate_id = 1
                if deals_with_stocks(fields['i_account_from'], fields['i_account_to']):
                    if fields['i_market_name'] != '':
                        market_id = dba.market_id_from_market(fields['i_market_name'])
                    if fields['i_stock_name'] != '':
                        stock_name_id = dba.stock_name_id_from_stock_name(
                                fields['i_stock_name'], market_id)
                    rate_id = dba.get_latest_rate_id()
                finance_record = dba.get_specific_finance_record(
                    fields['i_date'],
                    account_id,
                    category_id,
                    subcategory_id,
                    Decimal(fields['i_amount']),
                    fields['i_comment'],
                    stock_name_id,
                    int(fields['i_shares']),
                    Decimal(fields['i_shares']),
                    Decimal(fields['i_tax']),
                    Decimal(fields['i_commission']))
                amount_value = DEFAULT_DECIMAL
                if is_negative_amount(fields['i_amount']):
                    amount_value = -1.0 * Decimal(fields['i_amount'])
                else:
                    amount_value = Decimal(fields['i_amount'])
                if finance_record is None:
                        records = records + 1
                        statement_finance.add(
                            records,
                            {
                                'finance_id':None,
                                'date':fields['i_date'],
                                'year':fields['i_date'].year,
                                'month':fields['i_date'].month,
                                'day':fields['i_date'].day,
                                'account_id':account_id,
                                'amount': amount_value,
                                'comment':fields['i_comment'],
                                'stock_name_id':stock_name_id,
                                'shares':int(fields['i_shares']),
                                'price':Decimal(fields['i_price']),
                                'tax':Decimal(fields['i_tax']),
                                'commission':Decimal(fields['i_commission']),
                                'active':1,
                                'rate_id':rate_id,
                                'currency_exchange_id':currency_exchange_id,
                                'date_created':date_created,
                                'date_modified':date_modified
                            }
                        )
                        currency_exchange_id = currency_exchange_id + 1
            return statement_finance
        except Exception as ex:
            print(Error.CREATE_STATEMENTS_TABLE_FINANCE, ex)
        finally:
            dba = None
<commit_msg>Using a list instead of a dict.<commit_after>#! /usr/local/bin/python
"""
    See LICENSE file for copyright and license details.
"""
from datetime import datetime

from database.databaseaccess import DatabaseAccess
from modules.core_module import CoreModule
from modules.statement import Statement
from modules.constant import *
from modules.function import *
from generic.modules.function import *

class Finance(CoreModule):
    """
        Finance class.
    """

    def create_statements(self, input_fields):
        """
            Creates the records needed for Table.FINANCE
            and returns them as a Statement object.
        """
        #TODO: move all create statements tot their own module?
        #This code here is already way to big to my liking.
        try:
            dba = DatabaseAccess(self.config)
            date_created = current_date()
            date_modified = current_date()
            statement_finance = Statement(Table.FINANCE)
            records = 0
            currency_exchange_id = dba.first_currency_exchange_id_from_latest()
            rate_id = dba.first_rate_id_from_latest()
            for fields in input_fields:
                account_from_id = dba.account_id_from_account(fields[Input.ACCOUNT_FROM])
                account_to_id = dba.account_id_from_account(fields[Input.ACCOUNT_TO])

                #NOTE: in the database, the first values in the tables of the
                #below id's, are empty/dummy values, used for when we are not
                #dealing with stocks.
                market_id = 1
                stock_name_id = 1
                rate_id = 1
                if deals_with_stocks(fields[Input.ACCOUNT_FROM], fields[Input.ACCOUNT_TO]):
                    if fields[Input.MARKET] != '':
                        market_id = dba.market_id_from_market(fields[Input.MARKET])
                    if fields[Input.COMMODITY] != '':
                        commodity_name_id = dba.commodity_name_id_from_commodity_name(
                                fields[Input.COMMODITY], market_id)
                    rate_id = dba.get_latest_rate_id()
                finance_record = dba.get_specific_finance_record(
                    fields[Input.DATE],
                    account_from_id,
                    account_to_id,
                    fields[Input.AMOUNT],
                    fields[Input.COMMENT],
                    commodity_name_id,
                    fields[Input.QUANTITY],
                    fields[Input.TAX],
                    fields[Input.COMMISSION])
                amount_value = fields[Input.AMOUNT])
                if is_negative_amount(amount_value):
                    amount_value = -1.0 * amount_value

                if finance_record is None:
                        records = records + 1
                        statement_finance.add(
                            records,
                            {
                                'finance_id':None,
                                'date':fields[Input.DATE],
                                'year':fields[Input.DATE].year,
                                'month':fields[Input.DATE].month,
                                'day':fields[Input.DATE].day,
                                'account_from_id':account_from_id,
                                'account_to_id':account_to_id,
                                'amount': amount_value,
                                'comment':fields[Input.COMMENT],
                                'stock_name_id':stock_name_id,
                                'shares':fields[Input.QUANTITY],
                                'price':fields[Input.PRICE],
                                'tax':fields[Input.TAX],
                                'commission':fields[Input.COMMISSION],
                                'active':1,
                                'rate_id':rate_id,
                                'currency_exchange_id':currency_exchange_id,
                                'date_created':date_created,
                                'date_modified':date_modified
                            }
                        )
                        currency_exchange_id = currency_exchange_id + 1
            return statement_finance
        except Exception as ex:
            print(Error.CREATE_STATEMENTS_TABLE_FINANCE, ex)
        finally:
            dba = None
<|endoftext|>


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

import os
import subprocess
import sys

if __name__ == '__main__':

    ref_name = sys.argv[1]
    old_commit = sys.argv[2]
    new_commit = sys.argv[3]
    # check 1: allow branches other than master
    if ref_name != 'refs/heads/master':
        sys.exit()
    # no need to check at this point if old_commit or new_commit are 0, master can't be deleted or created
    # check 2: allow MarkUs through
    if os.environ.get('REMOTE_USER') is None:
        sys.exit()
    # check 3: forbid creating/deleting top-level files/directories (make an exception for .gitignore)
    old_ls = subprocess.run(['git', 'ls-tree', '--name-only', old_commit], stdout=subprocess.PIPE,
                            universal_newlines=True)
    new_ls = subprocess.run(['git', 'ls-tree', '--name-only', new_commit], stdout=subprocess.PIPE,
                            universal_newlines=True)
    old_ls = [line for line in old_ls.stdout.splitlines() if line.strip() != '.gitignore']
    new_ls = [line for line in new_ls.stdout.splitlines() if line.strip() != '.gitignore']
    if old_ls != new_ls:
        print('[MARKUS] Error: creating/deleting top level files and directories is not allowed on master!')
        sys.exit(1)
    # check 4: forbid modifying top-level files
    changes = subprocess.run(['git', 'diff', '--name-only', '--no-renames', old_commit, new_commit],
                             stdout=subprocess.PIPE, universal_newlines=True)
    if any(os.sep not in change for change in changes.stdout.splitlines()):
        print('[MARKUS] Error: modifying top level files is not allowed on master!')
        sys.exit(1)
<commit_msg>git-hooks: allow users to modify .gitignore at top level as well<commit_after>#!/usr/bin/env python3

import os
import subprocess
import sys

if __name__ == '__main__':

    ref_name = sys.argv[1]
    old_commit = sys.argv[2]
    new_commit = sys.argv[3]
    # check 1: allow branches other than master
    if ref_name != 'refs/heads/master':
        sys.exit()
    # no need to check at this point if old_commit or new_commit are 0, master can't be deleted or created
    # check 2: allow MarkUs through
    if os.environ.get('REMOTE_USER') is None:
        sys.exit()
    # check 3: forbid creating/deleting top-level files/directories (make an exception for .gitignore)
    old_ls = subprocess.run(['git', 'ls-tree', '--name-only', old_commit], stdout=subprocess.PIPE,
                            universal_newlines=True)
    new_ls = subprocess.run(['git', 'ls-tree', '--name-only', new_commit], stdout=subprocess.PIPE,
                            universal_newlines=True)
    old_ls = [line for line in old_ls.stdout.splitlines() if line != '.gitignore']
    new_ls = [line for line in new_ls.stdout.splitlines() if line != '.gitignore']
    if old_ls != new_ls:
        print('[MARKUS] Error: creating/deleting top level files and directories is not allowed on master!')
        sys.exit(1)
    # check 4: forbid modifying top-level files (make and exception for .gitignore)
    changes = subprocess.run(['git', 'diff', '--name-only', '--no-renames', old_commit, new_commit],
                             stdout=subprocess.PIPE, universal_newlines=True)
    if any(os.sep not in change for change in changes.stdout.splitlines() if change != '.gitignore'):
        print('[MARKUS] Error: modifying top level files is not allowed on master!')
        sys.exit(1)
<|endoftext|>


<commit_before>import sys
import os
import random
import datetime
import pygame
import pandas as pd
import PyQt4.QtCore as QtCore
import PyQt4.QtGui as QtGui

from designer import battery_window_qt
from interface import about_dialog, settings_window
from tasks import ant, mrt, sart, ravens, digitspan_backwards

class BatteryWindow(QtGui.QMainWindow, battery_window_qt.Ui_CognitiveBattery):
    # TODO move this class to a separate file
    def __init__(self, cur_directory, first_run, res_width, res_height):
        super(BatteryWindow, self).__init__()

        # Setup the main window UI
        self.setupUi(self)

        # Get screen resolution
        self.res_width = res_width
        self.res_height = res_height

        # Create/open settings file with no registry fallback
        self.settings = QtCore.QSettings("settings.ini",
                                         QtCore.QSettings.IniFormat)
        self.settings.setFallbacksEnabled(False)

        # If first run, store some default settings
        if first_run:
            # Main window size and position
            self.save_settings_window(self.size(), QtCore.QPoint(100, 100))

            # Settings - Task Windows
            self.settings.beginGroup("TaskWindows")
            self.settings.setValue('fullscreen', False)
            self.settings.setValue('width', 1280)
            self.settings.setValue('height', 1024)
            self.settings.endGroup()

        # Set initial window size/pos from saved settings
        self.settings.beginGroup("MainWindow")
        self.resize(self.settings.value("size").toSize())
        self.move(self.settings.value("pos").toPoint())
        self.settings.endGroup()

        # Initialize task settings
        self.task_fullscreen = None
        self.task_width = None
        self.task_height = None

        # Initialize the about and settings window objects
        self.about = None
        self.settings_window = None

        # Initialize pygame screen
        self.pygame_screen = None

        # Get current directory
        self.directory = cur_directory

        # Make data folder if it doesnt exist
        self.dataPath = self.directory + "\data\\"
        if not os.path.isdir(self.dataPath):
            os.makedirs(self.dataPath)

        # Get list of existing experiment names/IDs currently in use
        self.tempNames = []
        for fileName in os.listdir(self.dataPath):
            if fileName.endswith(".xls"):
                self.tempNames.append(fileName.split("_")[0])

        self.expNames = list(set(self.tempNames))

        # Autocomplete for experiment name/ID
        self.experiments = QtCore.QString(';'.join(self.expNames)).split(";")
        self.completer = QtGui.QCompleter(self.experiments,
                                          self.experimentIDBox)
        self.experimentIDBox.setCompleter(self.completer)

        # Handle menu bar item click events
        self.actionExit.triggered.connect(self.close)
        self.actionSettings.triggered.connect(self.show_settings)
        self.actionDocumentation.triggered.connect(self.show_documentation)
        self.actionLicense.triggered.connect(self.show_license)
        self.actionContribute.triggered.connect(self.show_contribute)
        self.actionBrowse_Issues.triggered.connect(self.show_browse_issues)
        self.actionReport_Bug.triggered.connect(self.show_new_issue)
        self.actionRequest_Feature.triggered.connect(self.show_new_issue)
        self.actionAbout.triggered.connect(self.show_about)

        # Bind button events
        self.cancelButton.clicked.connect(self.close)
        self.startButton.clicked.connect(self.start)
        self.randomOrderCheck.clicked.connect(self.random_order_selected)
        self.selectAllButton.clicked.connect(self.select_all)
        self.deselectAllButton.clicked.connect(self.deselect_all)
        self.upButton.clicked.connect(self.move_up)
        self.downButton.clicked.connect(self.move_down)

    # Open web browser to the documentation page
    def show_documentation(self):
        QtGui.QDesktopServices.openUrl(
            QtCore.QUrl("https://github.com/sho-87/cognitive-battery"))

    # Open web browser to the license page
    def show_license(self):
        QtGui.QDesktopServices.openUrl(
            QtCore.QUrl("https://github.com/sho-87/cognitive-battery/blob/master/LICENSE"))

    # Open web browser to the github develop branch for contribution
    def show_contribute(self):
        QtGui.QDesktopServices.openUrl(
            QtCore.QUrl("https://github.com/sho-87/cognitive-battery/tree/develop"))

    # Open web browser to the github issues page
    def show_browse_issues(self):
        QtGui.QDesktopServices.openUrl(
            QtCore.QUrl("https://github.com/sho-87/cognitive-battery/issues"))

    # Open web browser to the github new issue post
    def show_new_issue(self):
        QtGui.QDesktopServices.openUrl(
            QtCore.QUrl("https://github.com/sho-87/cognitive-battery/issues/new"))

    # Create a new SettingsWindow object and display it
    def show_settings(self):
        # If the settings window does not exist, create one
        if self.settings_window is None:
            self.settings_window = settings_window.SettingsWindow(self)
            self.settings_window.show()
            self.settings_window.finished.connect(
                lambda: setattr(self, 'settings_window', None))
        # If settings window exists, bring it to the front
        else:
            self.settings_window.activateWindow()
            self.settings_window.raise_()

    # Create a new AboutDialog object and display it
    def show_about(self):
        # If the about dialog does not exist, create one
        if self.about is None:
            self.about = about_dialog.AboutDialog(self)
            self.about.show()
            self.about.finished.connect(lambda: setattr(self, 'about', None))
        # If about dialog exists, bring it to the front
        else:
            self.about.activateWindow()
            self.about.raise_()

    def error_dialog(self, message):
        QtGui.QMessageBox.warning(self, 'Error', message)

    def random_order_selected(self):
        if self.randomOrderCheck.isChecked():
            self.upButton.setEnabled(False)
            self.downButton.setEnabled(False)
            return True
        else:
            self.upButton.setEnabled(True)
            self.downButton.setEnabled(True)
            return False

    def select_all(self):
        for index in range(self.taskList.count()):
            self.taskList.item(index).setCheckState(2)

    def deselect_all(self):
        for index in range(self.taskList.count()):
            self.taskList.item(index).setCheckState(0)

    def move_up(self):
        currentRow = self.taskList.currentRow()
        currentItem = self.taskList.takeItem(currentRow)
        self.taskList.insertItem(currentRow - 1, currentItem)
        self.taskList.setCurrentItem(currentItem)

    def move_down(self):
        currentRow = self.taskList.currentRow()
        currentItem = self.taskList.takeItem(currentRow)
        self.taskList.insertItem(currentRow + 1, currentItem)
        self.taskList.setCurrentItem(currentItem)

    # TODO add memory for previously selected tasks
    # Save window size/position to settings file
    def save_settings_window(self, size, pos):
        self.settings.beginGroup("MainWindow")
        self.settings.setValue('size', size)
        self.settings.setValue('pos', pos)
        self.settings.endGroup()

    # Get task window settings from file
    def get_task_settings(self):
        self.settings.beginGroup("TaskWindows")
        self.task_fullscreen = self.settings.value("fullscreen").toBool()
        self.task_width = self.settings.value("width").toInt()[0]
        self.task_height = self.settings.value("height").toInt()[0]

    # Redefine the closeEvent method
    def closeEvent(self, event):
        self.save_settings_window(self.size(), self.pos())

        event.accept()

        sys.exit(0)  # This closes any open pygame windows

    def start(self):
        # Store input values
        self.subNum = self.subNumBox.text()
        self.experimentID = self.experimentIDBox.text()
        self.condition = self.conditionBox.text()
        self.age = self.ageBox.text()
        self.ra = self.raBox.text()
        self.datetime = datetime.datetime.now().strftime("%Y-%m-%d %H:%M")

        if self.maleRadio.isChecked():
            self.sex = 'male'
        elif self.femaleRadio.isChecked():
            self.sex = 'female'

        # Get *selected* tasks and task order
        self.tasks = []
        for index in range(self.taskList.count()):
            # State 2 is set when item is selected
            if self.taskList.item(index).checkState() == 2:
                # Add selected task to task list
                self.tasks.append(str(self.taskList.item(index).text()))

        # Check to see if a random order is desired
        # If so, shuffle tasks
        if self.random_order_selected():
            random.shuffle(self.tasks)

        # Check for required inputs
        if not self.tasks:
            self.error_dialog('No tasks selected')
        elif not self.ra:
            self.error_dialog('Please enter RA name...')
        elif not self.subNum:
            self.error_dialog('Please enter a subject number...')
        elif not self.experimentID:
            self.error_dialog('Please enter an experiment ID...')
        elif not self.condition:
            self.error_dialog('Please enter a condition number...')
        elif not self.age:
            self.error_dialog('Please enter an age...')
        elif not self.maleRadio.isChecked() and not self.femaleRadio.isChecked():
            self.error_dialog('Please select a sex...')
        else:
            # Store subject info into a dataframe
            self.subjectInfo = pd.DataFrame(
                data=[(str(self.datetime), str(self.subNum),
                       str(self.experimentID), str(self.condition),
                       int(self.age), str(self.sex), str(self.ra),
                       ', '.join(self.tasks))],
                columns=['datetime', 'subNum', 'expID', 'condition',
                         'age', 'sex', 'RA', 'tasks']
            )

            # Set the output file name
            self.datafileName = "%s_%s_%s.xls" % (self.experimentID,
                                                  self.subNum, self.condition)

            # Check if file already exists
            # TODO use OS path join to create paths
            if os.path.isfile(self.dataPath + self.datafileName):
                self.error_dialog('Data file already exists!')
            else:
                # TODO save each experiment data to their own directory
                # Create the excel writer object and save the file
                self.writer = pd.ExcelWriter(self.dataPath + self.datafileName)
                self.subjectInfo.to_excel(self.writer, 'info', index=False)
                self.writer.save()

                # Minimize battery UI
                self.showMinimized()

                # Get most recent task window settings from file
                self.get_task_settings()

                # Center all pygame windows if not fullscreen
                if not self.task_fullscreen:
                    pos_x = str(self.res_width/2 - self.task_width/2)
                    pos_y = str(self.res_height/2 - self.task_height/2)

                    os.environ['SDL_VIDEO_WINDOW_POS'] = \
                        "%s, %s" % (pos_x, pos_y)

                # Initialize pygame
                pygame.init()

                # Create primary task window
                # pygame_screen is passed to each task as the display window
                if self.task_fullscreen:
                    self.pygame_screen = pygame.display.set_mode(
                        (0, 0), pygame.FULLSCREEN)
                else:
                    self.pygame_screen = pygame.display.set_mode(
                        (self.task_width, self.task_height))

                # Run each task
                # Return and save their output to dataframe/excel
                # TODO move data saving to end of each individual task module
                for task in self.tasks:
                    if task == "Attention Network Test (ANT)":
                        # Set number of blocks for ANT
                        antTask = ant.ANT(self.pygame_screen, blocks=3)
                        # Run ANT
                        self.antData = antTask.run()
                        # Save ANT data to excel
                        self.antData.to_excel(self.writer, 'ANT', index=False)
                    elif task == "Mental Rotation Task":
                        mrtTask = mrt.MRT(self.pygame_screen)
                        # Run MRT
                        self.mrtData = mrtTask.run()
                        # Save MRT data to excel
                        self.mrtData.to_excel(self.writer, 'MRT', index=False)
                    elif task == "Sustained Attention to Response Task (SART)":
                        sartTask = sart.SART(self.pygame_screen)
                        # Run SART
                        self.sartData = sartTask.run()
                        # Save SART data to excel
                        self.sartData.to_excel(self.writer, 'SART', index=False)
                    elif task == "Digit Span (backwards)":
                        digitspanBackwardsTask = \
                            digitspan_backwards.DigitspanBackwards(
                                self.pygame_screen)
                        # Run Digit span (Backwards)
                        self.digitspanBackwardsData = digitspanBackwardsTask.run()
                        # Save digit span (backwards) data to excel
                        self.digitspanBackwardsData.to_excel(self.writer, 'Digit span (backwards)', index=False)
                    elif task == "Raven's Progressive Matrices":
                        ravensTask = ravens.Ravens(self.pygame_screen,
                                                   start=9, numTrials=12)
                        # Run Raven's Matrices
                        self.ravensData = ravensTask.run()
                        # Save ravens data to excel
                        self.ravensData.to_excel(self.writer, 'Ravens Matrices', index=False)

                    # Save excel file
                    self.writer.save()

                # Quit pygame
                pygame.quit()

                print "--- Experiment complete"
                self.close()

def main():
    # Get current directory
    cur_directory = os.path.dirname(os.path.realpath(__file__))

    # Check if settings file exists. If not, this is a first run
    first_run = not os.path.isfile(os.path.join(cur_directory, "settings.ini"))

    # Create main app window
    app = QtGui.QApplication(sys.argv)
    screen_resolution = app.desktop().screenGeometry()

    battery_window = BatteryWindow(cur_directory,
                                   first_run,
                                   screen_resolution.width(),
                                   screen_resolution.height())
    battery_window.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()
<commit_msg>Import Sternberg task<commit_after>import sys
import os
import random
import datetime
import pygame
import pandas as pd
import PyQt4.QtCore as QtCore
import PyQt4.QtGui as QtGui

from designer import battery_window_qt
from interface import about_dialog, settings_window
from tasks import ant, mrt, sart, ravens, digitspan_backwards, sternberg

class BatteryWindow(QtGui.QMainWindow, battery_window_qt.Ui_CognitiveBattery):
    # TODO move this class to a separate file
    def __init__(self, cur_directory, first_run, res_width, res_height):
        super(BatteryWindow, self).__init__()

        # Setup the main window UI
        self.setupUi(self)

        # Get screen resolution
        self.res_width = res_width
        self.res_height = res_height

        # Create/open settings file with no registry fallback
        self.settings = QtCore.QSettings("settings.ini",
                                         QtCore.QSettings.IniFormat)
        self.settings.setFallbacksEnabled(False)

        # If first run, store some default settings
        if first_run:
            # Main window size and position
            self.save_settings_window(self.size(), QtCore.QPoint(100, 100))

            # Settings - Task Windows
            self.settings.beginGroup("TaskWindows")
            self.settings.setValue('fullscreen', False)
            self.settings.setValue('width', 1280)
            self.settings.setValue('height', 1024)
            self.settings.endGroup()

        # Set initial window size/pos from saved settings
        self.settings.beginGroup("MainWindow")
        self.resize(self.settings.value("size").toSize())
        self.move(self.settings.value("pos").toPoint())
        self.settings.endGroup()

        # Initialize task settings
        self.task_fullscreen = None
        self.task_width = None
        self.task_height = None

        # Initialize the about and settings window objects
        self.about = None
        self.settings_window = None

        # Initialize pygame screen
        self.pygame_screen = None

        # Get current directory
        self.directory = cur_directory

        # Make data folder if it doesnt exist
        self.dataPath = self.directory + "\data\\"
        if not os.path.isdir(self.dataPath):
            os.makedirs(self.dataPath)

        # Get list of existing experiment names/IDs currently in use
        self.tempNames = []
        for fileName in os.listdir(self.dataPath):
            if fileName.endswith(".xls"):
                self.tempNames.append(fileName.split("_")[0])

        self.expNames = list(set(self.tempNames))

        # Autocomplete for experiment name/ID
        self.experiments = QtCore.QString(';'.join(self.expNames)).split(";")
        self.completer = QtGui.QCompleter(self.experiments,
                                          self.experimentIDBox)
        self.experimentIDBox.setCompleter(self.completer)

        # Handle menu bar item click events
        self.actionExit.triggered.connect(self.close)
        self.actionSettings.triggered.connect(self.show_settings)
        self.actionDocumentation.triggered.connect(self.show_documentation)
        self.actionLicense.triggered.connect(self.show_license)
        self.actionContribute.triggered.connect(self.show_contribute)
        self.actionBrowse_Issues.triggered.connect(self.show_browse_issues)
        self.actionReport_Bug.triggered.connect(self.show_new_issue)
        self.actionRequest_Feature.triggered.connect(self.show_new_issue)
        self.actionAbout.triggered.connect(self.show_about)

        # Bind button events
        self.cancelButton.clicked.connect(self.close)
        self.startButton.clicked.connect(self.start)
        self.randomOrderCheck.clicked.connect(self.random_order_selected)
        self.selectAllButton.clicked.connect(self.select_all)
        self.deselectAllButton.clicked.connect(self.deselect_all)
        self.upButton.clicked.connect(self.move_up)
        self.downButton.clicked.connect(self.move_down)

    # Open web browser to the documentation page
    def show_documentation(self):
        QtGui.QDesktopServices.openUrl(
            QtCore.QUrl("https://github.com/sho-87/cognitive-battery"))

    # Open web browser to the license page
    def show_license(self):
        QtGui.QDesktopServices.openUrl(
            QtCore.QUrl("https://github.com/sho-87/cognitive-battery/blob/master/LICENSE"))

    # Open web browser to the github develop branch for contribution
    def show_contribute(self):
        QtGui.QDesktopServices.openUrl(
            QtCore.QUrl("https://github.com/sho-87/cognitive-battery/tree/develop"))

    # Open web browser to the github issues page
    def show_browse_issues(self):
        QtGui.QDesktopServices.openUrl(
            QtCore.QUrl("https://github.com/sho-87/cognitive-battery/issues"))

    # Open web browser to the github new issue post
    def show_new_issue(self):
        QtGui.QDesktopServices.openUrl(
            QtCore.QUrl("https://github.com/sho-87/cognitive-battery/issues/new"))

    # Create a new SettingsWindow object and display it
    def show_settings(self):
        # If the settings window does not exist, create one
        if self.settings_window is None:
            self.settings_window = settings_window.SettingsWindow(self)
            self.settings_window.show()
            self.settings_window.finished.connect(
                lambda: setattr(self, 'settings_window', None))
        # If settings window exists, bring it to the front
        else:
            self.settings_window.activateWindow()
            self.settings_window.raise_()

    # Create a new AboutDialog object and display it
    def show_about(self):
        # If the about dialog does not exist, create one
        if self.about is None:
            self.about = about_dialog.AboutDialog(self)
            self.about.show()
            self.about.finished.connect(lambda: setattr(self, 'about', None))
        # If about dialog exists, bring it to the front
        else:
            self.about.activateWindow()
            self.about.raise_()

    def error_dialog(self, message):
        QtGui.QMessageBox.warning(self, 'Error', message)

    def random_order_selected(self):
        if self.randomOrderCheck.isChecked():
            self.upButton.setEnabled(False)
            self.downButton.setEnabled(False)
            return True
        else:
            self.upButton.setEnabled(True)
            self.downButton.setEnabled(True)
            return False

    def select_all(self):
        for index in range(self.taskList.count()):
            self.taskList.item(index).setCheckState(2)

    def deselect_all(self):
        for index in range(self.taskList.count()):
            self.taskList.item(index).setCheckState(0)

    def move_up(self):
        currentRow = self.taskList.currentRow()
        currentItem = self.taskList.takeItem(currentRow)
        self.taskList.insertItem(currentRow - 1, currentItem)
        self.taskList.setCurrentItem(currentItem)

    def move_down(self):
        currentRow = self.taskList.currentRow()
        currentItem = self.taskList.takeItem(currentRow)
        self.taskList.insertItem(currentRow + 1, currentItem)
        self.taskList.setCurrentItem(currentItem)

    # TODO add memory for previously selected tasks
    # Save window size/position to settings file
    def save_settings_window(self, size, pos):
        self.settings.beginGroup("MainWindow")
        self.settings.setValue('size', size)
        self.settings.setValue('pos', pos)
        self.settings.endGroup()

    # Get task window settings from file
    def get_task_settings(self):
        self.settings.beginGroup("TaskWindows")
        self.task_fullscreen = self.settings.value("fullscreen").toBool()
        self.task_width = self.settings.value("width").toInt()[0]
        self.task_height = self.settings.value("height").toInt()[0]

    # Redefine the closeEvent method
    def closeEvent(self, event):
        self.save_settings_window(self.size(), self.pos())

        event.accept()

        sys.exit(0)  # This closes any open pygame windows

    def start(self):
        # Store input values
        self.subNum = self.subNumBox.text()
        self.experimentID = self.experimentIDBox.text()
        self.condition = self.conditionBox.text()
        self.age = self.ageBox.text()
        self.ra = self.raBox.text()
        self.datetime = datetime.datetime.now().strftime("%Y-%m-%d %H:%M")

        if self.maleRadio.isChecked():
            self.sex = 'male'
        elif self.femaleRadio.isChecked():
            self.sex = 'female'

        # Get *selected* tasks and task order
        self.tasks = []
        for index in range(self.taskList.count()):
            # State 2 is set when item is selected
            if self.taskList.item(index).checkState() == 2:
                # Add selected task to task list
                self.tasks.append(str(self.taskList.item(index).text()))

        # Check to see if a random order is desired
        # If so, shuffle tasks
        if self.random_order_selected():
            random.shuffle(self.tasks)

        # Check for required inputs
        if not self.tasks:
            self.error_dialog('No tasks selected')
        elif not self.ra:
            self.error_dialog('Please enter RA name...')
        elif not self.subNum:
            self.error_dialog('Please enter a subject number...')
        elif not self.experimentID:
            self.error_dialog('Please enter an experiment ID...')
        elif not self.condition:
            self.error_dialog('Please enter a condition number...')
        elif not self.age:
            self.error_dialog('Please enter an age...')
        elif not self.maleRadio.isChecked() and not self.femaleRadio.isChecked():
            self.error_dialog('Please select a sex...')
        else:
            # Store subject info into a dataframe
            self.subjectInfo = pd.DataFrame(
                data=[(str(self.datetime), str(self.subNum),
                       str(self.experimentID), str(self.condition),
                       int(self.age), str(self.sex), str(self.ra),
                       ', '.join(self.tasks))],
                columns=['datetime', 'subNum', 'expID', 'condition',
                         'age', 'sex', 'RA', 'tasks']
            )

            # Set the output file name
            self.datafileName = "%s_%s_%s.xls" % (self.experimentID,
                                                  self.subNum, self.condition)

            # Check if file already exists
            # TODO use OS path join to create paths
            if os.path.isfile(self.dataPath + self.datafileName):
                self.error_dialog('Data file already exists!')
            else:
                # TODO save each experiment data to their own directory
                # Create the excel writer object and save the file
                self.writer = pd.ExcelWriter(self.dataPath + self.datafileName)
                self.subjectInfo.to_excel(self.writer, 'info', index=False)
                self.writer.save()

                # Minimize battery UI
                self.showMinimized()

                # Get most recent task window settings from file
                self.get_task_settings()

                # Center all pygame windows if not fullscreen
                if not self.task_fullscreen:
                    pos_x = str(self.res_width/2 - self.task_width/2)
                    pos_y = str(self.res_height/2 - self.task_height/2)

                    os.environ['SDL_VIDEO_WINDOW_POS'] = \
                        "%s, %s" % (pos_x, pos_y)

                # Initialize pygame
                pygame.init()

                # Create primary task window
                # pygame_screen is passed to each task as the display window
                if self.task_fullscreen:
                    self.pygame_screen = pygame.display.set_mode(
                        (0, 0), pygame.FULLSCREEN)
                else:
                    self.pygame_screen = pygame.display.set_mode(
                        (self.task_width, self.task_height))

                # Run each task
                # Return and save their output to dataframe/excel
                # TODO move data saving to end of each individual task module
                for task in self.tasks:
                    if task == "Attention Network Test (ANT)":
                        # Set number of blocks for ANT
                        antTask = ant.ANT(self.pygame_screen, blocks=3)
                        # Run ANT
                        self.antData = antTask.run()
                        # Save ANT data to excel
                        self.antData.to_excel(self.writer, 'ANT', index=False)
                    elif task == "Mental Rotation Task":
                        mrtTask = mrt.MRT(self.pygame_screen)
                        # Run MRT
                        self.mrtData = mrtTask.run()
                        # Save MRT data to excel
                        self.mrtData.to_excel(self.writer, 'MRT', index=False)
                    elif task == "Sustained Attention to Response Task (SART)":
                        sartTask = sart.SART(self.pygame_screen)
                        # Run SART
                        self.sartData = sartTask.run()
                        # Save SART data to excel
                        self.sartData.to_excel(self.writer, 'SART', index=False)
                    elif task == "Digit Span (backwards)":
                        digitspanBackwardsTask = \
                            digitspan_backwards.DigitspanBackwards(
                                self.pygame_screen)
                        # Run Digit span (Backwards)
                        self.digitspanBackwardsData = digitspanBackwardsTask.run()
                        # Save digit span (backwards) data to excel
                        self.digitspanBackwardsData.to_excel(self.writer, 'Digit span (backwards)', index=False)
                    elif task == "Raven's Progressive Matrices":
                        ravensTask = ravens.Ravens(self.pygame_screen,
                                                   start=9, numTrials=12)
                        # Run Raven's Matrices
                        self.ravensData = ravensTask.run()
                        # Save ravens data to excel
                        self.ravensData.to_excel(self.writer, 'Ravens Matrices', index=False)
                    elif task == "Sternberg Task":
                        sternbergTask = sternberg.Sternberg(self.pygame_screen)
                        # Run Sternberg Task
                        self.sternbergData = sternbergTask.run()
                        # Save sternberg data to excel
                        self.sternbergData.to_excel(self.writer, 'Sternberg', index=False)

                    # Save excel file
                    self.writer.save()

                # Quit pygame
                pygame.quit()

                print "--- Experiment complete"
                self.close()

def main():
    # Get current directory
    cur_directory = os.path.dirname(os.path.realpath(__file__))

    # Check if settings file exists. If not, this is a first run
    first_run = not os.path.isfile(os.path.join(cur_directory, "settings.ini"))

    # Create main app window
    app = QtGui.QApplication(sys.argv)
    screen_resolution = app.desktop().screenGeometry()

    battery_window = BatteryWindow(cur_directory,
                                   first_run,
                                   screen_resolution.width(),
                                   screen_resolution.height())
    battery_window.show()
    sys.exit(app.exec_())

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