<commit_before>#     Copyright 2014, Kay Hayen, mailto:kay.hayen@gmail.com
#
#     Part of "Nuitka", an optimizing Python compiler that is compatible and
#     integrates with CPython, but also works on its own.
#
#     Licensed under the Apache License, Version 2.0 (the "License");
#     you may not use this file except in compliance with the License.
#     You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#     Unless required by applicable law or agreed to in writing, software
#     distributed under the License is distributed on an "AS IS" BASIS,
#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#     See the License for the specific language governing permissions and
#     limitations under the License.
#
""" Nodes for statements.

"""

from nuitka.Utils import python_version

from .NodeBases import NodeBase, StatementChildrenHavingBase

def checkStatements(value):
    """ Check that statements list value property.

    Must not be None, must not contain None, and of course only statements,
    may be empty.
    """

    assert value is not None
    assert None not in value

    for statement in value:
        assert statement.isStatement() or statement.isStatementsFrame(), \
          statement

    return tuple(value)

class StatementsSequence(StatementChildrenHavingBase):
    kind = "STATEMENTS_SEQUENCE"

    named_children = (
        "statements",
    )

    checkers = {
        "statements" : checkStatements
    }

    def __init__(self, statements, source_ref):
        StatementChildrenHavingBase.__init__(
            self,
            values     = {
                "statements" : statements
            },
            source_ref = source_ref
        )

    getStatements = StatementChildrenHavingBase.childGetter("statements")
    setStatements = StatementChildrenHavingBase.childSetter("statements")

    def getDetails(self):
        if self.getStatements():
            return {
                "statement_count" : len(self.getStatements())
            }
        else:
            return {
                "statement_count" : 0
            }

    # Overloading name based automatic check, so that derived ones know it too.
    def isStatementsSequence(self):
        # Virtual method, pylint: disable=R0201

        return True

    def trimStatements(self, statement):
        assert statement.parent is self

        old_statements = list(self.getStatements())
        assert statement in old_statements, \
          (statement, self)

        new_statements = old_statements[ : old_statements.index(statement)+1 ]

        self.setChild("statements", new_statements)

    def removeStatement(self, statement):
        assert statement.parent is self

        statements = list(self.getStatements())
        statements.remove(statement)

        self.setChild("statements", statements)

    def mergeStatementsSequence(self, statement_sequence):
        assert statement_sequence.parent is self

        old_statements = list(self.getStatements())
        assert statement_sequence in old_statements, \
          (statement_sequence, self)

        merge_index =  old_statements.index(statement_sequence)

        new_statements = tuple(old_statements[ : merge_index ])     + \
                         statement_sequence.getStatements()         + \
                         tuple(old_statements[ merge_index+1 : ])

        self.setChild("statements", new_statements)

    def mayHaveSideEffects(self):
        # Statement sequences have a side effect if one of the statements does.
        for statement in self.getStatements():
            if statement.mayHaveSideEffects():
                return True
        return False

    def mayRaiseException(self, exception_type):
        for statement in self.getStatements():
            if statement.mayRaiseException(exception_type):
                return True
        return False

    def needsFrame(self):
        for statement in self.getStatements():
            if statement.needsFrame():
                return True
        return False

    def mayReturn(self):
        for statement in self.getStatements():
            if statement.mayReturn():
                return True
        return False

    def mayBreak(self):
        for statement in self.getStatements():
            if statement.mayBreak():
                return True
        return False

    def mayContinue(self):
        for statement in self.getStatements():
            if statement.mayContinue():
                return True
        return False

    def mayRaiseExceptionOrAbort(self, exception_type):
        return self.mayRaiseException(exception_type) or \
               self.mayReturn() or \
               self.mayBreak() or \
               self.mayContinue()

    def isStatementAborting(self):
        return self.getStatements()[-1].isStatementAborting()

    def computeStatement(self, constraint_collection):
        # Don't want to be called like this.
        assert False, self

    def computeStatementsSequence(self, constraint_collection):
        new_statements = []

        statements = self.getStatements()
        assert statements, self

        for count, statement in enumerate(statements):
            # May be frames embedded.
            if statement.isStatementsFrame():
                new_statement = statement.computeStatementsSequence(
                    constraint_collection
                )
            else:
                new_statement = constraint_collection.onStatement(
                    statement = statement
                )

            if new_statement is not None:
                if new_statement.isStatementsSequence() and \
                   not new_statement.isStatementsFrame():
                    new_statements.extend(
                        new_statement.getStatements()
                    )
                else:
                    new_statements.append(
                        new_statement
                    )

                if statement is not statements[-1] and \
                   new_statement.isStatementAborting():
                    constraint_collection.signalChange(
                        "new_statements",
                        statements[count+1].getSourceReference(),
                        "Removed dead statements."
                    )

                    break

        if statements != new_statements:
            if new_statements:
                self.setStatements(new_statements)

                return self
            else:
                return None
        else:
            return self

def checkFrameStatements(value):
    """ Check that frames statements list value proper.

    Must not be None, must not contain None, and of course only statements
    sequences, or statements, may be empty.
    """

    assert value is not None
    assert None not in value

    for statement in value:
        assert statement.isStatement() or statement.isStatementsFrame(), \
          statement

    return tuple(value)

class StatementsFrame(StatementsSequence):
    # Frames got many details to store, pylint: disable=R0902

    kind = "STATEMENTS_FRAME"

    checkers = {
        "statements" : checkFrameStatements
    }

    def __init__(self, statements, guard_mode, code_name, var_names, arg_count,
                 kw_only_count, has_starlist, has_stardict, source_ref):
        StatementsSequence.__init__(
            self,
            statements = statements,
            source_ref = source_ref
        )

        self.var_names = tuple(var_names)
        self.code_name = code_name

        self.kw_only_count = kw_only_count
        self.arg_count = arg_count

        self.guard_mode = guard_mode

        self.has_starlist = has_starlist
        self.has_stardict = has_stardict

        self.needs_frame_exception_preserve = False

    def getDetails(self):
        result = {
            "code_name"  : self.code_name,
            "var_names"  : ", ".join(self.var_names),
            "guard_mode" : self.guard_mode
        }

        if python_version >= 300:
            result["kw_only_count"] = self.kw_only_count

        result.update(StatementsSequence.getDetails(self))

        return result

    def getGuardMode(self):
        return self.guard_mode

    def needsExceptionFramePreservation(self):
        if python_version < 300:
            preserving = ("full", "once")
        else:
            preserving = ("full", "once", "generator")

        return self.guard_mode in preserving

    def getVarNames(self):
        return self.var_names

    def getCodeObjectName(self):
        return self.code_name

    def getKwOnlyParameterCount(self):
        return self.kw_only_count

    def getArgumentCount(self):
        return self.arg_count

    def makeCloneAt(self, source_ref):
        assert False

    def markAsFrameExceptionPreserving(self):
        self.needs_frame_exception_preserve = True

    def needsFrameExceptionPreserving(self):
        return self.needs_frame_exception_preserve

    def getCodeObjectHandle(self, context):
        provider = self.getParentVariableProvider()

        # TODO: Why do this accessing a node, do this outside.
        return context.getCodeObjectHandle(
            filename      = self.source_ref.getFilename(),
            var_names     = self.getVarNames(),
            arg_count     = self.getArgumentCount(),
            kw_only_count = self.getKwOnlyParameterCount(),
            line_number   = 0
                              if provider.isPythonModule() else
                            self.source_ref.getLineNumber(),
            code_name     = self.getCodeObjectName(),
            is_generator  = provider.isExpressionFunctionBody() and \
                            provider.isGenerator(),
            is_optimized  = not provider.isPythonModule() and \
                            not provider.isClassDictCreation() and \
                            not provider.hasLocalsDict(),
            has_starlist  = self.has_starlist,
            has_stardict  = self.has_stardict,
            has_closure   = provider.isExpressionFunctionBody() and \
                            provider.getClosureVariables() != () and \
                            not provider.isClassDictCreation(),
            future_flags  = provider.getSourceReference().getFutureSpec().\
                              asFlags()
        )

    def computeStatementsSequence(self, constraint_collection):
        # The extraction of parts of the frame that can be moved before or after
        # the frame scope, takes it toll to complexity, pylint: disable=R0912

        new_statements = []

        statements = self.getStatements()

        for count, statement in enumerate(statements):
            # May be frames embedded.
            if statement.isStatementsFrame():
                new_statement = statement.computeStatementsSequence(
                    constraint_collection = constraint_collection
                )
            else:
                new_statement = constraint_collection.onStatement(
                    statement = statement
                )

            if new_statement is not None:
                if new_statement.isStatementsSequence() and \
                   not new_statement.isStatementsFrame():
                    new_statements.extend(new_statement.getStatements())
                else:
                    new_statements.append(new_statement)

                if statement is not statements[-1] and \
                   new_statement.isStatementAborting():
                    constraint_collection.signalChange(
                        "new_statements",
                        statements[count+1].getSourceReference(),
                        "Removed dead statements."
                    )

                    break

        if not new_statements:
            return None

        # Determine statements inside the frame, that need not be in a frame,
        # because they wouldn't raise an exception.
        outside_pre = []
        while new_statements and \
              not new_statements[0].needsFrame():
            outside_pre.append(new_statements[0])
            del new_statements[0]

        outside_post = []
        while new_statements and \
              not new_statements[-1].needsFrame():
            outside_post.insert(0, new_statements[-1])
            del new_statements[-1]

        if outside_pre or outside_post:
            from .NodeMakingHelpers import makeStatementsSequenceReplacementNode

            if new_statements:
                self.setStatements(new_statements)

                return makeStatementsSequenceReplacementNode(
                    statements = outside_pre + [self] + \
                                 outside_post,
                    node       = self
                )
            else:
                constraint_collection.signalChange(
                    "new_statements",
                    self.getSourceReference(),
                    "Removed useless frame"
                )

                return makeStatementsSequenceReplacementNode(
                    statements = outside_pre + outside_post,
                    node       = self
                )
        else:
            if statements != new_statements:
                self.setStatements(new_statements)

            return self

class StatementExpressionOnly(StatementChildrenHavingBase):
    kind = "STATEMENT_EXPRESSION_ONLY"

    named_children = ("expression",)

    def __init__(self, expression, source_ref):
        assert expression.isExpression()

        StatementChildrenHavingBase.__init__(
            self,
            values     = {
                "expression" : expression
            },
            source_ref = source_ref
        )

    def getDetail(self):
        return "expression %s" % self.getExpression()

    def mayHaveSideEffects(self):
        return self.getExpression().mayHaveSideEffects()

    getExpression = StatementChildrenHavingBase.childGetter(
        "expression"
    )

    def computeStatement(self, constraint_collection):
        constraint_collection.onExpression(
            expression = self.getExpression()
        )

        new_statement, change_tags, change_desc = \
          self.getExpression().computeExpressionDrop(
            statement             = self,
            constraint_collection = constraint_collection
        )

        if new_statement is None:
            return new_statement, change_tags, change_desc

        if new_statement is not self:
            return new_statement, change_tags, change_desc

        return self, None, None

class StatementGeneratorEntry(NodeBase):
    kind = "STATEMENT_GENERATOR_ENTRY"

    def __init__(self, source_ref):
        NodeBase.__init__(
            self,
            source_ref = source_ref
        )

    def mayRaiseException(self, exception_type):
        # Anything might be thrown into a generator, representing that is the
        # whole point of this statement.
        return True

    def computeStatement(self, constraint_collection):
        # Nothing we can about it.
        return self, None, None

class StatementPreserveFrameException(NodeBase):
    kind = "STATEMENT_PRESERVE_FRAME_EXCEPTION"

    def __init__(self, source_ref):
        NodeBase.__init__(
            self,
            source_ref = source_ref
        )

    def computeStatement(self, constraint_collection):
        # For Python2 generators, it's not necessary to preserve, the frame
        # decides it. TODO: This check makes only sense once.

        if self.getParentStatementsFrame().needsExceptionFramePreservation():
            return self, None, None
        else:
            return (
                None,
                "new_statements",
                "Removed frame preservation for generators."
            )

    def mayRaiseException(self, exception_type):
        return False

    def needsFrame(self):
        return True

class StatementRestoreFrameException(NodeBase):
    kind = "STATEMENT_RESTORE_FRAME_EXCEPTION"

    def __init__(self, source_ref):
        NodeBase.__init__(
            self,
            source_ref = source_ref
        )

    def computeStatement(self, constraint_collection):
        return self, None, None

    def mayRaiseException(self, exception_type):
        return False

class StatementReraiseFrameException(NodeBase):
    kind = "STATEMENT_RERAISE_FRAME_EXCEPTION"

    def __init__(self, source_ref):
        NodeBase.__init__(
            self,
            source_ref = source_ref
        )

    def computeStatement(self, constraint_collection):
        return self, None, None

    def mayRaiseException(self, exception_type):
        return True

class StatementPublishException(NodeBase):
    kind = "STATEMENT_PUBLISH_EXCEPTION"

    def __init__(self, source_ref):
        NodeBase.__init__(
            self,
            source_ref = source_ref
        )

    def computeStatement(self, constraint_collection):
        # TODO: Determine the need for it.
        return self, None, None

    def mayRaiseException(self, exception_type):
        return False
<commit_msg>Fix, expression only statements only raise when their expression does.<commit_after>#     Copyright 2014, Kay Hayen, mailto:kay.hayen@gmail.com
#
#     Part of "Nuitka", an optimizing Python compiler that is compatible and
#     integrates with CPython, but also works on its own.
#
#     Licensed under the Apache License, Version 2.0 (the "License");
#     you may not use this file except in compliance with the License.
#     You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#     Unless required by applicable law or agreed to in writing, software
#     distributed under the License is distributed on an "AS IS" BASIS,
#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#     See the License for the specific language governing permissions and
#     limitations under the License.
#
""" Nodes for statements.

"""

from nuitka.Utils import python_version

from .NodeBases import NodeBase, StatementChildrenHavingBase

def checkStatements(value):
    """ Check that statements list value property.

    Must not be None, must not contain None, and of course only statements,
    may be empty.
    """

    assert value is not None
    assert None not in value

    for statement in value:
        assert statement.isStatement() or statement.isStatementsFrame(), \
          statement

    return tuple(value)

class StatementsSequence(StatementChildrenHavingBase):
    kind = "STATEMENTS_SEQUENCE"

    named_children = (
        "statements",
    )

    checkers = {
        "statements" : checkStatements
    }

    def __init__(self, statements, source_ref):
        StatementChildrenHavingBase.__init__(
            self,
            values     = {
                "statements" : statements
            },
            source_ref = source_ref
        )

    getStatements = StatementChildrenHavingBase.childGetter("statements")
    setStatements = StatementChildrenHavingBase.childSetter("statements")

    def getDetails(self):
        if self.getStatements():
            return {
                "statement_count" : len(self.getStatements())
            }
        else:
            return {
                "statement_count" : 0
            }

    # Overloading name based automatic check, so that derived ones know it too.
    def isStatementsSequence(self):
        # Virtual method, pylint: disable=R0201

        return True

    def trimStatements(self, statement):
        assert statement.parent is self

        old_statements = list(self.getStatements())
        assert statement in old_statements, \
          (statement, self)

        new_statements = old_statements[ : old_statements.index(statement)+1 ]

        self.setChild("statements", new_statements)

    def removeStatement(self, statement):
        assert statement.parent is self

        statements = list(self.getStatements())
        statements.remove(statement)

        self.setChild("statements", statements)

    def mergeStatementsSequence(self, statement_sequence):
        assert statement_sequence.parent is self

        old_statements = list(self.getStatements())
        assert statement_sequence in old_statements, \
          (statement_sequence, self)

        merge_index =  old_statements.index(statement_sequence)

        new_statements = tuple(old_statements[ : merge_index ])     + \
                         statement_sequence.getStatements()         + \
                         tuple(old_statements[ merge_index+1 : ])

        self.setChild("statements", new_statements)

    def mayHaveSideEffects(self):
        # Statement sequences have a side effect if one of the statements does.
        for statement in self.getStatements():
            if statement.mayHaveSideEffects():
                return True
        return False

    def mayRaiseException(self, exception_type):
        for statement in self.getStatements():
            if statement.mayRaiseException(exception_type):
                return True
        return False

    def needsFrame(self):
        for statement in self.getStatements():
            if statement.needsFrame():
                return True
        return False

    def mayReturn(self):
        for statement in self.getStatements():
            if statement.mayReturn():
                return True
        return False

    def mayBreak(self):
        for statement in self.getStatements():
            if statement.mayBreak():
                return True
        return False

    def mayContinue(self):
        for statement in self.getStatements():
            if statement.mayContinue():
                return True
        return False

    def mayRaiseExceptionOrAbort(self, exception_type):
        return self.mayRaiseException(exception_type) or \
               self.mayReturn() or \
               self.mayBreak() or \
               self.mayContinue()

    def isStatementAborting(self):
        return self.getStatements()[-1].isStatementAborting()

    def computeStatement(self, constraint_collection):
        # Don't want to be called like this.
        assert False, self

    def computeStatementsSequence(self, constraint_collection):
        new_statements = []

        statements = self.getStatements()
        assert statements, self

        for count, statement in enumerate(statements):
            # May be frames embedded.
            if statement.isStatementsFrame():
                new_statement = statement.computeStatementsSequence(
                    constraint_collection
                )
            else:
                new_statement = constraint_collection.onStatement(
                    statement = statement
                )

            if new_statement is not None:
                if new_statement.isStatementsSequence() and \
                   not new_statement.isStatementsFrame():
                    new_statements.extend(
                        new_statement.getStatements()
                    )
                else:
                    new_statements.append(
                        new_statement
                    )

                if statement is not statements[-1] and \
                   new_statement.isStatementAborting():
                    constraint_collection.signalChange(
                        "new_statements",
                        statements[count+1].getSourceReference(),
                        "Removed dead statements."
                    )

                    break

        if statements != new_statements:
            if new_statements:
                self.setStatements(new_statements)

                return self
            else:
                return None
        else:
            return self

def checkFrameStatements(value):
    """ Check that frames statements list value proper.

    Must not be None, must not contain None, and of course only statements
    sequences, or statements, may be empty.
    """

    assert value is not None
    assert None not in value

    for statement in value:
        assert statement.isStatement() or statement.isStatementsFrame(), \
          statement

    return tuple(value)

class StatementsFrame(StatementsSequence):
    # Frames got many details to store, pylint: disable=R0902

    kind = "STATEMENTS_FRAME"

    checkers = {
        "statements" : checkFrameStatements
    }

    def __init__(self, statements, guard_mode, code_name, var_names, arg_count,
                 kw_only_count, has_starlist, has_stardict, source_ref):
        StatementsSequence.__init__(
            self,
            statements = statements,
            source_ref = source_ref
        )

        self.var_names = tuple(var_names)
        self.code_name = code_name

        self.kw_only_count = kw_only_count
        self.arg_count = arg_count

        self.guard_mode = guard_mode

        self.has_starlist = has_starlist
        self.has_stardict = has_stardict

        self.needs_frame_exception_preserve = False

    def getDetails(self):
        result = {
            "code_name"  : self.code_name,
            "var_names"  : ", ".join(self.var_names),
            "guard_mode" : self.guard_mode
        }

        if python_version >= 300:
            result["kw_only_count"] = self.kw_only_count

        result.update(StatementsSequence.getDetails(self))

        return result

    def getGuardMode(self):
        return self.guard_mode

    def needsExceptionFramePreservation(self):
        if python_version < 300:
            preserving = ("full", "once")
        else:
            preserving = ("full", "once", "generator")

        return self.guard_mode in preserving

    def getVarNames(self):
        return self.var_names

    def getCodeObjectName(self):
        return self.code_name

    def getKwOnlyParameterCount(self):
        return self.kw_only_count

    def getArgumentCount(self):
        return self.arg_count

    def makeCloneAt(self, source_ref):
        assert False

    def markAsFrameExceptionPreserving(self):
        self.needs_frame_exception_preserve = True

    def needsFrameExceptionPreserving(self):
        return self.needs_frame_exception_preserve

    def getCodeObjectHandle(self, context):
        provider = self.getParentVariableProvider()

        # TODO: Why do this accessing a node, do this outside.
        return context.getCodeObjectHandle(
            filename      = self.source_ref.getFilename(),
            var_names     = self.getVarNames(),
            arg_count     = self.getArgumentCount(),
            kw_only_count = self.getKwOnlyParameterCount(),
            line_number   = 0
                              if provider.isPythonModule() else
                            self.source_ref.getLineNumber(),
            code_name     = self.getCodeObjectName(),
            is_generator  = provider.isExpressionFunctionBody() and \
                            provider.isGenerator(),
            is_optimized  = not provider.isPythonModule() and \
                            not provider.isClassDictCreation() and \
                            not provider.hasLocalsDict(),
            has_starlist  = self.has_starlist,
            has_stardict  = self.has_stardict,
            has_closure   = provider.isExpressionFunctionBody() and \
                            provider.getClosureVariables() != () and \
                            not provider.isClassDictCreation(),
            future_flags  = provider.getSourceReference().getFutureSpec().\
                              asFlags()
        )

    def computeStatementsSequence(self, constraint_collection):
        # The extraction of parts of the frame that can be moved before or after
        # the frame scope, takes it toll to complexity, pylint: disable=R0912

        new_statements = []

        statements = self.getStatements()

        for count, statement in enumerate(statements):
            # May be frames embedded.
            if statement.isStatementsFrame():
                new_statement = statement.computeStatementsSequence(
                    constraint_collection = constraint_collection
                )
            else:
                new_statement = constraint_collection.onStatement(
                    statement = statement
                )

            if new_statement is not None:
                if new_statement.isStatementsSequence() and \
                   not new_statement.isStatementsFrame():
                    new_statements.extend(new_statement.getStatements())
                else:
                    new_statements.append(new_statement)

                if statement is not statements[-1] and \
                   new_statement.isStatementAborting():
                    constraint_collection.signalChange(
                        "new_statements",
                        statements[count+1].getSourceReference(),
                        "Removed dead statements."
                    )

                    break

        if not new_statements:
            return None

        # Determine statements inside the frame, that need not be in a frame,
        # because they wouldn't raise an exception.
        outside_pre = []
        while new_statements and \
              not new_statements[0].needsFrame():
            outside_pre.append(new_statements[0])
            del new_statements[0]

        outside_post = []
        while new_statements and \
              not new_statements[-1].needsFrame():
            outside_post.insert(0, new_statements[-1])
            del new_statements[-1]

        if outside_pre or outside_post:
            from .NodeMakingHelpers import makeStatementsSequenceReplacementNode

            if new_statements:
                self.setStatements(new_statements)

                return makeStatementsSequenceReplacementNode(
                    statements = outside_pre + [self] + \
                                 outside_post,
                    node       = self
                )
            else:
                constraint_collection.signalChange(
                    "new_statements",
                    self.getSourceReference(),
                    "Removed useless frame"
                )

                return makeStatementsSequenceReplacementNode(
                    statements = outside_pre + outside_post,
                    node       = self
                )
        else:
            if statements != new_statements:
                self.setStatements(new_statements)

            return self

class StatementExpressionOnly(StatementChildrenHavingBase):
    kind = "STATEMENT_EXPRESSION_ONLY"

    named_children = (
        "expression",
    )

    def __init__(self, expression, source_ref):
        assert expression.isExpression()

        StatementChildrenHavingBase.__init__(
            self,
            values     = {
                "expression" : expression
            },
            source_ref = source_ref
        )

    def getDetail(self):
        return "expression %s" % self.getExpression()

    def mayHaveSideEffects(self):
        return self.getExpression().mayHaveSideEffects()

    def mayRaiseException(self, exception_type):
        return self.getExpression().mayRaiseException(exception_type)

    getExpression = StatementChildrenHavingBase.childGetter(
        "expression"
    )

    def computeStatement(self, constraint_collection):
        constraint_collection.onExpression(
            expression = self.getExpression()
        )

        new_statement, change_tags, change_desc = \
          self.getExpression().computeExpressionDrop(
            statement             = self,
            constraint_collection = constraint_collection
        )

        if new_statement is None:
            return new_statement, change_tags, change_desc

        if new_statement is not self:
            return new_statement, change_tags, change_desc

        return self, None, None

class StatementGeneratorEntry(NodeBase):
    kind = "STATEMENT_GENERATOR_ENTRY"

    def __init__(self, source_ref):
        NodeBase.__init__(
            self,
            source_ref = source_ref
        )

    def mayRaiseException(self, exception_type):
        # Anything might be thrown into a generator, representing that is the
        # whole point of this statement.
        return True

    def computeStatement(self, constraint_collection):
        # Nothing we can about it.
        return self, None, None

class StatementPreserveFrameException(NodeBase):
    kind = "STATEMENT_PRESERVE_FRAME_EXCEPTION"

    def __init__(self, source_ref):
        NodeBase.__init__(
            self,
            source_ref = source_ref
        )

    def computeStatement(self, constraint_collection):
        # For Python2 generators, it's not necessary to preserve, the frame
        # decides it. TODO: This check makes only sense once.

        if self.getParentStatementsFrame().needsExceptionFramePreservation():
            return self, None, None
        else:
            return (
                None,
                "new_statements",
                "Removed frame preservation for generators."
            )

    def mayRaiseException(self, exception_type):
        return False

    def needsFrame(self):
        return True

class StatementRestoreFrameException(NodeBase):
    kind = "STATEMENT_RESTORE_FRAME_EXCEPTION"

    def __init__(self, source_ref):
        NodeBase.__init__(
            self,
            source_ref = source_ref
        )

    def computeStatement(self, constraint_collection):
        return self, None, None

    def mayRaiseException(self, exception_type):
        return False

class StatementReraiseFrameException(NodeBase):
    kind = "STATEMENT_RERAISE_FRAME_EXCEPTION"

    def __init__(self, source_ref):
        NodeBase.__init__(
            self,
            source_ref = source_ref
        )

    def computeStatement(self, constraint_collection):
        return self, None, None

    def mayRaiseException(self, exception_type):
        return True

class StatementPublishException(NodeBase):
    kind = "STATEMENT_PUBLISH_EXCEPTION"

    def __init__(self, source_ref):
        NodeBase.__init__(
            self,
            source_ref = source_ref
        )

    def computeStatement(self, constraint_collection):
        # TODO: Determine the need for it.
        return self, None, None

    def mayRaiseException(self, exception_type):
        return False
<|endoftext|>


<commit_before># Copyright 2022 The T5X Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for t5x.utils."""

import dataclasses
import os
import re

from absl import flags
from absl.testing import absltest
from absl.testing import parameterized

import flax.core
from flax.linen import partitioning as flax_partitioning
import jax
import numpy as np

import seqio
from t5x import partitioning
from t5x import test_utils
from t5x import utils
import tensorflow as tf

mock = absltest.mock
Evaluator = seqio.Evaluator
PartitionSpec = partitioning.PartitionSpec
AxisMetadata = flax_partitioning.AxisMetadata

# Parse absl flags test_srcdir and test_tmpdir.
jax.config.parse_flags_with_absl()

FLAGS = flags.FLAGS

def get_mock_train_state(params, param_states=None, step=0):
  """Returns a mock TrainState."""
  step = np.array(step) if step is not None else None
  state = mock.Mock(param_states=param_states, step=step)
  state_dict = dict(
      target=params, state=dict(param_states=param_states, step=step))
  return mock.Mock(
      params=params,
      param_states=param_states,
      step=step,
      state_dict=lambda: state_dict,
      optimizer=mock.Mock(
          target=params, state=state, state_dict=lambda: state_dict),
  )

class UtilsTest(parameterized.TestCase):

  def round_vocab_size_to_multiple(self):
    self.assertEqual(utils.round_vocab_size_to_multiple(1), 128)
    self.assertEqual(utils.round_vocab_size_to_multiple(128), 128)
    self.assertEqual(utils.round_vocab_size_to_multiple(129), 256)
    self.assertEqual(utils.round_vocab_size_to_multiple(129), 256)
    self.assertEqual(
        utils.round_vocab_size_to_multiple(25600, divisor=384), 256128)

  def test_get_zeros_batch_like_spec(self):
    test_utils.assert_same(
        utils.get_zeros_batch_like_spec({
            "i": jax.ShapeDtypeStruct((2, 5), dtype=np.int32),
            "j": jax.ShapeDtypeStruct((1,), dtype=np.float32),
        }), {
            "i": np.zeros((2, 5), dtype=np.int32),
            "j": np.zeros((1,), dtype=np.float32)
        })

  def test_get_zeros_batch_like_dataset(self):
    ds = tf.data.Dataset.from_tensors({
        "i": np.arange(10, dtype=np.int32).reshape((2, 5)),
        "j": np.ones((1,), dtype=np.float32)
    })

    test_utils.assert_same(
        utils.get_zeros_batch_like_dataset(ds), {
            "i": np.zeros((2, 5), dtype=np.int32),
            "j": np.zeros((1,), dtype=np.float32)
        })

    test_utils.assert_same(
        utils.get_zeros_batch_like_dataset(ds, batch_size=4), {
            "i": np.zeros((4, 5), dtype=np.int32),
            "j": np.zeros((4,), dtype=np.float32)
        })

  def test_log_model_info(self):
    log_file = self.create_tempfile()

    mock_train_state = get_mock_train_state(
        params={
            "a": {
                "aa": jax.ShapeDtypeStruct(shape=(2, 3), dtype=np.int32)
            },
            "c": jax.ShapeDtypeStruct(shape=(7, 8), dtype=np.int32)
        },
        param_states={
            "a": {
                "aa": {
                    "v_row": jax.ShapeDtypeStruct(shape=(2,), dtype=np.int32),
                    "v_col": jax.ShapeDtypeStruct(shape=(3,), dtype=np.int32)
                }
            },
            "c": {
                "v_row": jax.ShapeDtypeStruct(shape=(2, 4), dtype=np.int32),
                "v_col": None
            }
        })

    mock_logical_axes = get_mock_train_state(
        params={
            "a": {
                "aa": partitioning.AxisNames("a1", None)
            },
            "c": partitioning.AxisNames(None, "a1")
        },
        param_states={
            "a": {
                "aa": {
                    "v_row": partitioning.AxisNames(None,),
                    "v_col": partitioning.AxisNames(None,)
                }
            },
            "c": {
                "v_row": partitioning.AxisNames("a1",),
                "v_col": partitioning.AxisNames("a2",)
            }
        },
        step=None)

    mock_mesh_axes = get_mock_train_state(
        params={
            "a": {
                "aa": PartitionSpec("b1", None)
            },
            "c": PartitionSpec(None, "b1")
        },
        param_states={
            "a": {
                "aa": {
                    "v_row": partitioning.AxisNames(None,),
                    "v_col": partitioning.AxisNames(None,)
                }
            },
            "c": {
                "v_row": partitioning.AxisNames("b1",),
                "v_col": partitioning.AxisNames("b2",)
            }
        },
        step=None)

    partitioner = mock.Mock(
        get_logical_axes=lambda _: mock_logical_axes,
        get_mesh_axes=lambda _: mock_mesh_axes)

    utils.log_model_info(log_file.full_path, mock_train_state, partitioner)

    self.assertEqual(
        re.sub(r"\s+", " ", log_file.read_text()),
        "Variable a/aa size 6 shape (a1=2, None=3) partition spec ('b1', None) "
        "Variable c size 56 shape (None=7, a1=8) partition spec (None, 'b1') "
        "Total number of parameters: 62 "
        "Variable param_states/a/aa/v_col size 3 shape (None=3) partition spec (None,) "
        "Variable param_states/a/aa/v_row size 2 shape (None=2) partition spec (None,) "
        "Variable param_states/c/v_col None "
        "Variable param_states/c/v_row size 8 shape (2, 4) partition spec ('b1',) "
        "Variable step size 1 shape () partition spec None ")

  @mock.patch.object(utils, "get_dataset")
  def test_get_training_eval_datasets(self, mock_get_dataset):
    # Register a mock SeqIO mixture.
    task1 = mock.create_autospec(seqio.Task, instance=True)
    task1.name = "mock_task1"
    task1.splits = set(["train", "test"])
    task2 = mock.create_autospec(seqio.Task, instance=True)
    task2.name = "mock_task2"
    task2.splits = set(["train", "test"])
    seqio.TaskRegistry.add_provider("mock_task1", task1)
    seqio.TaskRegistry.add_provider("mock_task2", task2)
    mixture = seqio.Mixture(
        "mock_mix", ["mock_task1", "mock_task2"], default_rate=1.0)
    seqio.MixtureRegistry.add_provider("mock_mix", mixture)

    mock_get_dataset.return_value = tf.data.Dataset.from_tensor_slices(
        range(100))

    # Verify calls to utils.get_dataset
    cfg = utils.DatasetConfig(
        mixture_or_task_name="mock_mix",
        task_feature_lengths={},
        split="test",
        batch_size=2,
        shuffle=False,
        seed=23)

    _ = utils.get_training_eval_datasets(
        cfg,
        shard_id=0,
        num_shards=1,
        eval_steps=2,
        feature_converter_cls=seqio.FeatureConverter,
        get_dataset_fn=utils.get_dataset)

    expected_calls = [
        mock.call(
            dataclasses.replace(cfg, mixture_or_task_name="mock_task1"),
            0,
            1,
            seqio.FeatureConverter,
            continue_from_last_checkpoint=False,
            num_epochs=2),
        mock.call(
            dataclasses.replace(cfg, mixture_or_task_name="mock_task2"),
            0,
            1,
            seqio.FeatureConverter,
            continue_from_last_checkpoint=False,
            num_epochs=2),
        mock.call(
            dataclasses.replace(cfg, mixture_or_task_name="mock_mix"),
            0,
            1,
            seqio.FeatureConverter,
            continue_from_last_checkpoint=False,
            num_epochs=2)
    ]
    mock_get_dataset.assert_has_calls(expected_calls)

  def test_override_params_axes_names(self):
    model_variables = flax.core.freeze({
        "params": {
            "logits_dense": np.zeros((2, 4)),
            "mlp": {
                "wo": {
                    "kernel": np.zeros((4, 6)),
                    "bias": np.zeros(6),
                }
            }
        },
        "params_axes": {
            "logits_dense_axes": AxisMetadata(names=("vocab", "embed")),
            "mlp": {
                "wo": {
                    "kernel_axes": AxisMetadata(names=("embed", "mlp"))
                }
            }
        }
    })

    with self.assertRaisesWithLiteralMatch(
        ValueError,
        "Model variables do not contain a 'params_axes' collection to apply an "
        "override to."):
      utils.override_params_axes_names({"params": model_variables["params"]},
                                       [("mlp/wo/kernel", ("embed",))])

    with self.assertRaisesWithLiteralMatch(
        ValueError,
        "Provided axis name override for mlp/wo/kernel does not match param "
        "rank (2): ('embed',)"):
      utils.override_params_axes_names(model_variables,
                                       [("mlp/wo/kernel", ("embed",))])

    overridden_variables = utils.override_params_axes_names(
        model_variables,
        [
            ("wo/kernel", ("batch",)),  # unused since not a full match
            (".*/wo/kernel", ("batch", "embed")),  # this one is used
            ("mlp/wo/kernel", ("embed",)),  # unused since already matched
            ("mlp/wo/bias", ("embed",)),  # used
        ])

    jax.tree_multimap(
        np.testing.assert_equal, overridden_variables,
        flax.core.freeze({
            "params": {
                "logits_dense": np.zeros((2, 4)),
                "mlp": {
                    "wo": {
                        "kernel": np.zeros((4, 6)),
                        "bias": np.zeros(6),
                    }
                }
            },
            "params_axes": {
                "logits_dense_axes": AxisMetadata(names=("vocab", "embed")),
                "mlp": {
                    "wo": {
                        "kernel_axes": AxisMetadata(names=("batch", "embed")),
                        "bias_axes": AxisMetadata(names=("embed",)),
                    }
                }
            }
        }))

if __name__ == "__main__":
  absltest.main()
<commit_msg>Add basic test case for TrainStateInitializer, which covers restoration from checkpoints. The test primarily covers the logic that determines which checkpoints to restore given restore_cfg(s).<commit_after># Copyright 2022 The T5X Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for t5x.utils."""

import dataclasses
import os
import re
from typing import Optional

from absl import flags
from absl.testing import absltest
from absl.testing import parameterized
import flax.core
from flax.linen import partitioning as flax_partitioning
import jax
import numpy as np
import seqio
from t5x import checkpoints
from t5x import partitioning
from t5x import test_utils
from t5x import train_state as train_state_lib
from t5x import utils
import tensorflow as tf

mock = absltest.mock
Evaluator = seqio.Evaluator
PartitionSpec = partitioning.PartitionSpec
AxisMetadata = flax_partitioning.AxisMetadata

# Parse absl flags test_srcdir and test_tmpdir.
jax.config.parse_flags_with_absl()

FLAGS = flags.FLAGS

def get_mock_train_state(params, param_states=None, step=0):
  """Returns a mock TrainState."""
  step = np.array(step) if step is not None else None
  state = mock.Mock(param_states=param_states, step=step)
  state_dict = dict(
      target=params, state=dict(param_states=param_states, step=step))
  return mock.Mock(
      params=params,
      param_states=param_states,
      step=step,
      state_dict=lambda: state_dict,
      optimizer=mock.Mock(
          target=params, state=state, state_dict=lambda: state_dict),
  )

class UtilsTest(parameterized.TestCase):

  def round_vocab_size_to_multiple(self):
    self.assertEqual(utils.round_vocab_size_to_multiple(1), 128)
    self.assertEqual(utils.round_vocab_size_to_multiple(128), 128)
    self.assertEqual(utils.round_vocab_size_to_multiple(129), 256)
    self.assertEqual(utils.round_vocab_size_to_multiple(129), 256)
    self.assertEqual(
        utils.round_vocab_size_to_multiple(25600, divisor=384), 256128)

  def test_get_zeros_batch_like_spec(self):
    test_utils.assert_same(
        utils.get_zeros_batch_like_spec({
            "i": jax.ShapeDtypeStruct((2, 5), dtype=np.int32),
            "j": jax.ShapeDtypeStruct((1,), dtype=np.float32),
        }), {
            "i": np.zeros((2, 5), dtype=np.int32),
            "j": np.zeros((1,), dtype=np.float32)
        })

  def test_get_zeros_batch_like_dataset(self):
    ds = tf.data.Dataset.from_tensors({
        "i": np.arange(10, dtype=np.int32).reshape((2, 5)),
        "j": np.ones((1,), dtype=np.float32)
    })

    test_utils.assert_same(
        utils.get_zeros_batch_like_dataset(ds), {
            "i": np.zeros((2, 5), dtype=np.int32),
            "j": np.zeros((1,), dtype=np.float32)
        })

    test_utils.assert_same(
        utils.get_zeros_batch_like_dataset(ds, batch_size=4), {
            "i": np.zeros((4, 5), dtype=np.int32),
            "j": np.zeros((4,), dtype=np.float32)
        })

  def test_log_model_info(self):
    log_file = self.create_tempfile()

    mock_train_state = get_mock_train_state(
        params={
            "a": {
                "aa": jax.ShapeDtypeStruct(shape=(2, 3), dtype=np.int32)
            },
            "c": jax.ShapeDtypeStruct(shape=(7, 8), dtype=np.int32)
        },
        param_states={
            "a": {
                "aa": {
                    "v_row": jax.ShapeDtypeStruct(shape=(2,), dtype=np.int32),
                    "v_col": jax.ShapeDtypeStruct(shape=(3,), dtype=np.int32)
                }
            },
            "c": {
                "v_row": jax.ShapeDtypeStruct(shape=(2, 4), dtype=np.int32),
                "v_col": None
            }
        })

    mock_logical_axes = get_mock_train_state(
        params={
            "a": {
                "aa": partitioning.AxisNames("a1", None)
            },
            "c": partitioning.AxisNames(None, "a1")
        },
        param_states={
            "a": {
                "aa": {
                    "v_row": partitioning.AxisNames(None,),
                    "v_col": partitioning.AxisNames(None,)
                }
            },
            "c": {
                "v_row": partitioning.AxisNames("a1",),
                "v_col": partitioning.AxisNames("a2",)
            }
        },
        step=None)

    mock_mesh_axes = get_mock_train_state(
        params={
            "a": {
                "aa": PartitionSpec("b1", None)
            },
            "c": PartitionSpec(None, "b1")
        },
        param_states={
            "a": {
                "aa": {
                    "v_row": partitioning.AxisNames(None,),
                    "v_col": partitioning.AxisNames(None,)
                }
            },
            "c": {
                "v_row": partitioning.AxisNames("b1",),
                "v_col": partitioning.AxisNames("b2",)
            }
        },
        step=None)

    partitioner = mock.Mock(
        get_logical_axes=lambda _: mock_logical_axes,
        get_mesh_axes=lambda _: mock_mesh_axes)

    utils.log_model_info(log_file.full_path, mock_train_state, partitioner)

    self.assertEqual(
        re.sub(r"\s+", " ", log_file.read_text()),
        "Variable a/aa size 6 shape (a1=2, None=3) partition spec ('b1', None) "
        "Variable c size 56 shape (None=7, a1=8) partition spec (None, 'b1') "
        "Total number of parameters: 62 "
        "Variable param_states/a/aa/v_col size 3 shape (None=3) partition spec (None,) "
        "Variable param_states/a/aa/v_row size 2 shape (None=2) partition spec (None,) "
        "Variable param_states/c/v_col None "
        "Variable param_states/c/v_row size 8 shape (2, 4) partition spec ('b1',) "
        "Variable step size 1 shape () partition spec None ")

  @mock.patch.object(utils, "get_dataset")
  def test_get_training_eval_datasets(self, mock_get_dataset):
    # Register a mock SeqIO mixture.
    task1 = mock.create_autospec(seqio.Task, instance=True)
    task1.name = "mock_task1"
    task1.splits = set(["train", "test"])
    task2 = mock.create_autospec(seqio.Task, instance=True)
    task2.name = "mock_task2"
    task2.splits = set(["train", "test"])
    seqio.TaskRegistry.add_provider("mock_task1", task1)
    seqio.TaskRegistry.add_provider("mock_task2", task2)
    mixture = seqio.Mixture(
        "mock_mix", ["mock_task1", "mock_task2"], default_rate=1.0)
    seqio.MixtureRegistry.add_provider("mock_mix", mixture)

    mock_get_dataset.return_value = tf.data.Dataset.from_tensor_slices(
        range(100))

    # Verify calls to utils.get_dataset
    cfg = utils.DatasetConfig(
        mixture_or_task_name="mock_mix",
        task_feature_lengths={},
        split="test",
        batch_size=2,
        shuffle=False,
        seed=23)

    _ = utils.get_training_eval_datasets(
        cfg,
        shard_id=0,
        num_shards=1,
        eval_steps=2,
        feature_converter_cls=seqio.FeatureConverter,
        get_dataset_fn=utils.get_dataset)

    expected_calls = [
        mock.call(
            dataclasses.replace(cfg, mixture_or_task_name="mock_task1"),
            0,
            1,
            seqio.FeatureConverter,
            continue_from_last_checkpoint=False,
            num_epochs=2),
        mock.call(
            dataclasses.replace(cfg, mixture_or_task_name="mock_task2"),
            0,
            1,
            seqio.FeatureConverter,
            continue_from_last_checkpoint=False,
            num_epochs=2),
        mock.call(
            dataclasses.replace(cfg, mixture_or_task_name="mock_mix"),
            0,
            1,
            seqio.FeatureConverter,
            continue_from_last_checkpoint=False,
            num_epochs=2)
    ]
    mock_get_dataset.assert_has_calls(expected_calls)

  def test_override_params_axes_names(self):
    model_variables = flax.core.freeze({
        "params": {
            "logits_dense": np.zeros((2, 4)),
            "mlp": {
                "wo": {
                    "kernel": np.zeros((4, 6)),
                    "bias": np.zeros(6),
                }
            }
        },
        "params_axes": {
            "logits_dense_axes": AxisMetadata(names=("vocab", "embed")),
            "mlp": {
                "wo": {
                    "kernel_axes": AxisMetadata(names=("embed", "mlp"))
                }
            }
        }
    })

    with self.assertRaisesWithLiteralMatch(
        ValueError,
        "Model variables do not contain a 'params_axes' collection to apply an "
        "override to."):
      utils.override_params_axes_names({"params": model_variables["params"]},
                                       [("mlp/wo/kernel", ("embed",))])

    with self.assertRaisesWithLiteralMatch(
        ValueError,
        "Provided axis name override for mlp/wo/kernel does not match param "
        "rank (2): ('embed',)"):
      utils.override_params_axes_names(model_variables,
                                       [("mlp/wo/kernel", ("embed",))])

    overridden_variables = utils.override_params_axes_names(
        model_variables,
        [
            ("wo/kernel", ("batch",)),  # unused since not a full match
            (".*/wo/kernel", ("batch", "embed")),  # this one is used
            ("mlp/wo/kernel", ("embed",)),  # unused since already matched
            ("mlp/wo/bias", ("embed",)),  # used
        ])

    jax.tree_multimap(
        np.testing.assert_equal, overridden_variables,
        flax.core.freeze({
            "params": {
                "logits_dense": np.zeros((2, 4)),
                "mlp": {
                    "wo": {
                        "kernel": np.zeros((4, 6)),
                        "bias": np.zeros(6),
                    }
                }
            },
            "params_axes": {
                "logits_dense_axes": AxisMetadata(names=("vocab", "embed")),
                "mlp": {
                    "wo": {
                        "kernel_axes": AxisMetadata(names=("batch", "embed")),
                        "bias_axes": AxisMetadata(names=("embed",)),
                    }
                }
            }
        }))

@dataclasses.dataclass
class MockTrainState:
  path: Optional[str] = None
  from_scratch: Optional[bool] = None

class MockCheckpointer(checkpoints.Checkpointer):

  def __init__(self, *args, **kwargs):
    pass

  # restore should return TrainState, but we force it to return Mock with path
  # for simplicity.
  def restore(self, path, *args, **kwargs):
    return MockTrainState(path=path, from_scratch=False)

class TrainStateInitializerTest(parameterized.TestCase):

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

    def _partition(train_state, in_axis_resources, out_axis_resources):
      del train_state, in_axis_resources, out_axis_resources
      partitioned_fn = lambda _: MockTrainState(from_scratch=True)
      return partitioned_fn

    partitioner = mock.Mock(get_mesh_axes=lambda _: None, partition=_partition)
    mock_inference_state_create = self.enter_context(
        mock.patch.object(train_state_lib.InferenceState, "create"))
    mock_inference_state_create.return_value = None

    shapes = {
        "ones": (1, 1),
        "twos": (2, 2),
        "threes": (3, 3),
    }
    types = {
        "ones": int,
        "twos": float,
        "threes": int,
    }

    def _init_fn(rng, input_shapes, input_types):
      del rng
      return {
          "ones":
              np.ones(input_shapes["ones"], dtype=input_types["ones"]),
          "twos":
              np.ones(input_shapes["twos"], dtype=input_types["twos"]) * 2,
          "threes":
              np.ones(input_shapes["threes"], dtype=input_types["threes"]) * 3
      }

    init_fn = mock.Mock()
    init_fn.__call__ = _init_fn
    init_fn.__self__ = None

    self.train_state_init = utils.TrainStateInitializer(None, init_fn, shapes,
                                                        partitioner, types)

    self.ckptdir = self.create_tempdir(name="primary_checkpoints")
    steps = (2, 3)
    self.paths = []
    for s in steps:
      step_dir = self.ckptdir.mkdir(f"checkpoint_{s}")
      step_dir.create_file("checkpoint")
      self.paths += [step_dir.full_path]

  def test_from_checkpoints_specific(self):
    # multiple paths
    ckpt_cfg = utils.RestoreCheckpointConfig(
        path=self.paths, mode="specific", checkpointer_cls=MockCheckpointer)
    restored = self.train_state_init.from_checkpoints([ckpt_cfg])
    self.assertSequenceEqual(self.paths, [state.path for state in restored])
    with self.assertRaisesRegex(ValueError, r"^Expected at most 1 checkpoint"):
      self.train_state_init.from_checkpoint([ckpt_cfg])

  def test_from_checkpoints_latest(self):
    # only restore single latest
    ckpt_cfg = utils.RestoreCheckpointConfig(
        path=self.ckptdir.full_path,
        mode="latest",
        checkpointer_cls=MockCheckpointer)
    restored = list(self.train_state_init.from_checkpoints([ckpt_cfg]))
    assert len(restored) == 1
    self.assertEqual(self.paths[-1], restored[0].path)
    restored = self.train_state_init.from_checkpoint([ckpt_cfg])
    self.assertEqual(self.paths[-1], restored.path)

  def test_from_checkpoints_multiple_configs(self):
    # uses first checkpoint with files present.
    ckpt_cfg = utils.RestoreCheckpointConfig(
        path=self.ckptdir.full_path,
        mode="latest",
        checkpointer_cls=MockCheckpointer)
    secondary_ckptdir = self.create_tempdir(name="secondary_checkpoints")
    for s in (4, 5):
      step_dir = secondary_ckptdir.mkdir(f"checkpoint_{s}")
      step_dir.create_file("checkpoint")
    secondary_ckpt_cfg = utils.RestoreCheckpointConfig(
        path=secondary_ckptdir.full_path,
        mode="latest",
        checkpointer_cls=MockCheckpointer)
    restored = self.train_state_init.from_checkpoint(
        [ckpt_cfg, secondary_ckpt_cfg])
    self.assertEqual(self.paths[-1], restored.path)

  def test_from_checkpoints_multiple_configs_one_empty(self):
    # skips empty_checkpoints directory with no checkpoints present.
    ckpt_cfg = utils.RestoreCheckpointConfig(
        path=self.ckptdir.full_path,
        mode="latest",
        checkpointer_cls=MockCheckpointer)
    empty_ckptdir = self.create_tempdir(name="empty_checkpoints")
    empty_ckpt_cfg = utils.RestoreCheckpointConfig(
        path=empty_ckptdir.full_path,
        mode="latest",
        checkpointer_cls=MockCheckpointer)
    restored = self.train_state_init.from_checkpoint([empty_ckpt_cfg, ckpt_cfg])
    self.assertEqual(self.paths[-1], restored.path)

  def test_from_scratch(self):
    self.assertTrue(
        self.train_state_init.from_scratch(jax.random.PRNGKey(13)).from_scratch)

  def test_from_checkpoint_or_scratch(self):
    ckpt_cfg = utils.RestoreCheckpointConfig(
        path=self.ckptdir.full_path,
        mode="latest",
        checkpointer_cls=MockCheckpointer)
    empty_ckptdir = self.create_tempdir(name="empty_checkpoints")
    empty_ckpt_cfg = utils.RestoreCheckpointConfig(
        path=empty_ckptdir.full_path,
        mode="latest",
        checkpointer_cls=MockCheckpointer)

    init_rng = jax.random.PRNGKey(13)

    # ckpt_cfg has checkpoints, restore from there
    restored = self.train_state_init.from_checkpoint_or_scratch(
        [empty_ckpt_cfg, ckpt_cfg], init_rng=init_rng)
    self.assertEqual(self.paths[-1], restored.path)
    self.assertFalse(restored.from_scratch)

    # no checkpoints available, init from scratch
    initialized = self.train_state_init.from_checkpoint_or_scratch(
        [empty_ckpt_cfg], init_rng=init_rng)
    self.assertTrue(initialized.from_scratch)

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


<commit_before>from agora_site.agora_core.models import Agora
from agora_site.misc.generic_resource import GenericResource, GenericMeta
from agora_site.agora_core.resources.user import UserResource

from tastypie import fields

ELECTION_RESOURCE = 'agora_site.agora_core.resources.election.ElectionResource'

def open_elections(bundle):
    # bundle.obj is an Agora
    return bundle.obj.get_open_elections()

def tallied_elections(bundle):
    return bundle.obj.get_tallied_elections()

def all_elections(bundle):
    return bundle.obj.all_elections()

def active_delegates(bundle):
    return bundle.obj.active_delegates()

class AgoraResource(GenericResource):
    creator = fields.ForeignKey(UserResource, 'creator', full=True)
    members = fields.ManyToManyField(UserResource, 'members', full=True)
    admins = fields.ManyToManyField(UserResource, 'admins', full=True)

    open_elections = fields.ToManyField(ELECTION_RESOURCE,
                                        attribute=open_elections,
                                        null=True)

    tallied_elections = fields.ToManyField(ELECTION_RESOURCE,
                                        attribute=tallied_elections,
                                        null=True)

    all_elections = fields.ToManyField(ELECTION_RESOURCE,
                                       attribute=all_elections,
                                       null=True)

    active_delegates = fields.ToManyField(UserResource,
                                        attribute=active_delegates,
                                        null=True)

    class Meta:

        queryset = Agora.objects.all()
        #authentication = SessionAuthentication()
        list_allowed_methods = ['get', 'post']
        detail_allowed_methods = ['get']
<commit_msg>API: Implemented Agora creation<commit_after>from agora_site.agora_core.models import Agora
from agora_site.misc.generic_resource import GenericResource, GenericMeta
from agora_site.agora_core.resources.user import UserResource

from django.forms import ModelForm
from django.core.urlresolvers import reverse

from tastypie import fields
from tastypie.validation import Validation, CleanedDataFormValidation

ELECTION_RESOURCE = 'agora_site.agora_core.resources.election.ElectionResource'

class CreateAgoraForm(ModelForm):
    class Meta:
        model = Agora
        fields = ('pretty_name', 'short_description', 'is_vote_secret')

class AgoraValidation(Validation):
    def is_valid(self, bundle, request=None):
        if not bundle.data:
            return {'__all__': 'Not quite what I had in mind.'}

        if request.method == "POST":
            return self.validate_post(bundle, request)
        elif request.method == "PUT":
            return self.validate_put(bundle, request)

        return {}

    def validate_put(self, bundle, request):
        # TODO, validate PUT data
        return {}

    def validate_post(self, bundle, request):
        form = CleanedDataFormValidation(form_class=CreateAgoraForm)
        return form.is_valid(bundle, request)

def open_elections(bundle):
    # bundle.obj is an Agora
    return bundle.obj.get_open_elections()

def tallied_elections(bundle):
    return bundle.obj.get_tallied_elections()

def all_elections(bundle):
    return bundle.obj.all_elections()

def active_delegates(bundle):
    return bundle.obj.active_delegates()

class AgoraResource(GenericResource):
    creator = fields.ForeignKey(UserResource, 'creator', full=True)
    members = fields.ManyToManyField(UserResource, 'members', full=True)
    admins = fields.ManyToManyField(UserResource, 'admins', full=True)

    open_elections = fields.ToManyField(ELECTION_RESOURCE,
                                        attribute=open_elections,
                                        null=True)

    tallied_elections = fields.ToManyField(ELECTION_RESOURCE,
                                        attribute=tallied_elections,
                                        null=True)

    all_elections = fields.ToManyField(ELECTION_RESOURCE,
                                       attribute=all_elections,
                                       null=True)

    active_delegates = fields.ToManyField(UserResource,
                                        attribute=active_delegates,
                                        null=True)

    class Meta(GenericMeta):
        queryset = Agora.objects.all()
        list_allowed_methods = ['get', 'post']
        detail_allowed_methods = ['get', 'post']
        validation = AgoraValidation()

    def obj_create(self, bundle, request=None, **kwargs):
        pretty_name = bundle.data['pretty_name']
        short_description = bundle.data['short_description']
        is_vote_secret = bundle.data['is_vote_secret']

        agora = Agora(pretty_name=pretty_name,
                      short_description=short_description,
                      is_vote_secret=is_vote_secret)
        agora.create_name(request.user)
        agora.creator = request.user
        agora.url = request.build_absolute_uri(reverse('agora-view',
            kwargs=dict(username=agora.creator.username, agoraname=agora.name)))

        # we need to save before add members
        agora.save()

        agora.members.add(request.user)
        agora.admins.add(request.user)

        bundle = self.build_bundle(obj=agora, request=request)
        bundle = self.full_dehydrate(bundle)
        return bundle
<|endoftext|>


<commit_before># third party
from sqlalchemy import Column
from sqlalchemy import Integer
from sqlalchemy import String

# relative
from . import Base

# this table holds the list of known nodes usually peer domains
class Node(Base):
    __tablename__ = "node"

    id = Column(Integer(), primary_key=True, autoincrement=True)
    node_uid = Column(String(255))
    node_name = Column(String(255))
<commit_msg>ADD keep_connected column<commit_after># third party
from sqlalchemy import Boolean
from sqlalchemy import Column
from sqlalchemy import Integer
from sqlalchemy import String

# relative
from . import Base

# this table holds the list of known nodes usually peer domains
class Node(Base):
    __tablename__ = "node"

    id = Column(Integer(), primary_key=True, autoincrement=True)
    node_uid = Column(String(255))
    node_name = Column(String(255))
    keep_connected = Column(Boolean(), default=True)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
    pygments.lexers.clean
    ~~~~~~~~~~~~~~~~~~~~~

    Lexer for the Clean language.

    :copyright: Copyright 2006-2015 by the Pygments team, see AUTHORS.
    :license: BSD, see LICENSE for details.
"""

from pygments.lexer import ExtendedRegexLexer, words, include, bygroups
from pygments.token import Comment, Error, Keyword, Literal, Name, Number, \
    Operator, Punctuation, String, Whitespace

__all__ = ['CleanLexer']

class CleanLexer(ExtendedRegexLexer):
    """
    Lexer for the general purpose, state-of-the-art, pure and lazy functional
    programming language Clean (http://clean.cs.ru.nl/Clean).

    .. versionadded: 2.2
    """
    name = 'Clean'
    aliases = ['clean']
    filenames = ['*.icl', '*.dcl']

    keywords = (
        'case', 'ccall', 'class', 'code', 'derive', 'export', 'foreign',
        'from', 'generic', 'if', 'import', 'in', 'infix', 'infixl', 'infixr',
        'inline', 'instance', 'let', 'of', 'otherwise', 'qualified', 'special',
        'stdcall', 'where', 'with')

    modulewords = ('implementation', 'definition', 'system')

    lowerId = r'[a-z][\w\d`]*'
    upperId = r'[A-Z][\w\d`]*'
    funnyId = r'[~@#\$%\^?!+\-*<>\\/|&=:]+'
    scoreUpperId = r'_' + upperId
    scoreLowerId = r'_' + lowerId

    tokens = {
        'root': [
            include('comments'),
            include('keywords'),
            include('module'),
            include('whitespace'),
            include('literals'),
            include('operators'),
            include('delimiters'),
            include('names'),
        ],
        'whitespace': [
            (r'\s+', Whitespace),
        ],
        'comments': [
            (r'//.*\n', Comment.Single),
            (r'/\*', Comment.Multi, 'comments.in'),
            (r'/\*\*', Comment.Special, 'comments.in'),
        ],
        'comments.in': [
            (r'\*\/', Comment.Multi, '#pop'),
            (r'/\*', Comment.Multi, '#push'),
            (r'[^*/]+', Comment.Multi),
            (r'\*(?!/)', Comment.Multi),
            (r'/', Comment.Multi),
        ],
        'keywords': [
            (words(keywords, prefix=r'\b', suffix=r'\b'), Keyword),
        ],
        'module': [
            (words(modulewords, prefix=r'\b', suffix=r'\b'), Keyword.Namespace),
            (r'\bmodule\b', Keyword.Namespace, 'module.name'),
        ],
        'module.name': [
            include('whitespace'),
            (lowerId, Name.Class, '#pop'),
            (upperId, Name.Class, '#pop'),
            (funnyId, Name.Class, '#pop'),
            (scoreLowerId, Name.Class, '#pop'),
            (scoreUpperId, Name.Class, '#pop'),
        ],
        'literals': [
            (r'\'([^\'\\]|\\(x[\da-fA-F]+|\d+|.))\'', Literal.Char),
            (r'[+~-]?0[0-7]+\b', Number.Oct),
            (r'[+~-]?\d+\.\d+(E[+-]?\d+)?', Number.Float),
            (r'[+~-]?\d+\b', Number.Integer),
            (r'[+~-]?0x[\da-fA-F]+\b', Number.Hex),
            (r'True|False', Literal),
            (r'"', String.Double, 'literals.stringd'),
        ],
        'literals.stringd': [
            (r'[^\\"\n]+', String.Double),
            (r'"', String.Double, '#pop'),
            (r'\\.', String.Double),
            (r'[$\n]', Error, '#pop'),
        ],
        'operators': [
            (r'[-~@#\$%\^?!+*<>\\/|&=:\.]+', Operator),
            (r'\b_+\b', Operator),
        ],
        'delimiters': [
            (r'[,;(){}\[\]]', Punctuation),
            (r'(\')([\w`.]+)(\')',
                bygroups(Punctuation, Name.Class, Punctuation)),
        ],
        'names': [
            (r'\b' + lowerId, Name),
            (scoreLowerId, Name),
            (r'\b' + funnyId, Name.Function),
            (r'\b' + upperId, Name.Class),
            (scoreUpperId, Name.Class),
        ]
    }
<commit_msg>inline is only a keyword after code<commit_after># -*- coding: utf-8 -*-
"""
    pygments.lexers.clean
    ~~~~~~~~~~~~~~~~~~~~~

    Lexer for the Clean language.

    :copyright: Copyright 2006-2015 by the Pygments team, see AUTHORS.
    :license: BSD, see LICENSE for details.
"""

from pygments.lexer import ExtendedRegexLexer, words, include, bygroups
from pygments.token import Comment, Error, Keyword, Literal, Name, Number, \
    Operator, Punctuation, String, Whitespace

__all__ = ['CleanLexer']

class CleanLexer(ExtendedRegexLexer):
    """
    Lexer for the general purpose, state-of-the-art, pure and lazy functional
    programming language Clean (http://clean.cs.ru.nl/Clean).

    .. versionadded: 2.2
    """
    name = 'Clean'
    aliases = ['clean']
    filenames = ['*.icl', '*.dcl']

    keywords = (
        'case', 'ccall', 'class', 'code', 'code inline', 'derive', 'export',
        'foreign', 'from', 'generic', 'if', 'import', 'in', 'infix', 'infixl',
        'infixr', 'instance', 'let', 'of', 'otherwise', 'qualified', 'special',
        'stdcall', 'where', 'with')

    modulewords = ('implementation', 'definition', 'system')

    lowerId = r'[a-z][\w\d`]*'
    upperId = r'[A-Z][\w\d`]*'
    funnyId = r'[~@#\$%\^?!+\-*<>\\/|&=:]+'
    scoreUpperId = r'_' + upperId
    scoreLowerId = r'_' + lowerId

    tokens = {
        'root': [
            include('comments'),
            include('keywords'),
            include('module'),
            include('whitespace'),
            include('literals'),
            include('operators'),
            include('delimiters'),
            include('names'),
        ],
        'whitespace': [
            (r'\s+', Whitespace),
        ],
        'comments': [
            (r'//.*\n', Comment.Single),
            (r'/\*', Comment.Multi, 'comments.in'),
            (r'/\*\*', Comment.Special, 'comments.in'),
        ],
        'comments.in': [
            (r'\*\/', Comment.Multi, '#pop'),
            (r'/\*', Comment.Multi, '#push'),
            (r'[^*/]+', Comment.Multi),
            (r'\*(?!/)', Comment.Multi),
            (r'/', Comment.Multi),
        ],
        'keywords': [
            (words(keywords, prefix=r'\b', suffix=r'\b'), Keyword),
        ],
        'module': [
            (words(modulewords, prefix=r'\b', suffix=r'\b'), Keyword.Namespace),
            (r'\bmodule\b', Keyword.Namespace, 'module.name'),
        ],
        'module.name': [
            include('whitespace'),
            (lowerId, Name.Class, '#pop'),
            (upperId, Name.Class, '#pop'),
            (funnyId, Name.Class, '#pop'),
            (scoreLowerId, Name.Class, '#pop'),
            (scoreUpperId, Name.Class, '#pop'),
        ],
        'literals': [
            (r'\'([^\'\\]|\\(x[\da-fA-F]+|\d+|.))\'', Literal.Char),
            (r'[+~-]?0[0-7]+\b', Number.Oct),
            (r'[+~-]?\d+\.\d+(E[+-]?\d+)?', Number.Float),
            (r'[+~-]?\d+\b', Number.Integer),
            (r'[+~-]?0x[\da-fA-F]+\b', Number.Hex),
            (r'True|False', Literal),
            (r'"', String.Double, 'literals.stringd'),
        ],
        'literals.stringd': [
            (r'[^\\"\n]+', String.Double),
            (r'"', String.Double, '#pop'),
            (r'\\.', String.Double),
            (r'[$\n]', Error, '#pop'),
        ],
        'operators': [
            (r'[-~@#\$%\^?!+*<>\\/|&=:\.]+', Operator),
            (r'\b_+\b', Operator),
        ],
        'delimiters': [
            (r'[,;(){}\[\]]', Punctuation),
            (r'(\')([\w`.]+)(\')',
                bygroups(Punctuation, Name.Class, Punctuation)),
        ],
        'names': [
            (r'\b' + lowerId, Name),
            (scoreLowerId, Name),
            (r'\b' + funnyId, Name.Function),
            (r'\b' + upperId, Name.Class),
            (scoreUpperId, Name.Class),
        ]
    }
<|endoftext|>


<commit_before>import csv
from random import randint
from itertools import chain

class TableFormatError(Exception):
    '''Raised when the user tries to load an improperly formatted table
       file.'''

    def __init__(self, msg):
        self.err_msg = msg

def load_table(filepath):
    '''Return a dict representing a roll table loaded from filepath.

    Loads a roll table from the CSV file at filepath into a dict whose keys
    are ranges containing the range of rolls (min, max) associated with the
    event specified in that key's value (a string describing the event).
    It is assumed that the first row of the file contains some descriptive
    headers and the row is ignored.

    The first column of the CSV should be the numbers or ranges of numbers
    to roll in order to 'bring about' the associated event in the same row
    of the second column. Ranges should be specified with dashes e.g.
    a roll of 1 to 10 inclusive would be written as '1-10'. None of the
    intervals should overlap. If there is a gap in the table i.e. a roll
    within the bounds of the table which is not associated with an event, an
    IOError is raised.'''

    table = {}
    with open(filepath, newline='') as table_file:
        table_reader = csv.reader(table_file, strict=True)
        for row in table_reader:
            if len(row) != 2:
                # Tables should only have two columns.
                raise TableFormatError("Tables should only have two columns.")
            if table_reader.line_num == 1:
                # Ignore the first line.
                continue
            roll = row[0]
            event = row[1]
            if row[0].find("-") != -1:
                # A range of rolls has been specified for this table item.
                min_roll = int(roll[:roll.find("-")])
                max_roll = int(roll[roll.find("-")+1:])
                table[range(min_roll, max_roll+1)] = event
            else:
                # A single roll has been specified for this table item.
                roll_num = int(roll)
                table[range(roll_num, roll_num+1)] = event

    # Check if there is a gap in the table by comparing its keys to a range.
    rolls_in_table = sorted(list(chain(*table.keys())))
    max_in_table = max(rolls_in_table)
    min_in_table = min(rolls_in_table)
    gap = rolls_in_table != list(range(min_in_table, max_in_table+1))
    if gap:
        filename = filepath[filepath.rfind("/")+1:]
        err_msg = 'There is a gap in the table: {}'.format(filename)
        raise TableFormatError(err_msg)

    return table

def roll_against(table):
    '''Return a randomly chosen entry (a dict value) from a roll table.

    table must be a dict of the same format as those created by load_table.'''

    # Get the range of rolls that the table caters for.
    permitted = list(chain(*table.keys()))
    max_permitted = max(permitted)
    min_permitted = min(permitted)

    # Generated a random integer bounded by the maximum and minimum permitted
    # rolls.
    roll = randint(min_permitted, max_permitted)

    # Check which event the roll corresponds to and return the description of
    # that event.
    for entry in table.items():
        if roll in entry[0]:
            return entry[1]
<commit_msg>Remove file name from load_table gap error<commit_after>import csv
from random import randint
from itertools import chain

class TableFormatError(Exception):
    '''Raised when the user tries to load an improperly formatted table
       file.'''

    def __init__(self, msg):
        self.err_msg = msg

def load_table(filepath):
    '''Return a dict representing a roll table loaded from filepath.

    Loads a roll table from the CSV file at filepath into a dict whose keys
    are ranges containing the range of rolls (min, max) associated with the
    event specified in that key's value (a string describing the event).
    It is assumed that the first row of the file contains some descriptive
    headers and the row is ignored.

    The first column of the CSV should be the numbers or ranges of numbers
    to roll in order to 'bring about' the associated event in the same row
    of the second column. Ranges should be specified with dashes e.g.
    a roll of 1 to 10 inclusive would be written as '1-10'. None of the
    intervals should overlap. If there is a gap in the table i.e. a roll
    within the bounds of the table which is not associated with an event, an
    IOError is raised.'''

    table = {}
    with open(filepath, newline='') as table_file:
        table_reader = csv.reader(table_file, strict=True)
        for row in table_reader:
            if len(row) != 2:
                # Tables should only have two columns.
                raise TableFormatError("Tables should only have two columns.")
            if table_reader.line_num == 1:
                # Ignore the first line.
                continue
            roll = row[0]
            event = row[1]
            if row[0].find("-") != -1:
                # A range of rolls has been specified for this table item.
                min_roll = int(roll[:roll.find("-")])
                max_roll = int(roll[roll.find("-")+1:])
                table[range(min_roll, max_roll+1)] = event
            else:
                # A single roll has been specified for this table item.
                roll_num = int(roll)
                table[range(roll_num, roll_num+1)] = event

    # Check if there is a gap in the table by comparing its keys to a range.
    rolls_in_table = sorted(list(chain(*table.keys())))
    max_in_table = max(rolls_in_table)
    min_in_table = min(rolls_in_table)
    gap = rolls_in_table != list(range(min_in_table, max_in_table+1))
    if gap:
        filename = filepath[filepath.rfind("/")+1:]
        err_msg = 'There is a gap in the table.'.format(filename)
        raise TableFormatError(err_msg)

    return table

def roll_against(table):
    '''Return a randomly chosen entry (a dict value) from a roll table.

    table must be a dict of the same format as those created by load_table.'''

    # Get the range of rolls that the table caters for.
    permitted = list(chain(*table.keys()))
    max_permitted = max(permitted)
    min_permitted = min(permitted)

    # Generated a random integer bounded by the maximum and minimum permitted
    # rolls.
    roll = randint(min_permitted, max_permitted)

    # Check which event the roll corresponds to and return the description of
    # that event.
    for entry in table.items():
        if roll in entry[0]:
            return entry[1]
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- encoding: utf-8 -*-

# Copyright 2011-2014, Nigel Small
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from subprocess import check_output
from shutil import rmtree

from py2neo.env import NEO4J_DIST
from py2neo.packages.httpstream import NetworkAddressError
from py2neo.packages.httpstream.packages.urimagic import URI
from py2neo.server import download, dist_archive_name, GraphServer

__all__ = ["Warehouse", "Box"]

class Warehouse(object):
    """ Container for Neo4j :class:`.Box` installations.

    :arg home: Home directory for this warehouse (defaults to the path
               held in ``$NEOBOX_HOME`` if specified, otherwise ``$HOME/.neobox``)

    """

    __instances = {}

    def __new__(cls, home=None):
        if not home:
            home = os.getenv("NEOBOX_HOME") or os.path.expanduser("~/.neobox")
        key = home
        try:
            inst = cls.__instances[key]
        except KeyError:
            inst = super(Warehouse, cls).__new__(cls)
            inst.home = home
            cls.__instances[key] = inst
        return inst

    home = None

    def __repr__(self):
        return "<Warehouse home=%r>" % self.home

    def __hash__(self):
        return hash(self.home)

    def box(self, name):
        """ Fetch a named box. This box may or may not contain a Neo4j
        server installation.

        :arg name: Name of the box to fetch.
        :rtype: :class:`.Box`

        """
        return Box(self, name)

    def boxes(self):
        """ Fetch a list of all installed boxes.

        :rtype: generator

        """
        for name in os.listdir(os.path.join(self.home, "boxes")):
            yield Box(self, name)

    @property
    def _ports_path(self):
        return os.path.join(self.home, "ports")

    def ensure_downloaded(self, edition, version):
        dist_path = os.path.join(self.home, "dist")
        try:
            os.makedirs(dist_path)
        except FileExistsError:
            pass
        filename = os.path.join(dist_path, dist_archive_name(edition, version))
        if os.path.isfile(filename):
            return filename
        try:
            return download(edition, version, dist_path)
        except NetworkAddressError:
            uri = URI(NEO4J_DIST)
            raise RuntimeError("Not able to connect to %s" % uri.host)

    @property
    def _ports(self):
        ports_path = self._ports_path
        try:
            os.makedirs(ports_path)
        except FileExistsError:
            pass
        ports = [port for port in os.listdir(ports_path)]
        return {os.path.basename(os.readlink(os.path.join(ports_path, port))): int(port)
                for port in ports}

    def _assign_port(self, name, port=None):
        if not port:
            ports = self._ports.values()
            if ports:
                port = max(ports) + 2
            else:
                port = 47470
        os.symlink(os.path.join("..", "boxes", name), os.path.join(self._ports_path, str(port)))
        return port

    def _remove_port(self, name):
        try:
            port = self._ports[name]
        except KeyError:
            pass
        else:
            os.remove(os.path.join(self._ports_path, str(port)))

class Box(object):
    """ Named container for a single Neo4j installation.

    :arg warehouse: The warehouse to which this box belongs.
    :arg name: The name of this box.

    """

    __instances = {}

    def __new__(cls, warehouse, name):
        key = (warehouse, name)
        try:
            inst = cls.__instances[key]
        except KeyError:
            inst = super(Box, cls).__new__(cls)
            inst.warehouse = warehouse
            inst.name = name
            cls.__instances[key] = inst
        return inst

    warehouse = None
    name = None

    def __repr__(self):
        return "<Box warehouse=%r name=%r>" % (self.warehouse.home, self.name)

    def __hash__(self):
        return hash(self.warehouse) ^ hash(self.name)

    @property
    def home(self):
        """ Base directory for this box. This directory may not exist.
        """
        return os.path.join(self.warehouse.home, "boxes", self.name)

    @property
    def server(self):
        """ The :class:`py2neo.server.GraphServer` installed in this box or
        :const:`None` if this box is empty.
        """
        if os.path.exists(self.home):
            return GraphServer(os.path.join(self.home, "neo4j"))
        else:
            return None

    def install(self, edition, version):
        """ Install a Neo4j server in this box.

        :arg edition: Neo4j edition (`'community'` or `'enterprise'`)
        :arg version: Neo4j version (e.g. `'2.1.5'`)
        :rtype: :class:`py2neo.server.GraphServer`

        """
        inst_path = self.home
        if self.server is not None:
            raise ValueError("A box named %r already exists" % self.name)
        filename = self.warehouse.ensure_downloaded(edition, version)
        os.makedirs(inst_path)
        # The Python tarfile module doesn't seem to recognise the Neo4j tar format :-(
        check_output("tar -x -C \"%s\" -f \"%s\"" % (inst_path, filename), shell=True)
        os.symlink(os.listdir(inst_path)[0], self.server.home)
        port = self.warehouse._assign_port(self.name)
        server = self.server
        server.update_server_properties(webserver_port=port, webserver_https_port=(port + 1))
        return server

    def remove(self, force=False):
        """ Remove any installed Neo4j server from this box. Running
        servers cannot be removed unless forced.

        :arg force: :const:`True` to remove server even if running.

        """
        if self.server.pid and not force:
            raise RuntimeError("Cannot remove a running server instance")
        self.warehouse._remove_port(self.name)
        rmtree(self.home, ignore_errors=(not force))

    def rename(self, new_name):
        """ Rename this box.

        :arg new_name: The new name for this box.

        """
        inst_path = self.home
        if not os.path.isdir(inst_path):
            raise ValueError("No box named %r exists" % self.name)
        new_box = self.warehouse.box(new_name)
        new_inst_path = new_box.home
        if os.path.isdir(new_inst_path):
            raise ValueError("A server instance named %r already exists" % new_name)
        port = self.warehouse._ports[self.name]
        self.warehouse._remove_port(self.name)
        os.rename(inst_path, new_inst_path)
        self.warehouse._assign_port(new_name, port)
        self.name = new_name
<commit_msg>Swapped order of checks on removal to enhance forcing<commit_after>#!/usr/bin/env python
# -*- encoding: utf-8 -*-

# Copyright 2011-2014, Nigel Small
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from subprocess import check_output
from shutil import rmtree

from py2neo.env import NEO4J_DIST
from py2neo.packages.httpstream import NetworkAddressError
from py2neo.packages.httpstream.packages.urimagic import URI
from py2neo.server import download, dist_archive_name, GraphServer

__all__ = ["Warehouse", "Box"]

class Warehouse(object):
    """ Container for Neo4j :class:`.Box` installations.

    :arg home: Home directory for this warehouse (defaults to the path
               held in ``$NEOBOX_HOME`` if specified, otherwise ``$HOME/.neobox``)

    """

    __instances = {}

    def __new__(cls, home=None):
        if not home:
            home = os.getenv("NEOBOX_HOME") or os.path.expanduser("~/.neobox")
        key = home
        try:
            inst = cls.__instances[key]
        except KeyError:
            inst = super(Warehouse, cls).__new__(cls)
            inst.home = home
            cls.__instances[key] = inst
        return inst

    home = None

    def __repr__(self):
        return "<Warehouse home=%r>" % self.home

    def __hash__(self):
        return hash(self.home)

    def box(self, name):
        """ Fetch a named box. This box may or may not contain a Neo4j
        server installation.

        :arg name: Name of the box to fetch.
        :rtype: :class:`.Box`

        """
        return Box(self, name)

    def boxes(self):
        """ Fetch a list of all installed boxes.

        :rtype: generator

        """
        for name in os.listdir(os.path.join(self.home, "boxes")):
            yield Box(self, name)

    @property
    def _ports_path(self):
        return os.path.join(self.home, "ports")

    def ensure_downloaded(self, edition, version):
        dist_path = os.path.join(self.home, "dist")
        try:
            os.makedirs(dist_path)
        except FileExistsError:
            pass
        filename = os.path.join(dist_path, dist_archive_name(edition, version))
        if os.path.isfile(filename):
            return filename
        try:
            return download(edition, version, dist_path)
        except NetworkAddressError:
            uri = URI(NEO4J_DIST)
            raise RuntimeError("Not able to connect to %s" % uri.host)

    @property
    def _ports(self):
        ports_path = self._ports_path
        try:
            os.makedirs(ports_path)
        except FileExistsError:
            pass
        ports = [port for port in os.listdir(ports_path)]
        return {os.path.basename(os.readlink(os.path.join(ports_path, port))): int(port)
                for port in ports}

    def _assign_port(self, name, port=None):
        if not port:
            ports = self._ports.values()
            if ports:
                port = max(ports) + 2
            else:
                port = 47470
        os.symlink(os.path.join("..", "boxes", name), os.path.join(self._ports_path, str(port)))
        return port

    def _remove_port(self, name):
        try:
            port = self._ports[name]
        except KeyError:
            pass
        else:
            os.remove(os.path.join(self._ports_path, str(port)))

class Box(object):
    """ Named container for a single Neo4j installation.

    :arg warehouse: The warehouse to which this box belongs.
    :arg name: The name of this box.

    """

    __instances = {}

    def __new__(cls, warehouse, name):
        key = (warehouse, name)
        try:
            inst = cls.__instances[key]
        except KeyError:
            inst = super(Box, cls).__new__(cls)
            inst.warehouse = warehouse
            inst.name = name
            cls.__instances[key] = inst
        return inst

    warehouse = None
    name = None

    def __repr__(self):
        return "<Box warehouse=%r name=%r>" % (self.warehouse.home, self.name)

    def __hash__(self):
        return hash(self.warehouse) ^ hash(self.name)

    @property
    def home(self):
        """ Base directory for this box. This directory may not exist.
        """
        return os.path.join(self.warehouse.home, "boxes", self.name)

    @property
    def server(self):
        """ The :class:`py2neo.server.GraphServer` installed in this box or
        :const:`None` if this box is empty.
        """
        if os.path.exists(self.home):
            return GraphServer(os.path.join(self.home, "neo4j"))
        else:
            return None

    def install(self, edition, version):
        """ Install a Neo4j server in this box.

        :arg edition: Neo4j edition (`'community'` or `'enterprise'`)
        :arg version: Neo4j version (e.g. `'2.1.5'`)
        :rtype: :class:`py2neo.server.GraphServer`

        """
        inst_path = self.home
        if self.server is not None:
            raise ValueError("A box named %r already exists" % self.name)
        filename = self.warehouse.ensure_downloaded(edition, version)
        os.makedirs(inst_path)
        # The Python tarfile module doesn't seem to recognise the Neo4j tar format :-(
        check_output("tar -x -C \"%s\" -f \"%s\"" % (inst_path, filename), shell=True)
        os.symlink(os.listdir(inst_path)[0], self.server.home)
        port = self.warehouse._assign_port(self.name)
        server = self.server
        server.update_server_properties(webserver_port=port, webserver_https_port=(port + 1))
        return server

    def remove(self, force=False):
        """ Remove any installed Neo4j server from this box. Running
        servers cannot be removed unless forced.

        :arg force: :const:`True` to remove server even if running.

        """
        if not force and self.server.pid:
            raise RuntimeError("Cannot remove a running server instance")
        self.warehouse._remove_port(self.name)
        rmtree(self.home, ignore_errors=(not force))

    def rename(self, new_name):
        """ Rename this box.

        :arg new_name: The new name for this box.

        """
        inst_path = self.home
        if not os.path.isdir(inst_path):
            raise ValueError("No box named %r exists" % self.name)
        new_box = self.warehouse.box(new_name)
        new_inst_path = new_box.home
        if os.path.isdir(new_inst_path):
            raise ValueError("A server instance named %r already exists" % new_name)
        port = self.warehouse._ports[self.name]
        self.warehouse._remove_port(self.name)
        os.rename(inst_path, new_inst_path)
        self.warehouse._assign_port(new_name, port)
        self.name = new_name
<|endoftext|>