From c4e7e6b7078e5dfe8a68882d0d81d5617dcff523 Mon Sep 17 00:00:00 2001
From: "Ryan W. Moore" <2772216+ryanwmoore@users.noreply.github.com>
Date: Sun, 2 Mar 2025 23:08:46 -0500
Subject: [PATCH] Add pickle support to ActionShortcut

---
 ckanapi/common.py            | 23 +++++++++++++++++------
 ckanapi/tests/test_common.py | 20 ++++++++++++++++++++
 2 files changed, 37 insertions(+), 6 deletions(-)
 create mode 100644 ckanapi/tests/test_common.py

diff --git a/ckanapi/common.py b/ckanapi/common.py
index 87d2bd7..2eead76 100644
--- a/ckanapi/common.py
+++ b/ckanapi/common.py
@@ -4,9 +4,10 @@
 
 import json
 
-from ckanapi.errors import (CKANAPIError, NotAuthorized, NotFound,
-    ValidationError, SearchQueryError, SearchError, SearchIndexError,
-    ServerIncompatibleError)
+from ckanapi.errors import (CKANAPIError, NotAuthorized, NotFound, SearchError,
+                            SearchIndexError, SearchQueryError,
+                            ServerIncompatibleError, ValidationError)
+
 
 class ActionShortcut(object):
     """
@@ -33,9 +34,19 @@ class ActionShortcut(object):
             {'package_id': 'foo'}, files={'upload': open(..)})
 
     """
+
     def __init__(self, ckan):
         self._ckan = ckan
 
+    # __getstate__ and __setstate__ provide pickle support. If they didn't
+    # exist, then __getattr__ would result in a runtime pickle exception because
+    # it returns a local.
+    def __getstate__(self):
+        return {'_ckan': self._ckan}
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
     def __getattr__(self, name):
         def action(**kwargs):
             files = {}
@@ -44,10 +55,10 @@ def action(**kwargs):
                     files[k] = v
             if files:
                 nonfiles = dict((k, v) for k, v in kwargs.items()
-                    if k not in files)
+                                if k not in files)
                 return self._ckan.call_action(name,
-                    data_dict=nonfiles,
-                    files=files)
+                                              data_dict=nonfiles,
+                                              files=files)
             return self._ckan.call_action(name, data_dict=kwargs)
         return action
 
diff --git a/ckanapi/tests/test_common.py b/ckanapi/tests/test_common.py
new file mode 100644
index 0000000..7e17a20
--- /dev/null
+++ b/ckanapi/tests/test_common.py
@@ -0,0 +1,20 @@
+import pickle
+import unittest
+
+from ckanapi import RemoteCKAN
+from ckanapi.common import ActionShortcut
+
+
+class TestCommon(unittest.TestCase):
+    def test_pickling(self):
+        with RemoteCKAN('http://localhost:8901') as ckan:
+            action_shortcut = ActionShortcut(ckan)
+            # Verifies that pickling seems to work. Previously, this could
+            # result in errors like:
+            # TypeError: ActionShortcut.__getattr__.<locals>.action() takes 0
+            #   positional arguments but 1 was given
+            pickled = pickle.dumps(action_shortcut)
+            unpickled = pickle.loads(pickled)
+            # Partially check that the pickling+unpickling worked
+            self.assertEqual(action_shortcut._ckan.address,
+                             unpickled._ckan.address)