"""
Utility classes and functions for use with Qt objects (such as QObjects).
Note that contrary to schrodinger.ui.qt.utils, these utilities do not rely on
QtGui or QtWidgets. This allows these utilities to be used on headless servers
which shouldn't import QtGui or QtWidgets.
"""
import abc
import decorator
import sys
import traceback
from schrodinger.Qt import QtCore
LAST_EXCEPTION = None
[docs]class suppress_signals:
    """
    A context manager to prevent signals from being emitted from the specified
    widget(s).  All widgets to be suppressed should be passed as arguments.
    """
[docs]    def __init__(self, *args, suppress=True):
        """
        Create a suppress_signals instance to suppress signals for all the
        widgets given as arguments.
        :param suppress bool: If True, suppress signals. If False, don't.
        Allows constructs such as ::
            with suppress_signal(mywidget, suppress=self.resetting)
        """
        if suppress:
            self._widgets = args
        else:
            self._widgets = []
        self._block_status = [] 
    def __enter__(self):
        for cur_widget in self._widgets:
            prev_val = cur_widget.blockSignals(True)
            self._block_status.append(prev_val)
    def __exit__(self, *args):
        for cur_widget in self._widgets:
            cur_widget.blockSignals(self._block_status.pop(0)) 
[docs]class SignalTimeoutException(RuntimeError):
    pass 
[docs]def wait_for_signal(signal, timeout=None):
    """
    Uses an event loop to wait until a signal is emitted. If the signal is not
    emitted within a specified timeout duration, a SignalTimeoutException is
    raised.
    :param signal: the signal to wait for
    :param timeout: number of seconds to wait for the signal before timing out
    :type timeout: float
    :return: the args emitted with the signal, if any. If there is only one
        arg emitted, it will be returned directly. If there are more, they will
        be return as a tuple.
    """
    event_loop = QtCore.QEventLoop()
    signal_args = None
    if timeout is not None:
        QtCore.QTimer.singleShot(timeout * 1000, event_loop.quit)
    def on_signal(*args):
        nonlocal signal_args
        signal_args = args
        event_loop.quit()
    signal.connect(on_signal)
    event_loop.exec()
    if signal_args is None:
        raise SignalTimeoutException(f'Timeout while waiting for {signal}')
    if len(signal_args) == 0:
        return
    if len(signal_args) == 1:
        return signal_args[0]
    return signal_args 
[docs]def call_func_and_wait_for_signal(func, signal, timeout=None):
    """
    Calls the specified function and then waits for the signal. The function
    is called in such a way that the signal is guaranteed  to be caught, even
    if the signal is emitted before the function returns.
    :param func: the function to call
    See wait_for_signal for additional parameter documentation.
    """
    QtCore.QTimer.singleShot(0, func)
    return wait_for_signal(signal, timeout=timeout) 
[docs]def get_signals(source):
    """
    Utility method for iterating through the signals on a QObject.
    :param source: Any object or class with signals
    :type source: Type[QtCore.QObject] or QtCore.QObject
    :return: A dictionary of {name: signal}
    :rtype: dict[str, QtCore.pyqtSignal]
    """
    cls = source if isinstance(source, type) else type(source)
    signal = QtCore.pyqtSignal
    filtered_names = (name for name in dir(source) if name != "destroyed")
    names = (name for name in filtered_names
             if isinstance(getattr(cls, name, None), signal))
    return {name: getattr(source, name) for name in names} 
[docs]class SignalAndSlot:
    """
    A composite object to manage a single signal/slot pair.  Usage::
        class ClassName(QtWidgets.QWidget):
            fooChangedSignal = QtCore.pyqtSignal()
            def __init__(self, parent=None):
                super(ClassName, self).__init__(parent)
                self.fooChanged = qt_utils.SignalAndSlot(self.fooChangedSignal,
                                                         self.fooChangedSlot)
            def fooChangedSlot(self):
                pass
    """
[docs]    def __init__(self, signal, slot):
        """
        Create an object that acts as both a signal and a slot
        :param signal: The signal object
        :type signal: `PyQt5.QtCore.pyqtSignal`
        :param slot: The slot object
        :type slot: function
        """
        self.signal = signal
        self.slot = slot 
    def __call__(self, *args, **kwargs):
        return self.slot(*args, **kwargs)
[docs]    def emit(self, *args, **kwargs):
        self.signal.emit(*args, **kwargs) 
[docs]    def connect(self, *args, **kwargs):
        self.signal.connect(*args, **kwargs) 
[docs]    def disconnect(self, *args, **kwargs):
        self.signal.disconnect(*args, **kwargs) 
    def __getitem__(self, key):
        return self.signal[key] 
[docs]def add_enums_as_attributes(enum_):
    """
    A class decorator that takes in an enum and aliases its members on the
    decorated class. For example::
        Shape = enum.Enum('Shape', 'SQUARE TRIANGLE CIRCLE')
        @qt_utils.add_enums_as_attributes(Shape)
        class Foo:
            pass
        assert Foo.SQUARE is Shape.SQUARE
        assert Foo.TRIANGLE is Shape.TRIANGLE
        assert Foo.CIRCLE is Shape.CIRCLE
    """
    def cls_decorator(cls):
        for member in enum_:
            setattr(cls, member.name, member)
        return cls
    return cls_decorator 
[docs]@decorator.decorator
def exit_event_loop_on_exception(func, *args, **kwargs):
    """
    Decorates a function that passes an event_loop keyword so if func throws an
    exception, the event loop will exit. The exception is accesible in
    get_last_exception. Example usage::
        @exit_event_loop_on_exception
        def slot(event_loop=None):
            ...
            event_loop.quit()
        event_loop = schrodinger.QtCore.QEventLoop()
        timer = schrodinger.QtCore.QTimer()
        timer.timeout.connect(functools.partial(event_loop=event_loop))
        timer.start(1)
        event_loop.exec()
        exc = get_last_exception()
        if exc:
            raise exc
    """
    global LAST_EXCEPTION
    if "event_loop" not in kwargs:
        # must exit because this is called from a slot and we won't see this
        # raised as an exception
        print(f"event_loop is not a kwarg of {func} {kwargs}")
        sys.exit(1)
    event_loop = kwargs["event_loop"]
    try:
        return func(*args, **kwargs)
    except Exception as e:
        LAST_EXCEPTION = e
        event_loop.quit() 
[docs]def get_last_exception():
    """
    Returns an exception if one was thrown previously in
    exit_event_loop_on_exception. Returns None if no exception was thrown.
    Calling this function resets the exception state.
    """
    global LAST_EXCEPTION
    exc = LAST_EXCEPTION
    LAST_EXCEPTION = None
    return exc 
[docs]class EventLoop(QtCore.QEventLoop):
    """
    A modified QEventLoop that catches exceptions that occur in any slot while
    the event loop is running, stores that exception and optionally exits the
    event loop and/or re-raises that exception from EventLoop.exec() call.
    """
[docs]    def __init__(self,
                 parent=None,
                 exit_on_exception=True,
                 reraise_exception=True,
                 timeout=None):
        """
        :param exit_on_exception: Whether to exit the running event loop if an
            exception is raised in a slot
        :param reraise_exception: Whether to reraise the last detected exception
            from the exec() method (making it catchable in the calling code).
        :param timeout: if specified, the event loop will exit after this many
            seconds. This is useful as a failsafe so the event loop doesn't hang
            indefinitely, especially in unit tests
        """
        super().__init__(parent)
        self.exception_info = None
        self._exit_on_exception = exit_on_exception
        self._reraise_exception = reraise_exception
        self._original_excepthook = None
        self._timeout = timeout
        self._timed_out = False 
    def _exceptHook(self, typ, value, tb):
        if self.isRunning():
            self.handleException(typ, value, tb)
        else:
            self._original_excepthook(typ, value, tb)
    def _timeoutSlot(self):
        self.exit()
        self._timed_out = True
[docs]    def exec(self, *args, **kwargs):
        self._timed_out = False
        self.exception_info = None
        self._original_excepthook = sys.excepthook
        if self._timeout is not None:
            timer = QtCore.QTimer()
            timer.setSingleShot(True)
            timer.setInterval(self._timeout * 1000)
            timer.timeout.connect(self._timeoutSlot)
            timer.start()
        try:
            sys.excepthook = self._exceptHook
            super().exec(*args, **kwargs)
        finally:
            sys.excepthook = self._original_excepthook
            if self._timeout is not None:
                timer.stop()
        if self.exception_info and self._reraise_exception:
            raise self.exception_info[1]
        if self._timed_out:
            raise RuntimeError(
                f'{self} timed out after {self._timeout} seconds.') 
[docs]    def handleException(self, typ, value, tb):
        self.exception_info = (typ, value, tb)
        if self._exit_on_exception:
            self.exit() 
[docs]    def printException(self):
        if self.exception_info:
            traceback.print_exception(*self.exception_info)