@@ -536,6 +536,10 @@ class MapAt(Builtin):
536536 >> MapAt[f, {a, b, c, d}, 2]
537537 = {a, f[b], c, d}
538538
539+ Map $f$ onto multiple parts:
540+ >> MapAt[f, {a, b, c, d}, {{1}, {4}}]
541+ = {f[a], b, c, f[d]}
542+
539543 Map $f$ onto the at the end:
540544 >> MapAt[f, {a, b, c, d}, -1]
541545 = {a, b, c, f[d]}
@@ -557,28 +561,40 @@ class MapAt(Builtin):
557561 "MapAt[f_, pos_][expr_]" : "MapAt[f, expr, pos]" ,
558562 }
559563
560- def apply (self , f , expr , n , evaluation , options = {}):
561- "MapAt[f_, expr_, n_Integer ]"
562- i = n . get_int_value ()
564+ def apply (self , f , expr , args , evaluation , options = {}):
565+ "MapAt[f_, expr_, args_ ]"
566+
563567 m = len (expr .leaves )
564- if 1 <= i <= m :
565- j = i - 1
566- elif - m <= i <= - 1 :
567- j = m + i
568- else :
569- raise PartRangeError
570-
571- new_leaves = list (expr .leaves )
572- replace_leaf = new_leaves [j ]
573- if hasattr (replace_leaf , "head" ) and replace_leaf .head == Symbol ("System`Rule" ):
574- new_leaves [j ] = Expression (
575- "System`Rule" ,
576- replace_leaf .leaves [0 ],
577- Expression (f , replace_leaf .leaves [1 ]),
578- )
579- else :
580- new_leaves [j ] = Expression (f , replace_leaf )
581- return List (* new_leaves )
568+
569+ def map_at_one (i , leaves ):
570+ if 1 <= i <= m :
571+ j = i - 1
572+ elif - m <= i <= - 1 :
573+ j = m + i
574+ else :
575+ raise PartRangeError
576+ replace_leaf = new_leaves [j ]
577+ if hasattr (replace_leaf , "head" ) and replace_leaf .head == Symbol ("System`Rule" ):
578+ new_leaves [j ] = Expression (
579+ "System`Rule" ,
580+ replace_leaf .leaves [0 ],
581+ Expression (f , replace_leaf .leaves [1 ]),
582+ )
583+ else :
584+ new_leaves [j ] = Expression (f , replace_leaf )
585+ return new_leaves
586+
587+ a = args .to_python ()
588+ if isinstance (a , int ):
589+ new_leaves = list (expr .leaves )
590+ new_leaves = map_at_one (a , new_leaves )
591+ return List (* new_leaves )
592+ elif isinstance (a , list ):
593+ new_leaves = list (expr .leaves )
594+ for l in a :
595+ if len (l ) == 1 and isinstance (l [0 ], int ):
596+ new_leaves = map_at_one (l [0 ], new_leaves )
597+ return List (* new_leaves )
582598
583599
584600class Scan (Builtin ):
0 commit comments