Add service_account.Credentials.with_claims (#140)

Author: Jon Wayne Parrott

google/auth/jwt.py

@@ -42,6 +42,7 @@
 
 import base64
 import collections
+import copy
 import datetime
 import json
 
@@ -426,13 +427,15 @@ def with_claims(self, issuer=None, subject=None, audience=None,
         Returns:
             google.auth.jwt.Credentials: A new credentials instance.
         """
+        new_additional_claims = copy.deepcopy(self._additional_claims)
+        new_additional_claims.update(additional_claims or {})
+
         return Credentials(
             self._signer,
             issuer=issuer if issuer is not None else self._issuer,
             subject=subject if subject is not None else self._subject,
             audience=audience if audience is not None else self._audience,
-            additional_claims=self._additional_claims.copy().update(
-                additional_claims or {}))
+            additional_claims=new_additional_claims)
 
     def _make_jwt(self):
         """Make a signed JWT.

google/oauth2/service_account.py

@@ -70,6 +70,7 @@
 .. _RFC 7523: https://tools.ietf.org/html/rfc7523
 """
 
+import copy
 import datetime
 
 from google.auth import _helpers
@@ -246,6 +247,29 @@ def with_subject(self, subject):
             subject=subject,
             additional_claims=self._additional_claims.copy())
 
+    def with_claims(self, additional_claims):
+        """Returns a copy of these credentials with modified claims.
+
+        Args:
+            additional_claims (Mapping[str, str]): Any additional claims for
+                the JWT payload. This will be merged with the current
+                additional claims.
+
+        Returns:
+            google.auth.service_account.Credentials: A new credentials
+                instance.
+        """
+        new_additional_claims = copy.deepcopy(self._additional_claims)
+        new_additional_claims.update(additional_claims or {})
+
+        return Credentials(
+            self._signer,
+            service_account_email=self._service_account_email,
+            scopes=self._scopes,
+            token_uri=self._token_uri,
+            subject=self._subject,
+            additional_claims=new_additional_claims)
+
     def _make_authorization_grant_assertion(self):
         """Create the OAuth 2.0 assertion.
 

tests/oauth2/test_service_account.py

@@ -134,6 +134,10 @@ def test_create_scoped(self):
         credentials = self.credentials.with_scopes(scopes)
         assert credentials._scopes == scopes
 
+    def test_with_claims(self):
+        new_credentials = self.credentials.with_claims({'meep': 'moop'})
+        assert new_credentials._additional_claims == {'meep': 'moop'}
+
     def test__make_authorization_grant_assertion(self):
         token = self.credentials._make_authorization_grant_assertion()
         payload = jwt.decode(token, PUBLIC_CERT_BYTES)