JS Generator in Ruby

generator.rb

class EnhancedYielder < Enumerator::Yielder
  attr_writer :yielder

  def initialize
    @error = nil
    @yielder = nil
  end

  def throw(error)
    @error = error
  end

  def yield(*args)
    if @error
      error = @error
      @error = nil
      raise error
    end

    @yielder.yield(*args)
  end
end

class Generator
  def initialize(&block)
    @done = false
    @never_called = true
    @yielder = EnhancedYielder.new
    scoped_enum_proc = lambda do |yielder|
      @yielder.yielder = yielder
      block.yield(@yielder)
    end
    @enumerator = Enumerator.new(&scoped_enum_proc)
  end

  def next(yield_return_value = nil)
    @never_called = false
    return { done: true, value: nil } if @done

    value = next_value
    @enumerator.feed(yield_return_value) unless @done
    { done: @done, value: value }
  end

  def return(return_value = nil)
    @done = true
    @never_called = false
    { done: true, value: return_value }
  end

  def throw(error)
    raise error if @done || @never_called
    @yielder.throw(error)

    self.next
  end

  private

  def next_value
    @enumerator.next
  rescue StopIteration
    @done = true
    nil
  end
end

generator_test.rb

require 'minitest/autorun'
require './lib/generator'

class GeneratorTest < Minitest::Test
  class ExampleError < StandardError; end

  ## next method

  def test_next_returns_done_if_no_yields
    generator = Generator.new do
    end

    assert_equal({ done: true, value: nil }, generator.next)
  end

  def test_next_returns_done_false_if_one_yield
    generator = Generator.new do |y|
      y.yield 10
    end

    assert_equal({ done: false, value: 10 }, generator.next)
  end

  def test_next_returns_done_true_if_one_yield_and_next_called_twice
    generator = Generator.new do |y|
      y.yield 10
    end
    generator.next

    assert_equal({ done: true, value: nil }, generator.next)
  end

  def test_next_returns_done_false_if_two_yields
    generator = Generator.new do |y|
      y.yield 10
      y.yield 20
    end

    assert_equal({ done: false, value: 10 }, generator.next)
  end

  def test_not_passing_value_to_next_makes_yielder_return_nil
    generator = Generator.new do |y|
      first_value = y.yield 10
      y.yield first_value
    end
    generator.next

    assert_equal({ done: false, value: nil }, generator.next)
  end

  def test_passing_value_to_next_sets_yielder_return_value
    generator = Generator.new do |y|
      first_value = y.yield 10
      y.yield first_value
    end
    generator.next('hello')

    assert_equal({ done: false, value: 'hello' }, generator.next)
  end

  def test_passing_value_to_final_next_sets_yielder_return_value
    generator = Generator.new do |y|
      y.yield 10
      final_value = y.yield 20
      raise ExampleError if final_value != 'hello'
    end
    generator.next

    assert_equal({ done: false, value: 20 }, generator.next('hello'))
  end

  ## return method

  def test_calling_return_without_value_when_generator_already_done
    generator = Generator.new do |y|
    end
    expected = generator.next

    assert_equal(expected, generator.return)
  end

  def test_calling_return_with_value_when_generator_already_done
    generator = Generator.new do |y|
    end

    assert_equal({ done: true, value: 30 }, generator.return(30))
  end

  def test_calling_return_ends_generator
    generator = Generator.new do |y|
      y.yield 10
    end

    assert_equal({ done: true, value: nil }, generator.return)
    assert_equal({ done: true, value: nil }, generator.next)
  end

  def test_calling_return_ends_generator_and_returns_value
    generator = Generator.new do |y|
      y.yield 10
    end

    assert_equal({ done: true, value: 30 }, generator.return(30))
    assert_equal({ done: true, value: nil }, generator.next)
    assert_equal({ done: true, value: 20 }, generator.return(20))
  end

  ## throw method

  def test_raises_error_if_thrown_before_first_call_to_next
    generator = Generator.new do |y|
      begin
        y.yield 10
      rescue ExampleError
        assert false # should not reach the assertion
      end
    end

    assert_raises ExampleError do
      generator.throw(ExampleError.new)
    end
  end

  def test_raises_error_if_generator_done_when_error_thrown
    generator = Generator.new do |y|
      begin
        y.yield 10
      rescue ExampleError
        assert false # should not reach the assertion
      end
    end
    generator.next
    assert generator.next[:done]

    assert_raises ExampleError do
      generator.throw(ExampleError.new)
    end
  end

  def test_throw_resumes_generator_by_raising_an_error_in_it
    generator = Generator.new do |y|
      begin
        y.yield 10
        raise StandardError if y.yield(20).nil? # should not reach the assertion
      rescue ExampleError
        assert true
      end
    end
    generator.next

    assert_equal({ done: true, value: nil }, generator.throw(ExampleError.new))
  end

  def test_throw_returns_same_value_as_calling_next
    generator = Generator.new do |y|
      begin
        y.yield 10
        raise StandardError if y.yield(20).nil? # should not reach the assertion
      rescue ExampleError
        assert true
      end
      y.yield 30
    end
    generator.next

    assert_equal({ done: false, value: 30 }, generator.throw(ExampleError.new))
    assert_equal({ done: true, value: nil }, generator.next)
  end

  def test_throw_will_catch_error_if_another_yield_not_in_block
    generator = Generator.new do |y|
      begin
        y.yield 10
      rescue ExampleError
        assert true
      end
      y.yield 30
    end
    generator.next

    assert_equal({ done: false, value: 30 }, generator.throw(ExampleError.new))
    assert_equal({ done: true, value: nil }, generator.next)
  end
end