meh

betty/ancestry/__init__.py

@@ -56,6 +56,7 @@
 from betty.model.collections import (
     MultipleTypesEntityCollection,
 )
+from betty.model.graph import PickleableEntityGraph
 from betty.string import camel_case_to_kebab_case
 from typing_extensions import override
 
@@ -2143,6 +2144,13 @@ def __init__(self):
         super().__init__()
         self._check_graph = True
 
+    def __getstate__(self) -> PickleableEntityGraph:
+        return PickleableEntityGraph(*self)
+
+    def __setstate__(self, state: PickleableEntityGraph) -> None:
+        self._collections = {}
+        self.add_unchecked_graph(*state.build())
+
     def add_unchecked_graph(self, *entities: Entity) -> None:
         """
         Add entities to the ancestry but do not automatically add associates as well.

betty/generate/pool.py

@@ -22,11 +22,12 @@
     ParamSpec,
 )
 
-import dill
+import pickle
 from math import floor
 
 from betty.app import App
 from betty.asyncio import gather
+from betty.cache.memory import MemoryCache
 from betty.generate import GenerationContext
 from betty.project import Project
 
@@ -64,6 +65,7 @@ async def __aenter__(self) -> Self:
         except BaseException:
             self._cancel.set()
             await self._stop()
+            raise
         return self
 
     async def __aexit__(
@@ -76,15 +78,9 @@ async def __aexit__(
             self._finish.set()
         else:
             self._cancel.set()
-        try:
+        await self._stop()
+        if exc_val is None:
             await self._log_jobs()
-        except BaseException:
-            self._cancel.set()
-            raise
-        finally:
-            await self._stop()
-            if exc_val is None:
-                await self._log_jobs()
 
     async def _start(self) -> None:
         concurrency = os.cpu_count() or 2
@@ -98,21 +94,23 @@ async def _start(self) -> None:
             lambda: executor.shutdown(wait=False, cancel_futures=True)
         )
         # @todo Ensure we pass on the necessary dependencies
-        pickled_app_args = dill.dumps(
+        # @todo Make App pickleable?
+        # @todo
+        # @todo
+        pickled_app_args = pickle.dumps(
             (
                 self._project.app.configuration,
                 self._project.app._cache_directory_path,
             )
         )
-        pickled_app_kwargs = dill.dumps(
+        pickled_app_kwargs = pickle.dumps(
             {
                 # @todo Give it the actual cache.
-                "cache_factory": self._project.app._cache_factory,
+                "cache_factory": MemoryCache,
             }
         )
-        # @todo Can we do this without Dill?
-        pickled_project_args = dill.dumps((self._project.configuration,))
-        pickled_project_kwargs = dill.dumps({"ancestry": self._project.ancestry})
+        pickled_project_args = pickle.dumps((self._project.configuration.dump(),))
+        pickled_project_kwargs = pickle.dumps({"ancestry": self._project.ancestry})
         for _ in range(0, concurrency):
             self._workers.append(
                 executor.submit(
@@ -135,7 +133,6 @@ async def _start(self) -> None:
         self._exit_stack.callback(log_task.cancel)
 
     async def _stop(self) -> None:
-        # @todo NO TIMEOUT
         await to_thread(futures.wait, self._workers)
         await to_thread(self._exit_stack.close)
 
@@ -213,12 +210,12 @@ def __call__(self) -> None:
 
     async def _perform_tasks_concurrently(self) -> None:
         async with App(
-            *dill.loads(self._pickled_app_args),
-            **dill.loads(self._pickled_app_kwargs),
+            *pickle.loads(self._pickled_app_args),
+            **pickle.loads(self._pickled_app_kwargs),
         ) as app, Project(
             app,
-            *dill.loads(self._pickled_project_args),
-            **dill.loads(self._pickled_project_kwargs),
+            *pickle.loads(self._pickled_project_args),
+            **pickle.loads(self._pickled_project_kwargs),
         ) as project:
             job_context = GenerationContext(project)
             await gather(

betty/model/graph.py

@@ -8,7 +8,7 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from collections.abc import MutableSequence, Mapping, MutableMapping
+from collections.abc import MutableSequence, Mapping, MutableMapping, Iterable
 from typing import Iterator, TypeAlias
 
 from betty.model import (
@@ -126,3 +126,53 @@ def add_association(
         self._associations[owner_type][owner_attr_name][owner_id].append(
             (associate_type, associate_id)
         )
+
+
+class PickleableEntityGraph(_EntityGraphBuilder):
+    """
+    Allow an entity graph to be pickled.
+    """
+
+    def __init__(self, *entities: Entity) -> None:
+        super().__init__()
+        self._pickled = False
+        for entity in entities:
+            self._entities[entity.type][entity.id] = entity
+
+    def __getstate__(
+        self,
+    ) -> tuple[_EntityGraphBuilderEntities, _EntityGraphBuilderAssociations]:
+        self._flatten()
+        return self._entities, self._associations
+
+    def __setstate__(
+        self, state: tuple[_EntityGraphBuilderEntities, _EntityGraphBuilderAssociations]
+    ) -> None:
+        self._entities, self._associations = state
+        self._built = False
+        self._pickled = False
+
+    def _flatten(self) -> None:
+        if self._pickled:
+            raise RuntimeError("This entity graph has been pickled already.")
+        self._pickled = True
+
+        for owner in self._iter():
+            unaliased_entity = unalias(owner)
+            entity_type = unaliased_entity.type
+
+            for association in AssociationRegistry.get_all_associations(entity_type):
+                associates: Iterable[Entity]
+                if isinstance(association, ToOneAssociation):
+                    associate = association.get_attr(unaliased_entity)
+                    if associate is None:
+                        continue
+                    associates = [associate]
+                else:
+                    associates = association.get_attr(unaliased_entity)
+                for associate in associates:
+                    self._associations[entity_type][association.owner_attr_name][
+                        owner.id
+                    ].append(
+                        (associate.type, associate.id),
+                    )

mypy.ini

@@ -13,9 +13,6 @@ disallow_untyped_calls = False
 no_implicit_optional = True
 warn_unused_ignores = True
 
-[mypy-dill.*]
-ignore_missing_imports = True
-
 [mypy-docker.*]
 ignore_missing_imports = True
 

pyproject.toml

@@ -22,7 +22,6 @@ dependencies = [
     'babel ~= 2.15',
     'basedtyping ~= 0.1.4',
     'click ~= 8.1',
-    'dill ~= 0.3, >= 0.3.8',
     'docker ~= 7.1',
     'furo == 2024.8.6',
     'geopy ~= 2.4',