Advent of code 2023 day 2

def day2(s, *, part2=False):
  cubes = dict(red=12, green=13, blue=14)

  def drawn_cubes(game):
    for num_cube in game.split(', '):
      s_num, color = num_cube.split(' ')
      yield int(s_num), color

  def game_possible(game):
    return all(num <= cubes[color] for num, color in drawn_cubes(game))

  total = 0
  for game_id, line in enumerate(s.splitlines(), 1):
    games = line.split(': ')[1].split('; ')
    if part2:
      min_cubes = dict(red=0, green=0, blue=0)
      for game in games:
        for num, color in drawn_cubes(game):
          min_cubes[color] = max(min_cubes[color], num)
      total += math.prod(min_cubes.values())
    elif all(game_possible(game) for game in games):
      total += game_id

  return total