Skip to content

Commit

Permalink
Literal::Array#zip (#268)
Browse files Browse the repository at this point in the history
This allows you to zip a Literal::Array with other Arrays or
Literal::Arrays returning a Literal::Array of Literal::Tuples.

Closes #216
  • Loading branch information
joeldrapper authored Dec 19, 2024
1 parent ae16a55 commit bfc853b
Show file tree
Hide file tree
Showing 2 changed files with 160 additions and 0 deletions.
71 changes: 71 additions & 0 deletions lib/literal/array.rb
Original file line number Diff line number Diff line change
Expand Up @@ -645,4 +645,75 @@ def |(other)
def fetch(...)
@__value__.fetch(...)
end

def zip(*others)
other_types = others.map do |other|
case other
when Literal::Array
other.__type__
when Array
_Any?
else
raise ArgumentError
end
end

tuple = Literal::Tuple(
@__type__,
*other_types
)

my_length = length
max_length = [my_length, *others.map(&:length)].max

# Check we match the max length or our type is nilable
unless my_length == max_length || @__type__ === nil
raise ArgumentError.new(<<~MESSAGE)
The literal array could not be zipped becuase its type is not nilable and it has fewer items than the maximum number of items in the other arrays.
You can either make the type of this array nilable, or add more items so its length matches the others.
#{inspect}
MESSAGE
end

# Check others match the max length or their types is nilable
others.each_with_index do |other, index|
unless other.length == max_length || other_types[index] === nil
raise ArgumentError.new(<<~MESSAGE)
The literal array could not be zipped becuase its type is not nilable and it has fewer items than the maximum number of items in the other arrays.
You can either make the type of this array nilable, or add more items so its length matches the others.
#{inspect}
MESSAGE
end
end

i = 0

if block_given?
while i < max_length
yield tuple.new(
@__value__[i],
*others.map { |it| it[i] }
)
i += 1
end

nil
else
result_value = []

while i < max_length
result_value << tuple.new(
@__value__[i],
*others.map { |it| it[i] }
)
i += 1
end

__with__(result_value)
end
end
end
89 changes: 89 additions & 0 deletions test/array.test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -842,3 +842,92 @@
[2, 4],
]
end

test "#zip with other literal arrays when all the lengths match" do
a = Literal::Array(String).new("a", "b")
b = Literal::Array(Integer).new(1, 2)
c = Literal::Array(Symbol).new(:a, :b)

assert_equal a.zip(b, c), Literal::Array(
Literal::Tuple(String, Integer, Symbol)
).new(
Literal::Tuple(String, Integer, Symbol).new("a", 1, :a),
Literal::Tuple(String, Integer, Symbol).new("b", 2, :b),
)
end

test "#zip with other regular arrays" do
a = Literal::Array(String).new("a", "b")
b = [1, 2]
c = [:a, :b]

assert_equal a.zip(b, c), Literal::Array(
Literal::Tuple(String, _Any?, _Any?)
).new(
Literal::Tuple(String, _Any?, _Any?).new("a", 1, :a),
Literal::Tuple(String, _Any?, _Any?).new("b", 2, :b),
)
end

test "#zip with other literal arrays where one of the others length is not the max length and the type is not nilable" do
a = Literal::Array(String).new("a", "b")
b = Literal::Array(Integer).new(1)
c = Literal::Array(Symbol).new(:a, :b)

assert_raises ArgumentError do
a.zip(b, c)
end
end

test "#zip with literal arrays where our length is not the max length and the type is not nilable" do
a = Literal::Array(String).new("a")
b = Literal::Array(Integer).new(1, 2)
c = Literal::Array(Symbol).new(:a, :b)

assert_raises ArgumentError do
a.zip(b, c)
end
end

test "#zip when our length is not the max length but the type is nilable" do
a = Literal::Array(_Nilable(String)).new("a")
b = [1, 2]
c = [:a, :b]

assert_equal a.zip(b, c), Literal::Array(
Literal::Tuple(_Nilable(String), _Any, _Any)
).new(
Literal::Tuple(_Nilable(String), _Any, _Any).new("a", 1, :a),
Literal::Tuple(_Nilable(String), _Any, _Any).new(nil, 2, :b),
)
end

test "#zip with others length is not the max length but the types are nilable" do
a = Literal::Array(String).new("a", "b")
b = Literal::Array(_Nilable(Integer)).new(1)
c = [:a]

assert_equal a.zip(b, c), Literal::Array(
Literal::Tuple(String, _Nilable(Integer), _Any)
).new(
Literal::Tuple(String, _Nilable(Integer), _Any).new("a", 1, :a),
Literal::Tuple(String, _Nilable(Integer), _Any).new("b", nil, nil),
)
end

test "#zip with a block" do
a = Literal::Array(String).new("a", "b")
b = Literal::Array(Integer).new(1, 2)
c = [:a, :b]

results = []

return_value = a.zip(b, c) { |it| results << it }

assert_equal results, [
Literal::Tuple(String, Integer, _Any).new("a", 1, :a),
Literal::Tuple(String, Integer, _Any).new("b", 2, :b),
]

assert_equal return_value, nil
end

0 comments on commit bfc853b

Please sign in to comment.