chore: improve for mypy

src/joserfc/drafts/jwe_ecdh_1pu.py

@@ -47,7 +47,7 @@ def __init__(self, key_wrapping: t.Optional[JWEKeyWrapping]):
             self.key_size = key_wrapping.key_size
         self.key_wrapping = key_wrapping
 
-    def _check_enc(self, enc: JWEEncModel):
+    def _check_enc(self, enc: JWEEncModel) -> None:
         if self.key_wrapping and not isinstance(enc, CBCHS2EncModel):
             description = (
                 'In key agreement with key wrapping mode ECDH-1PU algorithm '

src/joserfc/jwk.py

@@ -60,8 +60,16 @@ class JWKRegistry(_JWKRegistry):
     }
 
 
+class GuestProtocol(t.Protocol):  # pragma: no cover
+    def headers(self) -> Header:
+        ...
+
+    def set_kid(self, kid: str):
+        ...
+
+
 Key = t.Union[OctKey, RSAKey, ECKey, OKPKey]
-KeyCallable = t.Callable[[t.Any], Key]
+KeyCallable = t.Callable[[GuestProtocol], Key]
 
 
 class KeySet(_KeySet):
@@ -71,14 +79,6 @@ class KeySet(_KeySet):
 KeyFlexible = t.Union[t.AnyStr, Key, KeySet, KeyCallable]
 
 
-class GuestProtocol(t.Protocol):  # pragma: no cover
-    def headers(self) -> Header:
-        ...
-
-    def set_kid(self, kid: str):
-        ...
-
-
 def guess_key(key: KeyFlexible, obj: GuestProtocol) -> Key:
     """Guess key from a various sources.
 

src/joserfc/jws.py

@@ -60,7 +60,7 @@
 
 
 # register supported alg models
-def __register():
+def __register() -> None:
     # register alg in RFC7518
     for _alg in JWS_ALGORITHMS:
         JWSRegistry.register(_alg)
@@ -193,7 +193,8 @@ def serialize_json(
         payload: t.AnyStr,
         private_key: KeyFlexible,
         algorithms: t.Optional[t.List[str]] = None,
-        registry: t.Optional[JWSRegistry] = None):
+        registry: t.Optional[JWSRegistry] = None,
+) -> t.Union[GeneralJSONSerialization, FlattenedJSONSerialization]:
     """Generate a JWS JSON Serialization (in dict). The JWS JSON Serialization
     represents digitally signed or MACed content as a JSON object. This representation
     is neither optimized for compactness nor URL-safe.
@@ -251,7 +252,8 @@ def deserialize_json(
         value: t.Union[GeneralJSONSerialization, FlattenedJSONSerialization],
         public_key: KeyFlexible,
         algorithms: t.Optional[t.List[str]] = None,
-        registry: t.Optional[JWSRegistry] = None):
+        registry: t.Optional[JWSRegistry] = None,
+) -> t.Union[GeneralJSONSignature, FlattenedJSONSignature]:
     """Extract and validate the JWS (in string) with the given key.
 
     :param value: a dict of the JSON signature

src/joserfc/rfc7518/derive_key.py

@@ -18,7 +18,7 @@ def derive_key_for_concat_kdf(
         header: Header,
         cek_size: int,
         key_size: t.Optional[int],
-        tag: t.Optional[bytes] = None):
+        tag: t.Optional[bytes] = None) -> bytes:
     # PartyUInfo
     apu_info = u32be_len_input(header.get("apu"), True)
     # PartyVInfo
@@ -48,7 +48,7 @@ def derive_key_for_concat_kdf(
     return ckdf.derive(shared_key)
 
 
-def u32be_len_input(s, use_base64=False):
+def u32be_len_input(s, use_base64=False) -> bytes:
     if not s:
         return b"\x00\x00\x00\x00"
     if use_base64: