deal with optioanal async type

mypy understands `isinstance(obj, Awaitable)`, but not `isawaitable(obj)`

Author: minrk

jupyter_server/auth/identity.py

@@ -13,7 +13,7 @@
 import sys
 import uuid
 from dataclasses import asdict, dataclass
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, Awaitable, cast
 
 from tornado import web
 from traitlets import Bool, Dict, Type, Unicode, default
@@ -175,7 +175,7 @@ def _token_default(self):
 
     need_token = Bool(True)
 
-    def get_user(self, handler: JupyterHandler) -> User | None:
+    async def get_user(self, handler: JupyterHandler) -> User | None:
         """Get the authenticated user for a request
 
         Must return a :class:`.jupyter_server.auth.User`,
@@ -188,11 +188,18 @@ def get_user(self, handler: JupyterHandler) -> User | None:
         if getattr(handler, "_jupyter_current_user", None):
             # already authenticated
             return handler._jupyter_current_user
-        token_user = self.get_user_token(handler)
-        cookie_user = self.get_user_cookie(handler)
+        _token_user: User | None | Awaitable[User | None] = self.get_user_token(handler)
+        if isinstance(_token_user, Awaitable):
+            _token_user = await _token_user
+        token_user: User | None = _token_user  # need second variable name to collapse type
+        _cookie_user = self.get_user_cookie(handler)
+        if isinstance(_cookie_user, Awaitable):
+            _cookie_user = await _cookie_user
+        cookie_user: User | None = _cookie_user
         # prefer token to cookie if both given,
         # because token is always explicit
         user = token_user or cookie_user
+
         if token_user:
             # if token-authenticated, persist user_id in cookie
             # if it hasn't already been stored there
@@ -258,16 +265,16 @@ def set_login_cookie(self, handler: JupyterHandler, user: User) -> None:
         cookie_options.setdefault("path", handler.base_url)
         handler.set_secure_cookie(handler.cookie_name, self.user_to_cookie(user), **cookie_options)
 
-    def get_user_cookie(self, handler: JupyterHandler) -> User | None:
+    def get_user_cookie(self, handler: JupyterHandler) -> User | None | Awaitable[User | None]:
         """Get user from a cookie
 
         Calls user_from_cookie to deserialize cookie value
         """
         get_secure_cookie_kwargs = handler.settings.get("get_secure_cookie_kwargs", {})
-        user_cookie = handler.get_secure_cookie(handler.cookie_name, **get_secure_cookie_kwargs)
-        if not user_cookie:
+        _user_cookie = handler.get_secure_cookie(handler.cookie_name, **get_secure_cookie_kwargs)
+        if not _user_cookie:
             return None
-        user_cookie = user_cookie.decode()
+        user_cookie = _user_cookie.decode()
         # TODO: try/catch in case of change in config?
         try:
             return self.user_from_cookie(user_cookie)
@@ -296,7 +303,7 @@ def get_token(self, handler: JupyterHandler) -> str | None:
                 user_token = m.group(2)
         return user_token
 
-    def get_user_token(self, handler: JupyterHandler) -> User | None:
+    async def get_user_token(self, handler: JupyterHandler) -> User | None:
         """Identify the user based on a token in the URL or Authorization header
 
         Returns:
@@ -321,7 +328,10 @@ def get_user_token(self, handler: JupyterHandler) -> User | None:
             # token does not correspond to user-id,
             # which is stored in a cookie.
             # still check the cookie for the user id
-            user = self.get_user_cookie(handler)
+            _user = self.get_user_cookie(handler)
+            if isinstance(_user, Awaitable):
+                _user = await _user
+            user: User | None = _user
             if user is None:
                 user = self.generate_anonymous_user(handler)
             return user

jupyter_server/base/handlers.py

@@ -1,6 +1,8 @@
 """Base Tornado handlers for the Jupyter server."""
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
+from __future__ import annotations
+
 import datetime
 import functools
 import inspect
@@ -14,6 +16,7 @@
 import warnings
 from http.client import responses
 from http.cookies import Morsel
+from typing import TYPE_CHECKING, Awaitable
 from urllib.parse import urlparse
 
 import prometheus_client
@@ -37,6 +40,9 @@
     urldecode_unix_socket_path,
 )
 
+if TYPE_CHECKING:
+    from jupyter_server.auth.identity import User
+
 # -----------------------------------------------------------------------------
 # Top-level handlers
 # -----------------------------------------------------------------------------
@@ -581,6 +587,7 @@ async def prepare(self):
 
         mod_obj = inspect.getmodule(self.get_current_user)
         assert mod_obj is not None
+        user: User | None = None
 
         if type(self.identity_provider) is IdentityProvider and mod_obj.__name__ != __name__:
             # check for overridden get_current_user + default IdentityProvider
@@ -594,10 +601,11 @@ async def prepare(self):
             )
             user = self.get_current_user()
         else:
-            user = self.identity_provider.get_user(self)
-            if inspect.isawaitable(user):
+            _user = self.identity_provider.get_user(self)
+            if isinstance(_user, Awaitable):
                 # IdentityProvider.get_user _may_ be async
-                user = await user
+                _user = await _user
+            user = _user
 
         # self.current_user for tornado's @web.authenticated
         # self._jupyter_current_user for backward-compat in deprecated get_current_user calls