Add more types

[Viicos]

A few fixes I had while trying to fix https://github.com/jpadilla/pyjwt/issues/602.

I was also wondering if Algorithm could be made an abstract base class? Afaik it is not instanciated as is and this would avoid having to explicitly raise NotImplementedError.

Regarding https://github.com/jpadilla/pyjwt/issues/602, I might have a few solutions, but that might be tricky to handle (regarding the fact that cryptography is optional, as stated in https://github.com/jpadilla/pyjwt/pull/605). The least that can be done:

• Add types to the following methods for each algorithm: prepare_key, sign and verify.
• For the Algorithm methods, use a conditional Union including all available types, depending on the availability of cryptography (not sure if this will work as excepted).
• Use this same Union type for the encode / decode / decode_complete public methods.

Additionally, overloads could be defined for the methods listed above:

@overload
def encode(
    self,
    payload: bytes,
    key: Union[str, bytes],
    algorithm: Literal["HS256", "HS384", "HS512"],
    headers: dict[str, Any] | None = None,
    json_encoder: Type[json.JSONEncoder] | None = None,
    is_payload_detached: bool = False,
    sort_headers: bool = True,
) -> str:

@overload
def encode(
    self,
    payload: bytes,
    key: Union[RSAPrivateKey, str, bytes],
    algorithm: Literal["RS256", "RS384", "RS512"],
    headers: dict[str, Any] | None = None,
    json_encoder: Type[json.JSONEncoder] | None = None,
    is_payload_detached: bool = False,
    sort_headers: bool = True,
) -> str:

# And so on

That way type checkers could attest that for instance when using encode(payload, key, "HS256"), key should be Union[str, bytes]. Although this improves type safety, it adds a lot of boilerplate code.

[Viicos]

I have an interesting issue with mypy, now that hash_alg is typed:

jwt/algorithms.py:592: error: Argument "salt_length" to "PSS" has incompatible type "Callable[[HashAlgorithm], int]"; expected "Union[int, _MaxLength, _Auto, _DigestLength]"  [arg-type]
jwt/algorithms.py:604: error: Argument "salt_length" to "PSS" has incompatible type "Callable[[HashAlgorithm], int]"; expected "Union[int, _MaxLength, _Auto, _DigestLength]"  [arg-type]

This is due to the fact that hash_alg is of type Type[HashAlgorithm], hence accessing the digest_size attribute makes mypy believe that it is a property (and not a int).

https://github.com/jpadilla/pyjwt/blob/701e8d9d3b8389e13ad6371a331e600b503c2550/jwt/algorithms.py#L589-L592

And because cryptography is doing a weird thing (defining an abstract property in HashAlgorithm but overriding it as a simple attribute in the hash implementations):

class HashAlgorithm(metaclass=abc.ABCMeta):
    @abc.abstractproperty
    def name(self) -> str:
        """
        A string naming this algorithm (e.g. "sha256", "md5").
        """

    @abc.abstractproperty
    def digest_size(self) -> int:
        """
        The size of the resulting digest in bytes.
        """

    @abc.abstractproperty
    def block_size(self) -> typing.Optional[int]:
        """
        The internal block size of the hash function, or None if the hash
        function does not use blocks internally (e.g. SHA3).
        """

class SHA256(HashAlgorithm):
    name = "sha256"
    digest_size = 32
    block_size = 64

So this still works at runtime, but I think this is something that should be sorted out in the cryptography package instead. As a workaround here, we can use an instance of hash_alg instead: salt_length=self.hash_alg().digest_size