From 6b747e62712570c837ea60257bd79c434fa76f3d Mon Sep 17 00:00:00 2001
From: Daniel Krenn <git@danielkrenn.at>
Date: Fri, 28 Aug 2015 15:57:40 +0200
Subject: [PATCH] mapping argument for .copy()

---
 src/sage/data_structures/mutable_poset.py | 29 ++++++++++++++---------
 1 file changed, 18 insertions(+), 11 deletions(-)

diff --git a/src/sage/data_structures/mutable_poset.py b/src/sage/data_structures/mutable_poset.py
index 1ebaeb1a2b9..3978e73ca5d 100644
--- a/src/sage/data_structures/mutable_poset.py
+++ b/src/sage/data_structures/mutable_poset.py
@@ -608,7 +608,7 @@ def eq(left, right):
     __eq__ = eq
 
 
-    def _copy_all_linked_(self, memo, poset):
+    def _copy_all_linked_(self, memo, poset, mapping):
         r"""
         Helper function for :meth:`MutablePoset.copy`.
 
@@ -619,6 +619,8 @@ def _copy_all_linked_(self, memo, poset):
 
         - ``poset`` -- the poset to which the newly created shells belongs.
 
+        - ``mapping`` -- a function which is applied on each of the elements.
+
         OUTPUT:
 
         A new shell.
@@ -629,7 +631,7 @@ def _copy_all_linked_(self, memo, poset):
             sage: P = MP()
             sage: Q = MP()
             sage: memo = {}
-            sage: z = P.null._copy_all_linked_(memo, Q)
+            sage: z = P.null._copy_all_linked_(memo, Q, lambda e: e)
             sage: z.poset is Q
             True
             sage: oo = z.successors().pop()
@@ -641,12 +643,13 @@ def _copy_all_linked_(self, memo, poset):
         except KeyError:
             pass
 
-        new = self.__class__(poset, self.element)
+        new = self.__class__(poset, mapping(self.element)
+                             if self.element is not None else None)
         memo[id(self)] = new
 
         for reverse in (False, True):
             for e in self.successors(reverse):
-                new.successors(reverse).add(e._copy_all_linked_(memo, poset))
+                new.successors(reverse).add(e._copy_all_linked_(memo, poset, mapping))
 
         return new
 
@@ -1231,7 +1234,7 @@ def __init__(self, data=None, key=None, merge=None, can_merge=None):
         if is_MutablePoset(data):
             if key is not None:
                 raise TypeError('Cannot use key when data is a poset.')
-            self._copy_shells_(data)
+            self._copy_shells_(data, lambda e: e)
 
         else:
             self.clear()
@@ -1435,7 +1438,7 @@ def get_key(self, element):
         return self._key_(element)
 
 
-    def _copy_shells_(self, other):
+    def _copy_shells_(self, other, mapping):
         r"""
         Helper function for copying shells.
 
@@ -1444,6 +1447,8 @@ def _copy_shells_(self, other):
         - ``other`` -- the mutable poset from which the shells
           should be copied this poset.
 
+        - ``mapping`` -- a function which is applied on each of the elements.
+
         OUTPUT:
 
         Nothing.
@@ -1461,7 +1466,7 @@ def _copy_shells_(self, other):
             sage: P.add(T((4, 4)))
             sage: P.add(T((1, 2)))
             sage: Q = MP()
-            sage: Q._copy_shells_(P)
+            sage: Q._copy_shells_(P, lambda e: e)
             sage: P.repr_full() == Q.repr_full()
             True
         """
@@ -1470,20 +1475,20 @@ def _copy_shells_(self, other):
         self._merge_ = copy(other._merge_)
         self._can_merge_ = copy(other._can_merge_)
         memo = {}
-        self._null_ = other._null_._copy_all_linked_(memo, self)
+        self._null_ = other._null_._copy_all_linked_(memo, self, mapping)
         self._oo_ = memo[id(other._oo_)]
         self._shells_ = dict((f.key, f) for f in
                               iter(memo[id(e)]
                                    for e in other._shells_.itervalues()))
 
 
-    def copy(self):
+    def copy(self, mapping=None):
         r"""
         Creates a shallow copy.
 
         INPUT:
 
-        Nothing.
+        - ``mapping`` -- a function which is applied on each of the elements.
 
         OUTPUT:
 
@@ -1505,8 +1510,10 @@ def copy(self):
             sage: P.repr_full() == Q.repr_full()
             True
         """
+        if mapping is None:
+            mapping = lambda element: element
         new = self.__class__()
-        new._copy_shells_(self)
+        new._copy_shells_(self, mapping)
         return new