-

cool-RR · cool-RR · commit 92edbc3cb072 · 2016-04-20T23:15:26.000+03:00
diff --git a/source_py3/python_toolbox/math_tools/misc.py b/source_py3/python_toolbox/math_tools/misc.py
@@ -3,6 +3,9 @@
 
 import numbers
 import math
+import random
+
+from python_toolbox import nifty_collections
 
 
 infinity = float('inf')
@@ -153,12 +156,19 @@ def is_integer(x):
     return inted_x == x
     
     
-def cute_round(x, step=1, up=False):
+class RoundMode(nifty_collections.CuteEnum):
+    CLOSEST_OR_DOWN = 0
+    CLOSEST_OR_UP = 1
+    ALWAYS_DOWN = 2
+    ALWAYS_UP = 3
+    PROBABILISTIC = 4
+
+def cute_round(x, round_mode=RoundMode.CLOSEST_OR_DOWN, *, step=1):
     '''
     Round with a chosen step.
     
     Examples:
-    
+    blocktododoc
         >>> cute_round(7.456)
         7
         >>> cute_round(7.456, up=True)
@@ -178,6 +188,17 @@ def cute_round(x, step=1, up=False):
         
     '''
     div, mod = divmod(x, step)
-    return (div + bool(mod and up)) * step
+    if round_mode == RoundMode.CLOSEST_OR_DOWN:
+        round_up = (mod > 0.5 * step)
+    elif round_mode == RoundMode.CLOSEST_OR_UP:
+        round_up = (mod >= 0.5 * step)
+    elif round_mode == RoundMode.ALWAYS_DOWN:
+        round_up = False
+    elif round_mode == RoundMode.ALWAYS_UP:
+        round_up = True
+    else:
+        assert round_mode == RoundMode.PROBABILISTIC
+        round_up = random.random() < mod / step
+    return (div + round_up) * step
     
     
diff --git a/source_py3/python_toolbox/math_tools/statistics.py b/source_py3/python_toolbox/math_tools/statistics.py
@@ -29,4 +29,4 @@ def get_mean(iterable):
         sum_ += value
     return sum_ / (i + 1)
 
-        
+        
diff --git a/source_py3/test_python_toolbox/test_math_tools/test_cute_round.py b/source_py3/test_python_toolbox/test_math_tools/test_cute_round.py
@@ -1,18 +1,22 @@
 # Copyright 2009-2015 Ram Rachum.
 # This program is distributed under the MIT license.
 
-from python_toolbox.math_tools import cute_round
+import inspect
+
+from python_toolbox.math_tools import cute_round, RoundMode
 
 def almost_equals(x, y):
     return (abs(1-(x / y)) < (10 ** -10))
     
 
-def test_cute_round():
-    assert almost_equals(cute_round(7.456), 7)
-    assert almost_equals(cute_round(7.456, up=True), 8)
-    assert almost_equals(cute_round(7.456, step=0.1), 7.4)
-    assert almost_equals(cute_round(7.456, step=0.1, up=True), 7.5)
-    assert almost_equals(cute_round(7.456, step=0.2), 7.4)
-    assert almost_equals(cute_round(7.456, step=0.2, up=True), 7.6)
-    assert almost_equals(cute_round(7.456, step=0.01), 7.45)
-    assert almost_equals(cute_round(7.456, step=0.01, up=True), 7.46)
+class CuteRoundTestCase:
+    def test_closest_or_down(self):
+        full_arg_spec = inspect.getfullargspec(cute_round)
+        assert RoundMode.CLOSEST_OR_DOWN in full_arg_spec.defaults
+                                                          
+        assert almost_equals(cute_round(7.456), 7)
+        assert almost_equals(cute_round(7.654), 8)
+        assert almost_equals(cute_round(7.5), 7)
+        assert almost_equals(cute_round(7.456, step=0.1), 7.5)
+        assert almost_equals(cute_round(7.456, step=0.2), 7.4)
+        assert almost_equals(cute_round(7.456, step=0.01), 7.46)
