add pickling support

.vscode/cspell.json

@@ -550,6 +550,12 @@
         "Jwcmlud"
       ]
     },
+    {
+      "filename": "sdk/identity/azure-identity/tests/*.py",
+      "words": [
+        "infile"
+      ]
+    },
     {
       "filename": "sdk/identity/test-resources*",
       "words": [

sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py

@@ -26,7 +26,12 @@ def __init__(
         identity_config: Optional[Dict] = None,
         **kwargs: Any
     ) -> None:
-        self._cache = kwargs.pop("_cache", None) or TokenCache()
+        self._custom_cache = False
+        self._cache = kwargs.pop("_cache", None)
+        if self._cache:
+            self._custom_cache = True
+        else:
+            self._cache = TokenCache()
         self._content_callback = kwargs.pop("_content_callback", None)
         self._identity_config = identity_config or {}
         if client_id:
@@ -91,6 +96,19 @@ def request_token(self, *scopes, **kwargs):
     def _build_pipeline(self, **kwargs):
         pass
 
+    def __getstate__(self) -> Dict[str, Any]:
+        state = self.__dict__.copy()
+        # Remove the non-picklable entries
+        if not self._custom_cache:
+            del state["_cache"]
+        return state
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        # Re-create the unpickable entries
+        if not self._custom_cache:
+            self._cache = TokenCache()
+
 
 class ManagedIdentityClient(ManagedIdentityClientBase):
     def __enter__(self) -> "ManagedIdentityClient":

sdk/identity/azure-identity/azure/identity/_internal/msal_client.py

@@ -131,3 +131,14 @@ def _store_auth_error(self, response: PipelineResponse) -> None:
             content = response.context.get(ContentDecodePolicy.CONTEXT_NAME)
             if content and "error" in content:
                 self._local.error = (content["error"], response.http_response)
+
+    def __getstate__(self) -> Dict[str, Any]:  # pylint:disable=client-method-name-no-double-underscore
+        state = self.__dict__.copy()
+        # Remove the non-picklable entries
+        del state["_local"]
+        return state
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:  # pylint:disable=client-method-name-no-double-underscore
+        self.__dict__.update(state)
+        # Re-create the unpickable entries
+        self._local = threading.local()

sdk/identity/azure-identity/azure/identity/_internal/msal_credentials.py

@@ -112,3 +112,16 @@ def _get_app(self, **kwargs: Any) -> msal.ClientApplication:
             )
 
         return client_applications_map[tenant_id]
+
+    def __getstate__(self) -> Dict[str, Any]:
+        state = self.__dict__.copy()
+        # Remove the non-picklable entries
+        del state["_client_applications"]
+        del state["_cae_client_applications"]
+        return state
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        # Re-create the unpickable entries
+        self._client_applications = {}
+        self._cae_client_applications = {}

sdk/identity/azure-identity/tests/test_pickling.py

@@ -0,0 +1,14 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import pickle
+from azure.identity import DefaultAzureCredential
+
+
+def test_pickle_dac():
+    cred = DefaultAzureCredential()
+    with open("data.pkl", "wb") as outfile:
+        pickle.dump(cred, outfile)
+    with open("data.pkl", "rb") as infile:
+        data_loaded = pickle.load(infile)

sdk/identity/azure-identity/tests/test_pickling_async.py

@@ -0,0 +1,14 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import pickle
+from azure.identity.aio import DefaultAzureCredential
+
+
+def test_pickle_dac():
+    cred = DefaultAzureCredential()
+    with open("data_aio.pkl", "wb") as outfile:
+        pickle.dump(cred, outfile)
+    with open("data_aio.pkl", "rb") as infile:
+        data_loaded = pickle.load(infile)