-

cool-RR · cool-RR · commit 7a463edfc88e · 2014-10-08T13:16:18.000+03:00
diff --git a/source_py3/python_toolbox/nifty_collections/bagging.py b/source_py3/python_toolbox/nifty_collections/bagging.py
@@ -8,6 +8,8 @@
 import collections
 import functools
 
+from python_toolbox import math_tools
+
 from .lazy_tuple import LazyTuple
 from .ordered_dict import OrderedDict
 from .frozen_dict_and_frozen_ordered_dict import FrozenDict, FrozenOrderedDict
@@ -21,6 +23,33 @@ def _count_elements(mapping, iterable):
         mapping_get = mapping.get
         for element in iterable:
             mapping[element] = mapping_get(element, 0) + 1
+            
+            
+class _NO_DEFAULT: 
+    pass
+            
+class _ZeroCountAttempted(Exception):
+    pass
+        
+
+def _process_count(count):
+    if not math_tools.is_integer(count):
+        raise TypeError(
+            'You passed %s as the count of %s, while a `Bag` can only handle '
+            'integer counts.' % (count, key)
+        )
+    if count < 0:
+        raise TypeError(
+            "You passed %s as the count of %s, while `Bag` doesn't support "
+            "negative amounts." % (count, key)
+        )
+    
+    if count == 0:
+        raise _ZeroCountAttempted
+    
+    return int(count)
+    
+    
 
 class _BaseBagMixin:
     '''Mixin for `FrozenBag` and `FrozenOrderedBag`.'''
@@ -32,23 +61,10 @@ def __init__(self, iterable={}):
         
         if isinstance(iterable, collections.Mapping):
             for key, value, in iterable.items():
-                if not math_tools.is_integer(value):
-                    raise TypeError(
-                        'You passed %s as the count of %s, while '
-                        'a `Bag` can only handle integer counts.' %
-                                                                   (value, key)
-                    )
-                if value < 0:
-                    raise TypeError(
-                        "You passed %s as the count of %s, while "
-                        "`Bag` doesn't support negative amounts." %
-                                                                   (value, key)
-                    )
-                    
-                if value == 0:
+                try:
+                    self._dict[key] = _process_count(value)
+                except _ZeroCountAttempted:
                     continue
-                
-                self._dict[key] = int(value)
         else:
             _count_elements(self._dict, iterable)
 
@@ -291,8 +307,52 @@ def __repr__(self):
 
 class _MutableBagMixin(_BaseBagMixin):
     # blocktodo: ensure all mutable methods, like __iadd__ and everything
-    pass
+    
+    def __setitem__(self, i, count):
+        try:
+            super().__setitem__(i, _process_count(count))
+        except _ZeroCountAttempted:
+            del self[i]
+        
+    
+    def __imod__(self, modulo):
+        for key in tuple(self.keys()):
+            self[key] %= modulo
+        return self
+    
+    def setdefault(self, key, default):
+        current_count = self[key]
+        if current_count > 0:
+            return current_count
+        else:
+            self[key] = default
+            return default
+
+    def __delitem__(self, key):
+        # We're making `__delitem__` not raise an exception on missing or
+        # zero-count elements because we're automatically deleting zero-count
+        # elements even though they seem to exist from the outside, so we're
+        # avoiding raising exceptions where someone would try to explicitly
+        # delete them.
+        try:
+            del self._dict[key]
+        except KeyError:
+            pass
         
+    def pop(self, key, default=_NO_DEFAULT):
+        '''D.pop(k[,d]) -> v, remove specified key and return the corresponding value.
+          If key is not found, d is returned if given, otherwise KeyError is raised.
+        '''
+        value = self[key]
+        if value == 0 and default is not _NO_DEFAULT:
+            return default
+        else:
+            del self[key]
+            return value
+
+
+    
+
 class _OrderedBagMixin(Ordered):
     def __repr__(self):
         if not self:
diff --git a/source_py3/test_python_toolbox/test_nifty_collections/test_bagging.py b/source_py3/test_python_toolbox/test_nifty_collections/test_bagging.py
@@ -209,7 +209,7 @@ def test_mutating(self):
 
         bag = self.bag_type('abracadabra')
         bag += bag
-        assert bag == self.bag_type('abracadabr' * 2)
+        assert bag == self.bag_type('abracadabra' * 2)
 
         bag = self.bag_type('abracadabra')
         bag -= bag
@@ -223,10 +223,6 @@ def test_mutating(self):
         bag['a'] = 7
         assert bag == self.bag_type('abracadabra' + 'aa')
 
-        bag = self.bag_type('abracadabra')
-        bag.set('a', 7)
-        assert bag == self.bag_type('abracadabra' + 'aa')
-
         bag = self.bag_type('abracadabra')
         assert bag.setdefault('a', 7) == 5
         assert bag == self.bag_type('abracadabra')
@@ -236,7 +232,9 @@ def test_mutating(self):
         assert bag == self.bag_type('abracadabra' + 'x' * 7)
 
         bag = self.bag_type('abracadabra')
-        assert bag.pop('a', 7) == 7
+        assert bag.pop('a', 7) == 5
+        assert bag == self.bag_type('brcdbr')
+        assert bag.pop('x', 7) == 7
         assert bag == self.bag_type('brcdbr')
 
         bag = self.bag_type('abracadabra')
@@ -289,8 +287,6 @@ def test_mutating(bag_type):
             bag %= bag
         with cute_testing.RaiseAssertor(TypeError):
             bag['a'] = 7
-        with cute_testing.RaiseAssertor(TypeError):
-            bag.set('a', 7)
         with cute_testing.RaiseAssertor(TypeError):
             bag.setdefault('a', 7)
         with cute_testing.RaiseAssertor(TypeError):
