<commit_before>"""
Tools for handling items on the canvas.

Although Gaphas has quite a few useful tools, some tools need to be extended:
 - PlacementTool: should perform undo
 - HandleTool: should support adapter based connection protocol
 - TextEditTool: should support adapter based edit protocol
"""

import logging
from zope import component

from gaphas.aspect import Connector, InMotion
from gaphas.guide import GuidedItemInMotion
from gaphas.tool import (
    Tool,
    PlacementTool as _PlacementTool,
    ToolChain,
    HoverTool,
    ItemTool,
    RubberbandTool,
    ConnectHandleTool,
)
from gi.repository import Gdk
from gi.repository import Gtk

from gaphor.core import Transaction, transactional
from gaphor.diagram.diagramline import DiagramLine
from gaphor.diagram.elementitem import ElementItem
from gaphor.diagram.interfaces import IEditor, IConnect, IGroup

# cursor to indicate grouping
IN_CURSOR = Gdk.Cursor.new(Gdk.CursorType.DIAMOND_CROSS)

# cursor to indicate ungrouping
OUT_CURSOR = Gdk.Cursor.new(Gdk.CursorType.SIZING)

log = logging.getLogger(__name__)

@Connector.when_type(DiagramLine)
class DiagramItemConnector(Connector.default):
    """
    Handle Tool (acts on item handles) that uses the IConnect protocol
    to connect items to one-another.

    It also adds handles to lines when a line is grabbed on the middle of
    a line segment (points are drawn by the LineSegmentPainter).
    """

    def allow(self, sink):
        adapter = component.queryMultiAdapter((sink.item, self.item), IConnect)
        return adapter and adapter.allow(self.handle, sink.port)

    @transactional
    def connect(self, sink):
        """
        Create connection at handle level and at model level.
        """
        handle = self.handle
        item = self.item
        cinfo = item.canvas.get_connection(handle)

        try:
            callback = DisconnectHandle(self.item, self.handle)
            if cinfo and cinfo.connected is sink.item:
                # reconnect only constraint - leave model intact
                log.debug("performing reconnect constraint")
                constraint = sink.port.constraint(item.canvas, item, handle, sink.item)
                item.canvas.reconnect_item(item, handle, constraint=constraint)
            elif cinfo:
                # first disconnect but disable disconnection handle as
                # reconnection is going to happen
                adapter = component.queryMultiAdapter((sink.item, item), IConnect)
                try:
                    connect = adapter.reconnect
                except AttributeError:
                    connect = adapter.connect
                else:
                    cinfo.callback.disable = True
                self.disconnect()

                # new connection
                self.connect_handle(sink, callback=callback)

                # adapter requires both ends to be connected.
                connect(handle, sink.port)
            else:
                # new connection
                adapter = component.queryMultiAdapter((sink.item, item), IConnect)
                self.connect_handle(sink, callback=callback)
                adapter.connect(handle, sink.port)
        except Exception as e:
            log.error("Error during connect", exc_info=True)

    @transactional
    def disconnect(self):
        super(DiagramItemConnector, self).disconnect()

class DisconnectHandle(object):
    """
    Callback for items disconnection using the adapters.

    This is an object so disconnection data can be serialized/deserialized
    using pickle.

    :Variables:
     item
        Connecting item.
     handle
        Handle of connecting item.
     disable
        If set, then disconnection is disabled.
    """

    def __init__(self, item, handle):
        self.item = item
        self.handle = handle
        self.disable = False

    def __call__(self):
        handle = self.handle
        item = self.item
        canvas = self.item.canvas
        cinfo = canvas.get_connection(handle)

        if self.disable:
            log.debug("Not disconnecting %s.%s (disabled)" % (item, handle))
        else:
            log.debug("Disconnecting %s.%s" % (item, handle))
            if cinfo:
                adapter = component.queryMultiAdapter((cinfo.connected, item), IConnect)
                adapter.disconnect(handle)

class TextEditTool(Tool):
    """
    Text edit tool. Allows for elements that can adapt to the
    IEditable interface to be edited.
    """

    def create_edit_window(self, x, y, text, editor):
        """
        Create a popup window with some editable text.
        """
        view = self.view
        window = Gtk.Window.new(Gtk.WindowType.TOPLEVEL)
        window.set_property("decorated", False)
        window.set_property("skip-taskbar-hint", True)
        window.set_modal(True)
        window.set_transient_for(view.get_toplevel())
        buffer = Gtk.TextBuffer()
        if text:
            buffer.set_text(text)
            startiter, enditer = buffer.get_bounds()
            buffer.move_mark_by_name("selection_bound", startiter)
            buffer.move_mark_by_name("insert", enditer)
        text_view = Gtk.TextView()
        text_view.set_buffer(buffer)
        text_view.set_left_margin(2)
        text_view.set_right_margin(2)

        frame = Gtk.Frame()
        frame.set_shadow_type(Gtk.ShadowType.IN)
        frame.add(text_view)

        window.add(frame)
        r = Gdk.Rectangle()
        r.x = 0
        r.y = 0
        r.width = 70
        r.height = 50
        window.size_allocate(r)
        window.move(int(x), int(y))

        def on_button_press(widget, event):
            if event.window == view.get_window():
                self.submit_text(widget, buffer, editor)

        def on_key_press_event(widget, event):
            if event.keyval == Gdk.KEY_Return and not event.get_state() & (
                Gdk.ModifierType.CONTROL_MASK | Gdk.ModifierType.SHIFT_MASK
            ):
                self.submit_text(widget, buffer, editor)
            elif event.keyval == Gdk.KEY_Escape:
                widget.get_toplevel().destroy()

        def on_focus_out_event(widget, event):
            print("Focus out event emitted")
            self.submit_text(widget, buffer, editor)

        window.add_events(Gdk.EventMask.FOCUS_CHANGE_MASK)
        window.connect("focus-out-event", on_focus_out_event)
        window.connect("button-press-event", on_button_press)
        text_view.connect("key-press-event", on_key_press_event)
        window.show_all()

    @transactional
    def submit_text(self, widget, buffer, editor):
        """
        Submit the final text to the edited item.
        """
        text = buffer.get_text(
            buffer.get_start_iter(), buffer.get_end_iter(), include_hidden_chars=True
        )
        editor.update_text(text)
        widget.get_toplevel().destroy()

    def on_double_click(self, event):
        view = self.view
        item = view.hovered_item
        if item:
            try:
                editor = IEditor(item)
            except TypeError:
                # Could not adapt to IEditor
                return False
            log.debug("Found editor %r" % editor)
            x, y = view.get_matrix_v2i(item).transform_point(event.x, event.y)
            if editor.is_editable(x, y):
                text = editor.get_text()
                root_coords = event.get_root_coords()
                self.create_edit_window(
                    root_coords.x_root, root_coords.y_root, text, editor
                )
                return True

class PlacementTool(_PlacementTool):
    """
    PlacementTool is used to place items on the canvas.
    """

    def __init__(self, view, item_factory, after_handler=None, handle_index=-1):
        """
        item_factory is a callable. It is used to create a CanvasItem
        that is displayed on the diagram.
        """
        _PlacementTool.__init__(
            self,
            view,
            factory=item_factory,
            handle_tool=ConnectHandleTool(),
            handle_index=handle_index,
        )
        self.after_handler = after_handler
        self._tx = None

    @transactional
    def create_item(self, pos):
        return self._create_item(pos)

    def on_button_press(self, event):
        assert not self._tx
        self._tx = Transaction()
        view = self.view
        view.unselect_all()
        if _PlacementTool.on_button_press(self, event):
            try:
                opposite = self.new_item.opposite(
                    self.new_item.handles()[self._handle_index]
                )
            except (KeyError, AttributeError):
                pass
            else:
                # Connect opposite handle first, using the HandleTool's
                # mechanisms

                # First make sure all matrices are updated:
                view.canvas.update_matrix(self.new_item)
                view.update_matrix(self.new_item)

                vpos = event.x, event.y

                item = self.handle_tool.glue(self.new_item, opposite, vpos)
                if item:
                    self.handle_tool.connect(self.new_item, opposite, vpos)
            return True
        return False

    def on_button_release(self, event):
        try:
            if self.after_handler:
                self.after_handler(self.new_item)
            return _PlacementTool.on_button_release(self, event)
        finally:
            self._tx.commit()
            self._tx = None

class GroupPlacementTool(PlacementTool):
    """
    Try to group items when placing them on diagram.
    """

    def __init__(self, view, item_factory, after_handler=None, handle_index=-1):
        super(GroupPlacementTool, self).__init__(
            view, item_factory, after_handler, handle_index
        )
        self._parent = None

    def on_motion_notify(self, event):
        """
        Change parent item to dropzone state if it can accept diagram item
        object to be created.
        """
        view = self.view

        if view.focused_item:
            view.unselect_item(view.focused_item)
            view.focused_item = None

        try:
            parent = view.get_item_at_point((event.x, event.y))
        except KeyError:
            parent = None

        if parent:
            # create dummy adapter
            adapter = component.queryMultiAdapter(
                (parent, self._factory.item_class()), IGroup
            )
            if adapter and adapter.can_contain():
                view.dropzone_item = parent
                view.window.set_cursor(IN_CURSOR)
                self._parent = parent
            else:
                view.dropzone_item = None
                view.window.set_cursor(None)
                self._parent = None
            parent.request_update(matrix=False)
        else:
            if view.dropzone_item:
                view.dropzone_item.request_update(matrix=False)
            view.dropzone_item = None
            view.window.set_cursor(None)

    def _create_item(self, pos, **kw):
        """
        Create diagram item and place it within parent's boundaries.
        """
        parent = self._parent
        view = self.view
        try:
            adapter = component.queryMultiAdapter(
                (parent, self._factory.item_class()), IGroup
            )
            if parent and adapter and adapter.can_contain():
                kw["parent"] = parent

            item = super(GroupPlacementTool, self)._create_item(pos, **kw)

            adapter = component.queryMultiAdapter((parent, item), IGroup)
            if parent and item and adapter:
                adapter.group()

                canvas = view.canvas
                parent.request_update(matrix=False)
        finally:
            self._parent = None
            view.dropzone_item = None
            view.window.set_cursor(None)
        return item

@InMotion.when_type(ElementItem)
class DropZoneInMotion(GuidedItemInMotion):
    def move(self, pos):
        """
        Move the item. x and y are in view coordinates.
        """
        super(DropZoneInMotion, self).move(pos)
        item = self.item
        view = self.view
        x, y = pos

        current_parent = view.canvas.get_parent(item)
        over_item = view.get_item_at_point((x, y), selected=False)

        if not over_item:
            view.dropzone_item = None
            view.window.set_cursor(None)
            return

        if current_parent and not over_item:  # are we going to remove from parent?
            adapter = component.queryMultiAdapter((current_parent, item), IGroup)
            if adapter:
                view.window.set_cursor(OUT_CURSOR)
                view.dropzone_item = current_parent
                current_parent.request_update(matrix=False)

        if over_item:
            # are we going to add to parent?
            adapter = component.queryMultiAdapter((over_item, item), IGroup)
            if adapter and adapter.can_contain():
                view.dropzone_item = over_item
                view.window.set_cursor(IN_CURSOR)
                over_item.request_update(matrix=False)

    def stop_move(self):
        """
        Motion stops: drop!
        """
        super(DropZoneInMotion, self).stop_move()
        item = self.item
        view = self.view
        canvas = view.canvas
        old_parent = view.canvas.get_parent(item)
        new_parent = view.dropzone_item
        try:

            if new_parent is old_parent:
                if old_parent is not None:
                    old_parent.request_update(matrix=False)
                return

            if old_parent:
                adapter = component.queryMultiAdapter((old_parent, item), IGroup)
                if adapter:
                    adapter.ungroup()

                canvas.reparent(item, None)
                m = canvas.get_matrix_i2c(old_parent)
                item.matrix *= m
                old_parent.request_update()

            if new_parent:
                adapter = component.queryMultiAdapter((new_parent, item), IGroup)
                if adapter and adapter.can_contain():
                    adapter.group()

                canvas.reparent(item, new_parent)
                m = canvas.get_matrix_c2i(new_parent)
                item.matrix *= m
                new_parent.request_update()
        finally:
            view.dropzone_item = None
            view.window.set_cursor(None)

class TransactionalToolChain(ToolChain):
    """
    In addition to a normal toolchain, this chain begins an undo-transaction
    at button-press and commits the transaction at button-release.
    """

    def __init__(self, view=None):
        super(TransactionalToolChain, self).__init__(view)
        self._tx = None

    def handle(self, event):
        # For double click: button_press, double_click, button_release
        # print 'event', self.EVENT_HANDLERS.get(event.type)
        if self.EVENT_HANDLERS.get(event.type) in ("on_button_press",):
            assert not self._tx
            self._tx = Transaction()

        try:
            super(TransactionalToolChain, self).handle(event)
        finally:
            if self._tx and self.EVENT_HANDLERS.get(event.type) in (
                "on_button_release",
                "on_double_click",
                "on_triple_click",
            ):
                self._tx.commit()
                self._tx = None

def DefaultTool():
    """
    The default tool chain build from HoverTool, ItemTool and HandleTool.
    """
    chain = TransactionalToolChain()
    chain.append(HoverTool())
    chain.append(ConnectHandleTool())
    chain.append(ItemTool())
    chain.append(TextEditTool())
    chain.append(RubberbandTool())
    return chain

# vim:sw=4:et:ai
<commit_msg>Fix Wayland Cursor TypeError: constructor returned NULL<commit_after>"""
Tools for handling items on the canvas.

Although Gaphas has quite a few useful tools, some tools need to be extended:
 - PlacementTool: should perform undo
 - HandleTool: should support adapter based connection protocol
 - TextEditTool: should support adapter based edit protocol
"""

import logging
from zope import component

from gaphas.aspect import Connector, InMotion
from gaphas.guide import GuidedItemInMotion
from gaphas.tool import (
    Tool,
    PlacementTool as _PlacementTool,
    ToolChain,
    HoverTool,
    ItemTool,
    RubberbandTool,
    ConnectHandleTool,
)
from gi.repository import Gdk
from gi.repository import Gtk

from gaphor.core import Transaction, transactional
from gaphor.diagram.diagramline import DiagramLine
from gaphor.diagram.elementitem import ElementItem
from gaphor.diagram.interfaces import IEditor, IConnect, IGroup

# cursor to indicate grouping
IN_CURSOR = Gdk.Cursor.new(Gdk.CursorType.DIAMOND_CROSS)

# cursor to indicate ungrouping
OUT_CURSOR = Gdk.Cursor.new(Gdk.CursorType.CROSSHAIR)

log = logging.getLogger(__name__)

@Connector.when_type(DiagramLine)
class DiagramItemConnector(Connector.default):
    """
    Handle Tool (acts on item handles) that uses the IConnect protocol
    to connect items to one-another.

    It also adds handles to lines when a line is grabbed on the middle of
    a line segment (points are drawn by the LineSegmentPainter).
    """

    def allow(self, sink):
        adapter = component.queryMultiAdapter((sink.item, self.item), IConnect)
        return adapter and adapter.allow(self.handle, sink.port)

    @transactional
    def connect(self, sink):
        """
        Create connection at handle level and at model level.
        """
        handle = self.handle
        item = self.item
        cinfo = item.canvas.get_connection(handle)

        try:
            callback = DisconnectHandle(self.item, self.handle)
            if cinfo and cinfo.connected is sink.item:
                # reconnect only constraint - leave model intact
                log.debug("performing reconnect constraint")
                constraint = sink.port.constraint(item.canvas, item, handle, sink.item)
                item.canvas.reconnect_item(item, handle, constraint=constraint)
            elif cinfo:
                # first disconnect but disable disconnection handle as
                # reconnection is going to happen
                adapter = component.queryMultiAdapter((sink.item, item), IConnect)
                try:
                    connect = adapter.reconnect
                except AttributeError:
                    connect = adapter.connect
                else:
                    cinfo.callback.disable = True
                self.disconnect()

                # new connection
                self.connect_handle(sink, callback=callback)

                # adapter requires both ends to be connected.
                connect(handle, sink.port)
            else:
                # new connection
                adapter = component.queryMultiAdapter((sink.item, item), IConnect)
                self.connect_handle(sink, callback=callback)
                adapter.connect(handle, sink.port)
        except Exception as e:
            log.error("Error during connect", exc_info=True)

    @transactional
    def disconnect(self):
        super(DiagramItemConnector, self).disconnect()

class DisconnectHandle(object):
    """
    Callback for items disconnection using the adapters.

    This is an object so disconnection data can be serialized/deserialized
    using pickle.

    :Variables:
     item
        Connecting item.
     handle
        Handle of connecting item.
     disable
        If set, then disconnection is disabled.
    """

    def __init__(self, item, handle):
        self.item = item
        self.handle = handle
        self.disable = False

    def __call__(self):
        handle = self.handle
        item = self.item
        canvas = self.item.canvas
        cinfo = canvas.get_connection(handle)

        if self.disable:
            log.debug("Not disconnecting %s.%s (disabled)" % (item, handle))
        else:
            log.debug("Disconnecting %s.%s" % (item, handle))
            if cinfo:
                adapter = component.queryMultiAdapter((cinfo.connected, item), IConnect)
                adapter.disconnect(handle)

class TextEditTool(Tool):
    """
    Text edit tool. Allows for elements that can adapt to the
    IEditable interface to be edited.
    """

    def create_edit_window(self, x, y, text, editor):
        """
        Create a popup window with some editable text.
        """
        view = self.view
        window = Gtk.Window.new(Gtk.WindowType.TOPLEVEL)
        window.set_property("decorated", False)
        window.set_property("skip-taskbar-hint", True)
        window.set_modal(True)
        window.set_transient_for(view.get_toplevel())
        buffer = Gtk.TextBuffer()
        if text:
            buffer.set_text(text)
            startiter, enditer = buffer.get_bounds()
            buffer.move_mark_by_name("selection_bound", startiter)
            buffer.move_mark_by_name("insert", enditer)
        text_view = Gtk.TextView()
        text_view.set_buffer(buffer)
        text_view.set_left_margin(2)
        text_view.set_right_margin(2)

        frame = Gtk.Frame()
        frame.set_shadow_type(Gtk.ShadowType.IN)
        frame.add(text_view)

        window.add(frame)
        r = Gdk.Rectangle()
        r.x = 0
        r.y = 0
        r.width = 70
        r.height = 50
        window.size_allocate(r)
        window.move(int(x), int(y))

        def on_button_press(widget, event):
            if event.window == view.get_window():
                self.submit_text(widget, buffer, editor)

        def on_key_press_event(widget, event):
            if event.keyval == Gdk.KEY_Return and not event.get_state() & (
                Gdk.ModifierType.CONTROL_MASK | Gdk.ModifierType.SHIFT_MASK
            ):
                self.submit_text(widget, buffer, editor)
            elif event.keyval == Gdk.KEY_Escape:
                widget.get_toplevel().destroy()

        def on_focus_out_event(widget, event):
            print("Focus out event emitted")
            self.submit_text(widget, buffer, editor)

        window.add_events(Gdk.EventMask.FOCUS_CHANGE_MASK)
        window.connect("focus-out-event", on_focus_out_event)
        window.connect("button-press-event", on_button_press)
        text_view.connect("key-press-event", on_key_press_event)
        window.show_all()

    @transactional
    def submit_text(self, widget, buffer, editor):
        """
        Submit the final text to the edited item.
        """
        text = buffer.get_text(
            buffer.get_start_iter(), buffer.get_end_iter(), include_hidden_chars=True
        )
        editor.update_text(text)
        widget.get_toplevel().destroy()

    def on_double_click(self, event):
        view = self.view
        item = view.hovered_item
        if item:
            try:
                editor = IEditor(item)
            except TypeError:
                # Could not adapt to IEditor
                return False
            log.debug("Found editor %r" % editor)
            x, y = view.get_matrix_v2i(item).transform_point(event.x, event.y)
            if editor.is_editable(x, y):
                text = editor.get_text()
                root_coords = event.get_root_coords()
                self.create_edit_window(
                    root_coords.x_root, root_coords.y_root, text, editor
                )
                return True

class PlacementTool(_PlacementTool):
    """
    PlacementTool is used to place items on the canvas.
    """

    def __init__(self, view, item_factory, after_handler=None, handle_index=-1):
        """
        item_factory is a callable. It is used to create a CanvasItem
        that is displayed on the diagram.
        """
        _PlacementTool.__init__(
            self,
            view,
            factory=item_factory,
            handle_tool=ConnectHandleTool(),
            handle_index=handle_index,
        )
        self.after_handler = after_handler
        self._tx = None

    @transactional
    def create_item(self, pos):
        return self._create_item(pos)

    def on_button_press(self, event):
        assert not self._tx
        self._tx = Transaction()
        view = self.view
        view.unselect_all()
        if _PlacementTool.on_button_press(self, event):
            try:
                opposite = self.new_item.opposite(
                    self.new_item.handles()[self._handle_index]
                )
            except (KeyError, AttributeError):
                pass
            else:
                # Connect opposite handle first, using the HandleTool's
                # mechanisms

                # First make sure all matrices are updated:
                view.canvas.update_matrix(self.new_item)
                view.update_matrix(self.new_item)

                vpos = event.x, event.y

                item = self.handle_tool.glue(self.new_item, opposite, vpos)
                if item:
                    self.handle_tool.connect(self.new_item, opposite, vpos)
            return True
        return False

    def on_button_release(self, event):
        try:
            if self.after_handler:
                self.after_handler(self.new_item)
            return _PlacementTool.on_button_release(self, event)
        finally:
            self._tx.commit()
            self._tx = None

class GroupPlacementTool(PlacementTool):
    """
    Try to group items when placing them on diagram.
    """

    def __init__(self, view, item_factory, after_handler=None, handle_index=-1):
        super(GroupPlacementTool, self).__init__(
            view, item_factory, after_handler, handle_index
        )
        self._parent = None

    def on_motion_notify(self, event):
        """
        Change parent item to dropzone state if it can accept diagram item
        object to be created.
        """
        view = self.view

        if view.focused_item:
            view.unselect_item(view.focused_item)
            view.focused_item = None

        try:
            parent = view.get_item_at_point((event.x, event.y))
        except KeyError:
            parent = None

        if parent:
            # create dummy adapter
            adapter = component.queryMultiAdapter(
                (parent, self._factory.item_class()), IGroup
            )
            if adapter and adapter.can_contain():
                view.dropzone_item = parent
                view.window.set_cursor(IN_CURSOR)
                self._parent = parent
            else:
                view.dropzone_item = None
                view.window.set_cursor(None)
                self._parent = None
            parent.request_update(matrix=False)
        else:
            if view.dropzone_item:
                view.dropzone_item.request_update(matrix=False)
            view.dropzone_item = None
            view.window.set_cursor(None)

    def _create_item(self, pos, **kw):
        """
        Create diagram item and place it within parent's boundaries.
        """
        parent = self._parent
        view = self.view
        try:
            adapter = component.queryMultiAdapter(
                (parent, self._factory.item_class()), IGroup
            )
            if parent and adapter and adapter.can_contain():
                kw["parent"] = parent

            item = super(GroupPlacementTool, self)._create_item(pos, **kw)

            adapter = component.queryMultiAdapter((parent, item), IGroup)
            if parent and item and adapter:
                adapter.group()

                canvas = view.canvas
                parent.request_update(matrix=False)
        finally:
            self._parent = None
            view.dropzone_item = None
            view.window.set_cursor(None)
        return item

@InMotion.when_type(ElementItem)
class DropZoneInMotion(GuidedItemInMotion):
    def move(self, pos):
        """
        Move the item. x and y are in view coordinates.
        """
        super(DropZoneInMotion, self).move(pos)
        item = self.item
        view = self.view
        x, y = pos

        current_parent = view.canvas.get_parent(item)
        over_item = view.get_item_at_point((x, y), selected=False)

        if not over_item:
            view.dropzone_item = None
            view.window.set_cursor(None)
            return

        if current_parent and not over_item:  # are we going to remove from parent?
            adapter = component.queryMultiAdapter((current_parent, item), IGroup)
            if adapter:
                view.window.set_cursor(OUT_CURSOR)
                view.dropzone_item = current_parent
                current_parent.request_update(matrix=False)

        if over_item:
            # are we going to add to parent?
            adapter = component.queryMultiAdapter((over_item, item), IGroup)
            if adapter and adapter.can_contain():
                view.dropzone_item = over_item
                view.window.set_cursor(IN_CURSOR)
                over_item.request_update(matrix=False)

    def stop_move(self):
        """
        Motion stops: drop!
        """
        super(DropZoneInMotion, self).stop_move()
        item = self.item
        view = self.view
        canvas = view.canvas
        old_parent = view.canvas.get_parent(item)
        new_parent = view.dropzone_item
        try:

            if new_parent is old_parent:
                if old_parent is not None:
                    old_parent.request_update(matrix=False)
                return

            if old_parent:
                adapter = component.queryMultiAdapter((old_parent, item), IGroup)
                if adapter:
                    adapter.ungroup()

                canvas.reparent(item, None)
                m = canvas.get_matrix_i2c(old_parent)
                item.matrix *= m
                old_parent.request_update()

            if new_parent:
                adapter = component.queryMultiAdapter((new_parent, item), IGroup)
                if adapter and adapter.can_contain():
                    adapter.group()

                canvas.reparent(item, new_parent)
                m = canvas.get_matrix_c2i(new_parent)
                item.matrix *= m
                new_parent.request_update()
        finally:
            view.dropzone_item = None
            view.window.set_cursor(None)

class TransactionalToolChain(ToolChain):
    """
    In addition to a normal toolchain, this chain begins an undo-transaction
    at button-press and commits the transaction at button-release.
    """

    def __init__(self, view=None):
        super(TransactionalToolChain, self).__init__(view)
        self._tx = None

    def handle(self, event):
        # For double click: button_press, double_click, button_release
        # print 'event', self.EVENT_HANDLERS.get(event.type)
        if self.EVENT_HANDLERS.get(event.type) in ("on_button_press",):
            assert not self._tx
            self._tx = Transaction()

        try:
            super(TransactionalToolChain, self).handle(event)
        finally:
            if self._tx and self.EVENT_HANDLERS.get(event.type) in (
                "on_button_release",
                "on_double_click",
                "on_triple_click",
            ):
                self._tx.commit()
                self._tx = None

def DefaultTool():
    """
    The default tool chain build from HoverTool, ItemTool and HandleTool.
    """
    chain = TransactionalToolChain()
    chain.append(HoverTool())
    chain.append(ConnectHandleTool())
    chain.append(ItemTool())
    chain.append(TextEditTool())
    chain.append(RubberbandTool())
    return chain

# vim:sw=4:et:ai
<|endoftext|>


<commit_before>from urlparse import urljoin

import logging
import requests

log = logging.getLogger('query')

class GDCIndexClient(object):

    def __init__(self, uri):
        self.uri = uri
        self.metadata = dict()

    def get_related_files(self, uuid):
        # type: str -> List[str]
        if uuid in self.metadata.keys():
            return self.metadata[uuid]['related_files']
        return []

    def get_annotations(self, uuid):
        # type: str -> List[str]
        if uuid in self.metadata.keys():
            return self.metadata[uuid]['annotations']
        return []

    def get_md5sum(self, uuid):
        # type: str -> str
        if uuid in self.metadata.keys():
            return self.metadata[uuid]['md5sum']

    def get_filesize(self, uuid):
        # type: str -> long
        if uuid in self.metadata.keys():
            return long(self.metadata[uuid]['file_size'])

    def get_access(self, uuid):
        # type: str -> long
        if uuid in self.metadata.keys():
            return self.metadata[uuid]['access']

    def _get_metadata(self, uuids):
        # type: List[str] -> Dict[str]str
        """ Capture the metadata of all the UUIDs while making
            as little open connections as possible.

            self.metadata = {
                str file_id: {
                    str       access
                    str       file_size
                    str       md5sum
                    List[str] annotations
                    List[str] related files
                }
            }
        """

        metadata_query = {
            'fields': 'file_id,file_size,md5sum,annotations.annotation_id,' \
                    'metadata_files.file_id,index_files.file_id,access',
            'filters': '{"op":"and","content":['
                       '{"op":"in","content":{'
                       '"field":"files.file_id","value":'
                       '["' + '","'.join(uuids) + '"]}}]}',
            'from': '0',
            'size': str(len(uuids)), # one big request
        }

        active_meta_url = urljoin(self.uri, 'v0/files')
        legacy_meta_url = urljoin(self.uri, 'v0/legacy/files')

        active_json_resp = dict()
        legacy_json_resp = dict()

        # using a POST request lets us avoid the MAX URL character length limit
        r_active = requests.post(active_meta_url, json=metadata_query, verify=False)
        r_legacy = requests.post(legacy_meta_url, json=metadata_query, verify=False)

        if r_active.status_code == requests.codes.ok:
            active_json_resp = r_active.json()

        if r_legacy.status_code == requests.codes.ok:
            legacy_json_resp = r_legacy.json()

        r_active.close()
        r_legacy.close()

        if not active_json_resp.get('data') and not legacy_json_resp.get('data'):
            log.debug('Unable to retrieve file metadata information. '
                        'continuing downloading as if they were large files')
            return self.metadata

        active_hits = active_json_resp['data']['hits']
        legacy_hits = legacy_json_resp['data']['hits']

        for h in active_hits + legacy_hits:
            related_returns = h.get('index_files', []) + h.get('metadata_files', [])
            related_files = [ r['file_id'] for r in related_returns ]

            annotations = [ a['annotation_id'] for a in h.get('annotations', []) ]

            # set the metadata as a class data member so that it can be
            # references as much as needed without needing to calculate
            # everything over again
            if h['id'] not in self.metadata.keys():
                # don't want to overwrite
                self.metadata[h['id']] = {
                    'access':        h['access'],
                    'file_size':     h['file_size'],
                    'md5sum':        h['md5sum'],
                    'annotations':   annotations,
                    'related_files': related_files,
                }

        return self.metadata

    def separate_small_files(self,
            ids,                    # type: Set[str]
            chunk_size,             # type: int
            related_files=False,    # type: bool
            annotations=False,      # type: bool
            ):
        # type: (...) -> (List[str], List[List[str]])
        """ Separate big and small files

        Separate the small files from the larger files in
        order to combine them into single grouped downloads. This will reduce
        the number of open connections needed to be made for many small files.

        On top of that, separate the small files by open and controlled access
        so that if a controlled grouping failed, you can handle it as the same
        edge case.
        """

        bigs = []
        smalls_open = []
        smalls_control = []
        potential_smalls = set()

        # go through all the UUIDs and pick out the ones with
        # relate and annotation files so they can be handled by parcel
        log.debug('Grouping ids by size')

        self._get_metadata(ids)
        for uuid in ids:
            if uuid not in self.metadata.keys():
                bigs.append(uuid)
                continue

            rf = self.get_related_files(uuid)
            af = self.get_annotations(uuid)

            # check for related files
            if related_files and rf and uuid not in bigs:
                bigs.append(uuid)

            # check for annotation files
            if annotations and af and uuid not in bigs:
                bigs.append(uuid)

            # if uuid has no related or annotation files
            # then proceed to the small file sorting with them
            if not af and not rf:
                potential_smalls |= set([uuid])

        # the following line is to trigger the first if statement
        # to start the process off properly
        bundle_open_size = chunk_size + 1
        bundle_control_size = chunk_size + 1

        i_open = -1
        i_control = -1

        for uuid in potential_smalls:
            # grouping of file exceeds chunk_size, create a new grouping
            if bundle_open_size > chunk_size:
                smalls_open.append([])
                i_open += 1
                bundle_open_size = 0

            if bundle_control_size > chunk_size:
                smalls_control.append([])
                i_control += 1
                bundle_control_size = 0

            # individual file is more than chunk_size, big file download
            if self.get_filesize(uuid) > chunk_size:
                bigs.append(uuid)

            # file size is less than chunk_size then group and tarfile it
            else:
                if self.get_access(uuid) == 'open':
                    smalls_open[i_open].append(uuid)
                    bundle_open_size += self.get_filesize(uuid)

                elif self.get_access(uuid) == 'controlled':
                    smalls_control[i_control].append(uuid)
                    bundle_control_size += self.get_filesize(uuid)

        # they are still small files to be downloaded in a group
        smalls = smalls_open + smalls_control

        # for logging/reporting purposes
        total_count = len(bigs) + sum([ len(s) for s in smalls ])
        if len(potential_smalls) > total_count:
            log.warning('There are less files to download than originally given')
            log.warning('Number of files originally given: {0}'\
                    .format(len(potential_smalls)))

        log.debug('{0} total number of files to download'.format(total_count))
        log.debug('{0} groupings of files'.format(len(smalls)))

        smalls = [ s for s in smalls if s != [] ]

        return bigs, smalls
<commit_msg>Turn SSL on for metadata-query<commit_after>from urlparse import urljoin

import logging
import requests

log = logging.getLogger('query')

class GDCIndexClient(object):

    def __init__(self, uri):
        self.uri = uri
        self.metadata = dict()

    def get_related_files(self, uuid):
        # type: str -> List[str]
        if uuid in self.metadata.keys():
            return self.metadata[uuid]['related_files']
        return []

    def get_annotations(self, uuid):
        # type: str -> List[str]
        if uuid in self.metadata.keys():
            return self.metadata[uuid]['annotations']
        return []

    def get_md5sum(self, uuid):
        # type: str -> str
        if uuid in self.metadata.keys():
            return self.metadata[uuid]['md5sum']

    def get_filesize(self, uuid):
        # type: str -> long
        if uuid in self.metadata.keys():
            return long(self.metadata[uuid]['file_size'])

    def get_access(self, uuid):
        # type: str -> long
        if uuid in self.metadata.keys():
            return self.metadata[uuid]['access']

    def _get_metadata(self, uuids):
        # type: List[str] -> Dict[str]str
        """ Capture the metadata of all the UUIDs while making
            as little open connections as possible.

            self.metadata = {
                str file_id: {
                    str       access
                    str       file_size
                    str       md5sum
                    List[str] annotations
                    List[str] related files
                }
            }
        """

        metadata_query = {
            'fields': 'file_id,file_size,md5sum,annotations.annotation_id,' \
                    'metadata_files.file_id,index_files.file_id,access',
            'filters': '{"op":"and","content":['
                       '{"op":"in","content":{'
                       '"field":"files.file_id","value":'
                       '["' + '","'.join(uuids) + '"]}}]}',
            'from': '0',
            'size': str(len(uuids)), # one big request
        }

        active_meta_url = urljoin(self.uri, 'v0/files')
        legacy_meta_url = urljoin(self.uri, 'v0/legacy/files')

        active_json_resp = dict()
        legacy_json_resp = dict()

        # using a POST request lets us avoid the MAX URL character length limit
        r_active = requests.post(active_meta_url, json=metadata_query, verify=True)
        r_legacy = requests.post(legacy_meta_url, json=metadata_query, verify=True)

        if r_active.status_code == requests.codes.ok:
            active_json_resp = r_active.json()

        if r_legacy.status_code == requests.codes.ok:
            legacy_json_resp = r_legacy.json()

        r_active.close()
        r_legacy.close()

        if not active_json_resp.get('data') and not legacy_json_resp.get('data'):
            log.debug('Unable to retrieve file metadata information. '
                        'continuing downloading as if they were large files')
            return self.metadata

        active_hits = active_json_resp['data']['hits']
        legacy_hits = legacy_json_resp['data']['hits']

        for h in active_hits + legacy_hits:
            related_returns = h.get('index_files', []) + h.get('metadata_files', [])
            related_files = [ r['file_id'] for r in related_returns ]

            annotations = [ a['annotation_id'] for a in h.get('annotations', []) ]

            # set the metadata as a class data member so that it can be
            # references as much as needed without needing to calculate
            # everything over again
            if h['id'] not in self.metadata.keys():
                # don't want to overwrite
                self.metadata[h['id']] = {
                    'access':        h['access'],
                    'file_size':     h['file_size'],
                    'md5sum':        h['md5sum'],
                    'annotations':   annotations,
                    'related_files': related_files,
                }

        return self.metadata

    def separate_small_files(self,
            ids,                    # type: Set[str]
            chunk_size,             # type: int
            related_files=False,    # type: bool
            annotations=False,      # type: bool
            ):
        # type: (...) -> (List[str], List[List[str]])
        """ Separate big and small files

        Separate the small files from the larger files in
        order to combine them into single grouped downloads. This will reduce
        the number of open connections needed to be made for many small files.

        On top of that, separate the small files by open and controlled access
        so that if a controlled grouping failed, you can handle it as the same
        edge case.
        """

        bigs = []
        smalls_open = []
        smalls_control = []
        potential_smalls = set()

        # go through all the UUIDs and pick out the ones with
        # relate and annotation files so they can be handled by parcel
        log.debug('Grouping ids by size')

        self._get_metadata(ids)
        for uuid in ids:
            if uuid not in self.metadata.keys():
                bigs.append(uuid)
                continue

            rf = self.get_related_files(uuid)
            af = self.get_annotations(uuid)

            # check for related files
            if related_files and rf and uuid not in bigs:
                bigs.append(uuid)

            # check for annotation files
            if annotations and af and uuid not in bigs:
                bigs.append(uuid)

            # if uuid has no related or annotation files
            # then proceed to the small file sorting with them
            if not af and not rf:
                potential_smalls |= set([uuid])

        # the following line is to trigger the first if statement
        # to start the process off properly
        bundle_open_size = chunk_size + 1
        bundle_control_size = chunk_size + 1

        i_open = -1
        i_control = -1

        for uuid in potential_smalls:
            # grouping of file exceeds chunk_size, create a new grouping
            if bundle_open_size > chunk_size:
                smalls_open.append([])
                i_open += 1
                bundle_open_size = 0

            if bundle_control_size > chunk_size:
                smalls_control.append([])
                i_control += 1
                bundle_control_size = 0

            # individual file is more than chunk_size, big file download
            if self.get_filesize(uuid) > chunk_size:
                bigs.append(uuid)

            # file size is less than chunk_size then group and tarfile it
            else:
                if self.get_access(uuid) == 'open':
                    smalls_open[i_open].append(uuid)
                    bundle_open_size += self.get_filesize(uuid)

                elif self.get_access(uuid) == 'controlled':
                    smalls_control[i_control].append(uuid)
                    bundle_control_size += self.get_filesize(uuid)

        # they are still small files to be downloaded in a group
        smalls = smalls_open + smalls_control

        # for logging/reporting purposes
        total_count = len(bigs) + sum([ len(s) for s in smalls ])
        if len(potential_smalls) > total_count:
            log.warning('There are less files to download than originally given')
            log.warning('Number of files originally given: {0}'\
                    .format(len(potential_smalls)))

        log.debug('{0} total number of files to download'.format(total_count))
        log.debug('{0} groupings of files'.format(len(smalls)))

        smalls = [ s for s in smalls if s != [] ]

        return bigs, smalls
<|endoftext|>


<commit_before>#!/usr/bin/env python3
"""
Example of aurora using Stan Solomon's GLOW Auroral model
code wrapping in Python by Michael Hirsch
"""
from __future__ import division,absolute_import
from itertools import chain
from matplotlib.pyplot import figure, subplots,tight_layout,draw
from matplotlib.ticker import MultipleLocator #LogFormatterMathtext,
from pandas import DataFrame,Panel
from numpy import hstack,asarray,rollaxis,degrees,zeros_like
from os import chdir
from os.path import join
try:
    import seaborn
    seaborn.set_set_context('poster')
except:
    pass
#
from histutils.fortrandates import datetime2yd
from histutils.findnearest import find_nearest
import glowaurora
from glowaurora import glowfort
#
glowpath=glowaurora.__path__[0]

dymaj=50
dymin=10
dpi = 100

def runglowaurora(eflux,e0,dt,glat,glon,f107a,f107,f107p,ap):
    chdir(glowpath)
    yd,utsec = datetime2yd(dt)[:2]
#%% (0) define altitude grid [km] (here i liberally copied that used by Stan)
    #z = range(80,110+1,1)
    z = list(range(30,110+1,1))
    z += (
         [111.5,113.,114.5,116.] +
         list(chain(range(118,150+2,2),range(153,168+3,3),range(172,180+4,4),
                    range(185,205+5,5),range(211,223+6,6),range(230,244+7,7),
                    range(252,300+8,8),range(309,345+9,9),range(355,395+10,10),
                    range(406,428+11,11))) +
         [440,453,467,482,498,515,533,551] +
         list(range(570,950+20,20))
         )

    z = asarray(z)
#%% (1) setup flux at top of ionosphere
    ener,dE = glowfort.egrid()

    if eflux:
        # maxwellian input PhiTop at top of ionosphere
        phitop = glowfort.maxt(eflux,e0,ener, dE, itail=0, fmono=0, emono=0)
    else: #eigenprofile generation
        e0ind = find_nearest(ener,e0)[0]
        phitop = zeros_like(ener)
        phitop[e0ind] = 1e6

    phi = hstack((ener[:,None],dE[:,None],phitop[:,None]))
#%% (2) msis,iri,glow model
    ion,ecalc,photI,ImpI,isr,prate,lrate = glowfort.aurora(z,yd,utsec,glat,glon%360,
                                             f107a,f107,f107p,ap,phi)
#%% handle the outputs including common blocks
    zeta=glowfort.cglow.zeta.T #columns 11:20 are identically zero

    ver = DataFrame(index=z,
                    data=zeta[:,:11],
                    columns=[3371, 4278, 5200, 5577, 6300,7320,10400,3466,7774, 8446,3726])
    photIon = DataFrame(index=z,
                   data=hstack((photI[:,None],ImpI[:,None],ecalc[:,None],ion)),
                    columns=['photoIoniz','eImpactIoniz','ne',
                    'nO+(2P)','nO+(2D)','nO+(4S)','nN+','nN2+','nO2+','nNO+',
                    'nO','nO2','nN2','nNO'])

    isrparam = DataFrame(index=z,
                         data=isr,
                         columns=['ne','Te','Ti'])

    phitop = DataFrame(index=phi[:,0], #eV
                       data=phi[:,2],  #diffnumflux
                       columns=['diffnumflux'])
    zceta = glowfort.cglow.zceta.T

    sza = degrees(glowfort.cglow.sza)

#%% production and loss rates
    prate = prate.T; lrate=lrate.T #fortran to C ordering 2x170x20, only first 12 columns are used

    #column labels by inspection of fortran/gchem.f staring after "DO 150 I=1,JMAX" (thanks Stan!)
    prateDF = Panel(items=['pre','final'],
                    major_axis=z,
                    minor_axis=['O+(2P)','O+(2D)','O+(4S)','N+','N2+','O2+','NO+',
                                 'N2(A)','N(2P)','N(2D)','O(1S)','O(1D)'],
                    data=prate[...,:12], #columns 12:20 are identically zero
                        )

    lrateDF = Panel(items=['pre','final'],
                    major_axis=z,
                    minor_axis=['O+(2P)','O+(2D)','O+(4S)','N+','N2+','O2+','NO+',
                                 'N2(A)','N(2P)','N(2D)','O(1S)','O(1D)'],
                    data=lrate[...,:12], #columns 12:20 are identically zero
                        )

    return ver,photIon,isrparam,phitop,zceta,sza,prateDF,lrateDF
#%% plot
def plotaurora(phitop,ver,zceta,photIon,isr,dtime,glat,glon,prate,lrate,
               E0=None,flux=None,sza=None,zminmax=(None,None),makeplot=None,odir=''):
    if makeplot is None:
        return

    if E0 and sza:
        titlend = '$E_0={:.0f}$ eV  SZA={:.1f}$^\circ$'.format(E0,sza)
    else:
        titlend = ''

    def _nicez(ax,zlim):
        ax.set_ylim(zlim)
        ax.yaxis.set_major_locator(MultipleLocator(dymaj))
        ax.yaxis.set_minor_locator(MultipleLocator(dymin))
        ax.grid(True,which='major',linewidth=1.)
        ax.grid(True,which='minor',linewidth=0.5)
        ax.tick_params(axis='both',which='major',labelsize='medium')

#%% neutral background (MSIS) and Te,Ti (IRI-90)
    if not 'eig' in makeplot:
        fg,axs = subplots(1,2,sharey=True,figsize=(15,8))
        fg.suptitle('{} ({},{}) '.format(dtime,glat,glon)+ titlend)

        ind = ['nO','nO2','nN2','nNO']
        ax = axs[0]
        ax.semilogx(photIon[ind], photIon.index)
        ax.set_xlabel('Number Density')
        ax.set_xscale('log')
        ax.set_xlim(left=1e1)
        ax.set_ylabel('Altitude [km]')
        _nicez(ax,zminmax)
        ax.legend(ind)
        ax.set_title('Neutral Number Density')

        ind=['Te','Ti']
        ax = axs[1]
        ax.semilogx(isr[ind], isr.index)
        ax.set_xlabel('Temperature [K]')
        ax.legend(ind)
        _nicez(ax,zminmax)
        ax.set_title('Background Temperature')

        writeplots(fg,'bg_',E0,makeplot,odir)
#%% volume emission rate
    fg,axs = subplots(1,3,sharey=False, figsize=(15,8))
    fg.suptitle('{} ({},{}) '.format(dtime,glat,glon) + titlend)
    tight_layout(pad=3.2, w_pad=0.6)

# incident flux at top of ionosphere
    ax = axs[0]
    ax.plot(phitop.index,phitop['diffnumflux'],marker='.')

    titxt='Incident Flux'
    if flux:
        titxt+='  Total Flux={:.1f},'.format(flux)
    ax.set_title(titxt)

    ax.set_xlabel('Beam Energy [eV]')
    ax.set_ylabel('Flux [erg sr$^{-1}$ s$^{-1}$]')
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_ylim(1e-4,1e6)
    ax.grid(True)
# ver visible
    ind= [4278, 5200, 5577, 6300]
    ax = axs[1]
    ax.plot(ver[ind].values,ver.index)
    ax.set_xlabel('Volume Emission Rate')
    ax.set_ylabel('altitude [km]')
    _nicez(ax,zminmax)
    ax.set_xscale('log')
    if not 'eig' in makeplot:
        ax.set_xlim(1e-5,1e3)
    ax.legend(ind,loc='best')
    ax.set_title('Volume Emission Rate: Visible')
# ver invisible
    ind = [3371,7320,10400,3466,7774, 8446,3726]
    ax = axs[2]
    ax.plot(ver[ind].values,ver.index)
    ax.set_xlabel('Volume Emission Rate')
    _nicez(ax,zminmax)
    ax.set_xscale('log')
    if not 'eig' in makeplot:
        ax.set_xlim(1e-5,1e3)
    ax.legend(ind,loc='best')
    ax.set_title('Volume Emission Rate: IR & UV')

    writeplots(fg,'ver_',E0,makeplot,odir)
#%% Ne, Ni
    if not 'eig' in makeplot:
        ind=['ne','nO+(2P)','nO+(2D)','nO+(4S)','nN+','nN2+','nO2+','nNO+']
        fg=figure()
        ax = fg.gca()
        ax.semilogx(photIon[ind], photIon.index)
        ax.set_xlabel('Density')
        ax.set_xscale('log')
        ax.set_xlim(left=1e-3)
        _nicez(ax,zminmax)
        ax.legend(ind)
        ax.set_title('Electron and Ion Densities')

        writeplots(fg,'effects_',E0,makeplot,odir)
#%% total energy deposition vs. altitude
    if not 'eig' in makeplot:
        fg,axs = subplots(1,3,sharey=True, figsize=(15,8))
        fg.suptitle('{} ({},{})  '.format(dtime,glat,glon) + titlend)
        tight_layout(pad=3.2, w_pad=0.3)

        ax = axs[2]
        tez = glowfort.cglow.tez
        ax.plot(tez,ver.index)
        ax.set_xscale('log')
        ax.set_xlim(1e-1,1e6)
        _nicez(ax,zminmax)
        ax.set_xlabel('Energy Deposited')
        ax.set_title('Total Energy Depostiion')

#%% e^- impact ionization rates from ETRANS
        ind=['photoIoniz','eImpactIoniz']
        ax = axs[1]
        ax.plot(photIon[ind],photIon.index)
        ax.set_xlabel('ionization')
        ax.set_xscale('log')
        ax.set_xlim(left=1e-1)
        _nicez(ax,zminmax)
        ax.legend(ind)
        ax.set_title('Photo and e$^-$ impact ionization')

        ind=['O','O2','N2']
        ax = axs[0]
        sion = glowfort.cglow.sion
        sion = DataFrame(index=ver.index,data=sion.T,columns=ind)
        ax.plot(sion,ver.index)
        ax.set_xscale('log')
        ax.set_xlim(1e-5,1e4)
        _nicez(ax,zminmax)
        ax.set_xlabel('e$^-$ impact ioniz. rate')
        ax.set_ylabel('Altitude [km]')
        ax.set_title('electron impact ioniz. rates')
        ax.legend(ind)

        writeplots(fg,'ioniz_',E0,makeplot,odir)
#%% constituants of per-wavelength VER
#    zcsum = zceta.sum(axis=-1)
    if not 'eig' in makeplot:
        fg = figure()
        ax = fg.gca()
        for zc in rollaxis(zceta,1):
            ax.plot(ver.index,zc)
        ax.set_xlabel('emission constituants  ' + titlend)
        #ax.legend(True)

        writeplots(fg,'constit_',E0,makeplot,odir)

#%%
def writeplots(fg,plotprefix,E0,method,odir):
    draw() #Must have this here or plot doesn't update in animation multiplot mode!
    #TIF was not faster and was 100 times the file size!
    #PGF is slow and big file,
    #RAW crashes
    #JPG no faster than PNG
    if 'png' in method:
        cn = join(odir,(plotprefix + 'beam{:.0f}.png'.format(E0)))
        print('write {}'.format(cn))
        fg.savefig(cn,bbox_inches='tight',format='png',dpi=dpi)  # this is slow and async
<commit_msg>plot volume prod/loss rates<commit_after>#!/usr/bin/env python3
"""
Example of aurora using Stan Solomon's GLOW Auroral model
code wrapping in Python by Michael Hirsch
"""
from __future__ import division,absolute_import
from itertools import chain
from matplotlib.pyplot import figure, subplots,tight_layout,draw
from matplotlib.ticker import MultipleLocator #LogFormatterMathtext,
from pandas import DataFrame,Panel
from numpy import hstack,asarray,rollaxis,degrees,zeros_like
from os import chdir
from os.path import join
try:
    import seaborn
    seaborn.set_set_context('poster')
except:
    pass
#
from histutils.fortrandates import datetime2yd
from histutils.findnearest import find_nearest
import glowaurora
from glowaurora import glowfort
#
glowpath=glowaurora.__path__[0]

dymaj=50
dymin=10
dpi = 100

def runglowaurora(eflux,e0,dt,glat,glon,f107a,f107,f107p,ap):
    chdir(glowpath)
    yd,utsec = datetime2yd(dt)[:2]
#%% (0) define altitude grid [km] (here i liberally copied that used by Stan)
    #z = range(80,110+1,1)
    z = list(range(30,110+1,1))
    z += (
         [111.5,113.,114.5,116.] +
         list(chain(range(118,150+2,2),range(153,168+3,3),range(172,180+4,4),
                    range(185,205+5,5),range(211,223+6,6),range(230,244+7,7),
                    range(252,300+8,8),range(309,345+9,9),range(355,395+10,10),
                    range(406,428+11,11))) +
         [440,453,467,482,498,515,533,551] +
         list(range(570,950+20,20))
         )

    z = asarray(z)
#%% (1) setup flux at top of ionosphere
    ener,dE = glowfort.egrid()

    if eflux:
        # maxwellian input PhiTop at top of ionosphere
        phitop = glowfort.maxt(eflux,e0,ener, dE, itail=0, fmono=0, emono=0)
    else: #eigenprofile generation
        e0ind = find_nearest(ener,e0)[0]
        phitop = zeros_like(ener)
        phitop[e0ind] = 1e6

    phi = hstack((ener[:,None],dE[:,None],phitop[:,None]))
#%% (2) msis,iri,glow model
    ion,ecalc,photI,ImpI,isr,prate,lrate = glowfort.aurora(z,yd,utsec,glat,glon%360,
                                             f107a,f107,f107p,ap,phi)
#%% handle the outputs including common blocks
    zeta=glowfort.cglow.zeta.T #columns 11:20 are identically zero

    ver = DataFrame(index=z,
                    data=zeta[:,:11],
                    columns=[3371, 4278, 5200, 5577, 6300,7320,10400,3466,7774, 8446,3726])
    photIon = DataFrame(index=z,
                   data=hstack((photI[:,None],ImpI[:,None],ecalc[:,None],ion)),
                    columns=['photoIoniz','eImpactIoniz','ne',
                    'nO+(2P)','nO+(2D)','nO+(4S)','nN+','nN2+','nO2+','nNO+',
                    'nO','nO2','nN2','nNO'])

    isrparam = DataFrame(index=z,
                         data=isr,
                         columns=['ne','Te','Ti'])

    phitop = DataFrame(index=phi[:,0], #eV
                       data=phi[:,2],  #diffnumflux
                       columns=['diffnumflux'])
    zceta = glowfort.cglow.zceta.T

    sza = degrees(glowfort.cglow.sza)

#%% production and loss rates
    prate = prate.T; lrate=lrate.T #fortran to C ordering 2x170x20, only first 12 columns are used

    #column labels by inspection of fortran/gchem.f staring after "DO 150 I=1,JMAX" (thanks Stan!)
    prateDF = Panel(items=['pre','final'],
                    major_axis=z,
                    minor_axis=['O+(2P)','O+(2D)','O+(4S)','N+','N2+','O2+','NO+',
                                 'N2(A)','N(2P)','N(2D)','O(1S)','O(1D)'],
                    data=prate[...,:12], #columns 12:20 are identically zero
                        )

    lrateDF = Panel(items=['pre','final'],
                    major_axis=z,
                    minor_axis=['O+(2P)','O+(2D)','O+(4S)','N+','N2+','O2+','NO+',
                                 'N2(A)','N(2P)','N(2D)','O(1S)','O(1D)'],
                    data=lrate[...,:12], #columns 12:20 are identically zero
                        )

    return ver,photIon,isrparam,phitop,zceta,sza,prateDF,lrateDF
#%% plot
def plotaurora(phitop,ver,zceta,photIon,isr,dtime,glat,glon,prate,lrate,
               E0=None,flux=None,sza=None,zlim=(None,None),makeplot=None,odir=''):
    if makeplot is None:
        return

    if E0 and sza:
        titlend = '$E_0={:.0f}$ eV  SZA={:.1f}$^\circ$'.format(E0,sza)
    else:
        titlend = ''

    def _nicez(ax,zlim):
        ax.set_ylim(zlim)
        ax.yaxis.set_major_locator(MultipleLocator(dymaj))
        ax.yaxis.set_minor_locator(MultipleLocator(dymin))
        ax.grid(True,which='major',linewidth=1.)
        ax.grid(True,which='minor',linewidth=0.5)
        ax.tick_params(axis='both',which='major',labelsize='medium')

    z=ver.index.values #[km], for convenience
#%% neutral background (MSIS) and Te,Ti (IRI-90)
    if not 'eig' in makeplot:
        fg,axs = subplots(1,2,sharey=True,figsize=(15,8))
        fg.suptitle('{} ({},{}) '.format(dtime,glat,glon)+ titlend)

        ind = ['nO','nO2','nN2','nNO']
        ax = axs[0]
        ax.semilogx(photIon[ind], z)
        ax.set_xlabel('Number Density')
        ax.set_xscale('log')
        ax.set_xlim(left=1e1)
        ax.set_ylabel('Altitude [km]')
        _nicez(ax,zlim)
        ax.legend(ind)
        ax.set_title('Neutral Number Density')

        ind=['Te','Ti']
        ax = axs[1]
        ax.semilogx(isr[ind], z)
        ax.set_xlabel('Temperature [K]')
        ax.legend(ind)
        _nicez(ax,zlim)
        ax.set_title('Background Temperature')

        writeplots(fg,'bg_',E0,makeplot,odir)
#%% production and loss rates for species
    if not 'eig' in makeplot:
        fg,ax = subplots(1,2,sharey=True,figsize=(15,8))
        fg.suptitle('Volume Production Rates   {} ({},{}) '.format(dtime,glat,glon)+ titlend)
        ax[0].plot(prate['pre'].values,z)
        ax[0].legend(prate.minor_axis,loc='best')
        ax[0].set_title('pre e^- density correction')
        _nicez(ax[0],zlim)
        ax[0].set_xlabel('Volume Production Rates [cm$^{-3}$ s$Y{-1}$')
        ax[0].set_xscale('log')
        ax[0].set_ylabel('altitude [km]')

        ax[1].plot(prate['final'].values,z)
        ax[1].legend(prate.minor_axis,loc='best')
        ax[1].set_title('post (final) e^- density correction')
        _nicez(ax[1],zlim)
        ax[1].set_xlabel('Volume Production Rates [cm$^{-3}$ s$Y{-1}$')
        ax[1].set_xscale('log')
        ax[1].set_ylabel('altitude [km]')

# loss rates
        fg,ax = subplots(1,2,sharey=True,figsize=(15,8))
        fg.suptitle('Volume Loss Rates   {} ({},{}) '.format(dtime,glat,glon)+ titlend)
        ax[0].plot(lrate['pre'].values,z)
        ax[0].legend(lrate.minor_axis,loc='best')
        ax[0].set_title('pre e^- density correction')
        _nicez(ax[0],zlim)
        ax[0].set_xlabel('Volume Loss Rates [cm$^{-3}$ s$Y{-1}$')
        ax[0].set_xscale('log')
        ax[0].set_ylabel('altitude [km]')

        ax[1].plot(lrate['final'].values,z)
        ax[1].legend(lrate.minor_axis,loc='best')
        ax[1].set_title('post (final) e^- density correction')
        _nicez(ax[1],zlim)
        ax[1].set_xlabel('Volume Loss Rates [cm$^{-3}$ s$Y{-1}$')
        ax[1].set_xscale('log')
        ax[1].set_ylabel('altitude [km]')

#%% volume emission rate
    fg,axs = subplots(1,3,sharey=False, figsize=(15,8))
    fg.suptitle('{} ({},{}) '.format(dtime,glat,glon) + titlend)
    tight_layout(pad=3.2, w_pad=0.6)

# incident flux at top of ionosphere
    ax = axs[0]
    ax.plot(phitop.index,phitop['diffnumflux'],marker='.')

    titxt='Incident Flux'
    if flux:
        titxt+='  Total Flux={:.1f},'.format(flux)
    ax.set_title(titxt)

    ax.set_xlabel('Beam Energy [eV]')
    ax.set_ylabel('Flux [erg sr$^{-1}$ s$^{-1}$]')
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.set_ylim(1e-4,1e6)
    ax.grid(True)
# ver visible
    ind= [4278, 5200, 5577, 6300]
    ax = axs[1]
    ax.plot(ver[ind].values,ver.index)
    ax.set_xlabel('Volume Emission Rate')
    ax.set_ylabel('altitude [km]')
    _nicez(ax,zlim)
    ax.set_xscale('log')
    if not 'eig' in makeplot:
        ax.set_xlim(1e-5,1e3)
    ax.legend(ind,loc='best')
    ax.set_title('Volume Emission Rate: Visible')
# ver invisible
    ind = [3371,7320,10400,3466,7774, 8446,3726]
    ax = axs[2]
    ax.plot(ver[ind].values,ver.index)
    ax.set_xlabel('Volume Emission Rate')
    _nicez(ax,zlim)
    ax.set_xscale('log')
    if not 'eig' in makeplot:
        ax.set_xlim(1e-5,1e3)
    ax.legend(ind,loc='best')
    ax.set_title('Volume Emission Rate: IR & UV')

    writeplots(fg,'ver_',E0,makeplot,odir)
#%% Ne, Ni
    if not 'eig' in makeplot:
        ind=['ne','nO+(2P)','nO+(2D)','nO+(4S)','nN+','nN2+','nO2+','nNO+']
        fg=figure()
        ax = fg.gca()
        ax.semilogx(photIon[ind], photIon.index)
        ax.set_xlabel('Density')
        ax.set_xscale('log')
        ax.set_xlim(left=1e-3)
        _nicez(ax,zlim)
        ax.legend(ind)
        ax.set_title('Electron and Ion Densities')

        writeplots(fg,'effects_',E0,makeplot,odir)
#%% total energy deposition vs. altitude
    if not 'eig' in makeplot:
        fg,axs = subplots(1,3,sharey=True, figsize=(15,8))
        fg.suptitle('{} ({},{})  '.format(dtime,glat,glon) + titlend)
        tight_layout(pad=3.2, w_pad=0.3)

        ax = axs[2]
        tez = glowfort.cglow.tez
        ax.plot(tez,ver.index)
        ax.set_xscale('log')
        ax.set_xlim(1e-1,1e6)
        _nicez(ax,zlim)
        ax.set_xlabel('Energy Deposited')
        ax.set_title('Total Energy Depostiion')

#%% e^- impact ionization rates from ETRANS
        ind=['photoIoniz','eImpactIoniz']
        ax = axs[1]
        ax.plot(photIon[ind],photIon.index)
        ax.set_xlabel('ionization')
        ax.set_xscale('log')
        ax.set_xlim(left=1e-1)
        _nicez(ax,zlim)
        ax.legend(ind)
        ax.set_title('Photo and e$^-$ impact ionization')

        ind=['O','O2','N2']
        ax = axs[0]
        sion = glowfort.cglow.sion
        sion = DataFrame(index=ver.index,data=sion.T,columns=ind)
        ax.plot(sion,ver.index)
        ax.set_xscale('log')
        ax.set_xlim(1e-5,1e4)
        _nicez(ax,zlim)
        ax.set_xlabel('e$^-$ impact ioniz. rate')
        ax.set_ylabel('Altitude [km]')
        ax.set_title('electron impact ioniz. rates')
        ax.legend(ind)

        writeplots(fg,'ioniz_',E0,makeplot,odir)
#%% constituants of per-wavelength VER
#    zcsum = zceta.sum(axis=-1)
    if not 'eig' in makeplot:
        fg = figure()
        ax = fg.gca()
        for zc in rollaxis(zceta,1):
            ax.plot(ver.index,zc)
        ax.set_xlabel('emission constituants  ' + titlend)
        #ax.legend(True)

        writeplots(fg,'constit_',E0,makeplot,odir)

#%%
def writeplots(fg,plotprefix,E0,method,odir):
    draw() #Must have this here or plot doesn't update in animation multiplot mode!
    #TIF was not faster and was 100 times the file size!
    #PGF is slow and big file,
    #RAW crashes
    #JPG no faster than PNG
    if 'png' in method:
        cn = join(odir,(plotprefix + 'beam{:.0f}.png'.format(E0)))
        print('write {}'.format(cn))
        fg.savefig(cn,bbox_inches='tight',format='png',dpi=dpi)  # this is slow and async
<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
import pytest

from django.test import TestCase

from production.models import Event
from diffusion.models import Place, MetaEvent

class CommandsTestCase(TestCase):
    """
        Tests Diffusion Commands
    """

    def setUp(self):
        # create place
        self.place = Place(name="Le Fresnoy", description="Le Fresnoy Studio National", address="22 rue du Fresnoy")
        self.place.save()
        # create event
        self.event = Event(title='PanoramaX',
                           starting_date="1970-01-01 00:00:00.0+00:00",
                           type="FEST",
                           place=self.place)
        self.event.save()
        # create meta_event
        self.main_event = Event(title='Panorama',
                                starting_date="1970-01-01 00:00:00.0+00:00",
                                type="EXIB",
                                main_event=True,
                                place=self.place)
        self.main_event.save()
        self.main_event.subevents.add(self.event)

    def tearDown(self):
        pass

    def test_meta_event(self):
        "simple TEST create meta_event"
        # create meta_event
        meta_event = MetaEvent(event=self.main_event, genres=['FILM'])
        meta_event.save()
        # get meta events
        meta_events = MetaEvent.objects.all()
        # test metaEvent created
        self.assertEqual(meta_events.count(), 1)
        parent_event = self.event.parent_event.first()
        # test parent event is main event
        self.assertEqual(parent_event.main_event, True)
        # test main event is meta event
        self.assertEqual(parent_event.meta_event.important, True)

@pytest.mark.django_db
class TestPlace:
    def test_str(self, place):
        place_str = str(place)
        assert place.name in place_str
        assert place.address[0:20] in place_str
        assert str(place.organization) in place_str

    def test_str_with_country(self, place):
        place.organization = None
        place.country = "FR"
        place.save()
        place_str = str(place)
        assert place.address[0:20] in place_str
        assert str(place.country) in place_str
        assert str(place.organization) not in place_str

    def test_str_without_address(self, place):
        # FIXME: conjonction peu probable, non ?
        place.name = place.address[0:20] + "..."
        place.save()
        place_str = str(place)

        assert place.name in place_str
        assert str(place.organization) in place_str

@pytest.mark.django_db
class TestMetaAward:
    def test_str(self, meta_award):
        meta_award_str = str(meta_award)
        assert meta_award.label in meta_award_str
        assert '(main event)' not in meta_award_str
        assert meta_award.event.title in meta_award_str
        assert str(meta_award.task) in meta_award_str

    def test_str_no_event(self, meta_award):
        meta_award.event = None
        meta_award.save()
        meta_award_str = str(meta_award)
        assert str(meta_award.task) in meta_award_str
        assert 'cat.' in meta_award_str

    def test_str_no_task(self, meta_award):
        meta_award.task = None
        meta_award.save()
        meta_award_str = str(meta_award)
        assert meta_award.event.title in meta_award_str
        assert 'cat.' not in meta_award_str

@pytest.mark.django_db
class TestAward:
    def test_str(self, award):
        award_str = str(award)
        assert str(award.date.year) in award_str
        assert str(award.meta_award) in award_str
        for artwork in award.artwork.all():
            assert str(artwork) in award_str

@pytest.mark.django_db
class TestMetaEvent:
    def test_str(self, meta_event):
        assert str(meta_event) == meta_event.event.title

@pytest.mark.django_db
class TestDiffusion:
    def test_str(self, diffusion):
        diffusion_str = str(diffusion)
        assert diffusion.artwork.title in diffusion_str
        assert diffusion.event.title in diffusion_str
<commit_msg>ref(tests): remove django TestCases<commit_after>import pytest

@pytest.mark.django_db
class TestPlace:
    def test_str(self, place):
        place_str = str(place)
        assert place.name in place_str
        assert place.address[0:20] in place_str
        assert str(place.organization) in place_str

    def test_str_with_country(self, place):
        place.organization = None
        place.country = "FR"
        place.save()
        place_str = str(place)
        assert place.address[0:20] in place_str
        assert str(place.country) in place_str
        assert str(place.organization) not in place_str

    def test_str_without_address(self, place):
        # FIXME: conjonction peu probable, non ?
        place.name = place.address[0:20] + "..."
        place.save()
        place_str = str(place)

        assert place.name in place_str
        assert str(place.organization) in place_str

@pytest.mark.django_db
class TestMetaAward:
    def test_str(self, meta_award):
        meta_award_str = str(meta_award)
        assert meta_award.label in meta_award_str
        assert '(main event)' not in meta_award_str
        assert meta_award.event.title in meta_award_str
        assert str(meta_award.task) in meta_award_str

    def test_str_no_event(self, meta_award):
        meta_award.event = None
        meta_award.save()
        meta_award_str = str(meta_award)
        assert str(meta_award.task) in meta_award_str
        assert 'cat.' in meta_award_str

    def test_str_no_task(self, meta_award):
        meta_award.task = None
        meta_award.save()
        meta_award_str = str(meta_award)
        assert meta_award.event.title in meta_award_str
        assert 'cat.' not in meta_award_str

@pytest.mark.django_db
class TestAward:
    def test_str(self, award):
        award_str = str(award)
        assert str(award.date.year) in award_str
        assert str(award.meta_award) in award_str
        for artwork in award.artwork.all():
            assert str(artwork) in award_str

@pytest.mark.django_db
class TestMetaEvent:
    def test_str(self, meta_event):
        assert str(meta_event) == meta_event.event.title

@pytest.mark.django_db
class TestDiffusion:
    def test_str(self, diffusion):
        diffusion_str = str(diffusion)
        assert diffusion.artwork.title in diffusion_str
        assert diffusion.event.title in diffusion_str
<|endoftext|>