From 007e63cf3b8fe6aee97ba6e53e897435512b74aa Mon Sep 17 00:00:00 2001
From: "Brandon T. Willard" <brandonwillard@users.noreply.github.com>
Date: Fri, 25 Mar 2022 21:23:31 -0500
Subject: [PATCH] Add rator and rands transform options to etuplize

---
 etuples/dispatch.py    | 36 +++++++++++++++++++++++++++++-------
 tests/test_dispatch.py | 18 ++++++++++++++++++
 2 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/etuples/dispatch.py b/etuples/dispatch.py
index 37bb0bc..4c6a435 100644
--- a/etuples/dispatch.py
+++ b/etuples/dispatch.py
@@ -102,8 +102,15 @@ def apply_ExpressionTuple(rator, rands):
 
 
 @dispatch(object)
-def etuplize(x, shallow=False, return_bad_args=False, convert_ConsPairs=True):
-    """Return an expression-tuple for an object (i.e. a tuple of rand and rators).
+def etuplize(
+    x,
+    shallow=False,
+    return_bad_args=False,
+    convert_ConsPairs=True,
+    rator_transform_fn=lambda x: x,
+    rands_transform_fn=lambda x: x,
+):
+    r"""Return an expression-tuple for an object (i.e. a tuple of rand and rators).
 
     When evaluated, the rand and rators should [re-]construct the object.  When
     the object cannot be given such a form, it is simply converted to an
@@ -112,12 +119,18 @@ def etuplize(x, shallow=False, return_bad_args=False, convert_ConsPairs=True):
     Parameters
     ----------
     x: object
-      Object to convert to expression-tuple form.
+        Object to convert to expression-tuple form.
     shallow: bool
-      Whether or not to do a shallow conversion.
+        Whether or not to do a shallow conversion.
     return_bad_args: bool
-      Return the passed argument when its type is not appropriate, instead
-      of raising an exception.
+        Return the passed argument when its type is not appropriate, instead
+        of raising an exception.
+    rator_transform_fn: callable
+        A function to be applied to each rator/CAR element of each constructed
+        `ExpressionTuple`.  The returned value is used in place of the input, and
+        the function is not applied to existing `ExpressionTuple`\s.
+    rands_transform_fn: callable
+        The same as `rator_transform_fn`, but for rands/CDR elements.
 
     """
 
@@ -133,7 +146,12 @@ def etuplize_step(
         elif (
             convert_ConsPairs and x is not None and isinstance(x, (ConsNull, ConsPair))
         ):
-            yield etuple(*x)
+            yield etuple(
+                *(
+                    (rator_transform_fn(rator(x)),)
+                    + tuple(rands_transform_fn(e) for e in rands(x))
+                )
+            )
             return
 
         try:
@@ -148,6 +166,9 @@ def etuplize_step(
             else:
                 raise TypeError(f"x is neither a non-str Sequence nor term: {type(x)}")
 
+        op = rator_transform_fn(op)
+        args = etuple(*tuple(rands_transform_fn(a) for a in args))
+
         if shallow:
             et_op = op
             et_args = args
@@ -155,6 +176,7 @@ def etuplize_step(
             et_op = yield etuplize_step(op, return_bad_args=True)
             et_args = []
             for a in args:
+
                 e = yield etuplize_step(
                     a, return_bad_args=True, convert_ConsPairs=False
                 )
diff --git a/tests/test_dispatch.py b/tests/test_dispatch.py
index 4f9d045..f77bb8f 100644
--- a/tests/test_dispatch.py
+++ b/tests/test_dispatch.py
@@ -93,6 +93,24 @@ def test_etuplize():
     assert type(etuplize(node_2)[-1]) == tuple
     assert etuplize(node_2, shallow=True) == etuple(op_1, node_1, 3, ())
 
+    def rands_transform(x):
+        if x == 1:
+            return 4
+        return x
+
+    assert etuplize(node_2, rands_transform_fn=rands_transform) == etuple(
+        op_1, etuple(op_2, 4, 2), 3, ()
+    )
+
+    def rator_transform(x):
+        if x == op_1:
+            return op_2
+        return x
+
+    assert etuplize(node_2, rator_transform_fn=rator_transform) == etuple(
+        op_2, etuple(op_2, 1, 2), 3, ()
+    )
+
 
 def test_unification():
     from cons import cons