Source code for schrodinger.utils.thread_utils

"""
Simplifies keeping track of threads, mainly for debugging. Allows threads to be
named and provides a print function that prepends which info on what thread it
is being called from. To use, either instantiate the class or use the global
singleton via the module-level functions.

"""
from schrodinger.Qt import QtCore


class ThreadTool:

    def __init__(self):
        self.thread_ids = {}

    def getThreadId(self, name=None):
        """
        Gets an integer id for a thread. If no name is provided, gets the id of
        the thread that is calling this method.

        :param name: a previously set name for a thread
        :type name: str

        :return: a thread id
        :rtype: int
        """
        if name is None:
            return int(QtCore.QThread.currentThreadId())
        for tid, tname in self.thread_ids.items():
            if name == tname:
                return tid
        raise ValueError(f'No thread named {name}')

    def getThreadName(self, tid=None):
        """
        Gets the name of a thread if it has been previously set, otherwise just
        return a string representation of the thread id.

        If no thread id is specified, use the id of the thread that is calling
        this method.

        :param tid: the id of the thread for which to retrieve the name
        :type tid: int or None

        :return: the assigned name for the thread
        :rtype: str
        """
        if tid is None:
            tid = self.getThreadId()
        return self.thread_ids.get(tid, str(tid))

    def setThreadName(self, name, tid=None):
        """
        Set a user-friendly name for a thread.

        If no thread id is specified, use the id of the thread that is calling
        this method.

        :param name: the name to assign (e.g. "main", "worker", "gui" etc.)
        :type name: str

        :param tid: the id of the thread for which to set a name
        :type tid: int or None
        """
        if tid is None:
            tid = self.getThreadId()

        if tid in self.thread_ids:
            print(f'Renaming thread {tid}: {name} -> {self.getThreadName(tid)}')
        self.thread_ids[tid] = name

    def threadPrint(self, *args):
        """
        Use in place of a normal print to prepend information about which thread
        is calling this print function.
        """
        print(f'<Thread: {self.getThreadName()}>\t', *args)


#===============================================================================
# Global functions - see corresponding ThreadTool methods for docs
#===============================================================================


def get_thread_id(name=None):
    return _thread_tool.getThreadId(name)


def get_thread_name(tid=None):
    return _thread_tool.getThreadName(tid)


def set_thread_name(name, tid=None):
    return _thread_tool.setThreadName(name, tid)


def thread_print(*args):
    _thread_tool.threadPrint(*args)


def in_main_thread():
    """
    Returns True when called from the main application thread. Otherwise returns
    False.
    """
    return (QtCore.QThread.currentThread() ==
            QtCore.QCoreApplication.instance().thread())


_thread_tool = ThreadTool()