diff --git a/src/sage/data_structures/mutable_poset.py b/src/sage/data_structures/mutable_poset.py index 2f30ccb4535..3848d59ddf7 100644 --- a/src/sage/data_structures/mutable_poset.py +++ b/src/sage/data_structures/mutable_poset.py @@ -3451,6 +3451,11 @@ def map(self, function, topological=False, reverse=False): Since this method works inplace, it is not allowed that ``function`` alters the key of an element. + .. NOTE:: + + If ``function`` returns ``None``, then the element is + removed after the mapping. + EXAMPLES:: sage: from sage.data_structures.mutable_poset import MutablePoset as MP @@ -3464,6 +3469,11 @@ def map(self, function, topological=False, reverse=False): sage: P poset((1, 2, 3), (1, 3, 4), (2, 1, 3), (2, 2, 4), (4, 4, 8)) + TESTS:: + + sage: P.map(lambda e: e if e[2] != 4 else None); P + poset((1, 2, 3), (2, 1, 3), (4, 4, 8)) + .. SEEALSO:: :meth:`copy`, @@ -3471,8 +3481,13 @@ def map(self, function, topological=False, reverse=False): """ shells = self.shells_topological(reverse=reverse) \ if topological else self.shells() + remove = [] for shell in shells: shell._element_ = function(shell._element_) + if shell._element_ is None: + remove.append(shell.key) + for key in remove: + self.remove(key) def mapped(self, function):