Skip to content

Instantly share code, notes, and snippets.

@yuki-teraoka

yuki-teraoka/enumerator_lazy_ruby19.rb

Created December 4, 2012 06:55
Show Gist options
  • Save yuki-teraoka/4201351 to your computer and use it in GitHub Desktop.
Save yuki-teraoka/4201351 to your computer and use it in GitHub Desktop.
Download ZIP
Enumerator::Lazy for ruby 1.9
Raw
enumerator_lazy_ruby19.rb
module Enumerable
def lazy
enum = self
Enumerator::Lazy.new(self).tap do |obj|
obj.instance_variable_set(:@enum , enum)
obj.instance_variable_set(:@method, false)
obj.instance_variable_set(:@args , [])
end
end
end
class Enumerator::Lazy < Enumerator
class << self
def _create(enum, method = :each, *args, &block)
Enumerator::Lazy.new(&block).tap do |obj|
obj.instance_variable_set(:@enum , enum)
obj.instance_variable_set(:@method, method)
obj.instance_variable_set(:@args , args)
end
end
end
def inspect
if instance_variable_defined?(:@enum)
result = "\#<#{self.class.name}: #{@enum.inspect}"
result << ":#{@method}" if @method
result << '(' << @args.map(&:inspect).join(',') << ')' unless @args.empty?
result << '>'
return result
else
return super
end
end
def map(&block)
enum = self
raise ArgumentError, "tried to call lazy map without a block" unless block
Enumerator::Lazy._create(self, :map) do |yielder|
enum.each {|v| yielder << block.call(v)}
end
end
def collect(&block)
enum = self
raise ArgumentError, "tried to call lazy collect without a block" unless block
Enumerator::Lazy._create(self, :collect) do |yielder|
enum.each {|v| yielder << block.call(v)}
end
end
def flat_map(&block)
enum = self
raise ArgumentError, "tried to call lazy flat_map without a block" unless block
Enumerator::Lazy._create(self, :flat_map) do |yielder|
enum.each do |v|
ary = block.call(v)
unless ary.kind_of?(Enumerable)
if ary.respond_to?(:to_ary)
ary = ary.to_ary
else
ary = [ary]
end
end
ary.each {|av| yielder << av}
end
end
end
def collect_concat(&block)
enum = self
raise ArgumentError, "tried to call lazy collect_concat without a block" unless block
Enumerator::Lazy._create(self, :collect_concat) do |yielder|
enum.each do |v|
ary = block.call(v)
unless ary.kind_of?(Enumerable)
if ary.respond_to?(:to_ary)
ary = ary.to_ary
else
ary = [ary]
end
end
ary.each {|av| yielder << av}
end
end
end
def select(&block)
enum = self
raise ArgumentError, "tried to call lazy select without a block" unless block
Enumerator::Lazy._create(self, :select) do |yielder|
_enum = enum.to_enum
loop do
begin
value = _enum.next
yielder << value if block.call(value)
rescue StopIteration
break
end
end
end
end
def find_all(&block)
enum = self
raise ArgumentError, "tried to call lazy find_all without a block" unless block
Enumerator::Lazy._create(self, :find_all) do |yielder|
_enum = enum.to_enum
loop do
begin
value = _enum.next
yielder << value if block.call(value)
rescue StopIteration
break
end
end
end
end
def reject(&block)
enum = self
raise ArgumentError, "tried to call lazy reject without a block" unless block
Enumerator::Lazy._create(self, :reject) do |yielder|
_enum = enum.to_enum
loop do
begin
value = _enum.next
yielder << value unless block.call(value)
rescue StopIteration
break
end
end
end
end
def grep(pattern)
enum = self
block = block_given? ? Proc.new : nil
Enumerator::Lazy._create(self, :grep, pattern) do |yielder|
_enum = enum.to_enum
loop do
begin
value = _enum.next
if pattern === value
yielder << (block ? block.call(value) : value)
end
rescue StopIteration
break
end
end
end
end
def zip(*lists)
return super if block_given?
enum = self
enum_lists = lists.map{|list| Enumerator.new(list)}
Enumerator::Lazy._create(self, :zip, *lists) do |yielder|
enum.each do |v|
yielder << enum_lists.map do |e|
begin
e.next
rescue StopIteration
nil
end
end.unshift(v)
end
end
end
def take(n)
enum = self
Enumerator::Lazy._create(self, :take, n) do |yielder|
i = 0
next if i >= n
enum.each do |v|
yielder << v
break if (i+=1) >= n
end
end
end
def take_while(&block)
enum = self
raise ArgumentError, "tried to call lazy take_while without a block" unless block
Enumerator::Lazy._create(self, :take_while) do |yielder|
enum.each do |v|
break unless block.call(v)
yielder << v
end
end
end
def drop(n)
enum = self
Enumerator::Lazy._create(self, :drop, n) do |yielder|
i = 0
enum.each do |v|
yielder << v unless (i+=1) <= n
end
end
end
def drop_while(&block)
enum = self
raise ArgumentError, "tried to call lazy drop_while without a block" unless block
Enumerator::Lazy._create(self, :drop_while) do |yielder|
droped = true
enum.each do |v|
yielder << v unless (droped and droped = block.call(v))
end
end
end
def cycle(n = nil)
return super if block_given?
enum = self
if n
return Enumerator::Lazy._create(self, :cycle, n) do |yielder|
n.times {enum.each {|v| yielder << v}}
end
else
return Enumerator::Lazy._create(self, :cycle) do |yielder|
loop {enum.each {|v| yielder << v}}
end
end
end
def lazy
self
end
def force
to_a
end
end
# --- test code -----------------------------------------------------------------
if __FILE__ == $0
require 'test/unit'
class TestLazyEnumerator < Test::Unit::TestCase
class Step
include Enumerable
attr_reader :current, :args
def initialize(enum)
@enum = enum
@current = nil
@args = nil
end
def each(*args)
@args = args
@enum.each {|i| @current = i; yield i}
end
end
def test_initialize
assert_equal([1, 2, 3], [1, 2, 3].lazy.to_a)
assert_equal([1, 2, 3], Enumerator::Lazy.new([1, 2, 3]).to_a)
end
# def test_each_args
# a = Step.new(1..3)
# assert_equal(1, a.lazy.each(4).first)
# assert_equal([4], a.args)
# end
# def test_each_line
# name = lineno = nil
# File.open(__FILE__) do |f|
# f.each("").map do |paragraph|
# paragraph[/\A\s*(.*)/, 1]
# end.find do |line|
# if name = line[/^class\s+(\S+)/, 1]
# lineno = f.lineno
# true
# end
# end
# end
# assert_equal(self.class.name, name)
# assert_operator(lineno, :>, 2)
#
# name = lineno = nil
# File.open(__FILE__) do |f|
# f.lazy.each("").map do |paragraph|
# paragraph[/\A\s*(.*)/, 1]
# end.find do |line|
# if name = line[/^class\s+(\S+)/, 1]
# lineno = f.lineno
# true
# end
# end
# end
# assert_equal(self.class.name, name)
# assert_equal(2, lineno)
# end
def test_select
a = Step.new(1..6)
assert_equal(4, a.select {|x| x > 3}.first)
assert_equal(6, a.current)
assert_equal(4, a.lazy.select {|x| x > 3}.first)
assert_equal(4, a.current)
a = Step.new(['word', nil, 1])
assert_raise(TypeError) {a.select {|x| "x"+x}.first}
assert_equal(nil, a.current)
assert_equal("word", a.lazy.select {|x| "x"+x}.first)
assert_equal("word", a.current)
end
def test_select_multiple_values
e = Enumerator.new { |yielder|
for i in 1..5
yielder.yield(i, i.to_s)
end
}
assert_equal([[2, "2"], [4, "4"]],
e.select {|x| x[0] % 2 == 0})
assert_equal([[2, "2"], [4, "4"]],
e.lazy.select {|x| x[0] % 2 == 0}.force)
end
def test_map
a = Step.new(1..3)
assert_equal(2, a.map {|x| x * 2}.first)
assert_equal(3, a.current)
assert_equal(2, a.lazy.map {|x| x * 2}.first)
assert_equal(1, a.current)
end
def test_flat_map
a = Step.new(1..3)
assert_equal(2, a.flat_map {|x| [x * 2]}.first)
assert_equal(3, a.current)
assert_equal(2, a.lazy.flat_map {|x| [x * 2]}.first)
assert_equal(1, a.current)
end
def test_flat_map_nested
a = Step.new(1..3)
assert_equal([1, "a"],
a.flat_map {|x| ("a".."c").map {|y| [x, y]}}.first)
assert_equal(3, a.current)
assert_equal([1, "a"],
a.lazy.flat_map {|x| ("a".."c").lazy.map {|y| [x, y]}}.first)
assert_equal(1, a.current)
end
def test_flat_map_to_ary
to_ary = Class.new {
def initialize(value)
@value = value
end
def to_ary
[:to_ary, @value]
end
}
assert_equal([:to_ary, 1, :to_ary, 2, :to_ary, 3],
[1, 2, 3].flat_map {|x| to_ary.new(x)})
assert_equal([:to_ary, 1, :to_ary, 2, :to_ary, 3],
[1, 2, 3].lazy.flat_map {|x| to_ary.new(x)}.force)
end
def test_flat_map_non_array
assert_equal(["1", "2", "3"], [1, 2, 3].flat_map {|x| x.to_s})
assert_equal(["1", "2", "3"], [1, 2, 3].lazy.flat_map {|x| x.to_s}.force)
end
def test_reject
a = Step.new(1..6)
assert_equal(4, a.reject {|x| x < 4}.first)
assert_equal(6, a.current)
assert_equal(4, a.lazy.reject {|x| x < 4}.first)
assert_equal(4, a.current)
a = Step.new(['word', nil, 1])
assert_equal(nil, a.reject {|x| x}.first)
assert_equal(1, a.current)
assert_equal(nil, a.lazy.reject {|x| x}.first)
assert_equal(nil, a.current)
end
def test_reject_multiple_values
e = Enumerator.new { |yielder|
for i in 1..5
yielder.yield(i, i.to_s)
end
}
assert_equal([[2, "2"], [4, "4"]],
e.reject {|x| x[0] % 2 != 0})
assert_equal([[2, "2"], [4, "4"]],
e.lazy.reject {|x| x[0] % 2 != 0}.force)
end
def test_grep
a = Step.new('a'..'f')
assert_equal('c', a.grep(/c/).first)
assert_equal('f', a.current)
assert_equal('c', a.lazy.grep(/c/).first)
assert_equal('c', a.current)
assert_equal(%w[a e], a.grep(proc {|x| /[aeiou]/ =~ x}))
assert_equal(%w[a e], a.lazy.grep(proc {|x| /[aeiou]/ =~ x}).to_a)
end
def test_grep_with_block
a = Step.new('a'..'f')
assert_equal('C', a.grep(/c/) {|i| i.upcase}.first)
assert_equal('C', a.lazy.grep(/c/) {|i| i.upcase}.first)
end
def test_grep_multiple_values
e = Enumerator.new { |yielder|
3.times { |i|
yielder.yield(i, i.to_s)
}
}
assert_equal([[2, "2"]], e.grep(proc {|x| x == [2, "2"]}))
assert_equal([[2, "2"]], e.lazy.grep(proc {|x| x == [2, "2"]}).force)
assert_equal(["22"],
e.lazy.grep(proc {|x| x == [2, "2"]}, &:join).force)
end
def test_zip
a = Step.new(1..3)
assert_equal([1, "a"], a.zip("a".."c").first)
assert_equal(3, a.current)
assert_equal([1, "a"], a.lazy.zip("a".."c").first)
assert_equal(1, a.current)
end
def test_zip_short_arg
a = Step.new(1..5)
assert_equal([5, nil], a.zip("a".."c").last)
assert_equal([5, nil], a.lazy.zip("a".."c").force.last)
end
def test_zip_without_arg
a = Step.new(1..3)
assert_equal([1], a.zip.first)
assert_equal(3, a.current)
assert_equal([1], a.lazy.zip.first)
assert_equal(1, a.current)
end
def test_zip_with_block
# zip should be eager when a block is given
a = Step.new(1..3)
ary = []
assert_equal(nil, a.lazy.zip("a".."c") {|x, y| ary << [x, y]})
assert_equal(a.zip("a".."c"), ary)
assert_equal(3, a.current)
end
def test_take
a = Step.new(1..10)
assert_equal(1, a.take(5).first)
assert_equal(5, a.current)
assert_equal(1, a.lazy.take(5).first)
assert_equal(1, a.current)
assert_equal((1..5).to_a, a.lazy.take(5).force)
assert_equal(5, a.current)
a = Step.new(1..10)
assert_equal([], a.lazy.take(0).force)
assert_equal(nil, a.current)
end
def test_take_recycle
bug6428 = '[ruby-dev:45634]'
a = Step.new(1..10)
take5 = a.lazy.take(5)
assert_equal((1..5).to_a, take5.force, bug6428)
assert_equal((1..5).to_a, take5.force, bug6428)
end
def test_take_while
a = Step.new(1..10)
assert_equal(1, a.take_while {|i| i < 5}.first)
assert_equal(5, a.current)
assert_equal(1, a.lazy.take_while {|i| i < 5}.first)
assert_equal(1, a.current)
assert_equal((1..4).to_a, a.lazy.take_while {|i| i < 5}.to_a)
end
def test_drop
a = Step.new(1..10)
assert_equal(6, a.drop(5).first)
assert_equal(10, a.current)
assert_equal(6, a.lazy.drop(5).first)
assert_equal(6, a.current)
assert_equal((6..10).to_a, a.lazy.drop(5).to_a)
end
def test_drop_while
a = Step.new(1..10)
assert_equal(5, a.drop_while {|i| i < 5}.first)
assert_equal(10, a.current)
assert_equal(5, a.lazy.drop_while {|i| i < 5}.first)
assert_equal(5, a.current)
assert_equal((5..10).to_a, a.lazy.drop_while {|i| i < 5}.to_a)
end
def test_drop_and_take
assert_equal([4, 5], (1..Float::INFINITY).lazy.drop(3).take(2).to_a)
end
def test_cycle
a = Step.new(1..3)
assert_equal("1", a.cycle(2).map(&:to_s).first)
assert_equal(3, a.current)
assert_equal("1", a.lazy.cycle(2).map(&:to_s).first)
assert_equal(1, a.current)
end
def test_cycle_with_block
# cycle should be eager when a block is given
a = Step.new(1..3)
ary = []
assert_equal(nil, a.lazy.cycle(2) {|i| ary << i})
assert_equal(a.cycle(2).to_a, ary)
assert_equal(3, a.current)
end
def test_cycle_chain
a = 1..3
assert_equal([1,2,3,1,2,3,1,2,3,1], a.lazy.cycle.take(10).force)
assert_equal([2,2,2,2,2,2,2,2,2,2], a.lazy.cycle.select {|x| x == 2}.take(10).force)
assert_equal([2,2,2,2,2,2,2,2,2,2], a.lazy.select {|x| x == 2}.cycle.take(10).force)
end
def test_force
assert_equal([1, 2, 3], (1..Float::INFINITY).lazy.take(3).force)
end
def test_inspect
assert_equal("#<Enumerator::Lazy: 1..10:each>",
Enumerator::Lazy.new(1..10).inspect)
assert_equal("#<Enumerator::Lazy: 1..10:cycle(2)>",
Enumerator::Lazy.new(1..10, :cycle, 2).inspect)
assert_equal("#<Enumerator::Lazy: 1..10>", (1..10).lazy.inspect)
assert_equal('#<Enumerator::Lazy: #<Enumerator: "foo":each_char>>',
"foo".each_char.lazy.inspect)
assert_equal("#<Enumerator::Lazy: #<Enumerator::Lazy: 1..10>:map>",
(1..10).lazy.map {}.inspect)
assert_equal("#<Enumerator::Lazy: #<Enumerator::Lazy: 1..10>:take(0)>",
(1..10).lazy.take(0).inspect)
assert_equal("#<Enumerator::Lazy: #<Enumerator::Lazy: 1..10>:take(3)>",
(1..10).lazy.take(3).inspect)
assert_equal('#<Enumerator::Lazy: #<Enumerator::Lazy: "a".."c">:grep(/b/)>',
("a".."c").lazy.grep(/b/).inspect)
assert_equal("#<Enumerator::Lazy: #<Enumerator::Lazy: 1..10>:cycle(3)>",
(1..10).lazy.cycle(3).inspect)
assert_equal("#<Enumerator::Lazy: #<Enumerator::Lazy: 1..10>:cycle>",
(1..10).lazy.cycle.inspect)
assert_equal("#<Enumerator::Lazy: #<Enumerator::Lazy: 1..10>:cycle(3)>",
(1..10).lazy.cycle(3).inspect)
l = (1..10).lazy.map {}.collect {}.flat_map {}.collect_concat {}.select {}.find_all {}.reject {}.grep(1).zip(?a..?c).take(10).take_while {}.drop(3).drop_while {}.cycle(3)
assert_equal(<<EOS.chomp, l.inspect)
#<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: 1..10>:map>:collect>:flat_map>:collect_concat>:select>:find_all>:reject>:grep(1)>:zip("a".."c")>:take(10)>:take_while>:drop(3)>:drop_while>:cycle(3)>
EOS
end
# def test_size
# lazy = [1, 2, 3].lazy
# assert_equal 3, lazy.size
# assert_equal 42, Enumerator.new(42){}.lazy.size
# %w[map collect flat_map collect_concat].map(&:to_sym).each do |m|
# assert_equal 3, lazy.send(m){}.size
# end
# assert_equal 3, lazy.zip([4]).size
# %w[select find_all reject take_while drop_while].map(&:to_sym).each do |m|
# assert_equal nil, lazy.send(m){}.size
# end
# assert_equal nil, lazy.grep(//).size
#
# assert_equal 2, lazy.take(2).size
# assert_equal 3, lazy.take(4).size
# assert_equal 4, loop.lazy.take(4).size
# assert_equal nil, lazy.select{}.take(4).size
#
# assert_equal 1, lazy.drop(2).size
# assert_equal 0, lazy.drop(4).size
# assert_equal Float::INFINITY, loop.lazy.drop(4).size
# assert_equal nil, lazy.select{}.drop(4).size
#
# assert_equal 0, lazy.cycle(0).size
# assert_equal 6, lazy.cycle(2).size
# assert_equal 3 << 80, 4.times.inject(lazy){|enum| enum.cycle(1 << 20)}.size
# assert_equal Float::INFINITY, lazy.cycle.size
# assert_equal Float::INFINITY, loop.lazy.cycle(4).size
# assert_equal Float::INFINITY, loop.lazy.cycle.size
# assert_equal nil, lazy.select{}.cycle(4).size
# assert_equal nil, lazy.select{}.cycle.size
# end
end
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment