Feat #1024: Add iterator for PyJWKSet (#1041)

* Feat #1024: Add iterator for PyJWKSet

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix mypy type

* changelog

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

CHANGELOG.rst

@@ -10,6 +10,8 @@ This project adheres to `Semantic Versioning <https://semver.org/>`__.
 Fixed
 ~~~~~
 - Validate key against allowed types for Algorithm family in `#964 <https://github.com/jpadilla/pyjwt/pull/964>`__
+- Add iterator for JWKSet in `#1041 <https://github.com/jpadilla/pyjwt/pull/1041>`__
+
 
 Added
 ~~~~~

jwt/api_jwk.py

@@ -2,7 +2,7 @@
 
 import json
 import time
-from typing import Any
+from typing import Any, Iterator
 
 from .algorithms import get_default_algorithms, has_crypto, requires_cryptography
 from .exceptions import (
@@ -131,6 +131,9 @@ def __getitem__(self, kid: str) -> PyJWK:
                 return key
         raise KeyError(f"keyset has no key for kid: {kid}")
 
+    def __iter__(self) -> Iterator[PyJWK]:
+        return iter(self.keys)
+
 
 class PyJWTSetWithTimestamp:
     def __init__(self, jwk_set: PyJWKSet):

tests/test_api_jwk.py

@@ -291,6 +291,20 @@ def test_keyset_should_index_by_kid(self):
         with pytest.raises(KeyError):
             _ = jwk_set["this-kid-does-not-exist"]
 
+    @crypto_required
+    def test_keyset_iterator(self):
+        algo = RSAAlgorithm(RSAAlgorithm.SHA256)
+
+        with open(key_path("jwk_rsa_pub.json")) as keyfile:
+            pub_key = algo.from_jwk(keyfile.read())
+
+        key_data_str = algo.to_jwk(pub_key)
+        key_data = json.loads(key_data_str)
+
+        jwk_set = PyJWKSet.from_dict({"keys": [key_data]})
+
+        assert jwk_set.keys == [jwk for jwk in jwk_set]
+
     @crypto_required
     def test_keyset_with_unknown_alg(self):
         # first keyset with unusable key and usable key