NSRunLoop integration with Python asyncio

nsrunloop_asyncio.py

"""
NSRunLoop integration with Python asyncio.
Created on 2020-06-25 by David Lechner <[email protected]>
"""

import asyncio
import selectors

import objc
from Foundation import NSDate, NSDefaultRunLoopMode, NSFileHandle, NSRunLoop


class NSRunLoopSelector(selectors.KqueueSelector):
    """Extension of standard library KqueueSelector that also pumps NSRunLoop
    events.
    """
    def __init__(self):
        super().__init__()
        # Wrap the asyincio kqueue file descriptor in a NSFileHandle so that it
        # can communicate with the NSRunLoop later.
        self._handle = NSFileHandle.alloc().initWithFileDescriptor_closeOnDealloc_(
            self.fileno(), False)

    def select(self, timeout=None):
        timeout = float("inf") if timeout is None else timeout
        # This registers the asyncio kqueue file handle with the current run
        # loop.
        self._handle.waitForDataInBackgroundAndNotify()
        date = NSDate.alloc().initWithTimeIntervalSinceNow_(timeout)
        # Then we process NSRunLoop events until asyncio kqueue notifies
        # or timeout, whichever comes first.
        NSRunLoop.currentRunLoop().runMode_beforeDate_(NSDefaultRunLoopMode, date)
        # At this point, we have already blocked until notification or timeout
        # so we always call with timeout of 0 to prevent deadlock.
        return super().select(0)


class NSRunLoopEventLoopPolicy(asyncio.DefaultEventLoopPolicy):
    """Extension of standard library DefaultEventLoopPolicy that ensures new
    event loops use the NSRunLoopSelector.
    """
    def new_event_loop(self):
        return asyncio.SelectorEventLoop(NSRunLoopSelector())