fix type annotations for TokenAuth (#310)

Author: dantownsend

piccolo_api/token_auth/endpoints.py

@@ -33,11 +33,11 @@ async def get_token(self, username: str, password: str) -> t.Optional[str]:
         if user:
             response = (
                 await TokenAuth.select(TokenAuth.token)
-                .first()
                 .where(TokenAuth.user == user)
-                .run()
+                .first()
             )
-            return response["token"]
+            if response:
+                return response["token"]
 
         return None
 

piccolo_api/token_auth/tables.py

@@ -8,9 +8,6 @@
 from piccolo.table import Table
 from piccolo.utils.sync import run_sync
 
-if t.TYPE_CHECKING:  # pragma: no cover
-    from piccolo.query.methods.select import First
-
 
 def generate_token() -> str:
     return str(uuid.uuid4())
@@ -56,23 +53,17 @@ def create_token_sync(cls, user_id: int) -> str:
         return run_sync(cls.create_token(user_id))
 
     @classmethod
-    async def authenticate(cls, token: str) -> First:
-        return cls.select(cls.user.id).where(cls.token == token).first()
+    async def authenticate(cls, token: str) -> t.Optional[t.Dict]:
+        return await cls.select(cls.user).where(cls.token == token).first()
 
     @classmethod
-    async def authenticate_sync(cls, token: str) -> First:
+    def authenticate_sync(cls, token: str) -> t.Optional[t.Dict]:
         return run_sync(cls.authenticate(token))
 
     @classmethod
     async def get_user_id(cls, token: str) -> t.Optional[int]:
         """
         Returns the user_id if the given token is valid, otherwise None.
         """
-        data = (
-            await cls.select(cls.user)
-            .where(cls.token == token)
-            .output(as_list=True)
-            .first()
-            .run()
-        )
+        data = await cls.authenticate(token=token)
         return data.get("user", None) if data else None