metafunctions 1.1.2__py3-none-any.whl

metafunctions/core/concurrent.py

@@ -0,0 +1,145 @@
+import functools
+import os
+import pickle
+from collections import namedtuple
+from multiprocessing import Queue
+from operator import itemgetter
+
+from metafunctions import exceptions
+from metafunctions.core.base import FunctionMerge
+from metafunctions.core.decorators import manage_call_state
+
+# Result tuple to be sent back from workers. Defined at module level for eas of pickling
+_ConcurrentResult = namedtuple(
+    "_ConcurrentResult", "index result call_state_data exception location"
+)
+
+
+class ConcurrentMerge(FunctionMerge):
+    def __init__(self, function_merge: FunctionMerge):
+        """A subclass of FunctionMerge that calls each of its component functions in parallel.
+
+        ConcurrentMerge takes a FunctionMerge object and upgrades it.
+        """
+        if not isinstance(function_merge, FunctionMerge):
+            # This check is necessary because functools.wraps will copy FunctionMerge attributes to
+            # objects that are not FunctionMerges, so this init will succeed, then result in errors
+            # at call time.
+            raise exceptions.CompositionError(
+                f"{type(self)} can only upgrade FunctionMerges"
+            )
+        if not hasattr(os, "fork"):
+            raise exceptions.CompositionError(
+                f"{type(self).__name__} requires os.fork, and thus is only available on unix"
+            )
+
+        super().__init__(
+            function_merge._merge_func,
+            function_merge._functions,
+            function_merge._function_join_str,
+        )
+        self._function_merge = function_merge
+
+    def __str__(self):
+        merge_name = str(self._function_merge)
+        return (
+            f"concurrent{merge_name}"
+            if merge_name.startswith("(")
+            else f"concurrent({merge_name})"
+        )
+
+    @manage_call_state
+    def __call__(self, *args, **kwargs):
+        """We fork here, and execute each function in a child process before joining the results
+        with _merge_func
+        """
+        arg_iter, func_iter = self._get_call_iterators(args)
+        enumerated_funcs = enumerate(func_iter)
+        result_q = Queue()
+
+        # spawn a child for each function
+        children = []
+        for arg, (i, f) in zip(arg_iter, enumerated_funcs):
+            child_pid = self._process_in_fork(i, f, result_q, (arg,), kwargs)
+            children.append(child_pid)
+
+        # iterate over any remaining functions for which we have no args
+        for i, f in enumerated_funcs:
+            child_pid = self._process_in_fork(i, f, result_q, (), kwargs)
+            children.append(child_pid)
+
+        # the parent waits for all children to complete
+        for pid in children:
+            os.waitpid(pid, 0)
+
+        # then retrieves the results
+        results = []
+        # iterate over the result q in sorted order
+        result_q.put(None)
+        for r in sorted(iter(result_q.get, None), key=itemgetter(0)):
+            if r.exception:
+                raise exceptions.ConcurrentException(
+                    "Caught exception in child process", location=r.location
+                ) from pickle.loads(r.exception)
+            kwargs["call_state"].data.update(pickle.loads(r.call_state_data))
+            results.append(pickle.loads(r.result))
+        return self._merge_func(*results)
+
+    def _get_call_iterators(self, args):
+        return self._function_merge._get_call_iterators(args)
+
+    def _call_function(self, f, args: tuple, kwargs: dict):
+        return self._function_merge._call_function(f, args, kwargs)
+
+    def _process_in_fork(self, idx, func, result_q, args, kwargs):
+        """Call self._call_function in a child process. This function returns the ID of the child
+        in the parent process, while the child process calls _call_function, puts the results in
+        the provided queues, then exits.
+        """
+        pid = os.fork()
+        if pid:
+            return pid
+
+        # here we are the child
+        make_result = functools.partial(
+            _ConcurrentResult,
+            result=None,
+            exception=None,
+            index=idx,
+            call_state_data=None,
+            location="",
+        )
+
+        result = None
+        try:
+            r = self._call_function(func, args, kwargs)
+
+            # pickle here, so that we can't crash with pickle errors in the finally clause
+            pickled_r = pickle.dumps(r)
+            data = pickle.dumps(kwargs["call_state"].data)
+            result = make_result(result=pickled_r, call_state_data=data)
+        except Exception as e:  # noqa: BLE001
+            try:
+                # In case func does something stupid like raising an unpicklable exception
+                pickled_exception = pickle.dumps(e)
+            except (AttributeError, pickle.PicklingError):  # fmt: skip
+                # Python < 3.14 raises AttributeError for local classes; 3.14+ raises PicklingError
+                pickled_exception = pickle.dumps(
+                    AttributeError(f"Unplicklable exception raised in {func}")
+                )
+            result = make_result(
+                exception=pickled_exception,
+                location=kwargs["call_state"].highlight_active_function(),
+            )
+        finally:
+            result_q.put(result)
+            # it's necessary to explicitly close the result_q and join its background thread here,
+            # because the below os._exit won't allow time for any cleanup.
+            result_q.close()
+            result_q.join_thread()
+
+            # This is the one place that the python docs say it's normal to use os._exit. Because
+            # this is executed in a child process, calling sys.exit can have unintended
+            # consequences. e.g., anything above this that catches the resulting SystemExit can
+            # cause the child process to stay alive. the unittest framework does this.
+            os._exit(0)

metafunctions/core/decorators.py

@@ -0,0 +1,40 @@
+"""
+Internal decorators that are applied to MetaFunction methods, not functions.
+"""
+
+from collections.abc import Callable
+from functools import wraps
+
+
+def binary_operation(method):
+    """Internal decorator to apply common type checking for binary operations"""
+
+    @wraps(method)
+    def new_operation(self, other):
+        if isinstance(other, Callable):
+            new_other = self.make_meta(other)
+        else:
+            new_other = self.defer_value(other)
+        return method(self, new_other)
+
+    return new_operation
+
+
+def manage_call_state(call_method):
+    """Decorates the call method to insure call_state is present in kwargs, or create a new one.
+    If the call state isn't active, we assume we are the meta entry point.
+    """
+
+    @wraps(call_method)
+    def with_call_state(self, *args, **kwargs):
+        try:
+            call_state = kwargs["call_state"]
+        except KeyError:
+            call_state = self.new_call_state()
+            kwargs["call_state"] = call_state
+        call_state.push(self)
+        r = call_method(self, *args, **kwargs)
+        call_state.pop()
+        return r
+
+    return with_call_state

metafunctions/exceptions.py

@@ -0,0 +1,25 @@
+class MetaFunctionError(Exception):
+    pass
+
+
+class CompositionError(MetaFunctionError, TypeError):
+    "An exception that occureds when MetaFunctions are composed incorrectly"
+
+
+class CallError(MetaFunctionError, TypeError):
+    "An exception that occures when a MetaFunction is called incorrectly"
+
+    def __init__(self, *args, location: str = ""):
+        self.location = location
+        super().__init__(*args)
+
+    def __repr__(self):
+        return "{}({}, {})".format(
+            self.__class__.__name__,
+            self.args,
+            f"location='{self.location}'" if self.location else "",
+        )
+
+
+class ConcurrentException(CallError):
+    "Concurrent specific call errors (e.g., things that aren't picklable)"

metafunctions/map.py

@@ -0,0 +1,36 @@
+import itertools
+import typing as tp
+
+from metafunctions.core.concurrent import FunctionMerge
+from metafunctions.operators import concat
+
+
+class MergeMap(FunctionMerge):
+    def __init__(self, function: tp.Callable, merge_function: tp.Callable = concat):
+        """
+        MergeMap is a FunctionMerge with only one function. When called, it behaves like the
+        builtin `map` function and calls its function once per item in the iterable(s) it receives.
+        """
+        super().__init__(merge_function, (function,))
+
+    def _get_call_iterators(self, args):
+        """
+        Each element in args is an iterable.
+        """
+        args_iter = zip(*args)
+
+        # Note that EVERY element in the func iter will be called, so we need to make sure the
+        # length of our iterator is the same as the shortest iterable we received.
+        shortest_arg = min(args, key=len)
+        func_iter = itertools.repeat(self.functions[0], len(shortest_arg))
+        return args_iter, func_iter
+
+    def _call_function(self, f, args: tuple, kwargs: dict):
+        """In MergeMap, args will be a single element tuple containing the args for this function."""
+        return f(*args[0], **kwargs)
+
+    def __str__(self):
+        return f"mmap({self.functions[0]!s})"
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.functions[0]}, merge_function={self._merge_func})"

metafunctions/operators.py

@@ -0,0 +1,14 @@
+"""
+Extra operators used by MetaFunctions
+"""
+
+# These are imported for use elsewhere
+from operator import add  # noqa: F401
+from operator import mul  # noqa: F401
+from operator import sub  # noqa: F401
+from operator import truediv  # noqa: F401
+
+
+def concat(*args):
+    "concat(1, 2, 3) -> (1, 2, 3)"
+    return args

metafunctions/tests/simple_nodes.py

@@ -0,0 +1,30 @@
+"""
+Some simple nodes to use in tests.
+"""
+
+from metafunctions.api import node
+
+
+@node
+def a(x):
+    return x + "a"
+
+
+@node()
+def b(x):
+    return x + "b"
+
+
+@node()
+def c(x):
+    return x + "c"
+
+
+@node
+def d(x):
+    return x + "d"
+
+
+@node
+def e(x):
+    return x + "e"

metafunctions/tests/test_api.py

@@ -0,0 +1,87 @@
+from metafunctions.api import bind_call_state
+from metafunctions.api import node
+from metafunctions.api import recall
+from metafunctions.api import store
+from metafunctions.core import CallState
+from metafunctions.core import SimpleFunction
+from metafunctions.tests.simple_nodes import *
+from metafunctions.tests.util import BaseTestCase
+
+
+class TestUnit(BaseTestCase):
+    def test_bind_call_state(self):
+        """
+        If decorated with bind_call_state, the function receives the call state dictionary as its
+        first argument.
+        """
+
+        @node
+        @bind_call_state
+        def a_(call_state, x):
+            self.assertIsInstance(call_state, CallState)
+            call_state.data["a"] = "b"
+            return x + "a"
+
+        @node
+        @bind_call_state
+        def f(call_state, x):
+            return x + call_state.data.get("a", "f")
+
+        self.assertEqual(a("_"), "_a")
+        self.assertEqual(f("_"), "_f")
+
+        cmp = a_ | f
+        self.assertEqual(cmp("_"), "_ab")
+        cmp = f | a_ | a_ | f + f
+        self.assertEqual(cmp("_"), "_faab_faab")
+
+    def test_node_bracketless(self):
+        """
+        I'm allowing the node decorator to be applied without calling because this is how both
+        celery and function_pipes work.
+        """
+
+        @node
+        def a(x):
+            return x + "a"
+
+        @node()
+        def b(x):
+            return x + "b"
+
+        self.assertIsInstance(a, SimpleFunction)
+        self.assertIsInstance(b, SimpleFunction)
+        self.assertEqual((b | a)("_"), "_ba")
+
+    def test_store(self):
+        state = CallState()
+        abc = a | b | store("output") | c
+        big = (
+            a
+            | b
+            | c + store("ab") + store("ab2")
+            | store("abc")
+            | recall("ab") + recall("ab2")
+            | c + recall("abc")
+        )
+
+        self.assertEqual(abc("_", call_state=state), "_abc")
+        self.assertEqual(state.data["output"], "_ab")
+        self.assertEqual(big("_"), "_ab_abc_abc_ab_ab")
+
+    def test_recall(self):
+        state = a.new_call_state()
+        state.data["k"] = "secret"
+
+        cmp = a + b | store("k") | c + recall("k")
+        self.assertEqual(cmp("_"), "_a_bc_a_b")
+
+        cmp = a + b | store("k") | c + recall("k") | recall("k", from_call_state=state)
+        self.assertEqual(cmp("_"), "secret")
+
+    def test_str_store(self):
+        # this should be possible
+        self.assertEqual(str(store("key")), "store('key')")
+
+    def test_str_recall(self):
+        self.assertEqual(str(recall("key")), "recall('key')")

metafunctions/tests/test_at_operator.py

@@ -0,0 +1,40 @@
+from metafunctions.api import node
+from metafunctions.tests.simple_nodes import *
+from metafunctions.tests.util import BaseTestCase
+
+
+class TestUnit(BaseTestCase):
+    def test_broadcast(self):
+        @node
+        def f(*args):
+            return args
+
+        cmp = (a | b) @ f
+        self.assertEqual(cmp("_"), ("_", "a", "b"))
+
+        rcmp = (1, 2, 3) @ f
+        self.assertEqual(rcmp(), (1, 2, 3))
+
+    def test_str_repr(self):
+        cmp = a @ b
+        self.assertEqual(str(cmp), "(a | star(b))")
+        self.assertEqual(repr(cmp), f"FunctionChain{(a, cmp.functions[1])}")
+
+        cmp = a | a @ b * 5 / 7 | b & b @ a
+        self.assertEqual(
+            str(cmp), "(a | (((a | star(b)) * 5) / 7) | (b & (b | star(a))))"
+        )
+
+        cmp = (1, 2, 3) @ a
+        self.assertEqual(str(cmp), "((1, 2, 3) | star(a))")
+
+        cmp = a @ (b & c)
+        self.assertEqual(str(cmp), "(a | star(b & c))")
+        self.assertEqual(
+            repr(cmp),
+            f"FunctionChain({a!r}, SimpleFunction({cmp._functions[1]._function}))",
+        )
+
+    def test_upgrade_merge(self):
+        aabbcc = (a & b & c) @ (a & b & c)
+        self.assertEqual(aabbcc("_"), ("_aa", "_bb", "_cc"))

metafunctions/tests/test_call_state.py

@@ -0,0 +1,251 @@
+from metafunctions.api import bind_call_state
+from metafunctions.api import locate_error
+from metafunctions.api import mmap
+from metafunctions.api import node
+from metafunctions.core import CallState
+from metafunctions.tests.simple_nodes import *
+from metafunctions.tests.util import BaseTestCase
+
+
+class TestUnit(BaseTestCase):
+    def test_bind_call_state(self):
+        # Bound functions receive call state
+        @node
+        @bind_call_state
+        def f(call_state, x):
+            self.assertIsInstance(call_state, CallState)
+            call_state.data["h"] = "b"
+            return x + "f"
+
+        @node
+        def g(x):
+            return x + "g"
+
+        @node
+        @bind_call_state
+        def h(call_state, x):
+            self.assertIsInstance(call_state, CallState)
+            return x + call_state.data["h"]
+
+        fg = f | g | h
+        self.assertEqual(fg("_"), "_fgb")
+
+    def test_provide_call_state(self):
+        # A call state you provide is passed to all functions
+        @node
+        @bind_call_state
+        def f(call_state, x):
+            self.assertIs(call_state, c)
+            call_state.data["h"] = "d"
+            return x + "f"
+
+        @node
+        def g(x):
+            return x + "g"
+
+        @node
+        @bind_call_state
+        def h(call_state, x):
+            self.assertIs(call_state, c)
+            return x + call_state.data["h"]
+
+        fg = f | g | h
+
+        c = CallState()
+        self.assertEqual(fg("_", call_state=c), "_fgd")
+
+    def test_add_remove(self):
+        # imagine these are functions
+        a, b, c, d = ["a", "b", "c", "d"]
+        Node = CallState.Node
+
+        tree = CallState()
+        tree.push(a)
+        tree.push(a)
+        self.assertEqual(tree.pop(), a)
+        self.assertDictEqual(tree._parents, {})
+
+        tree.push(b)
+        tree.push(a)
+        self.assertDictEqual(
+            tree._parents, {Node(b, 2): Node(a, 0), Node(a, 3): Node(b, 2)}
+        )
+        self.assertDictEqual(
+            tree._children,
+            {
+                Node(a, 0): [Node(a, 1), Node(b, 2)],
+                Node(b, 2): [Node(a, 3)],
+            },
+        )
+
+        tree.push(b)
+        self.assertDictEqual(
+            tree._children,
+            {
+                Node(a, 0): [Node(a, 1), Node(b, 2)],
+                Node(b, 2): [Node(a, 3)],
+                Node(a, 3): [Node(b, 4)],
+            },
+        )
+        self.assertDictEqual(
+            tree._parents,
+            {
+                Node(b, 4): Node(a, 3),
+                Node(a, 3): Node(b, 2),
+                Node(b, 2): Node(a, 0),
+            },
+        )
+
+        self.assertEqual(tree.pop(), b)
+        self.assertEqual(tree.pop(), a)
+        for f in a, b, c, d:
+            tree.push(f)
+            tree.pop()
+        self.assertDictEqual(
+            tree._children,
+            {
+                Node(a, 0): [Node(a, 1), Node(b, 2)],
+                Node(b, 2): [
+                    Node(a, 3),
+                    Node(a, 5),
+                    Node(b, 6),
+                    Node(c, 7),
+                    Node(d, 8),
+                ],
+            },
+        )
+        self.assertDictEqual(
+            tree._parents,
+            {
+                Node(b, 2): Node(a, 0),
+            },
+        )
+
+    def test_iter_parent_nodes(self):
+        @node
+        def l(*args):
+            return []
+
+        @node
+        @bind_call_state
+        def return_parent_nodes(cs, *args):
+            return list(cs.iter_parent_nodes(cs.active_node))
+
+        def nodes2str(nodes):
+            return [(str(p.function), p.insert_index) for p in nodes]
+
+        def get_parents(f):
+            cs = CallState()
+            return nodes2str(f("", call_state=cs))
+
+        self.assertEqual(return_parent_nodes(), [(return_parent_nodes, 0)])
+
+        # create some crazy compositions to get parents out of
+        simple_chain = a | b | c | return_parent_nodes
+        self.assertListEqual(get_parents(simple_chain), [(str(simple_chain), 0)])
+
+        merges = l + (l + (l + return_parent_nodes))
+        self.assertListEqual(
+            get_parents(merges),
+            [
+                ("(l + return_parent_nodes)", 4),
+                ("(l + (l + return_parent_nodes))", 2),
+                ("(l + (l + (l + return_parent_nodes)))", 0),
+            ],
+        )
+
+    def test_highlight_active_function(self):
+        fmt_index = 7
+
+        @node
+        def ff(x):
+            return x + "F"
+
+        @node
+        @bind_call_state
+        def f(call_state, x):
+            if call_state._nodes_visited == fmt_index:
+                location_string = call_state.highlight_active_function()
+                self.assertEqual(
+                    location_string, "(a | b | ff | f | f | ->f<- | f | f)"
+                )
+                self.assertEqual(x, "_abFff")
+            return x + "f"
+
+        pipe = a | b | ff | f | f | f | f | f
+        pipe("_")
+
+        state = CallState()
+        af = a + f
+        af("_", call_state=state)
+        with self.assertRaises(AttributeError):
+            state.highlight_active_function()
+
+    def test_hightlight_active_function_no_parents(self):
+        @node
+        @bind_call_state
+        def f(cs):
+            return cs.highlight_active_function()
+
+        self.assertEqual(f(), "->f<-")
+
+    def test_highlight_active_function_multichar(self):
+        # Don't fail on long named functions. This is a regression test
+        @node
+        def fail(x):
+            if not x:
+                raise ZeroDivisionError("division by zero")
+            return x - 1
+
+        cmp = fail | fail + a
+        color = locate_error(cmp, use_color=True)
+        no_color = locate_error(cmp, use_color=False)
+        with self.assertRaises(ZeroDivisionError) as e:
+            color(1)
+        self.assertTrue(
+            e.exception.args[0].endswith("(fail | (\x1b[31m->fail<-\x1b[0m + a))")
+        )
+        with self.assertRaises(ZeroDivisionError) as e:
+            no_color(1)
+        self.assertTrue(e.exception.args[0].endswith("(fail | (->fail<- + a))"))
+
+    def test_highlight_with_map(self):
+        @node
+        def no_fail(x):
+            return x
+
+        @node
+        def fail(*args):
+            raise ZeroDivisionError("division by zero")
+
+        cmp = a + b | (c & no_fail & fail)
+        mapper = locate_error(("aaaaa", "BBBBB") | mmap(cmp), use_color=False)
+        with self.assertRaises(ZeroDivisionError) as e:
+            mapper()
+        self.assertEqual(
+            str(e.exception),
+            "division by zero \n\nOccured in the following function: "
+            "(('aaaaa', 'BBBBB') | mmap(((a + b) | (c & no_fail & ->fail<-))))",
+        )
+
+    def test_tricky_highlight(self):
+        def wrapper(mf, new_name=None):
+            @node(name=new_name or str(mf))
+            @bind_call_state
+            def meta_wrapper(cs, *args, **kwargs):
+                return mf(*args, call_state=cs, **kwargs)
+
+            return meta_wrapper
+
+        @node
+        @bind_call_state
+        def locator_string(cs, *args):
+            return cs.highlight_active_function()
+
+        cmp = a | b + c | locator_string
+
+        tim = wrapper(cmp, new_name="tim")
+        atim = a | tim
+        self.assertEqual(str(tim), "tim")
+        self.assertEqual(str(atim), "(a | tim)")
+        self.assertEqual(atim("_"), "(a | ->tim<-)")