<commit_before>"""
Puzzles System - Provides a typeclass and commands for
objects that can be combined (i.e. 'use'd) to produce
new objects.

Evennia contribution - Henddher 2018

A Puzzle is a recipe of what objects (aka parts) must
be combined by a player so a new set of objects
(aka results) are automatically created.

Consider this simple Puzzle:

    orange, mango, yogurt, blender = fruit smoothie

As a Builder:

    @create/drop orange
    @create/drop mango
    @create/drop yogurt
    @create/drop blender

    @puzzle smoothie puzzle, orange, mango, yogurt, blender = fruit smoothie
    ...
    Puzzle smoothie puzzle (#1234) created successfuly.

    @destroy/force orange, mango, yogurt, blender

    @armpuzzle #1234
    Part orange is spawned at ...
    Part mango is spawned at ...
    ....
    Puzzle smoothie puzzle (#1234) has been armed successfully

As Player:

    use orange, mango, yogurt, blender
    ...
    Genius, you blended all fruits to create a yummy smoothie!

Details:

Puzzles are created from existing objects. The given
objects are introspected to create prototypes for the
puzzle parts. These prototypes become the puzzle recipe.
(See PuzzleRecipeObject and @puzzle command).

At a later time, a Builder or a Script can arm the puzzle
and spawn all puzzle parts (PuzzlePartObject) in their
respective locations (See @armpuzzle).

A regular player can collect the puzzle parts and combine
them (See use command). If player has specified
all pieces, the puzzle is considered solved and all
its puzzle parts are destroyed while the puzzle results
are spawened on their corresponding location.

Installation:

Add the PuzzleSystemCmdSet to all players.
Alternatively:

    @py self.cmdset.add('evennia.contrib.puzzles.PuzzleSystemCmdSet')

"""

import itertools
from random import choice
from django.conf import settings
from evennia import create_object
from evennia import CmdSet
from evennia import DefaultObject
from evennia import DefaultCharacter
from evennia import DefaultRoom
from evennia.commands.default.muxcommand import MuxCommand
from evennia.utils.utils import inherits_from
from evennia.utils import search, utils, logger
from evennia.utils.spawner import spawn

# ----------- UTILITY FUNCTIONS ------------

def proto_def(obj, with_tags=True):
    """
    Basic properties needed to spawn
    and compare recipe with candidate part
    """
    protodef = {
        'key': obj.key,
        'typeclass': 'evennia.contrib.puzzles.PuzzlePartObject',  # FIXME: what if obj is another typeclass
        'desc': obj.db.desc,
        'location': obj.location,
        # FIXME: Can tags be INVISIBLE? We don't want player to know an object belongs to a puzzle
        'tags': [(_PUZZLES_TAG_MEMBER, _PUZZLES_TAG_CATEGORY)],
    }
    if not with_tags:
        del(protodef['tags'])
    return protodef

# ------------------------------------------

# Tag used by puzzles
_PUZZLES_TAG_CATEGORY = 'puzzles'
_PUZZLES_TAG_RECIPE = 'puzzle_recipe'
# puzzle part and puzzle result
_PUZZLES_TAG_MEMBER = 'puzzle_member'

class PuzzlePartObject(DefaultObject):
    """
    Puzzle Part, typically used by @armpuzzle command
    """

    def mark_as_puzzle_member(self, puzzle_name):
        """
        Marks this object as a member of puzzle named
        puzzle_name
        """
        # FIXME: if multiple puzzles have the same
        # puzzle_name, their ingredients may be
        # combined but leave other parts orphan
        # Similarly, if a puzzle_name were changed,
        # its parts will become orphan
        # Perhaps we should use #dbref but that will
        # force specific parts to be combined
        self.db.puzzle_name = puzzle_name

class PuzzleRecipeObject(DefaultObject):
    """
    Definition of a Puzzle Recipe
    """

    def save_recipe(self, puzzle_name, parts, results):
        self.db.puzzle_name = puzzle_name
        self.db.parts = tuple(parts)
        self.db.results = tuple(results)
        self.tags.add(_PUZZLES_TAG_RECIPE, category=_PUZZLES_TAG_CATEGORY)

class CmdCreatePuzzleRecipe(MuxCommand):
    """
    Creates a puzzle recipe.

    Each part and result must exist and be placed in their corresponding location.
    All parts and results are left intact. Caller must explicitly
    destroy them.

    Usage:
        @puzzle name,<part1[,part2,...>] = <result1[,result2,...]>
    """

    key = '@puzzle'
    aliases = '@puzzlerecipe'
    locks = 'cmd:perm(puzzle) or perm(Builder)'
    help_category = 'Puzzles'

    def func(self):
        caller = self.caller

        if len(self.lhslist) < 2 \
            or not self.rhs:
            string = "Usage: @puzzle name,<part1[,...]> = <result1[,...]>"
            caller.msg(string)
            return

        puzzle_name = self.lhslist[0]

        def is_valid_obj_location(obj):
            valid = True
            # Valid locations are: room, ...
            # TODO: other valid locations must be added here
            # Certain locations can be handled accordingly: e.g,
            # a part is located in a character's inventory,
            # perhaps will translate into the player character
            # having the part in his/her inventory while being
            # located in the same room where the builder was
            # located.
            # Parts and results may have different valid locations
            # TODO: handle contents of a given part
            if not inherits_from(obj.location, settings.BASE_ROOM_TYPECLASS):
                caller.msg('Invalid location for %s' % (obj.key))
                valid = False
            return valid

        def is_valid_part_location(part):
            return is_valid_obj_location(part)

        def is_valid_result_location(part):
            return is_valid_obj_location(part)

        parts = []
        for objname in self.lhslist[1:]:
            obj = caller.search(objname)
            if not obj:
                return
            if not is_valid_part_location(obj):
                return
            parts.append(obj)

        results = []
        for objname in self.rhslist:
            obj = caller.search(objname)
            if not obj:
                return
            if not is_valid_result_location(obj):
                return
            results.append(obj)

        for part in parts:
            caller.msg('Part %s(%s)' % (part.name, part.dbref))

        for result in results:
            caller.msg('Result %s(%s)' % (result.name, result.dbref))

        proto_parts = [proto_def(obj) for obj in parts]
        proto_results = [proto_def(obj) for obj in results]

        puzzle = create_object(PuzzleRecipeObject, key=puzzle_name)
        puzzle.save_recipe(puzzle_name, proto_parts, proto_results)

        caller.msg(
            "Puzzle |y'%s' |w%s(%s)|n has been created |gsuccessfully|n."
            % (puzzle.db.puzzle_name, puzzle.name, puzzle.dbref))
        caller.msg(
            'You may now dispose all parts and results. '
            'Typically, results and parts are useless afterwards.\n'
            'You are now able to arm this puzzle using Builder command:\n'
            '    @armpuzzle <puzzle #dbref>\n\n'
            'Or programmatically.\n'
        )

        # FIXME: puzzle recipe object exists but it has no location
        # should we create a PuzzleLibrary where all puzzles are
        # kept and cannot be reached by players?

class CmdArmPuzzle(MuxCommand):
    """
    Arms a puzzle by spawning all its parts
    """

    key = '@armpuzzle'
    # FIXME: permissions for scripts?
    locks = 'cmd:perm(armpuzzle) or perm(Builder)'
    help_category = 'Puzzles'

    def func(self):
        caller = self.caller

        if self.args is None or not utils.dbref(self.args):
            caller.msg("A puzzle recipe's #dbref must be specified")
            return

        puzzle = caller.search(self.args, global_search=True)
        if not puzzle or not inherits_from(puzzle, PuzzleRecipeObject):
            return

        caller.msg(
            "Puzzle Recipe %s(%s) '%s' found.\nSpawning %d parts ..." % (
            puzzle.name, puzzle.dbref, puzzle.db.puzzle_name, len(puzzle.db.parts)))

        for proto_part in puzzle.db.parts:
            # caller.msg('Protopart %r %r' % (proto_part, type(proto_part)))
            part = spawn(proto_part)[0]
            caller.msg("Part %s(%s) spawned and placed at %s(%s)" % (part.name, part.dbref, part.location, part.location.dbref))
            part.mark_as_puzzle_member(puzzle.db.puzzle_name)

        caller.msg("Puzzle armed |gsuccessfully|n.")

class CmdUsePuzzleParts(MuxCommand):
    """
    Searches for all puzzles whose parts
    match the given set of objects. If
    there are matching puzzles, the result
    objects are spawned in their corresponding
    location if all parts have been passed in.

    Usage:
        use <part1[,part2,...>]
    """

    # TODO: consider allowing builder to provide
    # messages and "hooks" that can be displayed
    # and/or fired whenever the resolver of the puzzle
    # enters the location where a result was spawned

    key = 'use'
    aliases = 'combine'
    locks = 'cmd:pperm(use) or pperm(Player)'
    help_category = 'Puzzles'

    def func(self):
        caller = self.caller

        if not self.lhs:
            caller.msg('Use what?')
            return

        many = 'these' if len(self.lhslist) > 1 else 'this'

        # either all are parts, or abort finding matching puzzles
        parts = []
        partnames = self.lhslist[:]
        for partname in partnames:
            part = caller.search(
                partname,
                multimatch_string='Which %s. There are many.\n' % (partname),
                nofound_string='There is no %s around.' % (partname)
            )

            if not part:
                return

            if not part.tags.get(_PUZZLES_TAG_MEMBER, category=_PUZZLES_TAG_CATEGORY) \
                or not inherits_from(part, PuzzlePartObject):

                # not a puzzle part ... abort
                caller.msg('You have no idea how %s can be used' % (many))
                return

            # a valid part
            parts.append(part)

        # Create lookup dict
        parts_dict = dict((part.dbref, part) for part in parts)

        # Group parts by their puzzle name
        puzzle_ingredients = dict()
        for part in parts:
            puzzle_name = part.db.puzzle_name
            if puzzle_name not in puzzle_ingredients:
                puzzle_ingredients[puzzle_name] = []
            puzzle_ingredients[puzzle_name].append(
                (part.dbref, proto_def(part, with_tags=False))
            )

        # Find all puzzles by puzzle name
        # FIXME: we rely on obj.db.puzzle_name which is visible and may be cnaged afterwards. Can we lock it and hide it?
        puzzles = []
        for puzzle_name, parts in puzzle_ingredients.items():
            _puzzles = caller.search(
                    puzzle_name,
                    typeclass=[PuzzleRecipeObject],
                    attribute_name='puzzle_name',
                    quiet=True,
                    exact=True,
                    global_search=True)
            if not _puzzles:
                continue
            else:
                puzzles.extend(_puzzles)

        # Create lookup dict
        puzzles_dict = dict((puzzle.dbref, puzzle) for puzzle in puzzles)

        # Check if parts can be combined to solve a puzzle
        matched_puzzles = dict()
        for puzzle in puzzles:
            puzzleparts = puzzle.db.parts[:]
            parts = puzzle_ingredients[puzzle.db.puzzle_name][:]
            pz = 0
            p = 0
            matched_dbrefparts = set()
            while pz < len(puzzleparts) and p < len(parts):
                puzzlepart = puzzleparts[pz]
                if 'tags' in puzzlepart:
                    # remove 'tags' as they will prevent equality
                    del(puzzlepart['tags'])
                dbref, part = parts[p]
                if part == puzzlepart:
                    pz += 1
                    matched_dbrefparts.add(dbref)
                else:
                    pass
                p += 1
            else:
                if len(puzzleparts) == len(matched_dbrefparts):
                    matched_puzzles[puzzle.dbref] = matched_dbrefparts

        if len(matched_puzzles) == 0:
            # FIXME: Add more random messages
            #    random part falls and lands on your feet
            #    random part hits you square on the face
            caller.msg("As you try to utilize %s, nothing happens." % (many))
            return

        puzzletuples = sorted(matched_puzzles.items(), key=lambda t: len(t[1]), reverse=True)

        # sort all matched puzzles and pick largest one(s)
        puzzledbref, matched_dbrefparts = puzzletuples[0]
        nparts = len(matched_dbrefparts)
        largest_puzzles = list(itertools.takewhile(lambda t: len(t[1]) == nparts, puzzletuples))

        # if there are more than one, let user pick
        if len(largest_puzzles) > 1:
            # FIXME: pick a random one or let user choose?
            caller.msg(
                'Your gears start turning and a bunch of ideas come to your mind ...\n%s' % (
                ' ...\n'.join([lp.db.puzzle_name for lp in largest_puzzles]))
            )
            puzzle = choice(largest_puzzles)
            caller.msg("You try %s ..." % (puzzle.db.puzzle_name))

        # got one, spawn its results
        puzzle = puzzles_dict[puzzledbref]
        # FIXME: DRY with parts
        for proto_result in puzzle.db.results:
            result = spawn(proto_result)[0]
            result.mark_as_puzzle_member(puzzle.db.puzzle_name)
            # FIXME: add 'ramdon' messages:
            # Hmmm ... did I search result.location?
            # What was that? ... I heard something in result.location?
            # Eureka! you built a result

        # Destroy all parts used
        for dbref in matched_dbrefparts:
            parts_dict[dbref].delete()

        # FIXME: Add random messages
        #    You are a genius ... no matter what your 2nd grade teacher told you
        #    You hear thunders and a cloud of dust raises leaving
        caller.msg("Puzzle solved |gsuccessfully|n.")

class CmdListPuzzleRecipes(MuxCommand):
    """
    Searches for all puzzle recipes

    Usage:
        @lspuzzlerecipes
    """

    key = '@lspuzzlerecipes'
    locks = 'cmd:perm(lspuzzlerecipes) or perm(Builder)'
    help_category = 'Puzzles'

    def func(self):
        caller = self.caller
        # TODO: use @tags/search puzzle_recipe : puzzles

class CmdListArmedPuzzles(MuxCommand):
    """
    Searches for all armed puzzles

    Usage:
        @lsarmedpuzzles
    """

    key = '@lsarmedpuzzles'
    locks = 'cmd:perm(lsarmedpuzzles) or perm(Builder)'
    help_category = 'Puzzles'

    def func(self):
        caller = self.caller
        # TODO: use @tags/search puzzle_member : puzzles

class PuzzleSystemCmdSet(CmdSet):
    """
    CmdSet to create, arm and resolve Puzzles

    Add with @py self.cmdset.add("evennia.contrib.puzzles.PuzzlesCmdSet")
    """

    def at_cmdset_creation(self):
        super(PuzzleSystemCmdSetCmdSet, self).at_cmdset_creation()

        self.add(CmdCreatePuzzleRecipe())
        self.add(CmdArmPuzzle())
        self.add(CmdUsePuzzleParts())
<commit_msg>When multiple puzzles are matched, show their names to the caller and then randomly pick one<commit_after>"""
Puzzles System - Provides a typeclass and commands for
objects that can be combined (i.e. 'use'd) to produce
new objects.

Evennia contribution - Henddher 2018

A Puzzle is a recipe of what objects (aka parts) must
be combined by a player so a new set of objects
(aka results) are automatically created.

Consider this simple Puzzle:

    orange, mango, yogurt, blender = fruit smoothie

As a Builder:

    @create/drop orange
    @create/drop mango
    @create/drop yogurt
    @create/drop blender

    @puzzle smoothie puzzle, orange, mango, yogurt, blender = fruit smoothie
    ...
    Puzzle smoothie puzzle (#1234) created successfuly.

    @destroy/force orange, mango, yogurt, blender

    @armpuzzle #1234
    Part orange is spawned at ...
    Part mango is spawned at ...
    ....
    Puzzle smoothie puzzle (#1234) has been armed successfully

As Player:

    use orange, mango, yogurt, blender
    ...
    Genius, you blended all fruits to create a yummy smoothie!

Details:

Puzzles are created from existing objects. The given
objects are introspected to create prototypes for the
puzzle parts. These prototypes become the puzzle recipe.
(See PuzzleRecipeObject and @puzzle command).

At a later time, a Builder or a Script can arm the puzzle
and spawn all puzzle parts (PuzzlePartObject) in their
respective locations (See @armpuzzle).

A regular player can collect the puzzle parts and combine
them (See use command). If player has specified
all pieces, the puzzle is considered solved and all
its puzzle parts are destroyed while the puzzle results
are spawened on their corresponding location.

Installation:

Add the PuzzleSystemCmdSet to all players.
Alternatively:

    @py self.cmdset.add('evennia.contrib.puzzles.PuzzleSystemCmdSet')

"""

import itertools
from random import choice
from django.conf import settings
from evennia import create_object
from evennia import CmdSet
from evennia import DefaultObject
from evennia import DefaultCharacter
from evennia import DefaultRoom
from evennia.commands.default.muxcommand import MuxCommand
from evennia.utils.utils import inherits_from
from evennia.utils import search, utils, logger
from evennia.utils.spawner import spawn

# ----------- UTILITY FUNCTIONS ------------

def proto_def(obj, with_tags=True):
    """
    Basic properties needed to spawn
    and compare recipe with candidate part
    """
    protodef = {
        # FIXME: Don't we need to honor ALL properties? locks, perms, etc.
        'key': obj.key,
        'typeclass': 'evennia.contrib.puzzles.PuzzlePartObject',  # FIXME: what if obj is another typeclass
        'desc': obj.db.desc,
        'location': obj.location,
        # FIXME: Can tags be INVISIBLE? We don't want player to know an object belongs to a puzzle
        'tags': [(_PUZZLES_TAG_MEMBER, _PUZZLES_TAG_CATEGORY)],
    }
    if not with_tags:
        del(protodef['tags'])
    return protodef

# ------------------------------------------

# Tag used by puzzles
_PUZZLES_TAG_CATEGORY = 'puzzles'
_PUZZLES_TAG_RECIPE = 'puzzle_recipe'
# puzzle part and puzzle result
_PUZZLES_TAG_MEMBER = 'puzzle_member'

class PuzzlePartObject(DefaultObject):
    """
    Puzzle Part, typically used by @armpuzzle command
    """

    def mark_as_puzzle_member(self, puzzle_name):
        """
        Marks this object as a member of puzzle named
        puzzle_name
        """
        # FIXME: if multiple puzzles have the same
        # puzzle_name, their ingredients may be
        # combined but leave other parts orphan
        # Similarly, if a puzzle_name were changed,
        # its parts will become orphan
        # Perhaps we should use #dbref but that will
        # force specific parts to be combined
        self.db.puzzle_name = puzzle_name

class PuzzleRecipeObject(DefaultObject):
    """
    Definition of a Puzzle Recipe
    """

    def save_recipe(self, puzzle_name, parts, results):
        self.db.puzzle_name = puzzle_name
        self.db.parts = tuple(parts)
        self.db.results = tuple(results)
        self.tags.add(_PUZZLES_TAG_RECIPE, category=_PUZZLES_TAG_CATEGORY)

class CmdCreatePuzzleRecipe(MuxCommand):
    """
    Creates a puzzle recipe.

    Each part and result must exist and be placed in their corresponding location.
    All parts and results are left intact. Caller must explicitly
    destroy them.

    Usage:
        @puzzle name,<part1[,part2,...>] = <result1[,result2,...]>
    """

    key = '@puzzle'
    aliases = '@puzzlerecipe'
    locks = 'cmd:perm(puzzle) or perm(Builder)'
    help_category = 'Puzzles'

    def func(self):
        caller = self.caller

        if len(self.lhslist) < 2 \
            or not self.rhs:
            string = "Usage: @puzzle name,<part1[,...]> = <result1[,...]>"
            caller.msg(string)
            return

        puzzle_name = self.lhslist[0]

        def is_valid_obj_location(obj):
            valid = True
            # Valid locations are: room, ...
            # TODO: other valid locations must be added here
            # Certain locations can be handled accordingly: e.g,
            # a part is located in a character's inventory,
            # perhaps will translate into the player character
            # having the part in his/her inventory while being
            # located in the same room where the builder was
            # located.
            # Parts and results may have different valid locations
            # TODO: handle contents of a given part
            if not inherits_from(obj.location, settings.BASE_ROOM_TYPECLASS):
                caller.msg('Invalid location for %s' % (obj.key))
                valid = False
            return valid

        def is_valid_part_location(part):
            return is_valid_obj_location(part)

        def is_valid_result_location(part):
            return is_valid_obj_location(part)

        parts = []
        for objname in self.lhslist[1:]:
            obj = caller.search(objname)
            if not obj:
                return
            if not is_valid_part_location(obj):
                return
            parts.append(obj)

        results = []
        for objname in self.rhslist:
            obj = caller.search(objname)
            if not obj:
                return
            if not is_valid_result_location(obj):
                return
            results.append(obj)

        for part in parts:
            caller.msg('Part %s(%s)' % (part.name, part.dbref))

        for result in results:
            caller.msg('Result %s(%s)' % (result.name, result.dbref))

        proto_parts = [proto_def(obj) for obj in parts]
        proto_results = [proto_def(obj) for obj in results]

        puzzle = create_object(PuzzleRecipeObject, key=puzzle_name)
        puzzle.save_recipe(puzzle_name, proto_parts, proto_results)

        caller.msg(
            "Puzzle |y'%s' |w%s(%s)|n has been created |gsuccessfully|n."
            % (puzzle.db.puzzle_name, puzzle.name, puzzle.dbref))
        caller.msg(
            'You may now dispose all parts and results. '
            'Typically, results and parts are useless afterwards.\n'
            'You are now able to arm this puzzle using Builder command:\n'
            '    @armpuzzle <puzzle #dbref>\n\n'
            'Or programmatically.\n'
        )

        # FIXME: puzzle recipe object exists but it has no location
        # should we create a PuzzleLibrary where all puzzles are
        # kept and cannot be reached by players?

class CmdArmPuzzle(MuxCommand):
    """
    Arms a puzzle by spawning all its parts
    """

    key = '@armpuzzle'
    # FIXME: permissions for scripts?
    locks = 'cmd:perm(armpuzzle) or perm(Builder)'
    help_category = 'Puzzles'

    def func(self):
        caller = self.caller

        if self.args is None or not utils.dbref(self.args):
            caller.msg("A puzzle recipe's #dbref must be specified")
            return

        puzzle = caller.search(self.args, global_search=True)
        if not puzzle or not inherits_from(puzzle, PuzzleRecipeObject):
            return

        caller.msg(
            "Puzzle Recipe %s(%s) '%s' found.\nSpawning %d parts ..." % (
            puzzle.name, puzzle.dbref, puzzle.db.puzzle_name, len(puzzle.db.parts)))

        for proto_part in puzzle.db.parts:
            part = spawn(proto_part)[0]
            caller.msg("Part %s(%s) spawned and placed at %s(%s)" % (part.name, part.dbref, part.location, part.location.dbref))
            part.mark_as_puzzle_member(puzzle.db.puzzle_name)

        caller.msg("Puzzle armed |gsuccessfully|n.")

class CmdUsePuzzleParts(MuxCommand):
    """
    Searches for all puzzles whose parts
    match the given set of objects. If
    there are matching puzzles, the result
    objects are spawned in their corresponding
    location if all parts have been passed in.

    Usage:
        use <part1[,part2,...>]
    """

    # TODO: consider allowing builder to provide
    # messages and "hooks" that can be displayed
    # and/or fired whenever the resolver of the puzzle
    # enters the location where a result was spawned

    key = 'use'
    aliases = 'combine'
    locks = 'cmd:pperm(use) or pperm(Player)'
    help_category = 'Puzzles'

    def func(self):
        caller = self.caller

        if not self.lhs:
            caller.msg('Use what?')
            return

        many = 'these' if len(self.lhslist) > 1 else 'this'

        # either all are parts, or abort finding matching puzzles
        parts = []
        partnames = self.lhslist[:]
        for partname in partnames:
            part = caller.search(
                partname,
                multimatch_string='Which %s. There are many.\n' % (partname),
                nofound_string='There is no %s around.' % (partname)
            )

            if not part:
                return

            if not part.tags.get(_PUZZLES_TAG_MEMBER, category=_PUZZLES_TAG_CATEGORY) \
                or not inherits_from(part, PuzzlePartObject):

                # not a puzzle part ... abort
                caller.msg('You have no idea how %s can be used' % (many))
                return

            # a valid part
            parts.append(part)

        # Create lookup dict
        parts_dict = dict((part.dbref, part) for part in parts)

        # Group parts by their puzzle name
        puzzle_ingredients = dict()
        for part in parts:
            puzzle_name = part.db.puzzle_name
            if puzzle_name not in puzzle_ingredients:
                puzzle_ingredients[puzzle_name] = []
            puzzle_ingredients[puzzle_name].append(
                (part.dbref, proto_def(part, with_tags=False))
            )

        # Find all puzzles by puzzle name
        # FIXME: we rely on obj.db.puzzle_name which is visible and may be cnaged afterwards. Can we lock it and hide it?
        puzzles = []
        for puzzle_name, parts in puzzle_ingredients.items():
            _puzzles = caller.search(
                    puzzle_name,
                    typeclass=[PuzzleRecipeObject],
                    attribute_name='puzzle_name',
                    quiet=True,
                    exact=True,
                    global_search=True)
            if not _puzzles:
                continue
            else:
                puzzles.extend(_puzzles)

        logger.log_info("PUZZLES %r" % ([p.dbref for p in puzzles]))

        # Create lookup dict
        puzzles_dict = dict((puzzle.dbref, puzzle) for puzzle in puzzles)

        # Check if parts can be combined to solve a puzzle
        matched_puzzles = dict()
        for puzzle in puzzles:
            puzzleparts = puzzle.db.parts[:]
            parts = puzzle_ingredients[puzzle.db.puzzle_name][:]
            pz = 0
            p = 0
            matched_dbrefparts = set()
            while pz < len(puzzleparts) and p < len(parts):
                puzzlepart = puzzleparts[pz]
                if 'tags' in puzzlepart:
                    # remove 'tags' as they will prevent equality
                    del(puzzlepart['tags'])
                dbref, part = parts[p]
                if part == puzzlepart:
                    pz += 1
                    matched_dbrefparts.add(dbref)
                p += 1
            else:
                if len(puzzleparts) == len(matched_dbrefparts):
                    matched_puzzles[puzzle.dbref] = matched_dbrefparts

        if len(matched_puzzles) == 0:
            # FIXME: Add more random messages
            #    random part falls and lands on your feet
            #    random part hits you square on the face
            caller.msg("As you try to utilize %s, nothing happens." % (many))
            return

        puzzletuples = sorted(matched_puzzles.items(), key=lambda t: len(t[1]), reverse=True)

        logger.log_info("MATCHED PUZZLES %r" % (puzzletuples))

        # sort all matched puzzles and pick largest one(s)
        puzzledbref, matched_dbrefparts = puzzletuples[0]
        nparts = len(matched_dbrefparts)
        puzzle = puzzles_dict[puzzledbref]
        largest_puzzles = list(itertools.takewhile(lambda t: len(t[1]) == nparts, puzzletuples))

        # if there are more than one, ...
        if len(largest_puzzles) > 1:
            # FIXME: pick a random one or let user choose?
            # FIXME: do we show the puzzle name or something else?
            caller.msg(
                'Your gears start turning and a bunch of ideas come to your mind ...\n%s' % (
                ' ...\n'.join([puzzles_dict[lp[0]].db.puzzle_name for lp in largest_puzzles]))
            )
            puzzletuple = choice(largest_puzzles)
            puzzle = puzzles_dict[puzzletuple[0]]
            caller.msg("You try %s ..." % (puzzle.db.puzzle_name))

        # got one, spawn its results
        # FIXME: DRY with parts
        result_names = []
        for proto_result in puzzle.db.results:
            result = spawn(proto_result)[0]
            result.mark_as_puzzle_member(puzzle.db.puzzle_name)
            result_names.append(result.name)
            # FIXME: add 'ramdon' messages:
            # Hmmm ... did I search result.location?
            # What was that? ... I heard something in result.location?
            # Eureka! you built a result

        # Destroy all parts used
        for dbref in matched_dbrefparts:
            parts_dict[dbref].delete()

        # FIXME: Add random messages
        #    You are a genius ... no matter what your 2nd grade teacher told you
        #    You hear thunders and a cloud of dust raises leaving
        result_names = ', '.join(result_names)
        caller.msg(
            "You are a |wG|re|wn|ri|wu|rs|n!!!\nYou just created %s" % (
            result_names
        ))
        caller.location.msg_contents(
            "|c%s|n performs some kind of tribal dance"
            " and seems to create |y%s|n from thin air" % (
            caller, result_names), exclude=(caller,)
        )

class CmdListPuzzleRecipes(MuxCommand):
    """
    Searches for all puzzle recipes

    Usage:
        @lspuzzlerecipes
    """

    key = '@lspuzzlerecipes'
    locks = 'cmd:perm(lspuzzlerecipes) or perm(Builder)'
    help_category = 'Puzzles'

    def func(self):
        caller = self.caller
        # TODO: use @tags/search puzzle_recipe : puzzles

class CmdListArmedPuzzles(MuxCommand):
    """
    Searches for all armed puzzles

    Usage:
        @lsarmedpuzzles
    """

    key = '@lsarmedpuzzles'
    locks = 'cmd:perm(lsarmedpuzzles) or perm(Builder)'
    help_category = 'Puzzles'

    def func(self):
        caller = self.caller
        # TODO: use @tags/search puzzle_member : puzzles

class PuzzleSystemCmdSet(CmdSet):
    """
    CmdSet to create, arm and resolve Puzzles

    Add with @py self.cmdset.add("evennia.contrib.puzzles.PuzzlesCmdSet")
    """

    def at_cmdset_creation(self):
        super(PuzzleSystemCmdSet, self).at_cmdset_creation()

        self.add(CmdCreatePuzzleRecipe())
        self.add(CmdArmPuzzle())
        self.add(CmdUsePuzzleParts())
<|endoftext|>


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

# (c) 2018, Chris Houseknecht <@chouseknecht>
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)

from __future__ import absolute_import, division, print_function

__metaclass__ = type

ANSIBLE_METADATA = {'metadata_version': '1.1',
                    'status': ['preview'],
                    'supported_by': 'community'}

DOCUMENTATION = '''

module: k8s

short_description: Manage Kubernetes (K8s) objects

version_added: "2.6"

author:
    - "Chris Houseknecht (@chouseknecht)"
    - "Fabian von Feilitzsch (@fabianvf)"

description:
  - Use the OpenShift Python client to perform CRUD operations on K8s objects.
  - Pass the object definition from a source file or inline. See examples for reading
    files and using Jinja templates.
  - Access to the full range of K8s APIs.
  - Use the M(k8s_facts) module to obtain a list of items about an object of type C(kind)
  - Authenticate using either a config file, certificates, password or token.
  - Supports check mode.

extends_documentation_fragment:
  - k8s_state_options
  - k8s_name_options
  - k8s_resource_options
  - k8s_auth_options

options:
  merge_type:
    description:
    - Whether to override the default patch merge approach with a specific type. By the default, the strategic
      merge will typically be used.
    - For example, Custom Resource Definitions typically aren't updatable by the usual strategic merge. You may
      want to use C(merge) if you see "strategic merge patch format is not supported"
    - See U(https://kubernetes.io/docs/tasks/run-application/update-api-object-kubectl-patch/#use-a-json-merge-patch-to-update-a-deployment)
    - Requires openshift >= 0.6.2
    - If more than one merge_type is given, the merge_types will be tried in order
    - If openshift >= 0.6.2, this defaults to C(['strategic-merge', 'merge']), which is ideal for using the same parameters
      on resource kinds that combine Custom Resources and built-in resources. For openshift < 0.6.2, the default
      is simply C(strategic-merge).
    choices:
    - json
    - merge
    - strategic-merge
    type: list
    version_added: "2.7"

requirements:
  - "python >= 2.7"
  - "openshift >= 0.6"
  - "PyYAML >= 3.11"
'''

EXAMPLES = '''
- name: Create a k8s namespace
  k8s:
    name: testing
    api_version: v1
    kind: Namespace
    state: present

- name: Create a Service object from an inline definition
  k8s:
    state: present
    definition:
      apiVersion: v1
      kind: Service
      metadata:
        name: web
        namespace: testing
        labels:
          app: galaxy
          service: web
      spec:
        selector:
          app: galaxy
          service: web
        ports:
        - protocol: TCP
          targetPort: 8000
          name: port-8000-tcp
          port: 8000

- name: Create a Service object by reading the definition from a file
  k8s:
    state: present
    src: /testing/service.yml

- name: Remove an existing Service object
  k8s:
    state: absent
    api_version: v1
    kind: Service
    namespace: testing
    name: web

# Passing the object definition from a file

- name: Create a Deployment by reading the definition from a local file
  k8s:
    state: present
    src: /testing/deployment.yml

- name: Read definition file from the Ansible controller file system
  k8s:
    state: present
    definition: "{{ lookup('file', '/testing/deployment.yml') }}"

- name: Read definition file from the Ansible controller file system after Jinja templating
  k8s:
    state: present
    definition: "{{ lookup('template', '/testing/deployment.yml') }}"
'''

RETURN = '''
result:
  description:
  - The created, patched, or otherwise present object. Will be empty in the case of a deletion.
  returned: success
  type: complex
  contains:
     api_version:
       description: The versioned schema of this representation of an object.
       returned: success
       type: str
     kind:
       description: Represents the REST resource this object represents.
       returned: success
       type: str
     metadata:
       description: Standard object metadata. Includes name, namespace, annotations, labels, etc.
       returned: success
       type: complex
     spec:
       description: Specific attributes of the object. Will vary based on the I(api_version) and I(kind).
       returned: success
       type: complex
     status:
       description: Current status details for the object.
       returned: success
       type: complex
     items:
       description: Returned only when multiple yaml documents are passed to src or resource_definition
       returned: when resource_definition or src contains list of objects
       type: list
'''

from ansible.module_utils.k8s.raw import KubernetesRawModule

def main():
    KubernetesRawModule().execute_module()

if __name__ == '__main__':
    main()
<commit_msg>Correct wording in k8s module description (#47125)<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

# (c) 2018, Chris Houseknecht <@chouseknecht>
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)

from __future__ import absolute_import, division, print_function

__metaclass__ = type

ANSIBLE_METADATA = {'metadata_version': '1.1',
                    'status': ['preview'],
                    'supported_by': 'community'}

DOCUMENTATION = '''

module: k8s

short_description: Manage Kubernetes (K8s) objects

version_added: "2.6"

author:
    - "Chris Houseknecht (@chouseknecht)"
    - "Fabian von Feilitzsch (@fabianvf)"

description:
  - Use the OpenShift Python client to perform CRUD operations on K8s objects.
  - Pass the object definition from a source file or inline. See examples for reading
    files and using Jinja templates.
  - Access to the full range of K8s APIs.
  - Use the M(k8s_facts) module to obtain a list of items about an object of type C(kind)
  - Authenticate using either a config file, certificates, password or token.
  - Supports check mode.

extends_documentation_fragment:
  - k8s_state_options
  - k8s_name_options
  - k8s_resource_options
  - k8s_auth_options

options:
  merge_type:
    description:
    - Whether to override the default patch merge approach with a specific type. By default, the strategic
      merge will typically be used.
    - For example, Custom Resource Definitions typically aren't updatable by the usual strategic merge. You may
      want to use C(merge) if you see "strategic merge patch format is not supported"
    - See U(https://kubernetes.io/docs/tasks/run-application/update-api-object-kubectl-patch/#use-a-json-merge-patch-to-update-a-deployment)
    - Requires openshift >= 0.6.2
    - If more than one merge_type is given, the merge_types will be tried in order
    - If openshift >= 0.6.2, this defaults to C(['strategic-merge', 'merge']), which is ideal for using the same parameters
      on resource kinds that combine Custom Resources and built-in resources. For openshift < 0.6.2, the default
      is simply C(strategic-merge).
    choices:
    - json
    - merge
    - strategic-merge
    type: list
    version_added: "2.7"

requirements:
  - "python >= 2.7"
  - "openshift >= 0.6"
  - "PyYAML >= 3.11"
'''

EXAMPLES = '''
- name: Create a k8s namespace
  k8s:
    name: testing
    api_version: v1
    kind: Namespace
    state: present

- name: Create a Service object from an inline definition
  k8s:
    state: present
    definition:
      apiVersion: v1
      kind: Service
      metadata:
        name: web
        namespace: testing
        labels:
          app: galaxy
          service: web
      spec:
        selector:
          app: galaxy
          service: web
        ports:
        - protocol: TCP
          targetPort: 8000
          name: port-8000-tcp
          port: 8000

- name: Create a Service object by reading the definition from a file
  k8s:
    state: present
    src: /testing/service.yml

- name: Remove an existing Service object
  k8s:
    state: absent
    api_version: v1
    kind: Service
    namespace: testing
    name: web

# Passing the object definition from a file

- name: Create a Deployment by reading the definition from a local file
  k8s:
    state: present
    src: /testing/deployment.yml

- name: Read definition file from the Ansible controller file system
  k8s:
    state: present
    definition: "{{ lookup('file', '/testing/deployment.yml') }}"

- name: Read definition file from the Ansible controller file system after Jinja templating
  k8s:
    state: present
    definition: "{{ lookup('template', '/testing/deployment.yml') }}"
'''

RETURN = '''
result:
  description:
  - The created, patched, or otherwise present object. Will be empty in the case of a deletion.
  returned: success
  type: complex
  contains:
     api_version:
       description: The versioned schema of this representation of an object.
       returned: success
       type: str
     kind:
       description: Represents the REST resource this object represents.
       returned: success
       type: str
     metadata:
       description: Standard object metadata. Includes name, namespace, annotations, labels, etc.
       returned: success
       type: complex
     spec:
       description: Specific attributes of the object. Will vary based on the I(api_version) and I(kind).
       returned: success
       type: complex
     status:
       description: Current status details for the object.
       returned: success
       type: complex
     items:
       description: Returned only when multiple yaml documents are passed to src or resource_definition
       returned: when resource_definition or src contains list of objects
       type: list
'''

from ansible.module_utils.k8s.raw import KubernetesRawModule

def main():
    KubernetesRawModule().execute_module()

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


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

# (c) 2012, Luis Alberto Perez Lazaro <luisperlazaro@gmail.com>
# (c) 2015, Jakub Jirutka <jakub@jirutka.cz>
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

DOCUMENTATION = '''
---
module: patch
author: Luis Alberto Perez Lazaro, Jakub Jirutka
version_added: 1.9
description:
    - Apply patch files using the GNU patch tool.
short_description: Apply patch files using the GNU patch tool.
options:
  basedir:
    description:
      - Path of a base directory in which the patch file will be applied.
        May be ommitted when C(dest) option is specified, otherwise required.
    required: false
  dest:
    description:
      - Path of the file on the remote machine to be patched.
      - The names of the files to be patched are usually taken from the patch
        file, but if there's just one file to be patched it can specified with
        this option.
    required: false
    aliases: [ "originalfile" ]
  src:
    description:
      - Path of the patch file as accepted by the GNU patch tool.
    required: true
    aliases: [ "patchfile" ]
  remote_src:
    description:
      - If False, it will search for src at originating/master machine, if True it will
        go to the remote/target machine for the src. Default is False.
    choices: [ "True", "False" ]
    required: false
    default: "False"
  strip:
    description:
      - Number that indicates the smallest prefix containing leading slashes
        that will be stripped from each file name found in the patch file.
        For more information see the strip parameter of the GNU patch tool.
    required: false
    type: "int"
    default: "0"
note:
  - This module requires GNU I(patch) utility to be installed on the remote host.
'''

EXAMPLES = '''
- name: apply patch to one file
  patch: >
    src=/tmp/index.html.patch
    dest=/var/www/index.html

- name: apply patch to multiple files under basedir
  patch: >
    src=/tmp/customize.patch
    basedir=/var/www
    strip=1
'''

import os
from os import path, R_OK, W_OK

class PatchError(Exception):
    pass

def is_already_applied(patch_func, patch_file, basedir, dest_file=None, strip=0):
    opts = ['--quiet', '--reverse', '--forward', '--dry-run',
            "--strip=%s" % strip, "--directory='%s'" % basedir,
            "--input='%s'" % patch_file]
    if dest_file:
        opts.append("'%s'" % dest_file)

    (rc, _, _) = patch_func(opts)
    return rc == 0

def apply_patch(patch_func, patch_file, basedir, dest_file=None, strip=0, dry_run=False):
    opts = ['--quiet', '--forward', '--batch', '--reject-file=-',
            "--strip=%s" % strip, "--directory='%s'" % basedir,
            "--input='%s'" % patch_file]
    if dry_run:
        opts.append('--dry-run')
    if dest_file:
        opts.append("'%s'" % dest_file)

    (rc, out, err) = patch_func(opts)
    if rc != 0:
        msg = out if not err else err
        raise PatchError(msg)

def main():
    module = AnsibleModule(
        argument_spec={
            'src':     {'required': True, 'aliases': ['patchfile']},
            'dest':    {'aliases': ['originalfile']},
            'basedir': {},
            'strip':   {'default': 0, 'type': 'int'},
            'remote_src': {'default': False, 'type': 'bool'},
        },
        required_one_of=[['dest', 'basedir']],
        supports_check_mode=True
    )

    # Create type object as namespace for module params
    p = type('Params', (), module.params)

    if not os.access(p.src, R_OK):
        module.fail_json(msg="src %s doesn't exist or not readable" % (p.src))

    if p.dest and not os.access(p.dest, W_OK):
        module.fail_json(msg="dest %s doesn't exist or not writable" % (d.dest))

    if p.basedir and not path.exists(p.basedir):
        module.fail_json(msg="basedir %s doesn't exist" % (p.basedir))

    if not p.basedir:
        p.basedir = path.dirname(p.dest)

    patch_bin = module.get_bin_path('patch')
    patch_func = lambda opts: module.run_command("%s %s" % (patch_bin, ' '.join(opts)))

    changed = False
    if not is_already_applied(patch_func, p.src, p.basedir, dest_file=p.dest, strip=p.strip):
        try:
            apply_patch(patch_func, p.src, p.basedir, dest_file=p.dest, strip=p.strip,
                        dry_run=module.check_mode)
            changed = True
        except PatchError, e:
            module.fail_json(msg=str(e))

    module.exit_json(changed=changed)

# import module snippets
from ansible.module_utils.basic import *
main()
<commit_msg>Variable name typo<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

# (c) 2012, Luis Alberto Perez Lazaro <luisperlazaro@gmail.com>
# (c) 2015, Jakub Jirutka <jakub@jirutka.cz>
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

DOCUMENTATION = '''
---
module: patch
author: Luis Alberto Perez Lazaro, Jakub Jirutka
version_added: 1.9
description:
    - Apply patch files using the GNU patch tool.
short_description: Apply patch files using the GNU patch tool.
options:
  basedir:
    description:
      - Path of a base directory in which the patch file will be applied.
        May be ommitted when C(dest) option is specified, otherwise required.
    required: false
  dest:
    description:
      - Path of the file on the remote machine to be patched.
      - The names of the files to be patched are usually taken from the patch
        file, but if there's just one file to be patched it can specified with
        this option.
    required: false
    aliases: [ "originalfile" ]
  src:
    description:
      - Path of the patch file as accepted by the GNU patch tool.
    required: true
    aliases: [ "patchfile" ]
  remote_src:
    description:
      - If False, it will search for src at originating/master machine, if True it will
        go to the remote/target machine for the src. Default is False.
    choices: [ "True", "False" ]
    required: false
    default: "False"
  strip:
    description:
      - Number that indicates the smallest prefix containing leading slashes
        that will be stripped from each file name found in the patch file.
        For more information see the strip parameter of the GNU patch tool.
    required: false
    type: "int"
    default: "0"
note:
  - This module requires GNU I(patch) utility to be installed on the remote host.
'''

EXAMPLES = '''
- name: apply patch to one file
  patch: >
    src=/tmp/index.html.patch
    dest=/var/www/index.html

- name: apply patch to multiple files under basedir
  patch: >
    src=/tmp/customize.patch
    basedir=/var/www
    strip=1
'''

import os
from os import path, R_OK, W_OK

class PatchError(Exception):
    pass

def is_already_applied(patch_func, patch_file, basedir, dest_file=None, strip=0):
    opts = ['--quiet', '--reverse', '--forward', '--dry-run',
            "--strip=%s" % strip, "--directory='%s'" % basedir,
            "--input='%s'" % patch_file]
    if dest_file:
        opts.append("'%s'" % dest_file)

    (rc, _, _) = patch_func(opts)
    return rc == 0

def apply_patch(patch_func, patch_file, basedir, dest_file=None, strip=0, dry_run=False):
    opts = ['--quiet', '--forward', '--batch', '--reject-file=-',
            "--strip=%s" % strip, "--directory='%s'" % basedir,
            "--input='%s'" % patch_file]
    if dry_run:
        opts.append('--dry-run')
    if dest_file:
        opts.append("'%s'" % dest_file)

    (rc, out, err) = patch_func(opts)
    if rc != 0:
        msg = out if not err else err
        raise PatchError(msg)

def main():
    module = AnsibleModule(
        argument_spec={
            'src':     {'required': True, 'aliases': ['patchfile']},
            'dest':    {'aliases': ['originalfile']},
            'basedir': {},
            'strip':   {'default': 0, 'type': 'int'},
            'remote_src': {'default': False, 'type': 'bool'},
        },
        required_one_of=[['dest', 'basedir']],
        supports_check_mode=True
    )

    # Create type object as namespace for module params
    p = type('Params', (), module.params)

    if not os.access(p.src, R_OK):
        module.fail_json(msg="src %s doesn't exist or not readable" % (p.src))

    if p.dest and not os.access(p.dest, W_OK):
        module.fail_json(msg="dest %s doesn't exist or not writable" % (p.dest))

    if p.basedir and not path.exists(p.basedir):
        module.fail_json(msg="basedir %s doesn't exist" % (p.basedir))

    if not p.basedir:
        p.basedir = path.dirname(p.dest)

    patch_bin = module.get_bin_path('patch')
    patch_func = lambda opts: module.run_command("%s %s" % (patch_bin, ' '.join(opts)))

    changed = False
    if not is_already_applied(patch_func, p.src, p.basedir, dest_file=p.dest, strip=p.strip):
        try:
            apply_patch(patch_func, p.src, p.basedir, dest_file=p.dest, strip=p.strip,
                        dry_run=module.check_mode)
            changed = True
        except PatchError, e:
            module.fail_json(msg=str(e))

    module.exit_json(changed=changed)

# import module snippets
from ansible.module_utils.basic import *
main()
<|endoftext|>


<commit_before>import numpy
import matplotlib.mlab

# (c) Tom Wiesing 2015
# licensed under MIT license

def PCA(M,axis=1):
    """
        Applies an unscaled PCA to a matrix.

        M: Matrix to apply PCA to.
        axis: Axis to normalise against. Needs to be either 1 or 0. Defaults to 1.
    """

    # ensure axis is of the right type.
    if axis != 1 and axis != 0:
        raise ValueError("Axis must be 0 or 1")

    # create a matrix P.
    if axis == 1:
        P = M.T
    else:
        P = M

    # compute a PCA object
    pcaobj = matplotlib.mlab.PCA(P,standardize=False)

    # replace each row of the matrix P
    for r in range(P.shape[0]):
        P[r, :] = pcaobj.project(P[r, :])

    # and return the right version.
    if axis == 1:
        return P.T
    else:
        return P
def reorder(M,sigma,axis=1):
    """
        Reorders data items according to a permutation sigma.

        M: Matrix to apply PCA to.
        sigma: Permutation of the attributes of M.
        axis: Axis to permute against. Needs to be either 1 or 0. Defaults to 1.
    """

    if axis == 1:
        P = M.T
    else:
        P = M

    for r in range(P.shape[1]):
        Q[:, r] = Q[:, sigma[r]]

    if axis == 1:
        return Q.T
    else:
        return Q
<commit_msg>Fix a bug in the permuate transform, version bump 0.0.2<commit_after>import numpy
import matplotlib.mlab

# (c) Tom Wiesing 2015
# licensed under MIT license

def PCA(M,axis=1):
    """
        Applies an unscaled PCA to a matrix.

        M: Matrix to apply PCA to.
        axis: Axis to normalise against. Needs to be either 1 or 0. Defaults to 1.
    """

    # ensure axis is of the right type.
    if axis != 1 and axis != 0:
        raise ValueError("Axis must be 0 or 1")

    # create a matrix P.
    if axis == 1:
        P = M.T
    else:
        P = M

    # compute a PCA object
    pcaobj = matplotlib.mlab.PCA(P,standardize=False)

    # replace each row of the matrix P
    for r in range(P.shape[0]):
        P[r, :] = pcaobj.project(P[r, :])

    # and return the right version.
    if axis == 1:
        return P.T
    else:
        return P
def permute(M,sigma,axis=1):
    """
        Reorders data items according to a permutation sigma.

        M: Matrix to apply PCA to.
        sigma: Permutation of the attributes of M.
        axis: Axis to permute against. Needs to be either 1 or 0. Defaults to 1.
    """

    if axis == 1:
        P = M.T
    else:
        P = M

    Q = P

    for r in range(P.shape[1]):
        Q[:, r] = Q[:, sigma[r]]

    if axis == 1:
        return Q.T
    else:
        return Q
<|endoftext|>


<commit_before>from glob import glob
import json
from os.path import exists, join
from os import makedirs

from bdd_tester import bdd_tester
from lxml import etree
import requests
from werkzeug.utils import secure_filename

from DataQualityTester import app, celery, db

def _colorify(number):
    if number is None:
        return '#222'
    # Methodology colours:
    colours = (
        ((180, 96, 122), 0),
        ((232, 152, 77), 25),
        ((248, 204, 78), 50),
        ((159, 175, 52), 75),
        ((108, 137, 178), 100),
    )
    prev_colour, prev_perc = colours[0]
    for colour, perc in colours[1:]:
        if number <= perc:
            dist = (number - prev_perc) / (perc - prev_perc)
            rgb = (int(dist * (colour[d] - prev_colour[d]) + prev_colour[d])
                   for d in range(3))
            return 'rgb({}, {}, {})'.format(*rgb)
        prev_colour = colour
        prev_perc = perc

def _compute_score(results):
    scores = []
    for summary in results.values():
        total = summary['passed'] + summary['failed']
        if total != 0:
            score = summary['passed'] / total
            scores.append(score)
    if len(scores) > 0:
        return 100. * sum(scores) / len(scores)
    else:
        return None

def _load_codelists():
    all_codes = {}
    codelist_list = app.config.get('CODELIST_LIST')
    for codelist in codelist_list:
        codes = []
        for ver in ['1', '2']:
            with open(join('codelists', ver, codelist + '.json')) as f:
                codelist_data = json.load(f)['data']
                codes += [x['code'] for x in codelist_data]
        all_codes[codelist] = codes
    return all_codes

@celery.task(bind=True)
def test_file_task(self, path_to_file, feature_path, component_id,
                   output_path):
    results_path = '{}.json'.format(join(output_path, component_id))
    if exists(results_path):
        with open(results_path) as f:
            results = json.load(f)
        score = _compute_score(results)
        if score is not None:
            colour = _colorify(score)
        else:
            colour = '#222'
    else:
        # TODO: handle parse error
        xml = etree.parse(path_to_file)

        features = glob(join(feature_path, '*.feature'))
        feature_count = len(features)
        codelists = _load_codelists()

        results = {}
        score = None
        for idx, feature in enumerate(features):
            result = bdd_tester(
                join('steps', 'step_definitions.py'),
                features,
                etree=xml,
                output_path=output_path,
                codelists=codelists,
            )
            if not result:
                continue
            results.update(result)
            score = _compute_score(results)
            if score is not None:
                colour = _colorify(score)
            else:
                colour = '#222'
            self.update_state(
                state='RUNNING',
                meta={
                    'name': component_id,
                    'progress': 100. * (idx + 1) / feature_count,
                    'data': results,
                    'score': score,
                    'colour': colour,
                }
            )
        # cache the result
        with open(results_path, 'w') as f:
            json.dump(results, f)

    return {
        'name': component_id,
        'progress': 100,
        'data': results,
        'score': score,
        'colour': colour,
    }

@celery.task(bind=True)
def download_task(self, sd_uuid):
    from DataQualityTester import models
    supplied_data = models.SuppliedData.query.get_or_404(str(sd_uuid))
    url = supplied_data.source_url
    request_kwargs = {
        'headers': {'User-Agent': 'Publish What You Fund Data Quality Tester'},
        'stream': True,
    }
    try:
        resp = requests.get(url, **request_kwargs)
    except requests.exceptions.SSLError:
        resp = requests.get(url, verify=False, **request_kwargs)
    resp.raise_for_status()

    filename = resp.url.split('/')[-1].split('?')[0][:100]
    if filename == '':
        filename = 'file.xml'
    elif not filename.endswith('.xml'):
        filename += '.xml'
    filename = secure_filename(filename)
    makedirs(supplied_data.upload_dir(), exist_ok=True)

    filepath = join(supplied_data.upload_dir(), filename)
    total_length = resp.headers.get('content-length')
    if total_length:
        total_length = int(total_length)
    with open(filepath, 'wb') as f:
        dl = 0
        for block in resp.iter_content(1024):
            dl += len(block)
            f.write(block)
            if total_length:
                self.update_state(
                    state='RUNNING',
                    meta={
                        'progress': 100 * dl / total_length,
                    }
                )

    supplied_data.original_file = join(supplied_data.id, filename)
    supplied_data.downloaded = True
    db.session.add(supplied_data)
    db.session.commit()
    return {
        'progress': 100
    }
<commit_msg>Tiny refactor<commit_after>from glob import glob
import json
from os.path import exists, join
from os import makedirs

from bdd_tester import bdd_tester
from lxml import etree
import requests
from werkzeug.utils import secure_filename

from DataQualityTester import app, celery, db

def _colorify(number):
    if number is None:
        return '#222'
    # Methodology colours:
    colours = (
        ((180, 96, 122), 0),
        ((232, 152, 77), 25),
        ((248, 204, 78), 50),
        ((159, 175, 52), 75),
        ((108, 137, 178), 100),
    )
    prev_colour, prev_perc = colours[0]
    for colour, perc in colours[1:]:
        if number <= perc:
            dist = (number - prev_perc) / (perc - prev_perc)
            rgb = (int(dist * (colour[d] - prev_colour[d]) + prev_colour[d])
                   for d in range(3))
            return 'rgb({}, {}, {})'.format(*rgb)
        prev_colour = colour
        prev_perc = perc

def _compute_score(results):
    scores = []
    for summary in results.values():
        total = summary['passed'] + summary['failed']
        if total != 0:
            score = summary['passed'] / total
            scores.append(score)
    if len(scores) > 0:
        return 100. * sum(scores) / len(scores)
    else:
        return None

def _load_codelists():
    all_codes = {}
    codelist_list = app.config.get('CODELIST_LIST')
    for codelist in codelist_list:
        codes = []
        for ver in ['1', '2']:
            with open(join('codelists', ver, codelist + '.json')) as f:
                codelist_data = json.load(f)['data']
                codes += [x['code'] for x in codelist_data]
        all_codes[codelist] = codes
    return all_codes

@celery.task(bind=True)
def test_file_task(self, path_to_file, feature_path, component_id,
                   output_path):
    results_path = '{}.json'.format(join(output_path, component_id))
    if exists(results_path):
        with open(results_path) as f:
            results = json.load(f)
        score = _compute_score(results)
        if score is not None:
            colour = _colorify(score)
        else:
            colour = '#222'
    else:
        # TODO: handle parse error
        xml = etree.parse(path_to_file)

        step_definitions_path = join('steps', 'step_definitions.py')
        features = glob(join(feature_path, '*.feature'))
        feature_count = len(features)
        codelists = _load_codelists()

        results = {}
        score = None
        for idx, feature in enumerate(features):
            result = bdd_tester(
                step_definitions_path,
                features,
                etree=xml,
                output_path=output_path,
                codelists=codelists,
            )
            if not result:
                continue
            results.update(result)
            score = _compute_score(results)
            if score is not None:
                colour = _colorify(score)
            else:
                colour = '#222'
            self.update_state(
                state='RUNNING',
                meta={
                    'name': component_id,
                    'progress': 100. * (idx + 1) / feature_count,
                    'data': results,
                    'score': score,
                    'colour': colour,
                }
            )
        # cache the result
        with open(results_path, 'w') as f:
            json.dump(results, f)

    return {
        'name': component_id,
        'progress': 100,
        'data': results,
        'score': score,
        'colour': colour,
    }

@celery.task(bind=True)
def download_task(self, sd_uuid):
    from DataQualityTester import models
    supplied_data = models.SuppliedData.query.get_or_404(str(sd_uuid))
    url = supplied_data.source_url
    request_kwargs = {
        'headers': {'User-Agent': 'Publish What You Fund Data Quality Tester'},
        'stream': True,
    }
    try:
        resp = requests.get(url, **request_kwargs)
    except requests.exceptions.SSLError:
        resp = requests.get(url, verify=False, **request_kwargs)
    resp.raise_for_status()

    filename = resp.url.split('/')[-1].split('?')[0][:100]
    if filename == '':
        filename = 'file.xml'
    elif not filename.endswith('.xml'):
        filename += '.xml'
    filename = secure_filename(filename)
    makedirs(supplied_data.upload_dir(), exist_ok=True)

    filepath = join(supplied_data.upload_dir(), filename)
    total_length = resp.headers.get('content-length')
    if total_length:
        total_length = int(total_length)
    with open(filepath, 'wb') as f:
        dl = 0
        for block in resp.iter_content(1024):
            dl += len(block)
            f.write(block)
            if total_length:
                self.update_state(
                    state='RUNNING',
                    meta={
                        'progress': 100 * dl / total_length,
                    }
                )

    supplied_data.original_file = join(supplied_data.id, filename)
    supplied_data.downloaded = True
    db.session.add(supplied_data)
    db.session.commit()
    return {
        'progress': 100
    }
<|endoftext|>


<commit_before>import os
import random

class DataSets(object):
    def __init__(self, DIR='../../data', output_filename='all_videos.txt', **kwargs):
        self.DIR = DIR
        self.output_filename = output_filename
        self.flagged_activities = ['PlayingDaf', 'BodyWeightSquats', 'Nunchucks', 'ShavingBeard', 'SkyDiving']

    def videos_to_text_file(self):
        with open(self.output_filename, "w") as a:
            for path, subdirs, files in os.walk(self.DIR):
               for filename in files:
                 f = os.path.join(path, filename)
                 a.write(str(f) + os.linesep)

    def train_test_split(self, split_test_data, split_validation_data):
        """
        split_test_data : '%' of test data to split between 0 to 1
        """
        data = {}
        unseen = [line.rstrip('\n') for line in open(self.output_filename) if any(substring in line for substring in self.flagged_activities)]
        seen = [line.rstrip('\n') for line in open(self.output_filename) if any(substring in line for substring in self.flagged_activities)]
        validation_index = int(len(seen)* (split_test_data))
        seen = random.shuffle(seen)
        data['validation'] = seen[:validation_index]

        seen = seen[validation_index:]
        test_index = int(len(seen)* (split_test_data))
        data['train'] = seen[test_index:]
        data['test'] = seen[:test_index]
        data['unseen'] = unseen

        return data

    @staticmethod
    def batch_iterator(iterator, batch_size):
        """Returns lists of length batch_size.
        This is a generator function, and it returns lists of the
        entries from the supplied iterator.  Each list will have
        batch_size entries, although the final list may be shorter.
        """
        entry = True  # Make sure we loop once
        while entry:
            batch = []
            while len(batch) < batch_size:
                try:
                    entry = iterator.next()
                except StopIteration:
                    entry = None
                if entry is None:
                    # End of file
                    break
                batch.append(entry)
            if batch:
                yield batch

a = DataSets()
# a.videos_to_text_file()
videos = a.train_test_split(0.2, 0.1)
<commit_msg>Random Shuffle Bug<commit_after>import os
import random

class DataSets(object):
    def __init__(self, DIR='../../data', output_filename='all_videos.txt', **kwargs):
        self.DIR = DIR
        self.output_filename = output_filename
        self.flagged_activities = ['PlayingDaf', 'BodyWeightSquats', 'Nunchucks', 'ShavingBeard', 'SkyDiving']

    def videos_to_text_file(self):
        with open(self.output_filename, "w") as a:
            for path, subdirs, files in os.walk(self.DIR):
               for filename in files:
                 f = os.path.join(path, filename)
                 a.write(str(f) + os.linesep)

    def train_test_split(self, split_test_data, split_validation_data):
        """
        split_test_data : '%' of test data to split between 0 to 1
        """
        data = {}
        unseen = [line.rstrip('\n') for line in open(self.output_filename) if any(substring in line for substring in self.flagged_activities)]
        seen = [line.rstrip('\n') for line in open(self.output_filename) if any(substring in line for substring in self.flagged_activities)]
        validation_index = int(len(seen)* (split_test_data))

        #Random Shuffle
        random.shuffle(seen)
        data['validation'] = seen[:validation_index]

        seen = seen[validation_index:]
        test_index = int(len(seen)* (split_test_data))
        data['train'] = seen[test_index:]
        data['test'] = seen[:test_index]
        data['unseen'] = unseen

        return data

    @staticmethod
    def batch_iterator(iterator, batch_size):
        """Returns lists of length batch_size.
        This is a generator function, and it returns lists of the
        entries from the supplied iterator.  Each list will have
        batch_size entries, although the final list may be shorter.
        """
        entry = True  # Make sure we loop once
        while entry:
            batch = []
            while len(batch) < batch_size:
                try:
                    entry = iterator.next()
                except StopIteration:
                    entry = None
                if entry is None:
                    # End of file
                    break
                batch.append(entry)
            if batch:
                yield batch

a = DataSets()
# a.videos_to_text_file()
videos = a.train_test_split(0.2, 0.1)
<|endoftext|>


<commit_before># Anagram another word.
import collections

ANAGRAM_INDICATORS = frozenset([
  'about',
  'awful',
  'bad',
  'badly',
  'break',
  'round',
  'crash',
  'crazy',
  'confuse',
  'destruction',
  'downfall',
  'from',
  'kink',
  'make',
  'mess',
  'odd',
  'oddly',
  'off',
  'out',
  'put',
  'roll',
  'round',
  'shake',
  'transform',
  'upset',
  'wander',
  'wandering',
])
# Solution appears in clue.
EMBEDDED_INDICATORS = frozenset([
  'bit',
  'circling',
  'in',
  'part',
  'partly',
])
# Synonym sounds like solution.
HOMOPHONE_INDICATORS = frozenset([
  'hear',
  'say',
])
# First letter.
INITIAL_INDICATORS = frozenset([
  'first',
  'initial',
  'initially',
  'lead',
])
# Join words.
CONCATENATE_INDICATORS = frozenset([
  'and',
  'by',
  'put',
  'with',
])
# Insert a word into another word.
INSERT_INDICATORS = frozenset([
  'stuck',
  'in',
])
# Front and back letters.
EDGES_INDICATORS = frozenset([
  'edge',
])
# Reversing words.
REVERSAL_INDICATORS = frozenset([
  'turn',
  'turning',
  'back',
  'backward',
])
# Ambiguous definition (e.g., double entendre).
AMBIGUOUS_INDICATORS = frozenset([
  '?'
])
# Shorthand conversions.
# TODO: This doesn't scale. Need a synonym list.
SHORTHAND_CONVERSIONS = {
  'american': ['us', 'usa'],
  'attempt': ['try'],
  'father': ['pa'],
  'fathers': ['pas'],
  'foot': ['ft'],
  'good': ['g'],
  'right': ['r'],
  'left': ['l'],
  'microphone': ['mic'],
}

ALL_INDICATORS = collections.defaultdict(list)
for group in [
  ANAGRAM_INDICATORS, EMBEDDED_INDICATORS, HOMOPHONE_INDICATORS,
  INITIAL_INDICATORS, CONCATENATE_INDICATORS, INSERT_INDICATORS,
  EDGES_INDICATORS, REVERSAL_INDICATORS,
]:
  for indicator in group:
    ALL_INDICATORS[indicator].append(group)
<commit_msg>Note some new cryptic substitutions.<commit_after># Anagram another word.
import collections

ANAGRAM_INDICATORS = frozenset([
  'about',
  'awful',
  'bad',
  'badly',
  'break',
  'round',
  'crash',
  'crazy',
  'confuse',
  'destruction',
  'downfall',
  'from',
  'kink',
  'make',
  'mess',
  'odd',
  'oddly',
  'off',
  'out',
  'put',
  'roll',
  'round',
  'shake',
  'transform',
  'upset',
  'wander',
  'wandering',
])
# Solution appears in clue.
EMBEDDED_INDICATORS = frozenset([
  'bit',
  'circling',
  'in',
  'part',
  'partly',
])
# Synonym sounds like solution.
HOMOPHONE_INDICATORS = frozenset([
  'hear',
  'say',
])
# First letter.
INITIAL_INDICATORS = frozenset([
  'first',
  'initial',
  'initially',
  'lead',
])
# Join words.
CONCATENATE_INDICATORS = frozenset([
  'and',
  'by',
  'put',
  'with',
])
# Insert a word into another word.
INSERT_INDICATORS = frozenset([
  'stuck',
  'in',
])
# Front and back letters.
EDGES_INDICATORS = frozenset([
  'edge',
])
# Reversing words.
REVERSAL_INDICATORS = frozenset([
  'turn',
  'turning',
  'back',
  'backward',
])
# Ambiguous definition (e.g., double entendre).
AMBIGUOUS_INDICATORS = frozenset([
  '?'
])
# Shorthand conversions.
# TODO: This doesn't scale. Need a synonym list.
SHORTHAND_CONVERSIONS = {
  'american': ['us', 'usa'],
  'attempt': ['try'],
  'father': ['pa'],
  'fathers': ['pas'],
  'foot': ['ft'],
  'good': ['g'],
  'right': ['r'],
  'left': ['l'],
  'microphone': ['mic'],
}
# "Set of books" = OT (old testament) or NT (new testament).

ALL_INDICATORS = collections.defaultdict(list)
for group in [
  ANAGRAM_INDICATORS, EMBEDDED_INDICATORS, HOMOPHONE_INDICATORS,
  INITIAL_INDICATORS, CONCATENATE_INDICATORS, INSERT_INDICATORS,
  EDGES_INDICATORS, REVERSAL_INDICATORS,
]:
  for indicator in group:
    ALL_INDICATORS[indicator].append(group)
<|endoftext|>


<commit_before># Copyright 2020 Google LLC
#
# 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
#
#     https://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
import sys
import traceback
import types

from jax._src import util

_exclude_paths = [__file__, util.__file__]

def register_exclusion(path):
  _exclude_paths.append(path)

_jax_message_append = (
    'The stack trace above excludes JAX-internal frames.\n'
    'The following is the original exception that occurred, unmodified.\n'
    '\n--------------------')

def path_starts_with(path, path_prefix):
  path = os.path.abspath(path)
  path_prefix = os.path.abspath(path_prefix)
  if not os.path.exists(path_prefix):
    return False
  try:
    common = os.path.commonpath([path, path_prefix])
    return os.path.samefile(common, path_prefix)
  except ValueError:
    # path and path_prefix are both absolute, the only case will raise a
    # ValueError is different drives.
    # https://docs.python.org/3/library/os.path.html#os.path.commonpath
    return False

def include_frame(f):
  return not any(path_starts_with(f.f_code.co_filename, path)
                 for path in _exclude_paths)

# When scanning stack traces, we might encounter frames from cpython that are
# removed from printed stack traces, such as frames from parts of importlib. We
# ignore these frames heuristically based on source and name match.
def ignore_known_hidden_frame(f):
  return 'importlib._bootstrap' in f.f_code.co_filename

def filter_traceback_and_stack(e):
  out = None

  # Scan the traceback and collect relevant frames.

  for f, lineno in reversed(list(traceback.walk_tb(e.__traceback__))):
    if include_frame(f):
      out = types.TracebackType(out, f, f.f_lasti, lineno)  # pytype: disable=wrong-arg-count

  # Continue up the call stack.
  #
  # We would like to avoid stepping too far up, e.g. past the exec/eval point of
  # a REPL such as IPython. To that end, we stop past the first contiguous bunch
  # of module-level frames, if we reach any such frames at all. This is a
  # heuristic that might stop in advance of the REPL boundary. For example, if
  # the call stack includes module-level frames from the current module A, and
  # the current module A was imported from within a function F elsewhere, then
  # the stack trace we produce will be truncated at F's frame.

  reached_module_level = False
  for f, lineno in traceback.walk_stack(e.__traceback__.tb_frame):
    if ignore_known_hidden_frame(f):
      continue
    if reached_module_level and f.f_code.co_name != '<module>':
      break
    if include_frame(f):
      out = types.TracebackType(out, f, f.f_lasti, lineno)  # pytype: disable=wrong-arg-count
    if f.f_code.co_name == '<module>':
      reached_module_level = True

  return out

def is_reraiser_frame(f):
  return (f.filename == __file__ and
          f.name == 'reraise_with_filtered_traceback')

def is_under_reraiser(e):
  tb = traceback.extract_stack(e.__traceback__.tb_frame)
  return any(is_reraiser_frame(f) for f in tb[:-1])

def format_exception_only(e):
  return ''.join(traceback.format_exception_only(type(e), e)).strip()

def last_cause(e):
  prev, cur = e, e.__cause__
  while cur is not None:
    prev, cur = cur, cur.__cause__
  return prev

class FilteredStackTrace(Exception): pass

def filtered_tracebacks_supported():
  return sys.version_info >= (3, 7)

def api_boundary(fun):
  '''Wraps ``fun`` to form a boundary for filtering exception tracebacks.

  When an exception occurs below ``fun``, this appends to it a custom
  ``__cause__`` that carries a filtered traceback. The traceback imitates the
  stack trace of the original exception, but with JAX-internal frames removed.

  This boundary annotation works in composition with itself. The topmost frame
  corresponding to an ``api_boundary`` is the one below which stack traces are
  filtered. In other words, if ``api_boundary(f)`` calls ``api_boundary(g)``,
  directly or indirectly, the filtered stack trace provided is the same as if
  ``api_boundary(f)`` were to simply call ``g`` instead.

  This annotation is primarily useful in wrapping functions output by JAX's
  transformations. For example, consder ``g = jax.jit(f)``. When ``g`` is
  called, JAX's JIT compilation machinery is invoked, which in turn calls ``f``
  in order to trace and translate it. If the function ``f`` raises an exception,
  the stack unwinds through JAX's JIT internals up to the original call site of
  ``g``. Because the function returned by ``jax.jit`` is annotated as an
  ``api_boundary``, such an exception is accompanied by an additional traceback
  that excludes the frames specific to JAX's implementation.
  '''

  if not filtered_tracebacks_supported():
    return fun

  @util.wraps(fun)
  def reraise_with_filtered_traceback(*args, **kwargs):
    try:
      return fun(*args, **kwargs)
    except Exception as e:
      if not is_under_reraiser(e):
        filtered_tb, cause, filtered = None, None, None
        try:
          filtered_tb = filter_traceback_and_stack(e)
          if filtered_tb:
            msg = format_exception_only(e)
            msg = f'{msg}\n\n{_jax_message_append}'
            filtered = FilteredStackTrace(msg).with_traceback(filtered_tb)
            cause = last_cause(e)
            cause.__cause__ = filtered
            raise
          else:
            raise
        finally:
          del filtered_tb
          del cause
          del filtered
      else:
        raise
  return reraise_with_filtered_traceback
<commit_msg>[JAX] Enable filtered tracebacks on Python 3.6.<commit_after># Copyright 2020 Google LLC
#
# 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
#
#     https://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
import sys
import traceback
import types

from jax.lib import xla_extension
from jax._src import util

_exclude_paths = [__file__, util.__file__]

def register_exclusion(path):
  _exclude_paths.append(path)

_jax_message_append = (
    'The stack trace above excludes JAX-internal frames.\n'
    'The following is the original exception that occurred, unmodified.\n'
    '\n--------------------')

def path_starts_with(path, path_prefix):
  path = os.path.abspath(path)
  path_prefix = os.path.abspath(path_prefix)
  if not os.path.exists(path_prefix):
    return False
  try:
    common = os.path.commonpath([path, path_prefix])
    return os.path.samefile(common, path_prefix)
  except ValueError:
    # path and path_prefix are both absolute, the only case will raise a
    # ValueError is different drives.
    # https://docs.python.org/3/library/os.path.html#os.path.commonpath
    return False

def include_frame(f):
  return not any(path_starts_with(f.f_code.co_filename, path)
                 for path in _exclude_paths)

# When scanning stack traces, we might encounter frames from cpython that are
# removed from printed stack traces, such as frames from parts of importlib. We
# ignore these frames heuristically based on source and name match.
def ignore_known_hidden_frame(f):
  return 'importlib._bootstrap' in f.f_code.co_filename

def filter_traceback_and_stack(tb):
  out = None

  # Scan the traceback and collect relevant frames.

  for f, lineno in reversed(list(traceback.walk_tb(tb))):
    if include_frame(f):
      out = make_traceback(out, f, f.f_lasti, lineno)  # pytype: disable=wrong-arg-count

  # Continue up the call stack.
  #
  # We would like to avoid stepping too far up, e.g. past the exec/eval point of
  # a REPL such as IPython. To that end, we stop past the first contiguous bunch
  # of module-level frames, if we reach any such frames at all. This is a
  # heuristic that might stop in advance of the REPL boundary. For example, if
  # the call stack includes module-level frames from the current module A, and
  # the current module A was imported from within a function F elsewhere, then
  # the stack trace we produce will be truncated at F's frame.

  reached_module_level = False
  for f, lineno in traceback.walk_stack(tb.tb_frame):
    if ignore_known_hidden_frame(f):
      continue
    if reached_module_level and f.f_code.co_name != '<module>':
      break
    if include_frame(f):
      out = make_traceback(out, f, f.f_lasti, lineno)  # pytype: disable=wrong-arg-count
    if f.f_code.co_name == '<module>':
      reached_module_level = True

  return out

def is_reraiser_frame(f):
  return (f.filename == __file__ and
          f.name == 'reraise_with_filtered_traceback')

def is_under_reraiser(e):
  tb = traceback.extract_stack(e.__traceback__.tb_frame)
  return any(is_reraiser_frame(f) for f in tb[:-1])

def format_exception_only(e):
  return ''.join(traceback.format_exception_only(type(e), e)).strip()

def last_cause(e):
  prev, cur = e, e.__cause__
  while cur is not None:
    prev, cur = cur, cur.__cause__
  return prev

class FilteredStackTrace(Exception): pass

make_traceback = (types.TracebackType if sys.version_info >= (3, 7) else
                  getattr(xla_extension, "make_python_traceback", None))

def filtered_tracebacks_supported():
  return make_traceback is not None

def api_boundary(fun):
  '''Wraps ``fun`` to form a boundary for filtering exception tracebacks.

  When an exception occurs below ``fun``, this appends to it a custom
  ``__cause__`` that carries a filtered traceback. The traceback imitates the
  stack trace of the original exception, but with JAX-internal frames removed.

  This boundary annotation works in composition with itself. The topmost frame
  corresponding to an ``api_boundary`` is the one below which stack traces are
  filtered. In other words, if ``api_boundary(f)`` calls ``api_boundary(g)``,
  directly or indirectly, the filtered stack trace provided is the same as if
  ``api_boundary(f)`` were to simply call ``g`` instead.

  This annotation is primarily useful in wrapping functions output by JAX's
  transformations. For example, consder ``g = jax.jit(f)``. When ``g`` is
  called, JAX's JIT compilation machinery is invoked, which in turn calls ``f``
  in order to trace and translate it. If the function ``f`` raises an exception,
  the stack unwinds through JAX's JIT internals up to the original call site of
  ``g``. Because the function returned by ``jax.jit`` is annotated as an
  ``api_boundary``, such an exception is accompanied by an additional traceback
  that excludes the frames specific to JAX's implementation.
  '''

  if not filtered_tracebacks_supported():
    return fun

  @util.wraps(fun)
  def reraise_with_filtered_traceback(*args, **kwargs):
    try:
      return fun(*args, **kwargs)
    except Exception as e:
      if not is_under_reraiser(e):
        filtered_tb, cause, filtered = None, None, None
        try:
          filtered_tb = filter_traceback_and_stack(e.__traceback__)
          if filtered_tb:
            msg = format_exception_only(e)
            msg = f'{msg}\n\n{_jax_message_append}'
            filtered = FilteredStackTrace(msg).with_traceback(filtered_tb)
            cause = last_cause(e)
            cause.__cause__ = filtered
            raise
          else:
            raise
        finally:
          del filtered_tb
          del cause
          del filtered
      else:
        raise
  return reraise_with_filtered_traceback
<|endoftext|>