fix: improve type hints for jwk module

Author: lepture

src/joserfc/jwk.py

@@ -88,7 +88,25 @@ def guess_key(key: KeyFlexible, obj: GuestProtocol, use_random: bool = False) ->
         raise ValueError("Invalid key")
 
 
-def import_key(data: AnyKey, key_type: str | None = None, parameters: KeyParameters | None = None) -> Key:
+@t.overload
+def import_key(data: AnyKey, key_type: t.Literal["oct"], parameters: KeyParameters | None = None) -> OctKey: ...
+
+
+@t.overload
+def import_key(data: AnyKey, key_type: t.Literal["RSA"], parameters: KeyParameters | None = None) -> RSAKey: ...
+
+
+@t.overload
+def import_key(data: AnyKey, key_type: t.Literal["EC"], parameters: KeyParameters | None = None) -> ECKey: ...
+
+
+@t.overload
+def import_key(data: AnyKey, key_type: t.Literal["OKP"], parameters: KeyParameters | None = None) -> OKPKey: ...
+
+
+def import_key(
+    data: AnyKey, key_type: t.Literal["oct", "RSA", "EC", "OKP"] | None = None, parameters: KeyParameters | None = None
+) -> Key:
     """Importing a key from bytes, string, and dict. When ``value`` is a dict,
     this method can tell the key type automatically, otherwise, developers
     SHOULD pass the ``key_type`` themselves.
@@ -101,8 +119,48 @@ def import_key(data: AnyKey, key_type: str | None = None, parameters: KeyParamet
     return JWKRegistry.import_key(data, key_type, parameters)
 
 
+@t.overload
+def generate_key(
+    key_type: t.Literal["oct"],
+    crv_or_size: str | int | None = None,
+    parameters: KeyParameters | None = None,
+    private: bool = True,
+    auto_kid: bool = False,
+) -> OctKey: ...
+
+
+@t.overload
+def generate_key(
+    key_type: t.Literal["RSA"],
+    crv_or_size: str | int | None = None,
+    parameters: KeyParameters | None = None,
+    private: bool = True,
+    auto_kid: bool = False,
+) -> RSAKey: ...
+
+
+@t.overload
+def generate_key(
+    key_type: t.Literal["EC"],
+    crv_or_size: str | int | None = None,
+    parameters: KeyParameters | None = None,
+    private: bool = True,
+    auto_kid: bool = False,
+) -> ECKey: ...
+
+
+@t.overload
+def generate_key(
+    key_type: t.Literal["OKP"],
+    crv_or_size: str | int | None = None,
+    parameters: KeyParameters | None = None,
+    private: bool = True,
+    auto_kid: bool = False,
+) -> OKPKey: ...
+
+
 def generate_key(
-    key_type: str,
+    key_type: t.Literal["oct", "RSA", "EC", "OKP"],
     crv_or_size: str | int | None = None,
     parameters: KeyParameters | None = None,
     private: bool = True,

tests/jwt/test_jwt.py

@@ -8,6 +8,7 @@
     DecodeError,
 )
 
+
 def use_embedded_jwk(obj: GuestProtocol) -> Key:
     headers = obj.headers()
     return import_key(headers["jwk"])
@@ -82,13 +83,13 @@ def test_using_registry(self):
 
     def test_with_embedded_jwk(self):
         value = (
-            'eyJqd2siOnsiY3J2IjoiUC0yNTYiLCJ4IjoiVU05ZzVuS25aWFlvdldBbE'
-            '03NmNMejl2VG96UmpfX0NIVV9kT2wtZ09vRSIsInkiOiJkczhhZVF3MWwy'
-            'Y0RDQTdiQ2tPTnZ3REtwWEFidFhqdnFDbGVZSDhXc19VIiwia3R5IjoiRU'
-            'MifSwiYWxnIjoiRVMyNTYifQ.eyJpc3MiOiJ1cm46ZXhhbXBsZTppc3N1Z'
-            'XIiLCJhdWQiOiJ1cm46ZXhhbXBsZTphdWRpZW5jZSIsImlhdCI6MTYwNDU'
-            '4MDc5NH0.60boak3_dErnW47ZPty1C0nrjeVq86EN_eK0GOq6K8w2OA0th'
-            'KoBxFK4j-NuU9yZ_A9UKGxPT_G87DladBaV9g'
+            "eyJqd2siOnsiY3J2IjoiUC0yNTYiLCJ4IjoiVU05ZzVuS25aWFlvdldBbE"
+            "03NmNMejl2VG96UmpfX0NIVV9kT2wtZ09vRSIsInkiOiJkczhhZVF3MWwy"
+            "Y0RDQTdiQ2tPTnZ3REtwWEFidFhqdnFDbGVZSDhXc19VIiwia3R5IjoiRU"
+            "MifSwiYWxnIjoiRVMyNTYifQ.eyJpc3MiOiJ1cm46ZXhhbXBsZTppc3N1Z"
+            "XIiLCJhdWQiOiJ1cm46ZXhhbXBsZTphdWRpZW5jZSIsImlhdCI6MTYwNDU"
+            "4MDc5NH0.60boak3_dErnW47ZPty1C0nrjeVq86EN_eK0GOq6K8w2OA0th"
+            "KoBxFK4j-NuU9yZ_A9UKGxPT_G87DladBaV9g"
         )
         token = jwt.decode(value, use_embedded_jwk)
         self.assertEqual(token.claims["iss"], "urn:example:issuer")