From e6931b57fa36e7781538d45794363ce8ebc5dc8a Mon Sep 17 00:00:00 2001
From: rmcar17 <robert.mcarthur@anu.edu.au>
Date: Thu, 29 Aug 2024 04:00:00 +0000
Subject: [PATCH 1/2] DEV: make piqtree2 apps pickleable

---
 src/piqtree2/_app/__init__.py | 45 +++++++++++++++++++++++++++++++----
 1 file changed, 41 insertions(+), 4 deletions(-)

diff --git a/src/piqtree2/_app/__init__.py b/src/piqtree2/_app/__init__.py
index b94af4c9..c4d73fe3 100644
--- a/src/piqtree2/_app/__init__.py
+++ b/src/piqtree2/_app/__init__.py
@@ -1,9 +1,46 @@
 """cogent3 apps for piqtree2."""
 
+from typing import Optional, Union
+
+import cogent3
+import cogent3.app.typing
 from cogent3.app.composable import define_app
+from cogent3.util.misc import extend_docstring_from
+
+from piqtree2 import TreeGenMode, build_tree, fit_tree, random_trees
+from piqtree2.model import Model
+
+
+@define_app
+@extend_docstring_from(build_tree)
+def piqtree_phylo(
+    aln: Union[cogent3.Alignment, cogent3.ArrayAlignment],
+    model: Model,
+    rand_seed: Optional[int] = None,
+) -> Union[cogent3.PhyloNode, cogent3.app.typing.SerialisableType]:
+    return build_tree(aln, model, rand_seed)
+
+
+@define_app
+@extend_docstring_from(fit_tree)
+def piqtree_fit(
+    aln: Union[cogent3.Alignment, cogent3.ArrayAlignment],
+    tree: cogent3.PhyloNode,
+    model: Model,
+    rand_seed: Optional[int] = None,
+) -> Union[cogent3.PhyloNode, cogent3.app.typing.SerialisableType]:
+    return fit_tree(aln, tree, model, rand_seed)
+
+
+@define_app
+@extend_docstring_from(random_trees)
+def piqtree_random_trees(
+    num_taxa: int,
+    tree_mode: TreeGenMode,
+    num_trees: int,
+    rand_seed: Optional[int] = None,
+) -> tuple[cogent3.PhyloNode]:
+    return random_trees(num_taxa, tree_mode, num_trees, rand_seed)
 
-from piqtree2.iqtree import build_tree, fit_tree, random_trees
 
-piqtree_phylo = define_app(build_tree)
-piqtree_fit = define_app(fit_tree)
-piqtree_random_trees = define_app(random_trees)
+_ALL_APP_NAMES = ["piqtree_phylo", "piqtree_fit", "piqtree_random_trees"]

From 5d17fe7fbf4d90b8901d501beebab72792ca825f Mon Sep 17 00:00:00 2001
From: rmcar17 <robert.mcarthur@anu.edu.au>
Date: Thu, 29 Aug 2024 04:00:22 +0000
Subject: [PATCH 2/2] TST: test whether piqtree2 apps are pickleable

---
 tests/test_app/test_pickle.py | 11 +++++++++++
 1 file changed, 11 insertions(+)
 create mode 100644 tests/test_app/test_pickle.py

diff --git a/tests/test_app/test_pickle.py b/tests/test_app/test_pickle.py
new file mode 100644
index 00000000..587d575f
--- /dev/null
+++ b/tests/test_app/test_pickle.py
@@ -0,0 +1,11 @@
+"""Apps must be pickleable to be able to be run with parallel=True"""
+
+import pickle
+
+from cogent3 import get_app
+from piqtree2._app import _ALL_APP_NAMES
+
+
+def test_pickle():
+    for app_name in _ALL_APP_NAMES:
+        assert len(pickle.dumps(get_app(app_name))) > 0