Add service_account.Credentials.with_claims

google/auth/jwt.py

@@ -426,13 +426,15 @@ def with_claims(self, issuer=None, subject=None, audience=None,
         Returns:
             google.auth.jwt.Credentials: A new credentials instance.
         """
+        new_additional_claims = self._additional_claims.copy()
+        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

@@ -246,6 +246,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 = self._additional_claims.copy()
+        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)