lib/rumonade/monad.rb in rumonade-0.4.0 vs lib/rumonade/monad.rb in rumonade-0.4.1

@@ -53,12 +53,16 @@
     #
     # @example
     #   [Some(Some(1)), Some(Some(None))], [None]].flatten
     #   #=> [1]
     #
-    def flatten_with_monad
-      bind { |x| x.is_a?(Monad) ? x.flatten_with_monad : self.class.unit(x) }
+    def flatten_with_monad(depth=nil)
+      if depth.is_a? Integer
+        depth.times.inject(self) {|e, _| e.shallow_flatten }
+      else
+        bind { |x| x.is_a?(Monad) ? x.flatten_with_monad : self.class.unit(x) }
+      end
     end
 
     # Returns a monad whose elements are all those elements of this monad for which the given predicate returned true
     def select(lam = nil, &blk)
       bind { |x| (lam || blk).call(x) ? self.class.unit(x) : self.class.empty }
@@ -69,14 +73,12 @@
     #
     # This method is equivalent to the Scala flatten call (single-level flattening), whereas #flatten is in keeping
     # with the native Ruby flatten calls (multiple-level flattening).
     #
     # @example
-    #   [Some(Some(1)), Some(Some(None))], [None]].shallow_flatten
-    #   #=> [Some(Some(1)), Some(Some(None)), None]
-    #   [Some(Some(1)), Some(Some(None)), None].shallow_flatten
-    #   #=> [Some(1), Some(None)]
-    #   [Some(1), Some(None)].shallow_flatten
+    #   [Some(Some(1)), Some(Some(None)), [None]].shallow_flatten
+    #   #=> [Some(1), Some(None), None]
+    #   [Some(1), Some(None), None].shallow_flatten
     #   #=> [1, None]
     #   [1, None].shallow_flatten
     #   #=> [1]
     #
     def shallow_flatten