from unittest import mock
from schrodinger.Qt import QtGui
from schrodinger.Qt.QtCore import Qt
from schrodinger.ui.qt.utils import get_signals
[docs]class SlotMockManager:
    """
    A helper class to create slot mocks for every signal on an object.
    Slots will be attributes with the same name as the signal.
    Example::
        class SomeObject(QtCore.QObject):
            changed = QtCore.pyqtSignal()
            def change(self):
                self.changed.emit()
        my_obj = SomeObject()
        some_object_slot_mgr = SlotMockManager(my_obj)
        my_obj.change()
        some_object_slot_mgr.changed.assert_called_once()
    """
[docs]    def __init__(self, q_object):
        """
        Create and connect slot mocks for each signal on `q_object`
        :param q_object: The object with signals
        :type  q_object: QtCore.QObject
        """
        slot_dict = dict()
        for name, signal in get_signals(q_object).items():
            slot_mock = mock.MagicMock()
            signal.connect(slot_mock)
            setattr(self, name, slot_mock)
            slot_dict[name] = slot_mock
        self._slot_dict = slot_dict 
[docs]    def reset(self):
        """
        Reset all of the slot mocks
        """
        for slot in self._slot_dict.values():
            slot.reset_mock() 
[docs]    def assertSlotsCalledExclusive(self, *expected_slots):
        """
        Check that only the specified slots were called
        Example::
            mgr = SlotMockManager(my_obj)
            my_obj.emitChanged()  # emits "changed"
            my_obj.emitCleared()  # emits "cleared"
            mgr.assertSlotsCalledExclusive("changed", "cleared")
        :param expected_slots: Name(s) of slots. Call with no arguments to
            assert that no slots were called.
        :type  expected_slots: str
        """
        expected_slots = set(expected_slots)
        for name, slot in self._slot_dict.items():
            if name in expected_slots:
                assert slot.called, f"Slot {name} was not called"
            else:
                assert not slot.called, f"Slot {name} was unexpectedly called"  
def _update_pos(widget, pos):
    """
    If `pos` is None, set it to the center of `widget`.  Otherwise, return it
    unmodified.
    :param widget: A widget
    :type widget: QtWidgets.QWidget
    :param pos: The position to possibly modify.
    :type pos: QtCore.QPoint or None
    :return: The position.
    :type: QtCore.QPoint
    """
    if pos is None:
        pos = widget.rect().center()
    return pos
[docs]def mouse_press(widget, button=Qt.LeftButton, modifier=Qt.NoModifier, pos=None):
    """
    Press a mouse button over a widget.  (Note that this function only presses
    the mouse, it doesn't release it to complete a click.  Use `mouse_click`
    instead if you want a complete mouse click.)
    :param widget: The widget to press the mouse on.
    :type widget: QtWidgets.QWidget
    :param button: The mouse button to press.
    :type button: Qt.MouseButton
    :param modifier: Any keyboard modifier keys that should affect the mouse
        press.
    :type modifier: Qt.KeyboardModifier
    :param pos: Where to press the mouse button.  If `None`, the center of
        `widget` will be used.
    :type pos: QtCore.QPoint
    """
    pos = _update_pos(widget, pos)
    QME = QtGui.QMouseEvent
    event = QME(QME.MouseButtonPress, pos, button, button, modifier)
    widget.mousePressEvent(event) 
[docs]def mouse_release(widget,
                  button=Qt.LeftButton,
                  modifier=Qt.NoModifier,
                  pos=None):
    """
    Release a mouse button over a widget.
    :param widget: The widget to release the mouse on.
    :type widget: QtWidgets.QWidget
    :param button: The mouse button to release.
    :type button: Qt.MouseButton
    :param modifier: Any keyboard modifier keys that should affect the mouse
        release.
    :type modifier: Qt.KeyboardModifier
    :param pos: Where to release the mouse button.  If `None`, the center of
        `widget` will be used.
    :type pos: QtCore.QPoint
    """
    pos = _update_pos(widget, pos)
    QME = QtGui.QMouseEvent
    event = QME(QME.MouseButtonRelease, pos, button, button, modifier)
    widget.mouseReleaseEvent(event) 
[docs]def mouse_click(widget, button=Qt.LeftButton, modifier=Qt.NoModifier, pos=None):
    """
    Click a mouse button on a widget.
    :param widget: The widget to click the mouse on
    :type widget: QtWidgets.QWidget
    :param button: The mouse button to click
    :type button: Qt.MouseButton
    :param modifier: Any keyboard modifier keys that should affect the mouse
        click.
    :type modifier: Qt.KeyboardModifier
    :param pos: Where to click the mouse button.  If `None`, the center of
        `widget` will be used.
    :type pos: QtCore.QPoint
    """
    pos = _update_pos(widget, pos)
    mouse_press(widget, button, modifier, pos)
    mouse_release(widget, button, modifier, pos) 
[docs]def mouse_double_click(widget,
                       button=Qt.LeftButton,
                       modifier=Qt.NoModifier,
                       pos=None):
    """
    Double-click a mouse button on a widget.  Note that the widget will receive
    mouse press and release events for both clicks in addition to the double-
    click event.
    :param widget: The widget to click the mouse on
    :type widget: QtWidgets.QWidget
    :param button: The mouse button to click
    :type button: Qt.MouseButton
    :param modifier: Any keyboard modifier keys that should affect the mouse
        clicks.
    :type modifier: Qt.KeyboardModifier
    :param pos: Where to click the mouse button.  If `None`, the center of
        `widget` will be used.
    :type pos: QtCore.QPoint
    """
    pos = _update_pos(widget, pos)
    mouse_press(widget, button, modifier, pos)
    mouse_release(widget, button, modifier, pos)
    QME = QtGui.QMouseEvent
    event = QME(QME.MouseButtonDblClick, pos, button, button, modifier)
    widget.mouseDoubleClickEvent(event)
    # mouseDoubleClickEvent will trigger a call to mousePressEvent, but we still
    # need to manually call mouseReleaseEvent to mimic the release of the second
    # click
    mouse_release(widget, button, modifier, pos) 
[docs]def mouse_click_on_cell(view,
                        row=0,
                        column=0,
                        button=Qt.LeftButton,
                        modifier=Qt.NoModifier):
    """
    Click a mouse button on the specified cell of a view.
    :param widget: The view to click the mouse on
    :type widget: QtWidgets.QAbstractItemView
    :param row: The row to click on.
    :type row: int
    :param column: The column to click on.
    :type column: int
    :param button: The mouse button to click
    :type button: Qt.MouseButton
    :param modifier: Any keyboard modifier keys that should affect the mouse
        click.
    :type modifier: Qt.KeyboardModifier
    """
    model = view.model()
    index = model.index(row, column)
    mouse_click_on_index(view, index, button, modifier) 
[docs]def mouse_double_click_on_cell(view,
                               row=0,
                               column=0,
                               button=Qt.LeftButton,
                               modifier=Qt.NoModifier):
    """
    Double-click a mouse button on the specified cell of a view.
    :param widget: The view to click the mouse on
    :type widget: QtWidgets.QAbstractItemView
    :param row: The row to click on.
    :type row: int
    :param column: The column to click on.
    :type column: int
    :param button: The mouse button to click
    :type button: Qt.MouseButton
    :param modifier: Any keyboard modifier keys that should affect the mouse
        click.
    :type modifier: Qt.KeyboardModifier
    """
    model = view.model()
    index = model.index(row, column)
    mouse_double_click_on_index(view, index, button, modifier) 
[docs]def mouse_click_on_index(view,
                         index,
                         button=Qt.LeftButton,
                         modifier=Qt.NoModifier):
    """
    Click a mouse button on the specified index of a view.
    :param widget: The view to click the mouse on
    :type widget: QtWidgets.QAbstractItemView
    :param index: The index to click on.
    :type index: QtCore.QModelIndex
    :param button: The mouse button to click
    :type button: Qt.MouseButton
    :param modifier: Any keyboard modifier keys that should affect the mouse
        click.
    :type modifier: Qt.KeyboardModifier
    """
    rect = view.visualRect(index)
    pos = rect.center()
    mouse_click(view, button, modifier, pos) 
[docs]def mouse_double_click_on_index(view,
                                index,
                                button=Qt.LeftButton,
                                modifier=Qt.NoModifier):
    """
    Double-click a mouse button on the specified index of a view.
    :param widget: The view to click the mouse on
    :type widget: QtWidgets.QAbstractItemView
    :param index: The index to click on.
    :type index: QtCore.QModelIndex
    :param button: The mouse button to click
    :type button: Qt.MouseButton
    :param modifier: Any keyboard modifier keys that should affect the mouse
        click.
    :type modifier: Qt.KeyboardModifier
    """
    rect = view.visualRect(index)
    pos = rect.center()
    mouse_double_click(view, button, modifier, pos) 
[docs]def mouse_move(widget, pos=None, buttons=Qt.NoButton, modifier=Qt.NoModifier):
    """
    Move the mouse over a widget.
    :param widget: The widget to move the mouse over.
    :type widget: QtWidgets.QWidget
    :param pos: Where to move the mouse to.  If `None`, the center of
        `widget` will be used.
    :type pos: QtCore.QPoint
    :param buttons: Any mouse buttons to be held down during the move.
    :type buttons: Qt.MouseButtons
    :param modifier: Any keyboard modifier keys to be held down during the move.
    :type modifier: Qt.KeyboardModifier
    """
    pos = _update_pos(widget, pos)
    QME = QtGui.QMouseEvent
    event = QME(QME.MouseMove, pos, Qt.NoButton, buttons, modifier)
    widget.mouseMoveEvent(event) 
[docs]def mouse_drag(widget,
               from_pos,
               to_pos,
               button=Qt.LeftButton,
               modifier=Qt.NoModifier):
    """
    Click and drag the mouse on a widget.
    :param widget: The widget to click and drag on.
    :type widget: QtWidgets.QWidget
    :param from_pos: Where to press the mouse.  If `None`, the center of
        `widget` will be used.
    :type from_pos: QtCore.QPoint
    :param to_pos: Where to release the mouse.  If `None`, the center of
        `widget` will be used.
    :type to_pos: QtCore.QPoint
    :param button: The mouse button to click
    :type button: Qt.MouseButton
    :param modifier: Any keyboard modifier keys that should affect the click
        and drag.
    :type modifier: Qt.KeyboardModifier
    """
    mouse_press(widget, button, modifier, from_pos)
    mouse_move(widget, to_pos, button, modifier)
    mouse_release(widget, button, modifier, to_pos) 
[docs]def mouse_drag_indices(view,
                       from_index,
                       to_index,
                       button=Qt.LeftButton,
                       modifier=Qt.NoModifier):
    """
    Click and drag the mouse on a widget.
    :param widget: The widget to click and drag on.
    :type widget: QtWidgets.QWidget
    :param from_index: The index to press the mouse on.
    :type from_index: QtCore.QModelIndex
    :param to_index: The index to release the mouse on.
    :type to_index: QtCore.QModelIndex
    :param button: The mouse button to click
    :type button: Qt.MouseButton
    :param modifier: Any keyboard modifier keys that should affect the click
        and drag.
    :type modifier: Qt.KeyboardModifier
    """
    from_pos = view.visualRect(from_index).center()
    to_pos = view.visualRect(to_index).center()
    mouse_drag(view, from_pos, to_pos, button, modifier) 
def _get_key_event(event_type, key, modifier):
    """
    Create and return a QKeyEvent.
    :param event_type: The type of event to create.
    :type event_type: QtCore.QEvent.Type
    :param key: The key for the event.
    :type key: str or Qt.Key
    :param modifier: Any keyboard modifiers that should affect the event.
    :type modifier: Qt.KeyboardModifier
    :return: The newly-created event.
    :rtype: QtGui.QKeyEvent
    """
    ERR = "key must be a single-character string or a Qt.Key value"
    if isinstance(key, str):
        text = key
        if not len(text) == 1:
            raise ValueError(ERR)
        key = Qt.Key(QtGui.QKeySequence(text)[0])
        if modifier & Qt.ShiftModifier:
            text = text.upper()
        elif text.isupper():
            modifier |= Qt.ShiftModifier
    elif isinstance(key, Qt.Key):
        text = QtGui.QKeySequence(key).toString()
        if not modifier & Qt.ShiftModifier:
            text = text.lower()
    else:
        raise TypeError(ERR)
    return QtGui.QKeyEvent(event_type, key, modifier, text)
[docs]def key_press(widget, key, modifier=Qt.NoModifier):
    """
    Send a key press to the specified widget.  (Note that this function only
    presses the key, it doesn't release it to complete a click.  Use `key_click`
    instead if you want a complete key click.)
    :param widget: The widget to send the key press to.
    :type widget: QtWidgets.QWidget
    :param key: The key for the event.
    :type key: str or Qt.Key
    :param modifier: Any keyboard modifiers that should affect the event.
    :type modifier: Qt.KeyboardModifier
    """
    event = _get_key_event(QtGui.QKeyEvent.KeyPress, key, modifier)
    widget.keyPressEvent(event) 
[docs]def key_release(widget, key, modifier=Qt.NoModifier):
    """
    Send a key release to the specified widget.
    :param widget: The widget to send the key release to.
    :type widget: QtWidgets.QWidget
    :param key: The key for the event.
    :type key: str or Qt.Key
    :param modifier: Any keyboard modifiers that should affect the event.
    :type modifier: Qt.KeyboardModifier
    """
    event = _get_key_event(QtGui.QKeyEvent.KeyRelease, key, modifier)
    widget.keyReleaseEvent(event) 
[docs]def key_click(widget, key, modifier=Qt.NoModifier):
    """
    Send a key click to the specified widget.
    :param widget: The widget to send the key click to.
    :type widget: QtWidgets.QWidget
    :param key: The key for the event.
    :type key: str or Qt.Key
    :param modifier: Any keyboard modifiers that should affect the event.
    :type modifier: Qt.KeyboardModifier
    """
    key_press(widget, key, modifier)
    key_release(widget, key, modifier) 
[docs]def key_clicks(widget, keys, modifier=Qt.NoModifier):
    """
    Send multiple key clicks to the specified widget.
    :param widget: The widget to send the key clicks to.
    :type widget: QtWidgets.QWidget
    :param key: The keys to send.
    :type key: Iterable(str or Qt.Key)
    :param modifier: Any keyboard modifiers that should affect the events.
    :type modifier: Qt.KeyboardModifier
    """
    for cur_key in keys:
        key_click(widget, cur_key, modifier)