jwcrypto: type most of the rest of JWT and JWKSet.generate function (#13807)

Author: jkrejcha

stubs/jwcrypto/jwcrypto/jwk.pyi

@@ -1,7 +1,7 @@
 from collections.abc import Callable, Sequence
 from enum import Enum
 from typing import Any, Literal, NamedTuple, TypeVar, overload
-from typing_extensions import Self, deprecated
+from typing_extensions import Self, TypeAlias, deprecated
 
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.asymmetric import ec, rsa
@@ -46,7 +46,8 @@ class _X448_CURVE(NamedTuple):
     pubkey: UnimplementedOKPCurveKey
     privkey: UnimplementedOKPCurveKey
 
-JWKTypesRegistry: dict[str, str]
+_JWKKeyTypeSupported: TypeAlias = Literal["oct", "RSA", "EC", "OKP"]
+JWKTypesRegistry: dict[_JWKKeyTypeSupported, str]
 
 class ParmType(Enum):
     name = "A string with a name"  # pyright: ignore[reportAssignmentType]
@@ -63,8 +64,12 @@ class JWKParameter(NamedTuple):
 JWKValuesRegistry: dict[str, dict[str, JWKParameter]]
 JWKParamsRegistry: dict[str, JWKParameter]
 JWKEllipticCurveRegistry: dict[str, str]
-JWKUseRegistry: dict[str, str]
-JWKOperationsRegistry: dict[str, str]
+_JWKUseSupported: TypeAlias = Literal["sig", "enc"]
+JWKUseRegistry: dict[_JWKUseSupported, str]
+_JWKOperationSupported: TypeAlias = Literal[
+    "sign", "verify", "encrypt", "decrypt", "wrapKey", "unwrapKey", "deriveKey", "deriveBits"
+]
+JWKOperationsRegistry: dict[_JWKOperationSupported, str]
 JWKpycaCurveMap: dict[str, str]
 IANANamedInformationHashAlgorithmRegistry: dict[
     str,
@@ -98,9 +103,26 @@ class InvalidJWKValue(JWException): ...
 
 class JWK(dict[str, Any]):
     def __init__(self, **kwargs) -> None: ...
+    # `kty` and the other keyword arguments are passed as `params` to the called generator
+    # function. The possible arguments depend on the value of `kty`.
+    # TODO: Add overloads for the individual `kty` values.
+    @classmethod
+    @overload
+    def generate(
+        cls,
+        *,
+        kty: Literal["RSA"],
+        public_exponent: int | None = None,
+        size: int | None = None,
+        kid: str | None = None,
+        alg: str | None = None,
+        use: _JWKUseSupported | None = None,
+        key_ops: list[_JWKOperationSupported] | None = None,
+    ) -> Self: ...
     @classmethod
-    def generate(cls, **kwargs) -> Self: ...
-    def generate_key(self, **params) -> None: ...
+    @overload
+    def generate(cls, *, kty: _JWKKeyTypeSupported, **kwargs) -> Self: ...
+    def generate_key(self, *, kty: _JWKKeyTypeSupported, **kwargs) -> None: ...
     def import_key(self, **kwargs) -> None: ...
     @classmethod
     def from_json(cls, key) -> Self: ...

stubs/jwcrypto/jwcrypto/jwt.pyi

@@ -1,6 +1,6 @@
 from _typeshed import Incomplete
 from collections.abc import Mapping
-from typing import Any
+from typing import Any, SupportsInt
 from typing_extensions import deprecated
 
 from jwcrypto.common import JWException, JWKeyNotFound
@@ -49,31 +49,31 @@ class JWT:
     @header.setter
     def header(self, h: dict[str, Any] | str) -> None: ...
     @property
-    def claims(self): ...
+    def claims(self) -> str: ...
     @claims.setter
-    def claims(self, data) -> None: ...
+    def claims(self, data: str) -> None: ...
     @property
     def token(self): ...
     @token.setter
     def token(self, t) -> None: ...
     @property
-    def leeway(self): ...
+    def leeway(self) -> int: ...
     @leeway.setter
-    def leeway(self, lwy) -> None: ...
+    def leeway(self, lwy: SupportsInt) -> None: ...
     @property
-    def validity(self): ...
+    def validity(self) -> int: ...
     @validity.setter
-    def validity(self, v) -> None: ...
+    def validity(self, v: SupportsInt) -> None: ...
     @property
     def expected_type(self): ...
     @expected_type.setter
     def expected_type(self, v) -> None: ...
     def norm_typ(self, val): ...
-    def make_signed_token(self, key) -> None: ...
-    def make_encrypted_token(self, key) -> None: ...
-    def validate(self, key) -> None: ...
+    def make_signed_token(self, key: JWK) -> None: ...
+    def make_encrypted_token(self, key: JWK) -> None: ...
+    def validate(self, key: JWK | JWKSet) -> None: ...
     def deserialize(self, jwt, key: Incomplete | None = None) -> None: ...
-    def serialize(self, compact: bool = True): ...
+    def serialize(self, compact: bool = True) -> str: ...
     @classmethod
     def from_jose_token(cls, token): ...
     def __eq__(self, other: object) -> bool: ...