Skip to the code
cf.

Flatten nested iterables

Write a generator function flatten(iterable, depth = Infinity). It yields the values of iterable, and when a value is itself iterable (an array, a Set, a generator, or anything else with a [Symbol.iterator] method), it yields that value's values instead, at any depth of nesting.

Strings are iterable, but flatten yields them whole, never character by character. Values that aren't iterable, such as numbers, null, and plain objects, are yielded as they are.

depth is how many levels of nesting to open, as with Array.prototype.flat. With depth 1, only the values directly inside iterable are opened, and anything nested deeper is yielded as it is. With depth 0, it yields iterable's values unchanged.

It's lazy, so it works with endless iterables, nested or not. When the caller stops early, it closes the iterators it has open (calls their return()), so a generator's finally blocks run.

[...flatten([1, [2, [3, [4]]], 5])];             // [1, 2, 3, 4, 5]
[...flatten([new Set([1, 2]), ["ab", ["c"]]])]; // [1, 2, "ab", "c"]
[...flatten([1, [2, [3, [4]]]], 1)];             // [1, 2, [3, [4]]]
[...flatten([1, [2, [3, [4]]]], 2)];             // [1, 2, 3, [4]]

What the tests check

  1. flattens nested arrays all the way down
    expect([...flatten([1, [2, [3, [4]]], 5])]).toEqual([1, 2, 3, 4, 5]);
    expect([...flatten([[[]], []])]).toEqual([]);
    
  2. opens Sets, generators, and any other iterable
    function* pair() {
      yield 3;
      yield [4, 5];
    }
    const custom = {
      *[Symbol.iterator]() {
        yield 6;
        yield new Set([7]);
      },
    };
    expect([...flatten([new Set([1, 2]), pair(), [custom]])]).toEqual([1, 2, 3, 4, 5, 6, 7]);
    
  3. yields strings whole
    expect([...flatten(["ab", ["cd", ["e"]], ""])]).toEqual(["ab", "cd", "e", ""]);
    
  4. yields values that aren't iterable as they are
    const object = { a: 1 };
    const values = [...flatten([object, null, undefined, 0, false, [NaN]])];
    expect(values).toEqual([object, null, undefined, 0, false, NaN]);
    expect(values[0]).toBe(object);
    
  5. depth limits how many levels are opened
    expect([...flatten([1, [2, [3, [4]]]], 1)]).toEqual([1, 2, [3, [4]]]);
    expect([...flatten([1, [2, [3, [4]]]], 2)]).toEqual([1, 2, 3, [4]]);
    expect([...flatten([1, [2], new Set([3])], 0)]).toEqual([1, [2], new Set([3])]);
    
  6. is lazy, so endless iterables work
    // Endless for lazy code; eager code gets an error instead of a frozen page.
    function* naturals() {
      for (let n = 1; n < 1000; n++) yield n;
      throw new Error("read 1000 values from an endless generator: something isn't lazy");
    }
    const values = flatten([[0], naturals(), "never reached"]);
    const first = [];
    for (let i = 0; i < 4; i++) first.push(values.next().value);
    expect(first).toEqual([0, 1, 2, 3]);
    expect(flatten(naturals()).next()).toEqual({ value: 1, done: false });
    
  7. closes the iterators it has open when the caller stops early
    let closed = false;
    function* inner() {
      try {
        yield "a";
        yield "b";
        yield "c";
      } finally {
        closed = true;
      }
    }
    const seen = [];
    for (const value of flatten([[1], [inner()]])) {
      seen.push(value);
      if (value === "b") break;
    }
    expect(seen).toEqual([1, "a", "b"]);
    expect(closed).toBe(true);
    

Learn more

solution.jsSaved in this browser

Run tests runs every test on your code, here in your browser. The problem is solved when they all pass.

Ctrl + Enter runs the tests (⌘ + Enter on a Mac).