Skip to content

Instantly share code, notes, and snippets.

@clowder
Created November 8, 2012 12:44
Show Gist options
  • Save clowder/4038579 to your computer and use it in GitHub Desktop.
Save clowder/4038579 to your computer and use it in GitHub Desktop.
List Comprehension in Ruby
class ForComp
def initialize(*args)
if args.first.is_a? Hash
@context = Struct.new(*args.first.keys)
end
vars = args.first.is_a?(Hash) ? args.first.values : args
enum = init_enum(vars.shift)
@enumerator = vars.inject(enum) do |enum, var|
if var.respond_to? :call
enum.flat_map { |other| var.call.map { |x| other + [x] } }
else
enum.flat_map { |other| var.map { |x| other + [x] } }
end
end
end
def to_a
@enumerator.to_a
end
def map
if defined? @context
@enumerator.map { |values| @context.new(*values).instance_eval(&Proc.new) }
else
@enumerator.map { |values| yield(*values) }
end
end
private
def init_enum(var)
var = var.respond_to?(:call) ? var.call : var
var = var.lazy if var.respond_to? :lazy
var.map { |x| [x] }
end
end
module Kernel
def forcomp(*args)
comprehension = ForComp.new(*args)
if block_given?
comprehension.map(&Proc.new)
else
comprehension
end
end
end
describe ForComp do
it "works" do
forcomp(-> { [1,2,3] }, -> { [4,5,6] }).to_a.should == [
[1,4], [1,5], [1,6],
[2,4], [2,5], [2,6],
[3,4], [3,5], [3,6]
]
end
it "yields to a block if passed" do
forcomp(-> { [1,2,3] }, -> { [4,5,6] }) do |x,y|
[x,y]
end.to_a.should == [[1, 4], [1, 5], [1, 6],
[2, 4], [2, 5], [2, 6],
[3, 4], [3, 5], [3, 6]]
end
it "work with named arguments" do
forcomp(x: -> { [1,2,3] }, y: -> { [4,5,6] }) do
[x,y]
end.to_a.should == [[1, 4], [1, 5], [1, 6],
[2, 4], [2, 5], [2, 6],
[3, 4], [3, 5], [3, 6]]
end
it "is enumerable" do
forcomp(y: -> { [1,2,3] }, x: -> { [4,5,6] }) { x*y }.select { |i| i%6==0 }.to_a.should == [6,12,12,18]
end
it "works with plain enumerables" do
forcomp(x: 1..2, y: [3,4]).to_a.should == [[1,3], [1,4],
[2,3], [2,4]]
end
it "doesn't explode on more complex examples" do
expect {
forcomp(y: -> { (1..10).select { |i| i % 2 } },
x: -> { (1..10).select { |i| i % 2 >= 0 } },
z: 1..10) do
y+x+z
end.to_a
}.not_to raise_error
end
end
@clowder
Copy link
Author

clowder commented Nov 8, 2012

Simplified & gemified over at https://github.com/clowder/comprise.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment