@@ -25,10 +25,18 @@ def __init__(self, async_session):
2525 self ._request_cb = self ._notification_cb = None
2626 self ._pending_messages = deque ()
2727 self ._is_running = False
28+ self ._setup_exception = None
2829
29- def post (self , name , * args ):
30- """Simple wrapper around `AsyncSession.post`."""
31- self ._async_session .post (name , args )
30+ def threadsafe_call (self , fn , * args , ** kwargs ):
31+ """Wrapper around `AsyncSession.threadsafe_call`."""
32+ def handler ():
33+ fn (* args , ** kwargs )
34+
35+ def greenlet_wrapper ():
36+ gr = greenlet .greenlet (handler )
37+ gr .switch ()
38+
39+ self ._async_session .threadsafe_call (greenlet_wrapper )
3240
3341 def next_message (self ):
3442 """Block until a message(request or notification) is available.
@@ -42,7 +50,8 @@ def next_message(self):
4250 return self ._pending_messages .popleft ()
4351 self ._async_session .run (self ._enqueue_request_and_stop ,
4452 self ._enqueue_notification_and_stop )
45- return self ._pending_messages .popleft ()
53+ if self ._pending_messages :
54+ return self ._pending_messages .popleft ()
4655
4756 def request (self , method , * args ):
4857 """Send a msgpack-rpc request and block until as response is received.
@@ -67,7 +76,7 @@ def request(self, method, *args):
6776 raise self .error_wrapper (err )
6877 return rv
6978
70- def run (self , request_cb , notification_cb ):
79+ def run (self , request_cb , notification_cb , setup_cb = None ):
7180 """Run the event loop to receive requests and notifications from Nvim.
7281
7382 Like `AsyncSession.run()`, but `request_cb` and `notification_cb` are
@@ -76,6 +85,23 @@ def run(self, request_cb, notification_cb):
7685 self ._request_cb = request_cb
7786 self ._notification_cb = notification_cb
7887 self ._is_running = True
88+ self ._setup_exception = None
89+
90+ def on_setup ():
91+ try :
92+ setup_cb ()
93+ except Exception as e :
94+ self ._setup_exception = e
95+ self .stop ()
96+
97+ if setup_cb :
98+ # Create a new greenlet to handle the setup function
99+ gr = greenlet .greenlet (on_setup )
100+ gr .switch ()
101+
102+ if self ._setup_exception :
103+ raise self ._setup_exception
104+
79105 # Process all pending requests and notifications
80106 while self ._pending_messages :
81107 msg = self ._pending_messages .popleft ()
@@ -85,6 +111,9 @@ def run(self, request_cb, notification_cb):
85111 self ._request_cb = None
86112 self ._notification_cb = None
87113
114+ if self ._setup_exception :
115+ raise self ._setup_exception
116+
88117 def stop (self ):
89118 """Stop the event loop."""
90119 self ._async_session .stop ()
0 commit comments