Client: add session_callback for when a new session is created

...via createSession or refreshSession

Author: snarfed

README.md

@@ -234,7 +234,7 @@ Here's how to package, test, and ship a new release.
 * `Base`:
   * Expose lexicons in `defs` attribute.
 * `Client`:
-  * Add minimal auth support with `access_token` constructor kwarg and attribute. To send authenticated requests, call `createSession` or `refreshSession` to get an access token, then set it on a `Client`.
+  * Add minimal auth support with `access_token` and `refresh_token` constructor kwargs and `session` attribute. If you use a `Client` to call `com.atproto.server.createSession` or `com.atproto.server.refreshSession`, the returned tokens will be automatically stored and used in future requests.
   * Bug fix: handle trailing slash on server address, eg `http://ser.ver/` vs `http://ser.ver`.
   * Default server address to official `https://bsky.social` PDS.
   * Add default `User-Agent: lexrpc (https://lexrpc.readthedocs.io/)` request header.

lexrpc/client.py

@@ -68,14 +68,20 @@ class Client(Base):
     """
 
     def __init__(self, address=DEFAULT_PDS, access_token=None,
-                 refresh_token=None, headers=None, **kwargs):
+                 refresh_token=None, headers=None, session_callback=None,
+                 **kwargs):
         """Constructor.
 
         Args:
           address (str): base URL of XRPC server, eg ``https://bsky.social/``
           access_token (str): optional, will be sent in ``Authorization`` header
           refresh_token (str): optional; used to refresh access token
           headers (dict): optional, HTTP headers to include in every request
+          session_callback (callable, dict => None): called when a new session
+            is created with new access and refresh tokens. This callable is
+            passed one positional argument, the dict JSON output from
+            ``com.atproto.server.createSession`` or
+            ``com.atproto.server.refreshSession``.
           kwargs: passed through to :class:`Base`
 
         Raises:
@@ -93,6 +99,7 @@ def __init__(self, address=DEFAULT_PDS, access_token=None,
             'accessJwt': access_token,
             'refreshJwt': refresh_token,
         }
+        self.session_callback = session_callback
 
     def __getattr__(self, attr):
         if NSID_SEGMENT_RE.match(attr):
@@ -109,7 +116,9 @@ def call(self, nsid, input=None, **params):
           params: optional method parameters
 
         Returns:
-          dict or generator iterator: for queries and procedures, decoded JSON object, or None if the method has no output. For subscriptions, generator of messages from server.
+          dict or generator iterator: for queries and procedures, decoded JSON
+          object, or None if the method has no output. For subscriptions,
+          generator of messages from server.
 
         Raises:
           NotImplementedError: if the given NSID is not found in any of the
@@ -174,6 +183,8 @@ def call(self, nsid, input=None, **params):
         if output and nsid in (LOGIN_NSID, REFRESH_NSID):
             logger.info(f'Logged in as {output.get("did")}, storing session')
             self.session = output
+            if self.session_callback:
+                self.session_callback(output)
 
         self._maybe_validate(nsid, 'output', output)
         return output

lexrpc/tests/test_client.py

@@ -274,10 +274,17 @@ def test_refresh_token(self, mock_post, mock_get):
         ]
         mock_post.return_value = response(session)
 
-        client = Client(access_token='towkin', refresh_token='reephrush')
+        callback_got = []
+        def callback(session):
+            nonlocal callback_got
+            callback_got.append(session)
+
+        client = Client(access_token='towkin', refresh_token='reephrush',
+                        session_callback=callback)
         got = client.com.atproto.server.describeServer(x='y')
         self.assertEqual(output, got)
         self.assertEqual(session, client.session)
+        self.assertEqual([session], callback_got)
 
         mock_get.assert_any_call(
             'https://bsky.social/xrpc/com.atproto.server.describeServer?x=y',