Make PaginatedList pickleable

PyGithub · Nov 19, 2023 · 6abd3e5 · 6abd3e5
1 parent 5dae94e
commit 6abd3e5
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/github/PaginatedList.py b/github/PaginatedList.py
@@ -140,7 +140,7 @@ def __init__(
         list_item: str = "items",
         firstData: Optional[Any] = None,
         firstHeaders: Optional[Dict[str, Union[str, int]]] = None,
-        attributesTransformer: Callable[[Any], Any] = lambda x: x,
+        attributesTransformer: Optional[Callable[[Any], Any]] = None,
     ):
         self.__requester = requester
         self.__contentClass = contentClass
@@ -161,6 +161,11 @@ def __init__(
             first_page = self._getPage(firstData, firstHeaders)
         super().__init__(first_page)
 
+    def _transformAttribute(self, attribute: Any) -> Any:
+        if self._attributesTransformer is None:
+            return attribute
+        return self._attributesTransformer(attribute)
+
     @property
     def totalCount(self) -> int:
         if not self.__totalCount:
@@ -239,7 +244,7 @@ def _getPage(self, data: Any, headers: Dict[str, Any]) -> List[T]:
             self.__totalCount = data.get("total_count")
             data = data[self.__list_item]
         content = [
-            self.__contentClass(self.__requester, headers, self._attributesTransformer(element), completed=False)
+            self.__contentClass(self.__requester, headers, self._transformAttribute(element), completed=False)
             for element in data
             if element is not None
         ]
@@ -272,7 +277,7 @@ def get_page(self, page: int) -> List[T]:
             self.__totalCount = data.get("total_count")
             data = data[self.__list_item]
         return [
-            self.__contentClass(self.__requester, headers, self._attributesTransformer(element), completed=False)
+            self.__contentClass(self.__requester, headers, self._transformAttribute(element), completed=False)
             for element in data
         ]
 

diff --git a/tests/Pickle.py b/tests/Pickle.py
@@ -2,6 +2,7 @@
 import unittest
 
 import github
+from github.PaginatedList import PaginatedList
 from github.Repository import Repository
 
 REPO_NAME = "PyGithub/PyGithub"
@@ -24,3 +25,10 @@ def testPickleRepository(self):
         self.assertIsNotNone(repo2._requester._Requester__connection_lock)
         self.assertIsNone(repo2._requester._Requester__connection)
         self.assertEqual(len(repo2._requester._Requester__custom_connections), 0)
+
+    def testPicklePaginatedList(self):
+        gh = github.Github()
+        repo = gh.get_repo(REPO_NAME, lazy=True)
+        branches = repo.get_branches()
+        branches2 = pickle.loads(pickle.dumps(branches))
+        self.assertIsInstance(branches2, PaginatedList)