Skip to content

Commit

Permalink
Add more types (#843)
Browse files Browse the repository at this point in the history
* Add return types in `JWKSetCache`

* Add types for hash algorithms

* Add missing type annotation in `ECAlgorithm`
  • Loading branch information
Viicos authored Jan 3, 2023
1 parent 701e8d9 commit 45fff3a
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 17 deletions.
28 changes: 14 additions & 14 deletions jwt/algorithms.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
import hashlib
import hmac
import json
from typing import Any, Dict, Union
from typing import Any, ClassVar, Dict, Type, Union

from .exceptions import InvalidKeyError
from .types import JWKDict
from .types import HashlibHash, JWKDict
from .utils import (
base64url_decode,
base64url_encode,
Expand Down Expand Up @@ -212,11 +212,11 @@ class HMACAlgorithm(Algorithm):
and the specified hash function.
"""

SHA256 = hashlib.sha256
SHA384 = hashlib.sha384
SHA512 = hashlib.sha512
SHA256: ClassVar[HashlibHash] = hashlib.sha256
SHA384: ClassVar[HashlibHash] = hashlib.sha384
SHA512: ClassVar[HashlibHash] = hashlib.sha512

def __init__(self, hash_alg) -> None:
def __init__(self, hash_alg: HashlibHash) -> None:
self.hash_alg = hash_alg

def prepare_key(self, key):
Expand Down Expand Up @@ -271,11 +271,11 @@ class RSAAlgorithm(Algorithm):
RSASSA-PKCS-v1_5 and the specified hash function.
"""

SHA256 = hashes.SHA256
SHA384 = hashes.SHA384
SHA512 = hashes.SHA512
SHA256: ClassVar[Type[hashes.HashAlgorithm]] = hashes.SHA256
SHA384: ClassVar[Type[hashes.HashAlgorithm]] = hashes.SHA384
SHA512: ClassVar[Type[hashes.HashAlgorithm]] = hashes.SHA512

def __init__(self, hash_alg) -> None:
def __init__(self, hash_alg: Type[hashes.HashAlgorithm]) -> None:
self.hash_alg = hash_alg

def prepare_key(self, key):
Expand Down Expand Up @@ -419,11 +419,11 @@ class ECAlgorithm(Algorithm):
ECDSA and the specified hash function
"""

SHA256 = hashes.SHA256
SHA384 = hashes.SHA384
SHA512 = hashes.SHA512
SHA256: ClassVar[Type[hashes.HashAlgorithm]] = hashes.SHA256
SHA384: ClassVar[Type[hashes.HashAlgorithm]] = hashes.SHA384
SHA512: ClassVar[Type[hashes.HashAlgorithm]] = hashes.SHA512

def __init__(self, hash_alg) -> None:
def __init__(self, hash_alg: Type[hashes.HashAlgorithm]) -> None:
self.hash_alg = hash_alg

def prepare_key(self, key):
Expand Down
4 changes: 2 additions & 2 deletions jwt/jwk_set_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,11 @@


class JWKSetCache:
def __init__(self, lifespan: int):
def __init__(self, lifespan: int) -> None:
self.jwk_set_with_timestamp: Optional[PyJWTSetWithTimestamp] = None
self.lifespan = lifespan

def put(self, jwk_set: PyJWKSet):
def put(self, jwk_set: PyJWKSet) -> None:
if jwk_set is not None:
self.jwk_set_with_timestamp = PyJWTSetWithTimestamp(jwk_set)
else:
Expand Down
4 changes: 3 additions & 1 deletion jwt/types.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from typing import Any, Dict
from typing import Any, Callable, Dict

JWKDict = Dict[str, Any]

HashlibHash = Callable[..., Any]

0 comments on commit 45fff3a

Please sign in to comment.