Refactor connection/session context a bit

Add helper functions to trio_cdp.context. These can be used in the main
module but are also helpful in downstream projects where sessions might
be reused across multiple tasks.

Author: mehaase

trio_cdp/__init__.py

@@ -239,11 +239,8 @@ async def open_session(self, target_id: cdp.target.TargetID) -> \
         and it will execute on the current session automatically.
         '''
         session = await self.connect_session(target_id)
-        token = session_context.set(session)
-        try:
+        with session_context(session):
             yield session
-        finally:
-            session_context.reset(token)
 
     async def connect_session(self, target_id: cdp.target.TargetID) -> 'CdpSession':
         '''
@@ -369,11 +366,10 @@ async def open_cdp(url) -> typing.AsyncIterator[CdpConnection]:
     '''
     async with trio.open_nursery() as nursery:
         conn = await connect_cdp(nursery, url)
-        token = connection_context.set(conn)
         try:
-            yield conn
+            with connection_context(conn):
+                yield conn
         finally:
-            connection_context.reset(token)
             await conn.aclose()
 
 

trio_cdp/context.py

@@ -1,8 +1,11 @@
+from contextlib import contextmanager
 import contextvars
 
 
-connection_context: contextvars.ContextVar = contextvars.ContextVar('connection_context')
-session_context: contextvars.ContextVar = contextvars.ContextVar('session_context')
+_connection_context: contextvars.ContextVar = contextvars.ContextVar('connection_context')
+_session_context: contextvars.ContextVar = contextvars.ContextVar('session_context')
+
+
 
 
 def get_connection_context(fn_name):
@@ -11,7 +14,7 @@ def get_connection_context(fn_name):
     ``RuntimeError`` with a helpful message.
     '''
     try:
-        return connection_context.get()
+        return _connection_context.get()
     except LookupError:
         raise RuntimeError(f'{fn_name}() must be called in a connection context.')
 
@@ -22,19 +25,41 @@ def get_session_context(fn_name):
     ``RuntimeError`` with a helpful message.
     '''
     try:
-        return session_context.get()
+        return _session_context.get()
     except LookupError:
         raise RuntimeError(f'{fn_name}() must be called in a session context.')
 
 
+@contextmanager
+def connection_context(connection):
+    ''' This context manager installs ``connection`` as the session context for the current
+    Trio task. '''
+    token = _connection_context.set(connection)
+    try:
+        yield
+    finally:
+        _connection_context.reset(token)
+
+
+@contextmanager
+def session_context(session):
+    ''' This context manager installs ``session`` as the session context for the current
+    Trio task. '''
+    token = _session_context.set(session)
+    try:
+        yield
+    finally:
+        _session_context.reset(token)
+
+
 def set_global_connection(connection):
     '''
     Install ``connection`` in the root context so that it will become the default
     connection for all tasks. This is generally not recommended, except it may be
     necessary in certain use cases such as running inside Jupyter notebook.
     '''
-    global connection_context
-    connection_context = contextvars.ContextVar('connection_context',
+    global _connection_context
+    _connection_context = contextvars.ContextVar('_connection_context',
         default=connection)
 
 
@@ -44,5 +69,5 @@ def set_global_session(session):
     session for all tasks. This is generally not recommended, except it may be
     necessary in certain use cases such as running inside Jupyter notebook.
     '''
-    global session_context
-    session_context = contextvars.ContextVar('session_context', default=session)
+    global _session_context
+    _session_context = contextvars.ContextVar('_session_context', default=session)