Merge pull request #1055 from mathics/MapAt

Map at

Author: rocky

mathics/builtin/structure.py

@@ -536,6 +536,10 @@ class MapAt(Builtin):
     >> MapAt[f, {a, b, c, d}, 2]
      = {a, f[b], c, d}
 
+    Map $f$ onto multiple parts:
+    >> MapAt[f, {a, b, c, d}, {{1}, {4}}]
+     = {f[a], b, c, f[d]}
+
     Map $f$ onto the at the end:
     >> MapAt[f, {a, b, c, d}, -1]
      = {a, b, c, f[d]}
@@ -547,30 +551,50 @@ class MapAt(Builtin):
     Use negative position in an association:
     >> MapAt[f, <|"a" -> 1, "b" -> 2, "c" -> 3, "d" -> 4|>, -3]
      = {a -> 1, b -> f[2], c -> 3, d -> 4}
+
+    Use the operator form of MapAt:
+    >> MapAt[f, 1][{a, b, c, d}]
+     = {f[a], b, c, d}
     """
 
-    def apply(self, f, expr, n, evaluation, options={}):
-        "MapAt[f_, expr_, n_Integer]"
-        i = n.get_int_value()
+    rules = {
+        "MapAt[f_, pos_][expr_]": "MapAt[f, expr, pos]",
+    }
+
+    def apply(self, f, expr, args, evaluation, options={}):
+        "MapAt[f_, expr_, args_]"
+
         m = len(expr.leaves)
-        if 1 <= i <= m:
-            j = i - 1
-        elif -m <= i <= -1:
-            j = m + i
-        else:
-            raise PartRangeError
-
-        new_leaves = list(expr.leaves)
-        replace_leaf = new_leaves[j]
-        if hasattr(replace_leaf, "head") and replace_leaf.head == Symbol("System`Rule"):
-            new_leaves[j] = Expression(
-                "System`Rule",
-                replace_leaf.leaves[0],
-                Expression(f, replace_leaf.leaves[1]),
-            )
-        else:
-            new_leaves[j] = Expression(f, replace_leaf)
-        return List(*new_leaves)
+
+        def map_at_one(i, leaves):
+            if 1 <= i <= m:
+                j = i - 1
+            elif -m <= i <= -1:
+                j = m + i
+            else:
+                raise PartRangeError
+            replace_leaf = new_leaves[j]
+            if hasattr(replace_leaf, "head") and replace_leaf.head == Symbol("System`Rule"):
+                new_leaves[j] = Expression(
+                    "System`Rule",
+                    replace_leaf.leaves[0],
+                    Expression(f, replace_leaf.leaves[1]),
+                )
+            else:
+                new_leaves[j] = Expression(f, replace_leaf)
+            return new_leaves
+
+        a = args.to_python()
+        if isinstance(a, int):
+            new_leaves = list(expr.leaves)
+            new_leaves = map_at_one(a, new_leaves)
+            return List(*new_leaves)
+        elif isinstance(a, list):
+            new_leaves = list(expr.leaves)
+            for l in a:
+                if len(l) == 1 and isinstance(l[0], int):
+                    new_leaves = map_at_one(l[0], new_leaves)
+            return List(*new_leaves)
 
 
 class Scan(Builtin):

test/test_combinatorica.py

@@ -259,12 +259,18 @@ def test_combinatorica_permutations_1_2():
             "(4 m + 2 m ^ 2 + m ^ 4 + m ^ 8) / 8",
             "Polya counting resulting in polynomial 1.2.6, Page 25",
         ),
-        # MapAt not finished which is probably causing Poly to fail here...
-        # (
-        #     "Polya[Automorphisms[Cycle[8]], m]",
-        #     "(4 m + 2 m ^ 2 + m ^ 4 + m ^ 8) / 8",
-        #     "Polya counting resulting in polynomial 1.2.6, Page 26",
-        # ),
+        # Automorphism is slow. So we reduce Cycle[8] given as the example in the
+        # book to Cycle[3].
+        (
+            "Polya[Automorphisms[Cycle[3]], m]",
+            "(2 m + 3 m ^ 2 + m ^ 3) / 6",
+            "Polya counting resulting in polynomial 1.2.6, Page 26",
+        ),
+        (
+            "Factor[(2 m + 3 m ^ 2 + m ^ 3) / 6]",
+            "m (1 + m) (2 + m) / 6",
+            "Factor Polya polynomial 1.2.6, Page 26",
+        ),
         (
             "Factor[(4 m + 2 m^2 + 5m^4 + 4m^5 + m^8)/16]",
             "m (1 + m) (4 - 2 m + 2 m ^ 2 + 3 m ^ 3 + m ^ 4 - m ^ 5 + m ^ 6) / 16",