Advent of code 2023 day 3

def day3(s, *, part2=False):
  grid = np.pad([list(line) for line in s.splitlines()], 1, constant_values='.')

  def neighbors(y, x):
    for dy, dx in set(itertools.product((-1, 0, 1), repeat=2)) - {(0, 0)}:
      yield y + dy, x + dx

  def number_span(y, x):
    while grid[y, x - 1].isdigit():  # Move to first digit.
      x -= 1
    x_start = x
    while grid[y, x].isdigit():  # Move past last digit.
      x += 1
    return y, x_start, x

  def adjacent_number_spans(yx):
    return {number_span(*yx2) for yx2 in neighbors(*yx) if grid[yx2].isdigit()}

  def get_value(y, x_start, x_stop):
    return int(''.join(grid[y, x_start:x_stop]))

  if part2:
    yxs = (yx for yx, ch in np.ndenumerate(grid) if ch == '*')
    gears = (spans for yx in yxs if len(spans := adjacent_number_spans(yx)) == 2)
    return sum(math.prod(get_value(*span) for span in gear) for gear in gears)

  else:
    yxs = (yx for yx, ch in np.ndenumerate(grid) if ch not in '.0123456789')
    spans = set.union(*(adjacent_number_spans(yx) for yx in yxs))
    return sum(get_value(*span) for span in spans)