[patch] Allow _FactoryMade to explicitly define where its reduce should import from (#27)

* Allow `_FactoryMade` to explicitly define where its reduce should import from

* Give the same <locals> protection with the new syntax

* Add comment

* Extend tests to cover both syntaxes

* Edit docstring

* Format black

---------

Co-authored-by: pyiron-runner <pyiron@mpie.de>

Author: liamhuber

pyiron_snippets/factory.py

@@ -237,16 +237,22 @@ class _FactoryMade(ABC):
     """
     A mix-in to make class-factory-produced classes pickleable.
 
-    If the factory is used as a decorator for another function, it will conflict with
-    this function (i.e. the owned function will be the true function, and will mismatch
-    with imports from that location, which will return the post-decorator factory made
-    class). This can be resolved by setting the
+    If the factory is used as a decorator for another function (or class), it will
+    conflict with this function (i.e. the owned function will be the true function,
+    and will mismatch with imports from that location, which will return the
+    post-decorator factory made class). This can be resolved by setting the
+    :attr:`_reduce_imports_as` attribute to a tuple of the (module, qualname) obtained
+    from the decorated definition in order to manually specify where it should be
+    re-imported from. (DEPRECATED alternative: set
     :attr:`_class_returns_from_decorated_function` attribute to be the decorated
-    function in the decorator definition.
+    function in the decorator definition.)
     """
 
+    # DEPRECATED: Use _reduce_imports_as instead
     _class_returns_from_decorated_function: ClassVar[callable | None] = None
 
+    _reduce_imports_as: ClassVar[tuple[str, str] | None] = None  # Module and qualname
+
     def __init_subclass__(cls, /, class_factory, class_factory_args, **kwargs):
         super().__init_subclass__(**kwargs)
         cls._class_factory = class_factory
@@ -271,6 +277,19 @@ def __reduce__(self):
                 ),
                 self.__getstate__(),
             )
+        elif (
+            self._reduce_imports_as is not None
+            and "<locals>" not in self._reduce_imports_as[1]
+        ):
+            return (
+                _instantiate_from_decorated,
+                (
+                    self._reduce_imports_as[0],
+                    self._reduce_imports_as[1],
+                    self.__getnewargs_ex__(),
+                ),
+                self.__getstate__(),
+            )
         else:
             return (
                 _instantiate_from_factory,

tests/unit/test_factory.py

@@ -116,6 +116,20 @@ def add_to_function(self, *args, **kwargs):
 
 @classfactory
 def adder_factory(fnc, n, /):
+    return (
+        f"{AddsNandX.__name__}{fnc.__name__}",
+        (AddsNandX,),
+        {
+            "fnc": staticmethod(fnc),
+            "n": n,
+            "_reduce_imports_as": (fnc.__module__, fnc.__qualname__)
+        },
+        {},
+    )
+
+
+@classfactory
+def deprecated_adder_factory(fnc, n, /):
     return (
         f"{AddsNandX.__name__}{fnc.__name__}",
         (AddsNandX,),
@@ -131,7 +145,7 @@ def adder_factory(fnc, n, /):
 def add_to_this_decorator(n):
     def wrapped(fnc):
         factory_made = adder_factory(fnc, n)
-        factory_made._class_returns_from_decorated_function = fnc
+        factory_made._reduce_imports_as = (fnc.__module__, fnc.__qualname__)
         return factory_made
     return wrapped
 
@@ -141,6 +155,19 @@ def adds_5_plus_x(y: int):
     return y
 
 
+def deprecated_add_to_this_decorator(n):
+    def wrapped(fnc):
+        factory_made = adder_factory(fnc, n)
+        factory_made._class_returns_from_decorated_function = fnc
+        return factory_made
+    return wrapped
+
+
+@deprecated_add_to_this_decorator(5)
+def deprecated_adds_5_plus_x(y: int):
+    return y
+
+
 class TestClassfactory(unittest.TestCase):
 
     def test_factory_initialization(self):
@@ -474,21 +501,23 @@ def test_other_decorators(self):
         In case the factory-produced class itself comes from a decorator, we need to
         check that name conflicts between the class and decorated function are handled.
         """
-        a5 = adds_5_plus_x(2)
-        self.assertIsInstance(a5, AddsNandX)
-        self.assertIsInstance(a5, _FactoryMade)
-        self.assertEqual(5, a5.n)
-        self.assertEqual(2, a5.x)
-        self.assertEqual(
-            1 + 5 + 2,  # y + n=5 + x=2
-            a5.add_to_function(1),
-            msg="Should execute the function as part of call"
-        )
-
-        reloaded = pickle.loads(pickle.dumps(a5))
-        self.assertEqual(a5.n, reloaded.n)
-        self.assertIs(a5.fnc, reloaded.fnc)
-        self.assertEqual(a5.x, reloaded.x)
+        for fnc in [adds_5_plus_x, deprecated_adds_5_plus_x]:
+            with self.subTest(fnc.__name__):
+                a5 = fnc(2)
+                self.assertIsInstance(a5, AddsNandX)
+                self.assertIsInstance(a5, _FactoryMade)
+                self.assertEqual(5, a5.n)
+                self.assertEqual(2, a5.x)
+                self.assertEqual(
+                    1 + 5 + 2,  # y + n=5 + x=2
+                    a5.add_to_function(1),
+                    msg="Should execute the function as part of call"
+                )
+
+                reloaded = pickle.loads(pickle.dumps(a5))
+                self.assertEqual(a5.n, reloaded.n)
+                self.assertIs(a5.fnc, reloaded.fnc)
+                self.assertEqual(a5.x, reloaded.x)
 
     def test_other_decorators_inside_locals(self):
         @add_to_this_decorator(6)