Qubyte Codes

Tip: Find and type narrow an element from an array in ruby and sorbet

Published

Recently I had a problem where I had to find the first matching element of an array by type. Ruby provides a method to return the first matching (or nil) element in an array, but sorbet isn't smart enough to type narrow it when the match is related to the type of the element.

For illustrative purposes, here's a function which takes an array of strings and integers, and returns the first string element lowercased, or nil when no strings are in the array. My first try looked like this:

sig do
  params(
    maybe_strings: T::Array[T.any(String, Integer)]
  ).returns(
    T.nilable(String)
  )
end
def find_first_and_lower(maybe_strings)
  first = maybe_strings.find { |x| x.is_a?(String) }

  # Sorbet thinks that first is 
  # `T.nilable(T.any(String, Integer))`,
  # so this doesn't work because Integer
  # has no `downcase` method.
  first&.downcase
end

This unfortunately doesn't work. The code is fine, but sorbet doesn't (at the time of writing) understand that the type of first should be T.nilable(String), so it thinks the last line of the function is incorrect.

My next attempt was to manually iterate through the array:

sig do
  params(
    maybe_strings: T::Array[T.any(String, Integer)]
  ).returns(
    T.nilable(String)
  )
end
def find_first_and_lower(maybe_strings)
  maybe_strings.each do |s|
    return s.downcase if s.is_a?(String)
  end
  nil
end

This works! It's pretty ugly though. The only good thing going for it is that it stops iterating through the array once the string is found (like find). To folk new to ruby (me) the return applying to the function as a whole and not just the block was jarring too...

In the end I used filter_map again, which encodes most of the behaviour I want:

sig do
  params(
    maybe_strings: T::Array[T.any(String, Integer)]
  ).returns(
    T.nilable(String)
  )
end
def find_first_and_lower(maybe_strings)
  first = maybe_strings
    .filter_map { |x| x if x.is_a?(String) }
    .first

  first&.downcase
end

This works too! It's not quite ideal though, because the filter_map will build us a whole new array when we only want the first (if any) element. That's why I've kept the downcase outside the filter_map. The solution was to make the filter_map lazy:

sig do
  params(
    maybe_strings: T::Array[T.any(String, Integer)]
  ).returns(
    T.nilable(String)
  )
end
def find_first_and_lower(maybe_strings)
  first = maybe_strings
    .lazy
    .filter_map { |x| x.downcase if x.is_a?(String) }
    .first
end

I've moved the downcase call into the filter_map now because it looks a little cleaner, and it will only apply to the first found element.

Bonus: How about getting the last matching element? While there's a .last method I could use, I don't want to iterate over the whole array to get to it. It turns out that reverse_each returns an enumerator when it's called without a block:

sig do
  params(
    maybe_strings: T::Array[T.any(String, Integer)]
  ).returns(
    T.nilable(String)
  )
end
def find_first_and_lower(maybe_strings)
  first = maybe_strings
    .reverse_each
    .lazy
    .filter_map { |x| x.downcase if x.is_a?(String) }
    .first
end

Otherwise, the solution looks the same, and like the one for the first element it only looks at as many entries (from the end of the array this time) as it needs.