<commit_before>"""Activity Partition item."""

from typing import List

from cairo import Context as CairoContext
from gaphas.geometry import Rectangle
from generic.event import Event

from gaphor import UML
from gaphor.core.modeling import DrawContext
from gaphor.core.modeling.properties import attribute
from gaphor.core.styling import Style, VerticalAlign
from gaphor.diagram.presentation import ElementPresentation, Named
from gaphor.diagram.shapes import Box, cairo_state, draw_highlight, stroke
from gaphor.diagram.support import represents
from gaphor.diagram.text import Layout

INIT_NUM_PARTITIONS: int = 2

@represents(UML.ActivityPartition)
class PartitionItem(ElementPresentation, Named):
    def __init__(self, id=None, model=None):
        super().__init__(id, model)
        self.partitions: List[UML.ActivityPartition] = []
        self.min_width = 150
        self.min_height = 300
        self.style: Style = {
            "font-family": "sans",
            "font-size": 14,
            "line-width": 2.4,
            "padding": (2, 2, 2, 2),
            "vertical-align": VerticalAlign.TOP,
        }

        self.watch("subject[NamedElement].name")
        self.watch("subject.appliedStereotype.classifier.name")
        self.watch("num_partitions", self.update_shapes)
        self.watch("partitions_dirty", self.update_shapes)

    num_partitions: attribute[int] = attribute(
        "num_partitions", int, default=INIT_NUM_PARTITIONS
    )
    partitions_dirty: attribute[bool] = attribute(
        "partitions_dirty", bool, default=False
    )

    def pre_update(self, context: DrawContext) -> None:
        """Set the min width of all the swimlanes."""
        self.min_width = 150 * self.num_partitions

    def post_update(self, context: DrawContext) -> None:
        """Nest the Activity Partitions under the Activity.

        This needs to be done after the update because the activity's
        classifier doesn't always exist yet.
        """
        activity = self.diagram.namespace
        for partition in self.partitions:
            if not partition.activity and isinstance(
                activity.ownedClassifier[0], UML.Activity
            ):
                partition.activity = activity.ownedClassifier[0]

    def update_partitions(self) -> None:
        """Add and remove UML.ActivityPartitions."""
        if not len(self.partitions):
            self.partitions.append(self.subject)

        if self.num_partitions > len(self.partitions):
            partition = self.subject.model.create(UML.ActivityPartition)
            partition.name = "NewActivityPartition"
            self.partitions.append(partition)
        elif self.num_partitions < len(self.partitions):
            partition = self.partitions[-1]
            partition.unlink()
            self.partitions.pop()

    def update_shapes(self, event: Event = None) -> None:
        """Update the number of partitions and draw them."""
        self.update_partitions()
        self.shape = Box(
            style=self.style,
            draw=self.draw_swimlanes,
        )
        self.partitions_dirty = False

    def draw_swimlanes(
        self, box: Box, context: DrawContext, bounding_box: Rectangle
    ) -> None:
        """Draw the vertical partitions as connected swimlanes.

        The partitions are open on the bottom. We divide the total size
        by the total number of partitions and space them evenly.
        """
        assert self.canvas

        cr = context.cairo
        cr.set_line_width(context.style["line-width"])
        header_h = self.draw_outline(bounding_box, cr)
        self.draw_partitions(bounding_box, context, header_h)
        stroke(context)
        self.draw_hover(bounding_box, context)

    def draw_hover(self, bounding_box: Rectangle, context: DrawContext):
        """Add dashed line on bottom of swimlanes when hovered."""
        cr = context.cairo
        if context.hovered or context.dropzone:
            with cairo_state(cr):
                cr.set_dash((1.0, 5.0), 0)
                cr.set_line_width(1.0)
                cr.rectangle(0, 0, bounding_box.width, bounding_box.height)
                draw_highlight(context)
                cr.stroke()

    def draw_partitions(
        self, bounding_box: Rectangle, context: DrawContext, header_h: int
    ) -> None:
        """Draw partition separators and add the name."""
        cr = context.cairo
        partition_width = bounding_box.width / self.num_partitions
        layout = Layout()
        for num, partition in enumerate(self.partitions):
            cr.move_to(partition_width * num, 0)
            cr.line_to(partition_width * num, bounding_box.height)
            layout.set(text=partition.name, font=self.style)
            padding_top = context.style["padding"][0]
            cr.move_to(partition_width * num, padding_top * 3)
            layout.show_layout(
                cr,
                partition_width,
                default_size=(partition_width, header_h),
            )

    def draw_outline(self, bounding_box: Rectangle, cr: CairoContext) -> int:
        """Draw the outline and header of the swimlanes."""
        cr.move_to(0, bounding_box.height)
        cr.line_to(0, 0)
        cr.line_to(bounding_box.width, 0)
        cr.line_to(bounding_box.width, bounding_box.height)
        cr.move_to(0, bounding_box.height)
        header_h = 29
        cr.line_to(0, header_h)
        cr.line_to(0 + bounding_box.width, header_h)
        cr.line_to(0 + bounding_box.width, bounding_box.height)
        return header_h
<commit_msg>Make partition header height a constant<commit_after>"""Activity Partition item."""

from typing import List

from cairo import Context as CairoContext
from gaphas.geometry import Rectangle
from generic.event import Event

from gaphor import UML
from gaphor.core.modeling import DrawContext
from gaphor.core.modeling.properties import attribute
from gaphor.core.styling import Style, VerticalAlign
from gaphor.diagram.presentation import ElementPresentation, Named
from gaphor.diagram.shapes import Box, cairo_state, draw_highlight, stroke
from gaphor.diagram.support import represents
from gaphor.diagram.text import Layout

INIT_NUM_PARTITIONS: int = 2
HEADER_HEIGHT: int = 29

@represents(UML.ActivityPartition)
class PartitionItem(ElementPresentation, Named):
    def __init__(self, id=None, model=None):
        super().__init__(id, model)
        self.partitions: List[UML.ActivityPartition] = []
        self.min_width = 150
        self.min_height = 300
        self.style: Style = {
            "font-family": "sans",
            "font-size": 14,
            "line-width": 2.4,
            "padding": (2, 2, 2, 2),
            "vertical-align": VerticalAlign.TOP,
        }

        self.watch("subject[NamedElement].name")
        self.watch("subject.appliedStereotype.classifier.name")
        self.watch("num_partitions", self.update_shapes)
        self.watch("partitions_dirty", self.update_shapes)

    num_partitions: attribute[int] = attribute(
        "num_partitions", int, default=INIT_NUM_PARTITIONS
    )
    partitions_dirty: attribute[bool] = attribute(
        "partitions_dirty", bool, default=False
    )

    def pre_update(self, context: DrawContext) -> None:
        """Set the min width of all the swimlanes."""
        self.min_width = 150 * self.num_partitions

    def post_update(self, context: DrawContext) -> None:
        """Nest the Activity Partitions under the Activity.

        This needs to be done after the update because the activity's
        classifier doesn't always exist yet.
        """
        activity = self.diagram.namespace
        for partition in self.partitions:
            if not partition.activity and isinstance(
                activity.ownedClassifier[0], UML.Activity
            ):
                partition.activity = activity.ownedClassifier[0]

    def update_partitions(self) -> None:
        """Add and remove UML.ActivityPartitions."""
        if not len(self.partitions):
            self.partitions.append(self.subject)

        if self.num_partitions > len(self.partitions):
            partition = self.subject.model.create(UML.ActivityPartition)
            partition.name = "NewActivityPartition"
            self.partitions.append(partition)
        elif self.num_partitions < len(self.partitions):
            partition = self.partitions[-1]
            partition.unlink()
            self.partitions.pop()

    def update_shapes(self, event: Event = None) -> None:
        """Update the number of partitions and draw them."""
        self.update_partitions()
        self.shape = Box(
            style=self.style,
            draw=self.draw_swimlanes,
        )
        self.partitions_dirty = False

    def draw_swimlanes(
        self, box: Box, context: DrawContext, bounding_box: Rectangle
    ) -> None:
        """Draw the vertical partitions as connected swimlanes.

        The partitions are open on the bottom. We divide the total size
        by the total number of partitions and space them evenly.
        """
        assert self.canvas

        cr = context.cairo
        cr.set_line_width(context.style["line-width"])
        self.draw_outline(bounding_box, cr)
        self.draw_partitions(bounding_box, context)
        stroke(context)
        self.draw_hover(bounding_box, context)

    def draw_hover(self, bounding_box: Rectangle, context: DrawContext) -> None:
        """Add dashed line on bottom of swimlanes when hovered."""
        cr = context.cairo
        if context.hovered or context.dropzone:
            with cairo_state(cr):
                cr.set_dash((1.0, 5.0), 0)
                cr.set_line_width(1.0)
                cr.rectangle(0, 0, bounding_box.width, bounding_box.height)
                draw_highlight(context)
                cr.stroke()

    def draw_partitions(self, bounding_box: Rectangle, context: DrawContext) -> None:
        """Draw partition separators and add the name."""
        cr = context.cairo
        partition_width = bounding_box.width / self.num_partitions
        layout = Layout()
        for num, partition in enumerate(self.partitions):
            cr.move_to(partition_width * num, 0)
            cr.line_to(partition_width * num, bounding_box.height)
            layout.set(text=partition.name, font=self.style)
            padding_top = context.style["padding"][0]
            cr.move_to(partition_width * num, padding_top * 3)
            layout.show_layout(
                cr,
                partition_width,
                default_size=(partition_width, HEADER_HEIGHT),
            )

    def draw_outline(self, bounding_box: Rectangle, cr: CairoContext) -> None:
        """Draw the outline and header of the swimlanes."""
        cr.move_to(0, bounding_box.height)
        cr.line_to(0, 0)
        cr.line_to(bounding_box.width, 0)
        cr.line_to(bounding_box.width, bounding_box.height)
        cr.move_to(0, bounding_box.height)
        cr.line_to(0, HEADER_HEIGHT)
        cr.line_to(0 + bounding_box.width, HEADER_HEIGHT)
        cr.line_to(0 + bounding_box.width, bounding_box.height)
<|endoftext|>


<commit_before>"""Activity Partition item."""

from cairo import Context as CairoContext
from gaphas.geometry import Rectangle

from gaphor import UML
from gaphor.core.modeling import DrawContext
from gaphor.core.modeling.properties import association
from gaphor.core.styling import VerticalAlign
from gaphor.diagram.presentation import ElementPresentation
from gaphor.diagram.shapes import Box, cairo_state, draw_highlight, stroke
from gaphor.diagram.support import represents
from gaphor.diagram.text import Layout

HEADER_HEIGHT: int = 29

@represents(UML.ActivityPartition)
class PartitionItem(ElementPresentation):
    def __init__(self, diagram, id=None):
        super().__init__(diagram, id)
        self.min_height = 300

        self.shape = Box(
            style={
                "line-width": 2.4,
                "padding": (2, 2, 2, 2),
                "vertical-align": VerticalAlign.TOP,
            },
            draw=self.draw_swimlanes,
        )
        self.watch("subject[NamedElement].name")
        self.watch("subject.appliedStereotype.classifier.name")
        self.watch("partition.name")
        self.watch("partition")

    partition = association("partition", UML.ActivityPartition, composite=True)

    def postload(self):
        super().postload()
        if self.subject and self.subject not in self.partition:
            self.partition = self.subject

    def pre_update(self, context: DrawContext) -> None:
        """Set the min width of all the swimlanes."""
        self.min_width = 150 * len(self.partition)

    def draw_swimlanes(
        self, box: Box, context: DrawContext, bounding_box: Rectangle
    ) -> None:
        """Draw the vertical partitions as connected swimlanes.

        The partitions are open on the bottom. We divide the total size
        by the total number of partitions and space them evenly.
        """
        cr = context.cairo
        cr.set_line_width(context.style["line-width"])
        self.draw_outline(bounding_box, cr)
        self.draw_partitions(bounding_box, context)
        stroke(context)
        self.draw_hover(bounding_box, context)

    def draw_hover(self, bounding_box: Rectangle, context: DrawContext) -> None:
        """Add dashed line on bottom of swimlanes when hovered."""
        if context.hovered or context.dropzone:
            cr = context.cairo
            with cairo_state(cr):
                cr.set_dash((1.0, 5.0), 0)
                cr.set_line_width(1.0)
                cr.rectangle(0, 0, bounding_box.width, bounding_box.height)
                draw_highlight(context)
                cr.stroke()

    def draw_partitions(self, bounding_box: Rectangle, context: DrawContext) -> None:
        """Draw partition separators and add the name."""
        cr = context.cairo
        if self.partition:
            partition_width = bounding_box.width / len(self.partition)
        else:
            partition_width = bounding_box.width / 2
        layout = Layout()
        style = context.style
        padding_top = context.style["padding"][0]
        for num, partition in enumerate(self.partition):
            cr.move_to(partition_width * num, 0)
            cr.line_to(partition_width * num, bounding_box.height)
            layout.set(text=partition.name, font=style)
            cr.move_to(partition_width * num, padding_top * 3)
            layout.show_layout(
                cr,
                partition_width,
                default_size=(partition_width, HEADER_HEIGHT),
            )

    def draw_outline(self, bounding_box: Rectangle, cr: CairoContext) -> None:
        """Draw the outline and header of the swimlanes."""
        cr.move_to(0, bounding_box.height)
        cr.line_to(0, 0)
        cr.line_to(bounding_box.width, 0)
        cr.line_to(bounding_box.width, bounding_box.height)
        cr.move_to(0, bounding_box.height)
        cr.line_to(0, HEADER_HEIGHT)
        cr.line_to(0 + bounding_box.width, HEADER_HEIGHT)
        cr.line_to(0 + bounding_box.width, bounding_box.height)
<commit_msg>Add highlighting to Partition<commit_after>"""Activity Partition item."""

from gaphas.geometry import Rectangle

from gaphor import UML
from gaphor.core.modeling import DrawContext
from gaphor.core.modeling.properties import association
from gaphor.core.styling import VerticalAlign
from gaphor.diagram.presentation import ElementPresentation
from gaphor.diagram.shapes import Box, cairo_state, draw_highlight, stroke
from gaphor.diagram.support import represents
from gaphor.diagram.text import Layout

HEADER_HEIGHT: int = 29

@represents(UML.ActivityPartition)
class PartitionItem(ElementPresentation):
    def __init__(self, diagram, id=None):
        super().__init__(diagram, id)
        self.min_height = 300

        self.shape = Box(
            style={
                "line-width": 2.4,
                "padding": (2, 2, 2, 2),
                "vertical-align": VerticalAlign.TOP,
            },
            draw=self.draw_swimlanes,
        )
        self.watch("subject[NamedElement].name")
        self.watch("subject.appliedStereotype.classifier.name")
        self.watch("partition.name")
        self.watch("partition")

    partition = association("partition", UML.ActivityPartition, composite=True)

    def postload(self):
        super().postload()
        if self.subject and self.subject not in self.partition:
            self.partition = self.subject

    def pre_update(self, context: DrawContext) -> None:
        """Set the min width of all the swimlanes."""
        self.min_width = 150 * len(self.partition)

    def draw_swimlanes(
        self, box: Box, context: DrawContext, bounding_box: Rectangle
    ) -> None:
        """Draw the vertical partitions as connected swimlanes.

        The partitions are open on the bottom. We divide the total size
        by the total number of partitions and space them evenly.
        """
        cr = context.cairo
        cr.set_line_width(context.style["line-width"])
        self.draw_outline(bounding_box, context)
        self.draw_partitions(bounding_box, context)
        stroke(context)
        self.draw_hover(bounding_box, context)

    def draw_hover(self, bounding_box: Rectangle, context: DrawContext) -> None:
        """Add dashed line on bottom of swimlanes when hovered."""
        if context.hovered or context.dropzone:
            cr = context.cairo
            with cairo_state(cr):
                cr.set_dash((1.0, 5.0), 0)
                cr.set_line_width(1.0)
                cr.rectangle(0, 0, bounding_box.width, bounding_box.height)
                cr.stroke()

    def draw_partitions(self, bounding_box: Rectangle, context: DrawContext) -> None:
        """Draw partition separators and add the name."""
        cr = context.cairo
        if self.partition:
            partition_width = bounding_box.width / len(self.partition)
        else:
            partition_width = bounding_box.width / 2
        layout = Layout()
        style = context.style
        padding_top = context.style["padding"][0]
        for num, partition in enumerate(self.partition):
            cr.move_to(partition_width * num, 0)
            cr.line_to(partition_width * num, bounding_box.height)
            layout.set(text=partition.name, font=style)
            cr.move_to(partition_width * num, padding_top * 3)
            layout.show_layout(
                cr,
                partition_width,
                default_size=(partition_width, HEADER_HEIGHT),
            )

    def draw_outline(self, bounding_box: Rectangle, context: DrawContext) -> None:
        """Draw the outline and header of the swimlanes."""
        cr = context.cairo
        cr.move_to(0, bounding_box.height)
        cr.line_to(0, 0)
        cr.line_to(bounding_box.width, 0)
        cr.line_to(bounding_box.width, bounding_box.height)
        draw_highlight(context)
        cr.move_to(0, bounding_box.height)
        cr.line_to(0, HEADER_HEIGHT)
        cr.line_to(0 + bounding_box.width, HEADER_HEIGHT)
        cr.line_to(0 + bounding_box.width, bounding_box.height)
<|endoftext|>


<commit_before>import AppKit
NSApplication = AppKit.NSApplication

try:
    import nodebox
except ImportError:
    import sys, os
    nodebox_dir = os.path.dirname(os.path.abspath(__file__))
    sys.path.append(os.path.dirname(nodebox_dir))

import modebox.graphics
graphics = modebox.graphics

import nodebox.util
util = nodebox.util

#from nodebox import graphics
#from nodebox import util

class NodeBoxRunner(object):

    def __init__(self):
        # Force NSApp initialisation.
        NSApplication.sharedApplication().activateIgnoringOtherApps_(0)
        self.namespace = {}
        self.canvas = graphics.Canvas()
        self.context = graphics.Context(self.canvas, self.namespace)
        self.__doc__ = {}
        self._pageNumber = 1
        self.frame = 1

    def _check_animation(self):
        """Returns False if this is not an animation, True otherwise.
        Throws an expection if the animation is not correct (missing a draw method)."""
        if self.canvas.speed is not None:
            if not self.namespace.has_key('draw'):
                raise graphics.NodeBoxError('Not a correct animation: No draw() method.')
            return True
        return False

    def run(self, source_or_code):
        self._initNamespace()
        if isinstance(source_or_code, basestring):
            source_or_code = compile(source_or_code + "\n\n", "<Untitled>", "exec")
        exec source_or_code in self.namespace
        if self._check_animation():
            if self.namespace.has_key('setup'):
                self.namespace['setup']()
            self.namespace['draw']()

    def run_multiple(self, source_or_code, frames):
        if isinstance(source_or_code, basestring):
            source_or_code = compile(source_or_code + "\n\n", "<Untitled>", "exec")

        # First frame is special:
        self.run(source_or_code)
        yield 1
        animation = self._check_animation()

        for i in range(frames-1):
            self.canvas.clear()
            self.frame = i + 2
            self.namespace["PAGENUM"] = self.namespace["FRAME"] = self.frame
            if animation:
                self.namespace['draw']()
            else:
                exec source_or_code in self.namespace
            yield self.frame

    def _initNamespace(self, frame=1):
        self.canvas.clear()
        self.namespace.clear()
        # Add everything from the namespace
        for name in graphics.__all__:
            self.namespace[name] = getattr(graphics, name)
        for name in util.__all__:
            self.namespace[name] = getattr(util, name)
        # Add everything from the context object
        self.namespace["_ctx"] = self.context
        for attrName in dir(self.context):
            self.namespace[attrName] = getattr(self.context, attrName)
        # Add the document global
        self.namespace["__doc__"] = self.__doc__
        # Add the frame
        self.frame = frame
        self.namespace["PAGENUM"] = self.namespace["FRAME"] = self.frame

def make_image(source_or_code, outputfile):

    """Given a source string or code object, executes the scripts and saves the result as an image.
    Supported image extensions: pdf, tiff, png, jpg, gif"""

    runner = NodeBoxRunner()
    runner.run(source_or_code)
    runner.canvas.save(outputfile)

def make_movie(source_or_code, outputfile, frames, fps=30):

    """Given a source string or code object, executes the scripts and saves the result as a movie.
    You also have to specify the number of frames to render.
    Supported movie extension: mov"""

    from nodebox.util import QTSupport
    runner = NodeBoxRunner()
    movie = QTSupport.Movie(outputfile, fps)
    for frame in runner.run_multiple(source_or_code, frames):
        movie.add(runner.canvas)
    movie.save()

def usage(err=""):
    if len(err) > 0:
        err = '\n\nError: ' + str(err)
    print """NodeBox console runner
Usage: console.py sourcefile imagefile
   or: console.py sourcefile moviefile number_of_frames [fps]
Supported image extensions: pdf, tiff, png, jpg, gif
Supported movie extension:  mov""" + err

def main():
    import sys, os
    if len(sys.argv) < 2:
        usage()
    elif len(sys.argv) == 3: # Should be an image
        basename, ext = os.path.splitext(sys.argv[2])
        if ext not in ('.pdf', '.gif', '.jpg', '.jpeg', '.png', '.tiff'):
            return usage('This is not a supported image format.')
        make_image(open(sys.argv[1]).read(), sys.argv[2])
    elif len(sys.argv) == 4 or len(sys.argv) == 5: # Should be a movie
        basename, ext = os.path.splitext(sys.argv[2])
        if ext != '.mov':
            return usage('This is not a supported movie format.')
        if len(sys.argv) == 5:
            try:
                fps = int(sys.argv[4])
            except ValueError:
                return usage()
        else:
            fps = 30
        make_movie(open(sys.argv[1]).read(), sys.argv[2], int(sys.argv[3]), fps)

def test():
    # Creating the NodeBoxRunner class directly:
    runner = NodeBoxRunner()
    runner.run('size(500,500)\nfor i in range(400):\n  oval(random(WIDTH),random(HEIGHT),50,50, fill=(random(), 0,0,random()))')
    runner.canvas.save('console-test.pdf')
    runner.canvas.save('console-test.png')

    # Using the runner for animations:
    runner = NodeBoxRunner()
    for frame in runner.run_multiple('size(300, 300)\ntext(FRAME, 100, 100)', 10):
        runner.canvas.save('console-test-frame%02i.png' % frame)

    # Using the shortcut functions:
    make_image('size(200,200)\ntext(FRAME, 100, 100)', 'console-test.gif')
    make_movie('size(200,200)\ntext(FRAME, 100, 100)', 'console-test.mov', 10)

if __name__=='__main__':
    main()<commit_msg>Typo.<commit_after>import AppKit
NSApplication = AppKit.NSApplication

try:
    import nodebox
except ImportError:
    import sys, os
    nodebox_dir = os.path.dirname(os.path.abspath(__file__))
    sys.path.append(os.path.dirname(nodebox_dir))

import nodebox.graphics
graphics = nodebox.graphics

import nodebox.util
util = nodebox.util

#from nodebox import graphics
#from nodebox import util

class NodeBoxRunner(object):

    def __init__(self):
        # Force NSApp initialisation.
        NSApplication.sharedApplication().activateIgnoringOtherApps_(0)
        self.namespace = {}
        self.canvas = graphics.Canvas()
        self.context = graphics.Context(self.canvas, self.namespace)
        self.__doc__ = {}
        self._pageNumber = 1
        self.frame = 1

    def _check_animation(self):
        """Returns False if this is not an animation, True otherwise.
        Throws an expection if the animation is not correct (missing a draw method)."""
        if self.canvas.speed is not None:
            if not self.namespace.has_key('draw'):
                raise graphics.NodeBoxError('Not a correct animation: No draw() method.')
            return True
        return False

    def run(self, source_or_code):
        self._initNamespace()
        if isinstance(source_or_code, basestring):
            source_or_code = compile(source_or_code + "\n\n", "<Untitled>", "exec")
        exec source_or_code in self.namespace
        if self._check_animation():
            if self.namespace.has_key('setup'):
                self.namespace['setup']()
            self.namespace['draw']()

    def run_multiple(self, source_or_code, frames):
        if isinstance(source_or_code, basestring):
            source_or_code = compile(source_or_code + "\n\n", "<Untitled>", "exec")

        # First frame is special:
        self.run(source_or_code)
        yield 1
        animation = self._check_animation()

        for i in range(frames-1):
            self.canvas.clear()
            self.frame = i + 2
            self.namespace["PAGENUM"] = self.namespace["FRAME"] = self.frame
            if animation:
                self.namespace['draw']()
            else:
                exec source_or_code in self.namespace
            yield self.frame

    def _initNamespace(self, frame=1):
        self.canvas.clear()
        self.namespace.clear()
        # Add everything from the namespace
        for name in graphics.__all__:
            self.namespace[name] = getattr(graphics, name)
        for name in util.__all__:
            self.namespace[name] = getattr(util, name)
        # Add everything from the context object
        self.namespace["_ctx"] = self.context
        for attrName in dir(self.context):
            self.namespace[attrName] = getattr(self.context, attrName)
        # Add the document global
        self.namespace["__doc__"] = self.__doc__
        # Add the frame
        self.frame = frame
        self.namespace["PAGENUM"] = self.namespace["FRAME"] = self.frame

def make_image(source_or_code, outputfile):

    """Given a source string or code object, executes the scripts and saves the result as an image.
    Supported image extensions: pdf, tiff, png, jpg, gif"""

    runner = NodeBoxRunner()
    runner.run(source_or_code)
    runner.canvas.save(outputfile)

def make_movie(source_or_code, outputfile, frames, fps=30):

    """Given a source string or code object, executes the scripts and saves the result as a movie.
    You also have to specify the number of frames to render.
    Supported movie extension: mov"""

    from nodebox.util import QTSupport
    runner = NodeBoxRunner()
    movie = QTSupport.Movie(outputfile, fps)
    for frame in runner.run_multiple(source_or_code, frames):
        movie.add(runner.canvas)
    movie.save()

def usage(err=""):
    if len(err) > 0:
        err = '\n\nError: ' + str(err)
    print """NodeBox console runner
Usage: console.py sourcefile imagefile
   or: console.py sourcefile moviefile number_of_frames [fps]
Supported image extensions: pdf, tiff, png, jpg, gif
Supported movie extension:  mov""" + err

def main():
    import sys, os
    if len(sys.argv) < 2:
        usage()
    elif len(sys.argv) == 3: # Should be an image
        basename, ext = os.path.splitext(sys.argv[2])
        if ext not in ('.pdf', '.gif', '.jpg', '.jpeg', '.png', '.tiff'):
            return usage('This is not a supported image format.')
        make_image(open(sys.argv[1]).read(), sys.argv[2])
    elif len(sys.argv) == 4 or len(sys.argv) == 5: # Should be a movie
        basename, ext = os.path.splitext(sys.argv[2])
        if ext != '.mov':
            return usage('This is not a supported movie format.')
        if len(sys.argv) == 5:
            try:
                fps = int(sys.argv[4])
            except ValueError:
                return usage()
        else:
            fps = 30
        make_movie(open(sys.argv[1]).read(), sys.argv[2], int(sys.argv[3]), fps)

def test():
    # Creating the NodeBoxRunner class directly:
    runner = NodeBoxRunner()
    runner.run('size(500,500)\nfor i in range(400):\n  oval(random(WIDTH),random(HEIGHT),50,50, fill=(random(), 0,0,random()))')
    runner.canvas.save('console-test.pdf')
    runner.canvas.save('console-test.png')

    # Using the runner for animations:
    runner = NodeBoxRunner()
    for frame in runner.run_multiple('size(300, 300)\ntext(FRAME, 100, 100)', 10):
        runner.canvas.save('console-test-frame%02i.png' % frame)

    # Using the shortcut functions:
    make_image('size(200,200)\ntext(FRAME, 100, 100)', 'console-test.gif')
    make_movie('size(200,200)\ntext(FRAME, 100, 100)', 'console-test.mov', 10)

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


<commit_before>#!/usr/bin/python
# Copyright (c) Johannes Renner 2011, see LICENSE file.

import os
import sys
import datetime
import ConfigParser

VERSION = "0.02a"

CONFIG_FILE = "./config"
LOGLEVELS = {"DEBUG" : 0, "INFO" : 1, "NOTICE" : 2, "WARN" : 3, "ERROR" : 4}

# Override these in the config file
BACKUP_PATH = os.getcwd()
LOGLEVEL = "NOTICE"

class Target:
  def __init__(self, path):
    # Remember the path of the target file
    self.path = path
    # Set the filename as the name of the target
    self.name = os.path.split(path)[1]
    # Parse paths from the file
    self.paths = self.read_paths(self.path)
    plog("INFO", "Initialized target '" + self.name + "'")

  def read_paths(self, filename):
    """ Read a list of paths from a file """
    try:
      file = open(filename, 'r')
      lines = file.readlines()
      file.close()
    except IOError:
      plog("ERROR", "Targetfile '" + filename + "' is not existing/readable!")
      sys.exit(1)
    # Clear paths
    paths = []
    for l in lines:
      # Split to single values
      splitted = l.split()
      if len(splitted) > 0:
        p = splitted[0]
        paths.append(p)
        plog("DEBUG", "Added path: "+p)
    return paths

  def backup(self):
    """ Create a zipped tape archive """
    if len(self.paths) >= 1:
      # We have at least one path, generate the archive filename
      today = str(datetime.date.today())
      tarfile = BACKUP_PATH + "/" + self.name + "." + today + ".tar"
      for p in self.paths:
        if os.path.exists(p):
          plog("DEBUG", "Path exists: " + p)
          plog("DEBUG", "split 0: " + os.path.split(p)[0] + ", split 1: "
             + os.path.split(p)[1])
          # Change to parent directory in any case
          # FIXME: This will not work for target paths on the root level
          command = "cd " + os.path.split(p)[0] + ";"
          # Create archive or append
          if not os.path.exists(tarfile):
            command += "tar -cvf " + tarfile + " " + os.path.split(p)[1]
            plog("DEBUG", "Command: << " + command + " >>")
            plog("INFO", "Creating archive of " + p + " to " + tarfile)
          else:
            command += "tar -rvf " + tarfile + " " + os.path.split(p)[1]
            plog("DEBUG", "Command: << " + command + " >>")
            plog("INFO", "Appending " + p + " to " + tarfile)
          os.system(command)
        else:
          self.paths.remove(p)
          plog("WARN","Path doesn't exist! Remove '" + p + "' from target '"
             + self.name + "'?")
      # Zip the archive
      os.system("gzip " + tarfile)
      plog("INFO", "Backup finished on target '" + self.name + ".'")
      # DEBUG: Print out the return value of the last command
      os.system("echo '[Return value: '$?']'")

def plog(level, msg):
  """ Log messages to stdout """
  # TODO: Add timestamps
  if(LOGLEVELS[level] >= LOGLEVELS[LOGLEVEL]):
    print level + ": " + msg
    sys.stdout.flush()

def configure(config_file):
  """ Read configuration from a file """
  config = ConfigParser.SafeConfigParser()
  if os.path.exists(config_file):
    config.read(config_file)
  else:
    plog("NOTICE", "Config file '" + config_file
       + "' does not exist, running on default values.")
    return
  try:
    global LOGLEVEL
    LOGLEVEL = config.get("GENERAL", "LOGLEVEL")
  except Exception as e:
    plog("NOTICE", str(e) + " -> using default value: " + LOGLEVEL)
  try:
    global BACKUP_PATH
    BACKUP_PATH = config.get("GENERAL", "BACKUP_PATH")
  except Exception as e:
    plog("NOTICE", str(e) + " -> using cwd: " + BACKUP_PATH)
  finally:
    plog("DEBUG", "BACKUP_PATH is " + BACKUP_PATH)

if __name__ == '__main__':
  """ Program entry point """
  if len(sys.argv) != 2:
    plog("ERROR", "Usage is \"./backdub.py <targetfile>\"")
  else:
    configure(CONFIG_FILE)
    plog("INFO", "backdub v" + str(VERSION))
    # Create target and back it up
    target = Target(sys.argv[1])
    target.backup()

<commit_msg>Very basic support for commandline options<commit_after>#!/usr/bin/python
# Copyright (c) Johannes Renner 2011, see LICENSE file.

import os
import sys
import datetime
import ConfigParser

VERSION = "0.02a"

CONFIG_FILE = "./config"
LOGLEVELS = {"DEBUG" : 0, "INFO" : 1, "NOTICE" : 2, "WARN" : 3, "ERROR" : 4}
VERBOSE = False

# Override these in the config file
BACKUP_PATH = os.getcwd()
LOGLEVEL = "NOTICE"

class Target:
  def __init__(self, path):
    # Remember the path of the target file
    self.path = path
    # Set the filename as the name of the target
    self.name = os.path.split(path)[1]
    # Parse paths from the file
    self.paths = self.read_paths(self.path)
    plog("INFO", "Initialized target '" + self.name + "'")

  def read_paths(self, filename):
    """ Read a list of paths from a file """
    try:
      file = open(filename, 'r')
      lines = file.readlines()
      file.close()
    except IOError:
      plog("ERROR", "Targetfile '" + filename + "' is not existing/readable!")
      sys.exit(1)
    # Clear paths
    paths = []
    for l in lines:
      # Split to single values
      splitted = l.split()
      if len(splitted) > 0:
        p = splitted[0]
        paths.append(p)
        plog("DEBUG", "Added path: "+p)
    return paths

  def backup(self):
    """ Create a zipped tape archive """
    if len(self.paths) >= 1:
      # We have at least one path, generate the archive filename
      today = str(datetime.date.today())
      tarfile = BACKUP_PATH + "/" + self.name + "." + today + ".tar"
      for p in self.paths:
        if os.path.exists(p):
          plog("DEBUG", "Path exists: " + p)
          plog("DEBUG", "split 0: " + os.path.split(p)[0] + ", split 1: "
             + os.path.split(p)[1])
          # Change to parent directory in any case
          # FIXME: This will not work for target paths on the root level
          command = "cd " + os.path.split(p)[0] + ";"
          # Handle 'verbose'
          options = "-"
          if VERBOSE:
            options += "v"
          # Create archive or append
          if not os.path.exists(tarfile):
            options += "cf"
            command += "tar " + options + " " + tarfile + " " + os.path.split(p)[1]
            plog("DEBUG", "Command: << " + command + " >>")
            plog("INFO", "Creating archive of " + p + " to " + tarfile)
          else:
            options += "rf"
            command += "tar " + options + " " + tarfile + " " + os.path.split(p)[1]
            plog("DEBUG", "Command: << " + command + " >>")
            plog("INFO", "Appending " + p + " to " + tarfile)
          os.system(command)
        else:
          self.paths.remove(p)
          plog("WARN","Path doesn't exist! Remove '" + p + "' from target '"
             + self.name + "'?")
      # Zip the archive
      os.system("gzip " + tarfile)
      plog("INFO", "Backup finished on target '" + self.name + ".'")
      # DEBUG: Print out the return value of the last command
      os.system("echo '[Return value: '$?']'")

def plog(level, msg):
  """ Log messages to stdout """
  # TODO: Add timestamps
  if(LOGLEVELS[level] >= LOGLEVELS[LOGLEVEL]):
    print level + ": " + msg
    sys.stdout.flush()

def configure(config_file):
  """ Read configuration from a file """
  config = ConfigParser.SafeConfigParser()
  if os.path.exists(config_file):
    config.read(config_file)
  else:
    plog("NOTICE", "Config file '" + config_file
       + "' does not exist, running on default values.")
    return
  try:
    global LOGLEVEL
    LOGLEVEL = config.get("GENERAL", "LOGLEVEL")
  except Exception as e:
    plog("NOTICE", str(e) + " -> using default value: " + LOGLEVEL)
  try:
    global BACKUP_PATH
    BACKUP_PATH = config.get("GENERAL", "BACKUP_PATH")
  except Exception as e:
    plog("NOTICE", str(e) + " -> using cwd: " + BACKUP_PATH)
  finally:
    plog("DEBUG", "BACKUP_PATH is " + BACKUP_PATH)

if __name__ == '__main__':
  """ Program entry point """
  if len(sys.argv) < 2:
    plog("ERROR", "Usage is \"./backdub.py <targetfile> <options>\"")
  else:
    configure(CONFIG_FILE)
    plog("INFO", "backdub v" + str(VERSION))
    # Create target and back it up
    target = Target(sys.argv[1])

    # Handle options
    if len(sys.argv) > 2:
      # Init options
      options = sys.argv[2]
      if options and options == "-v":
        VERBOSE = True
      elif options:
        plog("ERROR", "Unsupported option: " + options)
        sys.exit()

    # Start the backup
    target.backup()

<|endoftext|>


<commit_before>from __future__ import absolute_import

import cPickle
import itertools
import operator
from time import time
import pdb

from lasagne.layers import get_all_layers
from lasagne.layers import get_all_params
from lasagne.objectives import mse
from lasagne.updates import nesterov_momentum
import numpy as np
from sklearn.base import BaseEstimator
from sklearn.cross_validation import KFold
from sklearn.cross_validation import StratifiedKFold
from sklearn.metrics import accuracy_score
from sklearn.metrics import mean_squared_error
from sklearn.preprocessing import LabelEncoder
import theano
from theano import tensor as T

class _list(list):
    pass

class _dict(dict):
    def __contains__(self, key):
        return True

class ansi:
    BLUE = '\033[94m'
    GREEN = '\033[32m'
    ENDC = '\033[0m'

def negative_log_likelihood(output, prediction):
    return -T.mean(T.log(output)[T.arange(prediction.shape[0]), prediction])

class BatchIterator(object):
    def __init__(self, batch_size):
        self.batch_size = batch_size

    def __call__(self, X, y=None, test=False):
        self.X, self.y = X, y
        self.test = test
        return self

    def __iter__(self):
        n_samples = self.X.shape[0]
        bs = self.batch_size
        for i in range((n_samples + bs - 1) / bs):
            sl = slice(i * bs, (i + 1) * bs)
            Xb = self.X[sl]
            if self.y is not None:
                yb = self.y[sl]
            else:
                yb = None
            yield self.transform(Xb, yb)

    def transform(self, Xb, yb):
        return Xb, yb

class NeuralNet(BaseEstimator):
    """A scikit-learn estimator based on Lasagne.
    """
    def __init__(
        self,
        layers,
        update=nesterov_momentum,
        loss=None,
        batch_iterator=BatchIterator(batch_size=128),
        regression=False,
        max_epochs=100,
        eval_size=0.2,
        X_tensor_type=None,
        y_tensor_type=None,
        use_label_encoder=False,
        on_epoch_finished=(),
        on_training_finished=(),
        more_params=None,
        verbose=0,
        **kwargs
        ):
        if loss is None:
            loss = mse if regression else negative_log_likelihood
        if X_tensor_type is None:
            types = {
                2: T.matrix,
                3: T.tensor3,
                4: T.tensor4,
                }
            X_tensor_type = types[len(kwargs['input_shape'])]
        if y_tensor_type is None:
            y_tensor_type = T.fmatrix if regression else T.ivector

        self.layers = layers
        self.update = update
        self.loss = loss
        self.batch_iterator = batch_iterator
        self.regression = regression
        self.max_epochs = max_epochs
        self.eval_size = eval_size
        self.X_tensor_type = X_tensor_type
        self.y_tensor_type = y_tensor_type
        self.use_label_encoder = use_label_encoder
        self.on_epoch_finished = on_epoch_finished
        self.on_training_finished = on_training_finished
        self.more_params = more_params or {}
        self.verbose = verbose

        for key in kwargs.keys():
            assert not hasattr(self, key)
        vars(self).update(kwargs)
        self._kwarg_keys = kwargs.keys()

        self.train_history_ = []

    def fit(self, X, y):
        if not self.regression and self.use_label_encoder:
            self.enc_ = LabelEncoder()
            y = self.enc_.fit_transform(y).astype(np.int32)
            self.classes_ = self.enc_.classes_

        out = getattr(self, '_output_layer', None)
        if out is None:
            out = self._output_layer = self.initialize_layers()
        if self.verbose:
            self._print_layer_info(self.get_all_layers())

        iter_funcs = self._create_iter_funcs(
            out, self.loss, self.update,
            self.X_tensor_type,
            self.y_tensor_type,
            )
        self.train_iter_, self.eval_iter_, self.predict_iter_ = iter_funcs

        try:
            self.train_loop(X, y)
        except KeyboardInterrupt:
            pdb.set_trace()
        return self

    def train_loop(self, X, y):
        X_train, X_valid, y_train, y_valid = self.train_test_split(
            X, y, self.eval_size)

        on_epoch_finished = self.on_epoch_finished
        if not isinstance(on_epoch_finished, (list, tuple)):
            on_epoch_finished = [on_epoch_finished]

        on_training_finished = self.on_training_finished
        if not isinstance(on_training_finished, (list, tuple)):
            on_training_finished = [on_training_finished]

        epoch = 0
        info = None
        best_valid_loss = np.inf
        best_train_loss = np.inf

        if self.verbose:
            print("""
 Epoch  |  Train loss  |  Valid loss  |  Train / Val  |  Valid acc  |  Dur
--------|--------------|--------------|---------------|-------------|-------\
""")

        while epoch < self.max_epochs:
            epoch += 1

            train_losses = []
            valid_losses = []
            valid_accuracies = []

            t0 = time()

            for Xb, yb in self.batch_iterator(X_train, y_train):
                batch_train_loss = self.train_iter_(Xb, yb)
                train_losses.append(batch_train_loss)

            for Xb, yb in self.batch_iterator(X_valid, y_valid, test=True):
                batch_valid_loss, accuracy = self.eval_iter_(Xb, yb)
                valid_losses.append(batch_valid_loss)
                valid_accuracies.append(accuracy)

            avg_train_loss = np.mean(train_losses)
            avg_valid_loss = np.mean(valid_losses)
            avg_valid_accuracy = np.mean(valid_accuracies)

            if avg_train_loss < best_train_loss:
                best_train_loss = avg_train_loss
            if avg_valid_loss < best_valid_loss:
                best_valid_loss = avg_valid_loss

            if self.verbose:
                best_train = best_train_loss == avg_train_loss
                best_valid = best_valid_loss == avg_valid_loss
                print(" {:>5}  |  {}{:>10.6f}{}  |  {}{:>10.6f}{}  "
                      "|  {:>11.6f}  |  {:>9}  |  {:>3.1f}s".format(
                          epoch,
                          ansi.BLUE if best_train else "",
                          avg_train_loss,
                          ansi.ENDC if best_train else "",
                          ansi.GREEN if best_valid else "",
                          avg_valid_loss,
                          ansi.ENDC if best_valid else "",
                          avg_train_loss / avg_valid_loss,
                          "{:.2f}%".format(avg_valid_accuracy * 100)
                          if not self.regression else "",
                          time() - t0,
                          ))

            info = dict(
                epoch=epoch,
                train_loss=avg_train_loss,
                valid_loss=avg_valid_loss,
                valid_accuracy=avg_valid_accuracy,
                )
            self.train_history_.append(info)
            try:
                for func in on_epoch_finished:
                    func(self, self.train_history_)
            except StopIteration:
                break

        for func in on_training_finished:
            func(self, self.train_history_)

    def predict_proba(self, X):
        probas = []
        for Xb, yb in self.batch_iterator(X, test=True):
            probas.append(self.predict_iter_(Xb))
        return np.vstack(probas)

    def predict(self, X):
        if self.regression:
            return self.predict_proba(X)
        else:
            y_pred = np.argmax(self.predict_proba(X), axis=1)
            if self.use_label_encoder:
                y_pred = self.enc_.inverse_transform(y_pred)
            return y_pred

    def score(self, X, y):
        score = mean_squared_error if self.regression else accuracy_score
        return float(score(self.predict(X), y))

    def train_test_split(self, X, y, eval_size):
        if self.regression:
            kf = KFold(y.shape[0], 1. / eval_size)
        else:
            kf = StratifiedKFold(y, 1. / eval_size)

        train_indices, valid_indices = iter(kf).next()
        X_train, y_train = X[train_indices], y[train_indices]
        X_valid, y_valid = X[valid_indices], y[valid_indices]
        return X_train, X_valid, y_train, y_valid

    def get_all_layers(self):
        return get_all_layers(self._output_layer)[::-1]

    def get_all_params(self):
        return get_all_params(self._output_layer)[::-1]

    def _create_iter_funcs(self, output_layer, loss_func, update, input_type,
                           output_type):
        X = input_type('x')
        y = output_type('y')
        X_batch = input_type('x_batch')
        y_batch = output_type('y_batch')

        loss_train = loss_func(
            output_layer.get_output(X_batch), y_batch)
        loss_eval = loss_func(
            output_layer.get_output(X_batch, deterministic=True), y_batch)
        predict_proba = output_layer.get_output(X_batch, deterministic=True)
        if not self.regression:
            predict = predict_proba.argmax(axis=1)
            accuracy = T.mean(T.eq(predict, y_batch))
        else:
            accuracy = loss_eval

        all_params = get_all_params(output_layer)
        update_params = self._get_params_for('update')
        updates = update(loss_train, all_params, **update_params)

        train_iter = theano.function(
            inputs=[theano.Param(X_batch), theano.Param(y_batch)],
            outputs=[loss_train],
            updates=updates,
            givens={
                X: X_batch,
                y: y_batch,
                },
            )
        eval_iter = theano.function(
            inputs=[theano.Param(X_batch), theano.Param(y_batch)],
            outputs=[loss_eval, accuracy],
            givens={
                X: X_batch,
                y: y_batch,
                },
            )
        predict_iter = theano.function(
            inputs=[theano.Param(X_batch)],
            outputs=predict_proba,
            givens={
                X: X_batch,
                },
            )

        return train_iter, eval_iter, predict_iter

    def _get_params_for(self, name):
        collected = {}
        prefix = '{}_'.format(name)

        params = vars(self)
        more_params = self.more_params

        for key, value in itertools.chain(params.items(), more_params.items()):
            if key.startswith(prefix):
                collected[key[len(prefix):]] = value

        return collected

    def load_weights_from(self, source):
        self._output_layer = self.initialize_layers()

        if isinstance(source, str):
            source = np.load(source)

        if isinstance(source, NeuralNet):
            source = source.get_all_params()

        source_weights = [
            w.get_value() if hasattr(w, 'get_value') else w for w in source]

        for w1, w2 in zip(source_weights, self.get_all_params()):
            if w1.shape != w2.get_value().shape:
                continue
            w2.set_value(w1)

    def save_weights_to(self, fname):
        weights = [w.get_value() for w in self.get_all_params()]
        with open(fname, 'wb') as f:
            cPickle.dump(weights, f, -1)

    def initialize_layers(self, layers=None):
        if layers is not None:
            self.layers = layers

        input_layer_name, input_layer_factory = self.layers[0]
        input_layer_params = self._get_params_for(input_layer_name)
        layer = input_layer_factory(**input_layer_params)

        for (layer_name, layer_factory) in self.layers[1:]:
            layer_params = self._get_params_for(layer_name)
            layer = layer_factory(layer, **layer_params)

        self._output_layer = layer
        return layer

    def _print_layer_info(self, layers):
        for layer in layers:
            output_shape = layer.get_output_shape()
            print("  {:<18}\t{:<20}\tproduces {:>7} outputs".format(
                layer.__class__.__name__,
                output_shape,
                reduce(operator.mul, output_shape[1:]),
                ))

    def get_params(self, deep=True):
        params = super(NeuralNet, self).get_params(deep=deep)

        # Incidentally, Lasagne layers have a 'get_params' too, which
        # for sklearn's 'clone' means it would treat it in a special
        # way when cloning.  Wrapping the list of layers in a custom
        # list type does the trick here, but of course it's crazy:
        params['layers'] = _list(params['layers'])
        return _dict(params)

    def _get_param_names(self):
        # This allows us to have **kwargs in __init__ (woot!):
        param_names = super(NeuralNet, self)._get_param_names()
        return param_names + self._kwarg_keys
<commit_msg>Added forced_even to BatchIterator ensure equal batch size<commit_after>from __future__ import absolute_import

import cPickle
import itertools
import operator
from time import time
import pdb

from lasagne.layers import get_all_layers
from lasagne.layers import get_all_params
from lasagne.objectives import mse
from lasagne.updates import nesterov_momentum
import numpy as np
from sklearn.base import BaseEstimator
from sklearn.cross_validation import KFold
from sklearn.cross_validation import StratifiedKFold
from sklearn.metrics import accuracy_score
from sklearn.metrics import mean_squared_error
from sklearn.preprocessing import LabelEncoder
import theano
from theano import tensor as T

class _list(list):
    pass

class _dict(dict):
    def __contains__(self, key):
        return True

class ansi:
    BLUE = '\033[94m'
    GREEN = '\033[32m'
    ENDC = '\033[0m'

def negative_log_likelihood(output, prediction):
    return -T.mean(T.log(output)[T.arange(prediction.shape[0]), prediction])

class BatchIterator(object):
    def __init__(self, batch_size, forced_even=False):
        self.batch_size = batch_size
        self.forced_even = forced_even

    def __call__(self, X, y=None, test=False):
        self.X, self.y = X, y
        self.test = test
        return self

    def __iter__(self):
        n_samples = self.X.shape[0]
        bs = self.batch_size
        for i in range((n_samples + bs - 1) / bs):
            sl = slice(i * bs, (i + 1) * bs)
            Xb = self.X[sl]
            if self.forced_even and len(Xb) != bs:
                continue
            if self.y is not None:
                yb = self.y[sl]
            else:
                yb = None
            yield self.transform(Xb, yb)

    def transform(self, Xb, yb):
        return Xb, yb

class NeuralNet(BaseEstimator):
    """A scikit-learn estimator based on Lasagne.
    """
    def __init__(
        self,
        layers,
        update=nesterov_momentum,
        loss=None,
        batch_iterator=BatchIterator(batch_size=128),
        regression=False,
        max_epochs=100,
        eval_size=0.2,
        X_tensor_type=None,
        y_tensor_type=None,
        use_label_encoder=False,
        on_epoch_finished=(),
        on_training_finished=(),
        more_params=None,
        verbose=0,
        **kwargs
        ):
        if loss is None:
            loss = mse if regression else negative_log_likelihood
        if X_tensor_type is None:
            types = {
                2: T.matrix,
                3: T.tensor3,
                4: T.tensor4,
                }
            X_tensor_type = types[len(kwargs['input_shape'])]
        if y_tensor_type is None:
            y_tensor_type = T.fmatrix if regression else T.ivector

        self.layers = layers
        self.update = update
        self.loss = loss
        self.batch_iterator = batch_iterator
        self.regression = regression
        self.max_epochs = max_epochs
        self.eval_size = eval_size
        self.X_tensor_type = X_tensor_type
        self.y_tensor_type = y_tensor_type
        self.use_label_encoder = use_label_encoder
        self.on_epoch_finished = on_epoch_finished
        self.on_training_finished = on_training_finished
        self.more_params = more_params or {}
        self.verbose = verbose

        for key in kwargs.keys():
            assert not hasattr(self, key)
        vars(self).update(kwargs)
        self._kwarg_keys = kwargs.keys()

        self.train_history_ = []

    def fit(self, X, y):
        if not self.regression and self.use_label_encoder:
            self.enc_ = LabelEncoder()
            y = self.enc_.fit_transform(y).astype(np.int32)
            self.classes_ = self.enc_.classes_

        out = getattr(self, '_output_layer', None)
        if out is None:
            out = self._output_layer = self.initialize_layers()
        if self.verbose:
            self._print_layer_info(self.get_all_layers())

        iter_funcs = self._create_iter_funcs(
            out, self.loss, self.update,
            self.X_tensor_type,
            self.y_tensor_type,
            )
        self.train_iter_, self.eval_iter_, self.predict_iter_ = iter_funcs

        try:
            self.train_loop(X, y)
        except KeyboardInterrupt:
            pdb.set_trace()
        return self

    def train_loop(self, X, y):
        X_train, X_valid, y_train, y_valid = self.train_test_split(
            X, y, self.eval_size)

        on_epoch_finished = self.on_epoch_finished
        if not isinstance(on_epoch_finished, (list, tuple)):
            on_epoch_finished = [on_epoch_finished]

        on_training_finished = self.on_training_finished
        if not isinstance(on_training_finished, (list, tuple)):
            on_training_finished = [on_training_finished]

        epoch = 0
        info = None
        best_valid_loss = np.inf
        best_train_loss = np.inf

        if self.verbose:
            print("""
 Epoch  |  Train loss  |  Valid loss  |  Train / Val  |  Valid acc  |  Dur
--------|--------------|--------------|---------------|-------------|-------\
""")

        while epoch < self.max_epochs:
            epoch += 1

            train_losses = []
            valid_losses = []
            valid_accuracies = []

            t0 = time()

            for Xb, yb in self.batch_iterator(X_train, y_train):
                batch_train_loss = self.train_iter_(Xb, yb)
                train_losses.append(batch_train_loss)

            for Xb, yb in self.batch_iterator(X_valid, y_valid, test=True):
                batch_valid_loss, accuracy = self.eval_iter_(Xb, yb)
                valid_losses.append(batch_valid_loss)
                valid_accuracies.append(accuracy)

            avg_train_loss = np.mean(train_losses)
            avg_valid_loss = np.mean(valid_losses)
            avg_valid_accuracy = np.mean(valid_accuracies)

            if avg_train_loss < best_train_loss:
                best_train_loss = avg_train_loss
            if avg_valid_loss < best_valid_loss:
                best_valid_loss = avg_valid_loss

            if self.verbose:
                best_train = best_train_loss == avg_train_loss
                best_valid = best_valid_loss == avg_valid_loss
                print(" {:>5}  |  {}{:>10.6f}{}  |  {}{:>10.6f}{}  "
                      "|  {:>11.6f}  |  {:>9}  |  {:>3.1f}s".format(
                          epoch,
                          ansi.BLUE if best_train else "",
                          avg_train_loss,
                          ansi.ENDC if best_train else "",
                          ansi.GREEN if best_valid else "",
                          avg_valid_loss,
                          ansi.ENDC if best_valid else "",
                          avg_train_loss / avg_valid_loss,
                          "{:.2f}%".format(avg_valid_accuracy * 100)
                          if not self.regression else "",
                          time() - t0,
                          ))

            info = dict(
                epoch=epoch,
                train_loss=avg_train_loss,
                valid_loss=avg_valid_loss,
                valid_accuracy=avg_valid_accuracy,
                )
            self.train_history_.append(info)
            try:
                for func in on_epoch_finished:
                    func(self, self.train_history_)
            except StopIteration:
                break

        for func in on_training_finished:
            func(self, self.train_history_)

    def predict_proba(self, X):
        probas = []
        for Xb, yb in self.batch_iterator(X, test=True):
            probas.append(self.predict_iter_(Xb))
        return np.vstack(probas)

    def predict(self, X):
        if self.regression:
            return self.predict_proba(X)
        else:
            y_pred = np.argmax(self.predict_proba(X), axis=1)
            if self.use_label_encoder:
                y_pred = self.enc_.inverse_transform(y_pred)
            return y_pred

    def score(self, X, y):
        score = mean_squared_error if self.regression else accuracy_score
        return float(score(self.predict(X), y))

    def train_test_split(self, X, y, eval_size):
        if self.regression:
            kf = KFold(y.shape[0], 1. / eval_size)
        else:
            kf = StratifiedKFold(y, 1. / eval_size)

        train_indices, valid_indices = iter(kf).next()
        X_train, y_train = X[train_indices], y[train_indices]
        X_valid, y_valid = X[valid_indices], y[valid_indices]
        return X_train, X_valid, y_train, y_valid

    def get_all_layers(self):
        return get_all_layers(self._output_layer)[::-1]

    def get_all_params(self):
        return get_all_params(self._output_layer)[::-1]

    def _create_iter_funcs(self, output_layer, loss_func, update, input_type,
                           output_type):
        X = input_type('x')
        y = output_type('y')
        X_batch = input_type('x_batch')
        y_batch = output_type('y_batch')

        loss_train = loss_func(
            output_layer.get_output(X_batch), y_batch)
        loss_eval = loss_func(
            output_layer.get_output(X_batch, deterministic=True), y_batch)
        predict_proba = output_layer.get_output(X_batch, deterministic=True)
        if not self.regression:
            predict = predict_proba.argmax(axis=1)
            accuracy = T.mean(T.eq(predict, y_batch))
        else:
            accuracy = loss_eval

        all_params = get_all_params(output_layer)
        update_params = self._get_params_for('update')
        updates = update(loss_train, all_params, **update_params)

        train_iter = theano.function(
            inputs=[theano.Param(X_batch), theano.Param(y_batch)],
            outputs=[loss_train],
            updates=updates,
            givens={
                X: X_batch,
                y: y_batch,
                },
            )
        eval_iter = theano.function(
            inputs=[theano.Param(X_batch), theano.Param(y_batch)],
            outputs=[loss_eval, accuracy],
            givens={
                X: X_batch,
                y: y_batch,
                },
            )
        predict_iter = theano.function(
            inputs=[theano.Param(X_batch)],
            outputs=predict_proba,
            givens={
                X: X_batch,
                },
            )

        return train_iter, eval_iter, predict_iter

    def _get_params_for(self, name):
        collected = {}
        prefix = '{}_'.format(name)

        params = vars(self)
        more_params = self.more_params

        for key, value in itertools.chain(params.items(), more_params.items()):
            if key.startswith(prefix):
                collected[key[len(prefix):]] = value

        return collected

    def load_weights_from(self, source):
        self._output_layer = self.initialize_layers()

        if isinstance(source, str):
            source = np.load(source)

        if isinstance(source, NeuralNet):
            source = source.get_all_params()

        source_weights = [
            w.get_value() if hasattr(w, 'get_value') else w for w in source]

        for w1, w2 in zip(source_weights, self.get_all_params()):
            if w1.shape != w2.get_value().shape:
                continue
            w2.set_value(w1)

    def save_weights_to(self, fname):
        weights = [w.get_value() for w in self.get_all_params()]
        with open(fname, 'wb') as f:
            cPickle.dump(weights, f, -1)

    def initialize_layers(self, layers=None):
        if layers is not None:
            self.layers = layers

        input_layer_name, input_layer_factory = self.layers[0]
        input_layer_params = self._get_params_for(input_layer_name)
        layer = input_layer_factory(**input_layer_params)

        for (layer_name, layer_factory) in self.layers[1:]:
            layer_params = self._get_params_for(layer_name)
            layer = layer_factory(layer, **layer_params)

        self._output_layer = layer
        return layer

    def _print_layer_info(self, layers):
        for layer in layers:
            output_shape = layer.get_output_shape()
            print("  {:<18}\t{:<20}\tproduces {:>7} outputs".format(
                layer.__class__.__name__,
                output_shape,
                reduce(operator.mul, output_shape[1:]),
                ))

    def get_params(self, deep=True):
        params = super(NeuralNet, self).get_params(deep=deep)

        # Incidentally, Lasagne layers have a 'get_params' too, which
        # for sklearn's 'clone' means it would treat it in a special
        # way when cloning.  Wrapping the list of layers in a custom
        # list type does the trick here, but of course it's crazy:
        params['layers'] = _list(params['layers'])
        return _dict(params)

    def _get_param_names(self):
        # This allows us to have **kwargs in __init__ (woot!):
        param_names = super(NeuralNet, self)._get_param_names()
        return param_names + self._kwarg_keys
<|endoftext|>


<commit_before>from types import IntType, StringType, UnicodeType
import time
import datetime
import calendar
import os
import sys

class Config(dict):
    '''
    based on http://stackoverflow.com/questions/4984647/accessing-dict-keys-like-an-attribute-in-python
    create a config object based on an imported module.
    with the module, you have nice config.key attrib access, but you can't config.copy() it
    so we convert it into a dict... but then we loose the nice attrib access again.
    so this class gives an object that supports both config.key and config.copy() basically
    '''
    def __init__(self, module_or_dict):
        try:
            # it's a module
            d = module_or_dict.__dict__
        except:
            # it's a dict:
            d = module_or_dict
        for k, v in d.items():
            if not k.startswith('__'):
                self[k] = v

    def __getattr__(self, attr):
        return self[attr]

    def __setattr__(self, attr, value):
        self[attr] = value

    def copy(self):
        return Config(self)

class Event():
    '''
    timestamp must be a unix timestamp (int)
    desc is a string in whatever markup you want (html usually)
    tags is a list of strings (usally simple words)
    event_id is optional, it's the elasticsearch _id field
    '''

    def __init__(self, timestamp=None, desc=None, tags=[], event_id=None, extra_attributes={}):
        assert type(timestamp) is IntType, "timestamp must be an integer: %r" % timestamp
        assert type(desc) in (StringType, UnicodeType), "desc must be a non-empty string: %r" % desc
        assert desc, "desc must be a non-empty string: %r" % desc
        self.timestamp = timestamp
        self.desc = desc
        self.tags = tags  # just a list of strings
        self.event_id = event_id
        self.extra_attributes = extra_attributes

    def __str__(self):
        pretty_desc = self.desc
        if "\n" in self.desc:
            pretty_desc = "%s..." % self.desc[:self.desc.find('\n')]
        return "Event object. event_id=%s, ts=%i, tags=%s, desc=%s" % (self.event_id, self.timestamp, ','.join(self.tags), pretty_desc)

    def __getattr__(self, nm):
        if nm == 'outage':
            for tag in self.tags:
                if tag.startswith('outage='):
                    return tag.replace('outage=', '')
            return None
        if nm == 'impact':
            for tag in self.tags:
                if tag.startswith('impact='):
                    return tag.replace('impact=', '')
            return None
        raise AttributeError("no attribute %s" % nm)

class Reportpoint():

    def __init__(self, event, outages, muptime, ttf, tttf, ttd, tttd, ttr, tttr):
        self.event = event
        self.outages = outages  # number of outages occured until now (including this one, if appropriate)
        self.muptime = muptime
        self.ttf = ttf
        self.tttf = tttf
        self.ttd = ttd
        self.tttd = tttd
        self.ttr = ttr
        self.tttr = tttr

    def __getattr__(self, nm):
        divisor = self.outages
        if divisor == 0:
            divisor = 1
        if nm == 'mttf':
            return self.tttf / divisor
        if nm == 'mttd':
            return self.tttd / divisor
        if nm == 'mttr':
            return self.tttr / divisor
        raise AttributeError("no attribute %s" % nm)

class Backend():

    def __init__(self):
        sys.path.append("%s/%s" % (os.getcwd(), 'python-dateutil'))
        sys.path.append("%s/%s" % (os.getcwd(), 'requests'))
        sys.path.append("%s/%s" % (os.getcwd(), 'rawes'))
        import rawes
        import requests
        from rawes.elastic_exception import ElasticException
        # pyflakes doesn't like globals()['ElasticException'] = ElasticException  so:
        self.ElasticException = ElasticException
        self.es = rawes.Elastic('localhost:9200', except_on_error=True)
        # make sure the index exists
        try:
            # to explain the custom mapping:
            # * _source enabled is maybe not really needed, but it's easiest at
            # least. we just need to be able to reconstruct the original document.
            # * tags are not analyzed so that when we want to get a list of all
            # tags (a facet search) it returns the original tags, not the
            # tokenized terms.
            self.es.post('anthracite', data={
                "mappings": {
                    "event": {
                        "_source": {
                            "enabled": True
                        },
                        "properties": {
                            "tags": {
                                "type": "string",
                                "index": "not_analyzed"
                            }
                        }
                    }
                }
            })
            print "created new ElasticSearch Index"
        except ElasticException as e:
            if e.result['error'] == 'IndexAlreadyExistsException[[anthracite] Already exists]':
                pass
            else:
                raise
        except requests.exceptions.ConnectionError as e:
            sys.stderr.write("Could not connect to ElasticSearch: %s" % e)
            sys.exit(2)

    def object_to_dict(self, event):
        iso = self.unix_timestamp_to_iso8601(event.timestamp)
        data = {
            'date': iso,
            'tags': event.tags,
            'desc': event.desc
        }
        data.update(event.extra_attributes)
        return data

    def unix_timestamp_to_iso8601(self, unix_timestamp):
        return datetime.datetime.utcfromtimestamp(unix_timestamp).isoformat()

    def iso8601_to_unix_timestamp(self, iso8601):
        '''
            elasticsearch returns something like 2013-03-20T20:41:16

        '''
        unix = calendar.timegm(datetime.datetime.strptime(iso8601, "%Y-%m-%dT%H:%M:%S").timetuple())
        return unix

    def hit_to_object(self, hit):
        event_id = hit['_id']
        hit = hit['_source']
        unix = self.iso8601_to_unix_timestamp(hit['date'])
        extra_attributes= {}
        for (k, v) in hit.items():
            if k not in ('desc', 'tags', 'date'):
                extra_attributes[k] = v
        return Event(timestamp=unix, desc=hit['desc'], tags=hit['tags'], event_id=event_id, extra_attributes=extra_attributes)

    def add_event(self, event):
        ret = self.es.post('anthracite/event', data=self.object_to_dict(event))
        return ret['_id']

    def delete_event(self, event_id):
        try:
            self.es.delete('anthracite/event/%s' % event_id)
        except self.ElasticException as e:
            if 'found' in e.result and not e.result['found']:
                raise Exception("Document %s can't be found" % event_id)
            else:
                raise

    def edit_event(self, event):
        self.es.post('anthracite/event/%s/_update' % event.event_id, data={'doc': self.object_to_dict(event)})

    def es_get_events(self):
        return self.es.get('anthracite/event/_search?size=1000', data={
            "query": {
                "query_string": {
                    "query": "*"
                }
            },
            "sort": [
                {
                    "date": {
                        "order": "desc",
                        "ignore_unmapped": True  # avoid 'No mapping found for [date] in order to sort on' when we don't have data yet
                    }
                }
            ]
        })

    def get_events_raw(self):
        '''
        return format that's optimized for elasticsearch
        '''
        hits = self.es_get_events()
        events = hits['hits']['hits']
        for (i, event) in enumerate(events):
            event_id = event['_id']
            events[i] = event['_source']
            events[i]['id'] = event_id
            events[i]['date'] = self.iso8601_to_unix_timestamp(events[i]['date'])
        return events

    def get_events_objects(self):
        # retuns a list of event objects
        hits = self.es_get_events()
        return [self.hit_to_object(event_hit) for event_hit in hits['hits']['hits']]

    def get_event(self, event_id):
        # http://localhost:9200/dieterfoobarbaz/event/PZ1su5w5Stmln_c2Kc4B2g
        event_hit = self.es.get('anthracite/event/%s' % event_id)
        event_obj = self.hit_to_object(event_hit)
        return event_obj

    def get_tags(self):
        # get all different tags
        # curl -X POST "http://localhost:9200/anthracite/_search?pretty=true&size=0" -d '{  "query" : {"query_string" : {"query" : "*"}}, "facets":{"tags" : { "terms" : {"field" : "tags"} }}}"'
        tags = self.es.post('anthracite/_search?size=0', data={
            'query': {
                'query_string': {
                    'query': '*'
                }
            },
            'facets': {
                'tags': {
                    'terms': {
                        'field': 'tags'
                    }
                }
            }
        })
        tags = tags['facets']['tags']['terms']
        tags = [t['term'] for t in tags]
        return tags

    def get_events_range(self):
        low = self.es.post('anthracite/_search?size=1', data={
            "query": {
                "field": {
                    "date": {
                        "query": "*"
                    }
                }
            },
            "sort": [
                {
                    "date": {
                        "order": "asc",
                        "ignore_unmapped": True  # avoid 'No mapping found for [date] in order to sort on' when we don't have data yet
                    }
                }
            ]
        })
        # if there's not a single record in the database:
        if not len(low['hits']['hits']):
            return (0, time.time())
        high = self.es.post('anthracite/_search?size=1', data={
            "query": {
                "field": {
                    "date": {
                        "query": "*"
                    }
                }
            },
            "sort": [
                {
                    "date": {
                        "order": "desc",
                        "ignore_unmapped": True  # avoid 'No mapping found for [date] in order to sort on' when we don't have data yet
                    }
                }
            ]
        })
        low = self.iso8601_to_unix_timestamp(low['hits']['hits'][0]['_source']['date'])
        high = self.iso8601_to_unix_timestamp(high['hits']['hits'][0]['_source']['date'])
        return (low, high)

    def get_events_count(self):
        count = 0
        events = self.es.get('anthracite/event/_search')
        count = events['hits']['total']
        return count

    def get_outage_events(self):
        # TODO sanity checking (order of detected, resolved tags, etc)
        hits = self.es.get('anthracite/event/_search', data={
            'query': {
                'query_string': {
                    'query': 'tag like outage=_%'
                }
            },
            "sort": [
                {
                    "date": {
                        "order": "asc",
                        "ignore_unmapped": True  # avoid 'No mapping found for [date] in order to sort on' when we don't have data yet
                    }
                }
            ]
        })
        events = []
        for event_hit in hits['hits']['hits']:
            event_obj = self.hit_to_object(event_hit)
            events.append(event_obj)
        return events

class PluginError(Exception):

    def __init__(self, plugin, msg, underlying_error):
        self.plugin = plugin
        self.msg = msg
        self.underlying_error = underlying_error

    def __str__(self):
        return "%s -> %s (%s)" % (self.plugin, self.msg, self.underlying_error)

def load_plugins(plugins_to_load, config):
    '''
    loads all the plugins sub-modules
    returns encountered errors, doesn't raise them because
    whoever calls this function defines how any errors are
    handled. meanwhile, loading must continue
    '''
    import plugins
    errors = []
    add_urls = {}
    remove_urls = []
    handlers = {}  # not namespaced by plugin to keep it easy. include plugin name in name.
    plugins_dir = os.path.dirname(plugins.__file__)
    wd = os.getcwd()
    os.chdir(plugins_dir)
    for module in plugins_to_load:
        try:
            print "importing plugin '%s'" % module
            imp = __import__('plugins.' + module, {}, {}, ['*'])
            try:
                add_urls[module] = imp.add_urls
            except Exception:
                pass
            try:
                remove_urls.extend(imp.remove_urls)
            except Exception:
                pass
            try:
                handlers.update(imp.handlers)
            except Exception:
                pass
        except Exception, e:
            errors.append(PluginError(module, "Failed to add plugin '%s'" % module, e))
            continue

    os.chdir(wd)
    state = {
        'add_urls': add_urls,
        'remove_urls': remove_urls,
        'handlers': handlers
    }
    #  make some vars accessible for all imported plugins
    __builtins__['state'] = state
    __builtins__['config'] = config
    return (state, errors)
<commit_msg>support specifying custom elasticsearch query<commit_after>from types import IntType, StringType, UnicodeType
import time
import datetime
import calendar
import os
import sys

class Config(dict):
    '''
    based on http://stackoverflow.com/questions/4984647/accessing-dict-keys-like-an-attribute-in-python
    create a config object based on an imported module.
    with the module, you have nice config.key attrib access, but you can't config.copy() it
    so we convert it into a dict... but then we loose the nice attrib access again.
    so this class gives an object that supports both config.key and config.copy() basically
    '''
    def __init__(self, module_or_dict):
        try:
            # it's a module
            d = module_or_dict.__dict__
        except:
            # it's a dict:
            d = module_or_dict
        for k, v in d.items():
            if not k.startswith('__'):
                self[k] = v

    def __getattr__(self, attr):
        return self[attr]

    def __setattr__(self, attr, value):
        self[attr] = value

    def copy(self):
        return Config(self)

class Event():
    '''
    timestamp must be a unix timestamp (int)
    desc is a string in whatever markup you want (html usually)
    tags is a list of strings (usally simple words)
    event_id is optional, it's the elasticsearch _id field
    '''

    def __init__(self, timestamp=None, desc=None, tags=[], event_id=None, extra_attributes={}):
        assert type(timestamp) is IntType, "timestamp must be an integer: %r" % timestamp
        assert type(desc) in (StringType, UnicodeType), "desc must be a non-empty string: %r" % desc
        assert desc, "desc must be a non-empty string: %r" % desc
        self.timestamp = timestamp
        self.desc = desc
        self.tags = tags  # just a list of strings
        self.event_id = event_id
        self.extra_attributes = extra_attributes

    def __str__(self):
        pretty_desc = self.desc
        if "\n" in self.desc:
            pretty_desc = "%s..." % self.desc[:self.desc.find('\n')]
        return "Event object. event_id=%s, ts=%i, tags=%s, desc=%s" % (self.event_id, self.timestamp, ','.join(self.tags), pretty_desc)

    def __getattr__(self, nm):
        if nm == 'outage':
            for tag in self.tags:
                if tag.startswith('outage='):
                    return tag.replace('outage=', '')
            return None
        if nm == 'impact':
            for tag in self.tags:
                if tag.startswith('impact='):
                    return tag.replace('impact=', '')
            return None
        raise AttributeError("no attribute %s" % nm)

class Reportpoint():

    def __init__(self, event, outages, muptime, ttf, tttf, ttd, tttd, ttr, tttr):
        self.event = event
        self.outages = outages  # number of outages occured until now (including this one, if appropriate)
        self.muptime = muptime
        self.ttf = ttf
        self.tttf = tttf
        self.ttd = ttd
        self.tttd = tttd
        self.ttr = ttr
        self.tttr = tttr

    def __getattr__(self, nm):
        divisor = self.outages
        if divisor == 0:
            divisor = 1
        if nm == 'mttf':
            return self.tttf / divisor
        if nm == 'mttd':
            return self.tttd / divisor
        if nm == 'mttr':
            return self.tttr / divisor
        raise AttributeError("no attribute %s" % nm)

class Backend():

    def __init__(self):
        sys.path.append("%s/%s" % (os.getcwd(), 'python-dateutil'))
        sys.path.append("%s/%s" % (os.getcwd(), 'requests'))
        sys.path.append("%s/%s" % (os.getcwd(), 'rawes'))
        import rawes
        import requests
        from rawes.elastic_exception import ElasticException
        # pyflakes doesn't like globals()['ElasticException'] = ElasticException  so:
        self.ElasticException = ElasticException
        self.es = rawes.Elastic('localhost:9200', except_on_error=True)
        # make sure the index exists
        try:
            # to explain the custom mapping:
            # * _source enabled is maybe not really needed, but it's easiest at
            # least. we just need to be able to reconstruct the original document.
            # * tags are not analyzed so that when we want to get a list of all
            # tags (a facet search) it returns the original tags, not the
            # tokenized terms.
            self.es.post('anthracite', data={
                "mappings": {
                    "event": {
                        "_source": {
                            "enabled": True
                        },
                        "properties": {
                            "tags": {
                                "type": "string",
                                "index": "not_analyzed"
                            }
                        }
                    }
                }
            })
            print "created new ElasticSearch Index"
        except ElasticException as e:
            if e.result['error'] == 'IndexAlreadyExistsException[[anthracite] Already exists]':
                pass
            else:
                raise
        except requests.exceptions.ConnectionError as e:
            sys.stderr.write("Could not connect to ElasticSearch: %s" % e)
            sys.exit(2)

    def object_to_dict(self, event):
        iso = self.unix_timestamp_to_iso8601(event.timestamp)
        data = {
            'date': iso,
            'tags': event.tags,
            'desc': event.desc
        }
        data.update(event.extra_attributes)
        return data

    def unix_timestamp_to_iso8601(self, unix_timestamp):
        return datetime.datetime.utcfromtimestamp(unix_timestamp).isoformat()

    def iso8601_to_unix_timestamp(self, iso8601):
        '''
            elasticsearch returns something like 2013-03-20T20:41:16

        '''
        unix = calendar.timegm(datetime.datetime.strptime(iso8601, "%Y-%m-%dT%H:%M:%S").timetuple())
        return unix

    def hit_to_object(self, hit):
        event_id = hit['_id']
        hit = hit['_source']
        unix = self.iso8601_to_unix_timestamp(hit['date'])
        extra_attributes= {}
        for (k, v) in hit.items():
            if k not in ('desc', 'tags', 'date'):
                extra_attributes[k] = v
        return Event(timestamp=unix, desc=hit['desc'], tags=hit['tags'], event_id=event_id, extra_attributes=extra_attributes)

    def add_event(self, event):
        ret = self.es.post('anthracite/event', data=self.object_to_dict(event))
        return ret['_id']

    def delete_event(self, event_id):
        try:
            self.es.delete('anthracite/event/%s' % event_id)
        except self.ElasticException as e:
            if 'found' in e.result and not e.result['found']:
                raise Exception("Document %s can't be found" % event_id)
            else:
                raise

    def edit_event(self, event):
        self.es.post('anthracite/event/%s/_update' % event.event_id, data={'doc': self.object_to_dict(event)})

    def es_get_events(self, query = None):
        if query is None:
            query = {
                "query_string": {
                    "query": "*"
                }
            }
        return self.es.get('anthracite/event/_search?size=1000', data={
            "query": query,
            "sort": [
                {
                    "date": {
                        "order": "desc",
                        "ignore_unmapped": True  # avoid 'No mapping found for [date] in order to sort on' when we don't have data yet
                    }
                }
            ]
        })

    def get_events_raw(self, query=None):
        '''
        return format that's optimized for elasticsearch
        '''
        hits = self.es_get_events(query)
        events = hits['hits']['hits']
        for (i, event) in enumerate(events):
            event_id = event['_id']
            events[i] = event['_source']
            events[i]['id'] = event_id
            events[i]['date'] = self.iso8601_to_unix_timestamp(events[i]['date'])
        return events

    def get_events_objects(self):
        # retuns a list of event objects
        hits = self.es_get_events()
        return [self.hit_to_object(event_hit) for event_hit in hits['hits']['hits']]

    def get_event(self, event_id):
        # http://localhost:9200/dieterfoobarbaz/event/PZ1su5w5Stmln_c2Kc4B2g
        event_hit = self.es.get('anthracite/event/%s' % event_id)
        event_obj = self.hit_to_object(event_hit)
        return event_obj

    def get_tags(self):
        # get all different tags
        # curl -X POST "http://localhost:9200/anthracite/_search?pretty=true&size=0" -d '{  "query" : {"query_string" : {"query" : "*"}}, "facets":{"tags" : { "terms" : {"field" : "tags"} }}}"'
        tags = self.es.post('anthracite/_search?size=0', data={
            'query': {
                'query_string': {
                    'query': '*'
                }
            },
            'facets': {
                'tags': {
                    'terms': {
                        'field': 'tags'
                    }
                }
            }
        })
        tags = tags['facets']['tags']['terms']
        tags = [t['term'] for t in tags]
        return tags

    def get_events_range(self):
        low = self.es.post('anthracite/_search?size=1', data={
            "query": {
                "field": {
                    "date": {
                        "query": "*"
                    }
                }
            },
            "sort": [
                {
                    "date": {
                        "order": "asc",
                        "ignore_unmapped": True  # avoid 'No mapping found for [date] in order to sort on' when we don't have data yet
                    }
                }
            ]
        })
        # if there's not a single record in the database:
        if not len(low['hits']['hits']):
            return (0, time.time())
        high = self.es.post('anthracite/_search?size=1', data={
            "query": {
                "field": {
                    "date": {
                        "query": "*"
                    }
                }
            },
            "sort": [
                {
                    "date": {
                        "order": "desc",
                        "ignore_unmapped": True  # avoid 'No mapping found for [date] in order to sort on' when we don't have data yet
                    }
                }
            ]
        })
        low = self.iso8601_to_unix_timestamp(low['hits']['hits'][0]['_source']['date'])
        high = self.iso8601_to_unix_timestamp(high['hits']['hits'][0]['_source']['date'])
        return (low, high)

    def get_events_count(self):
        count = 0
        events = self.es.get('anthracite/event/_search')
        count = events['hits']['total']
        return count

    def get_outage_events(self):
        # TODO sanity checking (order of detected, resolved tags, etc)
        hits = self.es.get('anthracite/event/_search', data={
            'query': {
                'query_string': {
                    'query': 'tag like outage=_%'
                }
            },
            "sort": [
                {
                    "date": {
                        "order": "asc",
                        "ignore_unmapped": True  # avoid 'No mapping found for [date] in order to sort on' when we don't have data yet
                    }
                }
            ]
        })
        events = []
        for event_hit in hits['hits']['hits']:
            event_obj = self.hit_to_object(event_hit)
            events.append(event_obj)
        return events

class PluginError(Exception):

    def __init__(self, plugin, msg, underlying_error):
        self.plugin = plugin
        self.msg = msg
        self.underlying_error = underlying_error

    def __str__(self):
        return "%s -> %s (%s)" % (self.plugin, self.msg, self.underlying_error)

def load_plugins(plugins_to_load, config):
    '''
    loads all the plugins sub-modules
    returns encountered errors, doesn't raise them because
    whoever calls this function defines how any errors are
    handled. meanwhile, loading must continue
    '''
    import plugins
    errors = []
    add_urls = {}
    remove_urls = []
    handlers = {}  # not namespaced by plugin to keep it easy. include plugin name in name.
    plugins_dir = os.path.dirname(plugins.__file__)
    wd = os.getcwd()
    os.chdir(plugins_dir)
    for module in plugins_to_load:
        try:
            print "importing plugin '%s'" % module
            imp = __import__('plugins.' + module, {}, {}, ['*'])
            try:
                add_urls[module] = imp.add_urls
            except Exception:
                pass
            try:
                remove_urls.extend(imp.remove_urls)
            except Exception:
                pass
            try:
                handlers.update(imp.handlers)
            except Exception:
                pass
        except Exception, e:
            errors.append(PluginError(module, "Failed to add plugin '%s'" % module, e))
            continue

    os.chdir(wd)
    state = {
        'add_urls': add_urls,
        'remove_urls': remove_urls,
        'handlers': handlers
    }
    #  make some vars accessible for all imported plugins
    __builtins__['state'] = state
    __builtins__['config'] = config
    return (state, errors)
<|endoftext|>