Add more types (#843)

* Add return types in `JWKSetCache`

* Add types for hash algorithms

* Add missing type annotation in `ECAlgorithm`

3 files changed, 19 insertions, 17 deletions

diff --git a/jwt/algorithms.py b/jwt/algorithms.py
index c809f15..80d303c 100644
--- a/jwt/algorithms.py
+++ b/jwt/algorithms.py
@@ -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,
@@ -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):
@@ -271,11 +271,11 @@ if has_crypto:
         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):
@@ -419,11 +419,11 @@ if has_crypto:
         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):
diff --git a/jwt/jwk_set_cache.py b/jwt/jwk_set_cache.py
index e8c2a7e..1b2465c 100644
--- a/jwt/jwk_set_cache.py
+++ b/jwt/jwk_set_cache.py
@@ -5,11 +5,11 @@ from .api_jwk import PyJWKSet, PyJWTSetWithTimestamp
 
 
 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:
diff --git a/jwt/types.py b/jwt/types.py
index 830b185..7d99352 100644
--- a/jwt/types.py
+++ b/jwt/types.py
@@ -1,3 +1,5 @@
-from typing import Any, Dict
+from typing import Any, Callable, Dict
 
 JWKDict = Dict[str, Any]
+
+HashlibHash = Callable[..., Any]