From c39b9e5bd557ac53b4e17001ffd8fccabd1ce50c Mon Sep 17 00:00:00 2001
From: hasier <hasier@users.noreply.github.com>
Date: Wed, 6 Jul 2022 18:06:08 +0100
Subject: [PATCH] Improve python-jose jws and jwt modules (#8238)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
---
 stubs/python-jose/jose/jws.pyi | 27 ++++++++++++++++++++-----
 stubs/python-jose/jose/jwt.pyi | 36 ++++++++++++++++++++++------------
 2 files changed, 45 insertions(+), 18 deletions(-)

diff --git a/stubs/python-jose/jose/jws.pyi b/stubs/python-jose/jose/jws.pyi
index 4703aee13..40e944a1a 100644
--- a/stubs/python-jose/jose/jws.pyi
+++ b/stubs/python-jose/jose/jws.pyi
@@ -1,7 +1,24 @@
+from collections.abc import Container, Mapping
 from typing import Any
 
-def sign(payload, key, headers: Any | None = ..., algorithm=...): ...
-def verify(token, key, algorithms, verify: bool = ...): ...
-def get_unverified_header(token): ...
-def get_unverified_headers(token): ...
-def get_unverified_claims(token): ...
+from .backends.base import Key
+
+def sign(
+    payload: str | Mapping[str, Any],
+    # Internally it's passed down to jwk.construct(), which explicitly checks for
+    # key as dict instance, instead of a Mapping
+    key: str | dict[str, Any] | Key,
+    headers: Mapping[str, Any] | None = ...,
+    algorithm: str = ...,
+) -> str: ...
+def verify(
+    token: str,
+    key: str | Mapping[str, Any] | Key,
+    # Callers of this function, like jwt.decode(), and functions called internally,
+    # like jws._verify_signature(), use and accept algorithms=None
+    algorithms: str | Container[str] | None,
+    verify: bool = ...,
+) -> str: ...
+def get_unverified_header(token: str) -> dict[str, Any]: ...
+def get_unverified_headers(token: str) -> dict[str, Any]: ...
+def get_unverified_claims(token: str) -> str: ...
diff --git a/stubs/python-jose/jose/jwt.pyi b/stubs/python-jose/jose/jwt.pyi
index 4778dcfcb..cf9dcbc74 100644
--- a/stubs/python-jose/jose/jwt.pyi
+++ b/stubs/python-jose/jose/jwt.pyi
@@ -1,16 +1,26 @@
+from collections.abc import Container, Iterable, Mapping, MutableMapping
 from typing import Any
 
-def encode(claims, key, algorithm=..., headers: Any | None = ..., access_token: Any | None = ...): ...
+from .backends.base import Key
+
+def encode(
+    claims: MutableMapping[str, Any],
+    # Internally it calls jws.sign() that expects a key dict instance instead of Mapping
+    key: str | dict[str, Any] | Key,
+    algorithm: str = ...,
+    headers: Mapping[str, Any] | None = ...,
+    access_token: str | None = ...,
+) -> str: ...
 def decode(
-    token,
-    key,
-    algorithms: Any | None = ...,
-    options: Any | None = ...,
-    audience: Any | None = ...,
-    issuer: Any | None = ...,
-    subject: Any | None = ...,
-    access_token: Any | None = ...,
-): ...
-def get_unverified_header(token): ...
-def get_unverified_headers(token): ...
-def get_unverified_claims(token): ...
+    token: str,
+    key: str | Mapping[str, Any] | Key,
+    algorithms: str | Container[str] | None = ...,
+    options: Mapping[str, Any] | None = ...,
+    audience: str | None = ...,
+    issuer: str | Iterable[str] | None = ...,
+    subject: str | None = ...,
+    access_token: str | None = ...,
+) -> dict[str, Any]: ...
+def get_unverified_header(token: str) -> dict[str, Any]: ...
+def get_unverified_headers(token: str) -> dict[str, Any]: ...
+def get_unverified_claims(token: str) -> dict[str, Any]: ...