fix: improve type hints for import_key

Author: lepture

src/joserfc/_rfc7517/models.py

@@ -42,15 +42,15 @@ def import_from_dict(cls, value: DictKey) -> t.Any:
 
     @classmethod
     @abstractmethod
-    def import_from_bytes(cls, value: bytes, password: t.Optional[t.Any] = None) -> t.Any:
+    def import_from_bytes(cls, value: bytes, password: t.Any = None) -> t.Any:
         pass
 
     @staticmethod
     def as_bytes(
         key: GenericKey,
-        encoding: t.Optional[t.Literal["PEM", "DER"]] = None,
-        private: t.Optional[bool] = None,
-        password: t.Optional[str] = None,
+        encoding: t.Literal["PEM", "DER"] | None = None,
+        private: bool | None = None,
+        password: str | None = None,
     ) -> bytes:
         raise NotImplementedError()
 
@@ -88,7 +88,7 @@ def __init__(
         self,
         raw_value: NativePrivateKey | NativePublicKey,
         original_value: t.Any,
-        parameters: t.Optional[KeyParameters] = None,
+        parameters: KeyParameters | None = None,
     ):
         self._raw_value = raw_value
         self.original_value = original_value
@@ -160,7 +160,7 @@ def public_key(self) -> NativePublicKey:
         raise NotImplementedError()
 
     @property
-    def private_key(self) -> t.Optional[NativePrivateKey]:
+    def private_key(self) -> NativePrivateKey | None:
         raise NotImplementedError()
 
     def thumbprint(self) -> str:
@@ -177,7 +177,7 @@ def thumbprint_uri(self) -> str:
         value = self.thumbprint()
         return concat_thumbprint_uri(value, self.thumbprint_digest_method)
 
-    def as_dict(self, private: t.Optional[bool] = None, **params: t.Any) -> DictKey:
+    def as_dict(self, private: bool | None = None, **params: t.Any) -> DictKey:
         """Output this key to a JWK format (in dict). By default, it will return
         the ``dict_value`` of this key.
 
@@ -270,8 +270,8 @@ def validate_dict_key(cls, data: DictKey) -> None:
     def import_key(
         cls: t.Type[GenericKey],
         value: AnyKey,
-        parameters: t.Optional[KeyParameters] = None,
-        password: t.Optional[t.Any] = None,
+        parameters: KeyParameters | None = None,
+        password: t.Any = None,
     ) -> GenericKey:
         if isinstance(value, dict):
             cls.validate_dict_key(value)
@@ -285,7 +285,7 @@ def import_key(
     def generate_key(
         cls: t.Type[GenericKey],
         size_or_crv: t.Any,
-        parameters: t.Optional[KeyParameters] = None,
+        parameters: KeyParameters | None = None,
         private: bool = True,
         auto_kid: bool = False,
     ) -> GenericKey:
@@ -321,16 +321,16 @@ def raw_value(self) -> t.Union[NativePublicKey, NativePrivateKey]:
 
     def as_bytes(
         self,
-        encoding: t.Optional[t.Literal["PEM", "DER"]] = None,
-        private: t.Optional[bool] = None,
-        password: t.Optional[str] = None,
+        encoding: t.Literal["PEM", "DER"] | None = None,
+        private: bool | None = None,
+        password: str | None = None,
     ) -> bytes:
         return self.binding.as_bytes(self, encoding, private, password)
 
-    def as_pem(self, private: t.Optional[bool] = None, password: t.Optional[str] = None) -> bytes:
+    def as_pem(self, private: bool | None = None, password: str | None = None) -> bytes:
         return self.as_bytes(private=private, password=password)
 
-    def as_der(self, private: t.Optional[bool] = None, password: t.Optional[str] = None) -> bytes:
+    def as_der(self, private: bool | None = None, password: str | None = None) -> bytes:
         return self.as_bytes(encoding="DER", private=private, password=password)
 
 

src/joserfc/_rfc7518/ec_key.py

@@ -17,7 +17,7 @@
 from ..errors import InvalidExchangeKeyError
 from .._rfc7517.models import CurveKey
 from .._rfc7517.pem import CryptographyBinding
-from .._rfc7517.types import KeyParameters
+from .._rfc7517.types import KeyParameters, AnyKey
 from ..util import base64_to_int, int_to_base64
 from ..registry import KeyParameter
 
@@ -142,6 +142,15 @@ def curve_name(self) -> str:
     def curve_key_size(self) -> int:
         return self.raw_value.curve.key_size
 
+    @classmethod
+    def import_key(
+        cls: t.Any,
+        value: AnyKey,
+        parameters: KeyParameters | None = None,
+        password: t.Any = None,
+    ) -> "ECKey":
+        return super(ECKey, cls).import_key(value, parameters, password)
+
     @classmethod
     def generate_key(
         cls,

src/joserfc/_rfc7518/oct_key.py

@@ -10,7 +10,7 @@
 )
 from ..registry import KeyParameter
 from .._rfc7517.models import SymmetricKey, NativeKeyBinding
-from .._rfc7517.types import KeyParameters, DictKey
+from .._rfc7517.types import KeyParameters, DictKey, AnyKey
 
 
 POSSIBLE_UNSAFE_KEYS = (
@@ -50,6 +50,15 @@ class OctKey(SymmetricKey):
     #: https://www.rfc-editor.org/rfc/rfc7518#section-6.4
     value_registry = {"k": KeyParameter("Key Value", "str", True, True)}
 
+    @classmethod
+    def import_key(
+        cls: Any,
+        value: AnyKey,
+        parameters: KeyParameters | None = None,
+        password: Any = None,
+    ) -> "OctKey":
+        return super(OctKey, cls).import_key(value, parameters, password)
+
     @classmethod
     def generate_key(
         cls,

src/joserfc/_rfc7518/rsa_key.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import warnings
-from typing import TypedDict
+from typing import TypedDict, Any
 from functools import cached_property
 from cryptography.hazmat.primitives.asymmetric.rsa import (
     generate_private_key,
@@ -18,7 +18,7 @@
 from ..errors import SecurityWarning
 from .._rfc7517.models import AsymmetricKey
 from .._rfc7517.pem import CryptographyBinding
-from .._rfc7517.types import KeyParameters
+from .._rfc7517.types import KeyParameters, AnyKey
 from ..util import int_to_base64, base64_to_int
 
 
@@ -132,6 +132,15 @@ def private_key(self) -> RSAPrivateKey | None:
             return self.raw_value
         return None
 
+    @classmethod
+    def import_key(
+        cls: Any,
+        value: AnyKey,
+        parameters: KeyParameters | None = None,
+        password: Any = None,
+    ) -> "RSAKey":
+        return super(RSAKey, cls).import_key(value, parameters, password)
+
     @classmethod
     def generate_key(
         cls,

src/joserfc/_rfc8037/okp_key.py

@@ -13,7 +13,7 @@
     NoEncryption,
 )
 from .._rfc7517.models import CurveKey
-from .._rfc7517.types import KeyParameters
+from .._rfc7517.types import KeyParameters, AnyKey
 from .._rfc7517.pem import CryptographyBinding
 from ..errors import InvalidExchangeKeyError
 from ..util import to_bytes, urlsafe_b64decode, urlsafe_b64encode
@@ -120,6 +120,15 @@ def private_key(self) -> PrivateOKPKey | None:
     def curve_name(self) -> str:
         return get_key_curve(self.raw_value)
 
+    @classmethod
+    def import_key(
+        cls: t.Any,
+        value: AnyKey,
+        parameters: KeyParameters | None = None,
+        password: t.Any = None,
+    ) -> "OKPKey":
+        return super(OKPKey, cls).import_key(value, parameters, password)
+
     @classmethod
     def generate_key(
         cls,

src/joserfc/jwk.py

@@ -115,6 +115,10 @@ def import_key(data: AnyKey, key_type: t.Literal["EC"], parameters: KeyParameter
 def import_key(data: AnyKey, key_type: t.Literal["OKP"], parameters: KeyParameters | None = None) -> OKPKey: ...
 
 
+@t.overload
+def import_key(data: DictKey, key_type: None = None, parameters: KeyParameters | None = None) -> Key: ...
+
+
 def import_key(
     data: AnyKey,
     key_type: t.Literal["oct", "RSA", "EC", "OKP"] | None = None,

tests/jwk/test_key_methods.py

@@ -142,8 +142,8 @@ def test_find_correct_key_with_use(self):
         key = OctKey.generate_key()
         dict_key = key.as_dict()
 
-        key1: OctKey = OctKey.import_key(dict_key, {"use": "enc"})
-        key2: OctKey = OctKey.import_key(dict_key, {"use": "sig"})
+        key1 = OctKey.import_key(dict_key, {"use": "enc"})
+        key2 = OctKey.import_key(dict_key, {"use": "sig"})
         self.assertEqual(key1.kid, key2.kid)
 
         key_set = KeySet([key1, key2])
@@ -165,8 +165,8 @@ def test_find_correct_key_with_alg(self):
         key = OctKey.generate_key()
         dict_key = key.as_dict()
 
-        key1: OctKey = OctKey.import_key(dict_key, {"alg": "HS256"})
-        key2: OctKey = OctKey.import_key(dict_key, {"alg": "dir"})
+        key1 = OctKey.import_key(dict_key, {"alg": "HS256"})
+        key2 = OctKey.import_key(dict_key, {"alg": "dir"})
 
         self.assertEqual(key1.kid, key2.kid)