Repository: cogent3/cogent3
Pull request #975 (2021-10-31 02:40:01)

ENH: refactor of appify so decorated functiuons can be pickled
[CHANGED] converted appify from a class to 3-nested functions. This was done
    because it allowed usage of functools.wraps to wrap the user function
    and we can store a reference to that wrapped function for provision to
    user_function. This was key to enabling pickling support. Failure to do this
    resulted in PicklingError, function is not the same object as __main__.foo
    because the latter was the decorated function.


Patch:
diff --git a/src/cogent3/app/composable.py b/src/cogent3/app/composable.py
index 3f8ac8f7d..bb782dd6d 100644
--- a/src/cogent3/app/composable.py
+++ b/src/cogent3/app/composable.py
@@ -8,6 +8,7 @@ import time
 import traceback
 
 from copy import deepcopy
+from functools import wraps
 
 import scitrack
 
@@ -721,7 +722,10 @@ class user_function(Composable):
         args = self._args + args
         kwargs_ = deepcopy(self._kwargs)
         kwargs_.update(kwargs)
-        return self._user_func(*args, **kwargs_)
+        # the following enables a decorated user function (via @appify())
+        # or directly passed user function
+        func = getattr(self._user_func, "__wrapped__", self._user_func)
+        return func(*args, **kwargs_)
 
     def __str__(self):
         txt = "" if not self.input else str(self.input)
@@ -737,45 +741,45 @@ class user_function(Composable):
         return str(self)
 
 
-class appify:
+@extend_docstring_from(ComposableType.__init__, pre=True)
+def appify(input_types, output_types, data_types=None):
     """function decorator for generating user apps. Simplifies creation of
-    user_function() instancese, e.g.
+    user_function() instances, e.g.
 
     >>> @appify(SEQUENCE_TYPE, SEQUENCE_TYPE, data_types="SequenceCollection")
     ... def omit_seqs(seqs, quantile=None, gap_fraction=1, moltype="dna"):
     ...     return seqs.omit_bad_seqs(quantile=quantile, gap_fraction=gap_fraction, moltype="dna")
     ...
 
-    `omit_seqs()` is now an app factory, allowing creating variants of the app.
+    ``omit_seqs()`` is now an app factory, allowing creating variants of the app.
 
     >>> omit_bad = omit_seqs(quantile=0.95)
 
-    omit_bad is now a composable user_function app. Calling with different
+    ``omit_bad`` is now a composable ``user_function`` app. Calling with different
     args/kwargs values returns a variant app, as per the behaviour of builtin
     apps.
+
     """
+    # the 3 nested functions are required to allow setting decorator arguments
+    # to allow using functools.wraps so the decorated function has the correct
+    # docstring, name etc... And, the final inner one gets to pass the
+    # reference to the wrapped function (wrapped_ref) to user_function. This
+    # latter is required to enable pickling of the user_function instance.
+    def enclosed(func):
+        @wraps(func)
+        def maker(*args, **kwargs):
+            # construct the user_function app
+            return user_function(
+                wrapped_ref,
+                input_types,
+                output_types,
+                *args,
+                data_types=data_types,
+                **kwargs,
+            )
 
-    @extend_docstring_from(ComposableType.__init__)
-    def __init__(self, input_types, output_types, data_types=None) -> None:
-        self._it = input_types
-        self._ot = output_types
-        self._dt = data_types
-        self._func = None
-
-    def __call__(self, func):
-        # executed on use as decorator
-        self._func = func
-        # makes the returned reference have the name, docs etc.
-        # of original function
-        self._make_app.__func__.__doc__ = f"appify: {func.__doc__}"
-        self._make_app.__func__.__repr__ = lambda x: repr(func)
-        self._make_app.__func__.__name__ = func.__name__
-        self._make_app.__func__.__module__ = func.__module__
-
-        return self._make_app
-
-    def _make_app(self, *args, **kwargs):
-        # construct the user_function app
-        return user_function(
-            self._func, self._it, self._ot, *args, data_types=self._dt, **kwargs
-        )
+        wrapped_ref = maker
+
+        return maker
+
+    return enclosed


Test patch:
diff --git a/tests/test_app/test_composable.py b/tests/test_app/test_composable.py
index 766f08f65..fa736951d 100644
--- a/tests/test_app/test_composable.py
+++ b/tests/test_app/test_composable.py
@@ -1,6 +1,7 @@
 import os
 import pathlib
 
+from pickle import dumps, loads
 from tempfile import TemporaryDirectory
 from unittest import TestCase, main
 from unittest.mock import Mock
@@ -276,7 +277,6 @@ class TestNotCompletedResult(TestCase):
 class TestPicklable(TestCase):
     def test_composite_pickleable(self):
         """composable functions should be pickleable"""
-        from pickle import dumps
 
         from cogent3.app import align, evo, io, sample, translate, tree
 
@@ -299,8 +299,6 @@ class TestPicklable(TestCase):
 
     def test_not_completed_result(self):
         """should survive roundtripping pickle"""
-        from pickle import dumps, loads
-
         err = NotCompleted("FAIL", "mytest", "can we roundtrip")
         p = dumps(err)
         new = loads(p)
@@ -325,6 +323,13 @@ def _demo(ctx, expect):
     return ctx.frame_start == expect
 
 
+# for testing appify
+@appify(SERIALISABLE_TYPE, SERIALISABLE_TYPE)
+def slicer(val, index=2):
+    """my docstring"""
+    return val[:index]
+
+
 class TestUserFunction(TestCase):
     def foo(self, val, *args, **kwargs):
         return val[:4]
@@ -371,13 +376,7 @@ class TestUserFunction(TestCase):
 
     def test_appify(self):
         """acts like a decorator should!"""
-
-        @appify(SERIALISABLE_TYPE, SERIALISABLE_TYPE)
-        def slicer(val, index=2):
-            """my docstring"""
-            return val[:index]
-
-        self.assertEqual(slicer.__doc__, "appify: my docstring")
+        self.assertEqual(slicer.__doc__, "my docstring")
         self.assertEqual(slicer.__name__, "slicer")
         app = slicer()
         self.assertTrue(SERIALISABLE_TYPE in app._input_types)
@@ -386,6 +385,13 @@ class TestUserFunction(TestCase):
         app2 = slicer(index=3)
         self.assertEqual(app2(list(range(4))), [0, 1, 2])
 
+    def test_appify_pickle(self):
+        """appified function should be pickleable"""
+        app = slicer(index=6)
+        dumped = dumps(app)
+        loaded = loads(dumped)
+        self.assertEqual(loaded(list(range(10))), list(range(6)))
+
     def test_user_function_repr(self):
         u_function_1 = user_function(self.foo, "aligned", "aligned")
         u_function_2 = user_function(self.bar, "aligned", "pairwise_distances")