session: provide a way for library user to properly cleanup the event loop

Also close the loop automatically on connection error.

Author: bfredl

neovim/__init__.py

@@ -96,6 +96,18 @@ def attach(session_type, address=None, port=None,
         nvim = attach('socket', path=<path>)
         nvim = attach('child', argv=<argv>)
         nvim = attach('stdio')
+
+    When the session is not needed anymore, it is recommended to explicitly
+    close it:
+       nvim.close()
+    It is also possible to use the session as a context mangager:
+       with attach('socket', path=thepath) as nvim:
+           print(nvim.funcs.getpid())
+           print(nvim.current.line)
+    This will automatically close the session when you're done with it, or
+    when an error occured.
+
+
     """
     session = (tcp_session(address, port) if session_type == 'tcp' else
                socket_session(path) if session_type == 'socket' else

neovim/api/nvim.py

@@ -1,7 +1,7 @@
 """Main Nvim interface."""
-import functools
 import os
 import sys
+from functools import partial
 from traceback import format_stack
 
 from msgpack import ExtType
@@ -180,6 +180,21 @@ def stop_loop(self):
         """Stop the event loop being started with `run_loop`."""
         self._session.stop()
 
+    def close(self):
+        """Close the nvim session and release its resources."""
+        self._session.close()
+
+    def __enter__(self):
+        """"Enter nvim session as a context manager."""
+        return self
+
+    def __exit__(self, *exc_info):
+        """"Exict nvim session as context manager.
+
+        Closes the event loop.
+        """
+        self.close()
+
     def with_decode(self, decode=True):
         """Initialize a new Nvim instance."""
         return Nvim(self._session, self.channel_id,
@@ -439,7 +454,7 @@ def __init__(self, nvim):
         self._nvim = nvim
 
     def __getattr__(self, name):
-        return functools.partial(self._nvim.call, name)
+        return partial(self._nvim.call, name)
 
 
 class NvimError(Exception):

neovim/msgpack_rpc/async_session.py

@@ -70,6 +70,10 @@ def stop(self):
         """Stop the event loop."""
         self._msgpack_stream.stop()
 
+    def close(self):
+        """Close the event loop."""
+        self._msgpack_stream.close()
+
     def _on_message(self, msg):
         try:
             self._handlers.get(msg[0], self._on_invalid_message)(msg)

neovim/msgpack_rpc/event_loop/asyncio.py

@@ -39,6 +39,7 @@ class AsyncioEventLoop(BaseEventLoop, asyncio.Protocol,
     def connection_made(self, transport):
         """Used to signal `asyncio.Protocol` of a successful connection."""
         self._transport = transport
+        self._raw_transport = transport
         if isinstance(transport, asyncio.SubprocessTransport):
             self._transport = transport.get_pipe_transport(0)
 
@@ -74,6 +75,7 @@ def _init(self):
         self._loop = loop_cls()
         self._queued_data = deque()
         self._fact = lambda: self
+        self._raw_transport = None
 
     def _connect_tcp(self, address, port):
         coroutine = self._loop.create_connection(self._fact, address, port)
@@ -112,6 +114,11 @@ def _run(self):
     def _stop(self):
         self._loop.stop()
 
+    def _close(self):
+        if self._raw_transport is not None:
+            self._raw_transport.close()
+        self._loop.close()
+
     def _threadsafe_call(self, fn):
         self._loop.call_soon_threadsafe(fn)
 

neovim/msgpack_rpc/event_loop/base.py

@@ -85,7 +85,11 @@ def __init__(self, transport_type, *args):
         self._on_data = None
         self._error = None
         self._init()
-        getattr(self, '_connect_{}'.format(transport_type))(*args)
+        try:
+            getattr(self, '_connect_{}'.format(transport_type))(*args)
+        except Exception as e:
+            self.close()
+            raise e
         self._start_reading()
 
     def connect_tcp(self, address, port):
@@ -148,6 +152,11 @@ def stop(self):
         self._stop()
         debug('Stopped event loop')
 
+    def close(self):
+        """Stop the event loop."""
+        self._close()
+        debug('Closed event loop')
+
     def _on_signal(self, signum):
         msg = 'Received {}'.format(self._signames[signum])
         debug(msg)

neovim/msgpack_rpc/event_loop/uv.py

@@ -97,6 +97,9 @@ def _run(self):
     def _stop(self):
         self._loop.stop()
 
+    def _close(self):
+        pass
+
     def _threadsafe_call(self, fn):
         self._callbacks.append(fn)
         self._async.send()

neovim/msgpack_rpc/msgpack_stream.py

@@ -48,6 +48,10 @@ def stop(self):
         """Stop the event loop."""
         self._event_loop.stop()
 
+    def close(self):
+        """Close the event loop."""
+        self._event_loop.close()
+
     def _on_data(self, data):
         self._unpacker.feed(data)
         while True:

neovim/msgpack_rpc/session.py

@@ -140,6 +140,10 @@ def stop(self):
         """Stop the event loop."""
         self._async_session.stop()
 
+    def close(self):
+        """Close the event loop."""
+        self._async_session.close()
+
     def _yielding_request(self, method, args):
         gr = greenlet.getcurrent()
         parent = gr.parent