Add experimental partitionN() function

...which is like `partition()` but allows partitioning items into more
than two buckets. The reason it's experimental is that the type
signature is a lot more complicated to understand than the `partition()`
one.

CHANGELOG.md

@@ -1,5 +1,15 @@
 ## [Unreleased]
 
+- Add new `partitionN()` function, which is like `partition()` but
+  allows partitioning items into more than two buckets.
+  The reason the `partition()` isn't extended to support more than one
+  predicate is that the type signature is significantly more complex,
+  and I'm not sure yet if I want to pay that price for the most common case.
+  ```ts
+  const [matches1, nonMatches] = partitionN(items, pred1);
+  const [matches1, matches2, rest] = partitionN(items, pred1, pred2);
+  const [matches1, matches2, matches3, rest] = partitionN(items, pred1, pred2, pred3);
+  ```
 - Add second param `index` to all predicates. This will make operations like
   partitioning a list based on the element position as easy as partitioning
   based on the element value, for example:

src/index.ts

@@ -53,6 +53,7 @@ export {
   itake,
   pairwise,
   partition,
+  partitionN,
   roundrobin,
   take,
   uniqueEverseen,

src/more-itertools.ts

@@ -1,4 +1,4 @@
-import { enumerate, iter, map } from "./builtins";
+import { iter, map } from "./builtins";
 import { izip, repeat } from "./itertools";
 import type { Predicate, Primitive } from "./types";
 import { primitiveIdentity } from "./utils";
@@ -102,9 +102,8 @@ export function* pairwise<T>(iterable: Iterable<T>): IterableIterator<[T, T]> {
 }
 
 /**
- * Returns a 2-tuple of arrays.  Splits the elements in the input iterable into
- * either of the two arrays.  Will fully exhaust the input iterable.  The first
- * array contains all items that match the predicate, the second the rest:
+ * Splits an input array's elements into two arrays: ones that match the
+ * predicate, and ones that don't. Will fully exhaust the input iterable.
  *
  *     >>> const isOdd = x => x % 2 !== 0;
  *     >>> const iterable = range(10);
@@ -117,7 +116,7 @@ export function* pairwise<T>(iterable: Iterable<T>): IterableIterator<[T, T]> {
  */
 export function partition<T, N extends T>(
   iterable: Iterable<T>,
-  predicate: (item: T, index: number) => item is N,
+  predicate: TypePredicate<T, N>,
 ): [N[], Exclude<T, N>[]];
 export function partition<T>(iterable: Iterable<T>, predicate: Predicate<T>): [T[], T[]];
 export function partition<T>(iterable: Iterable<T>, predicate: Predicate<T>): [T[], T[]] {
@@ -136,6 +135,53 @@ export function partition<T>(iterable: Iterable<T>, predicate: Predicate<T>): [T
   return [good, bad];
 }
 
+type TypePredicate<T, N extends T> = (item: T, index: number) => item is N;
+type PredicateUnion<T, P extends (TypePredicate<T, T> | Predicate<T>)[]> = P extends TypePredicate<T, infer U>[]
+  ? U
+  : never;
+
+/**
+ * Like partition(), but takes N predicates and returns N + 1 "buckets".
+ *
+ *     >>> const isNegative = x => x < 0;
+ *     >>> const isOdd = x => x % 2 !== 0;
+ *     >>> const iterable = range(-5, 10);
+ *     >>> const [negs, odds, evens] = partition(iterable, isNegative, isOdd);
+ *     >>> negs
+ *     [-5, -4, -3, -2, -1]
+ *     >>> odds
+ *     [1, 3, 5, 7, 9]
+ *     >>> evens
+ *     [0, 2, 4, 6, 8]
+ *
+ */
+export function partitionN<T, Ps extends (TypePredicate<T, T> | Predicate<T>)[]>(
+  iterable: Iterable<T>,
+  ...predicates: [...Ps]
+): {
+  [K in keyof Ps]: Ps[K] extends TypePredicate<T, infer U> ? U[] : T[];
+} & { [K in Ps["length"]]: Exclude<T, PredicateUnion<T, Ps>>[] } {
+  if (predicates.length === 0) throw new Error("Must provide at least one predicate");
+
+  const buckets: T[][] = predicates.map(() => []);
+  const rest: T[] = [];
+
+  let index = 0;
+  outer: for (const item of iterable) {
+    for (let i = 0; i < predicates.length; i++) {
+      if (predicates[i](item, index)) {
+        buckets[i].push(item);
+        continue outer;
+      }
+    }
+    rest.push(item);
+    index++;
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  return [...buckets, rest] as any;
+}
+
 /**
  * Yields the next item from each iterable in turn, alternating between them.
  * Continues until all items are exhausted.

test-d/inference.test-d.ts

@@ -1,14 +1,40 @@
-import { partition } from "../dist";
+import { partition, partitionN } from "../dist";
 import { expectType } from "tsd";
 
 function isStr(x: unknown): x is string {
   return typeof x === "string";
 }
 
+function isNum(x: unknown): x is number {
+  return typeof x === "number";
+}
+
 {
   // partition with type predicate
   const items: unknown[] = [1, 2, null, true, 0, "hi", false, -1];
   const [strings, others] = partition(items, isStr);
   expectType<string[]>(strings);
   expectType<unknown[]>(others);
 }
+
+{
+  // partitionN with 2 (type) predicates
+  const items: unknown[] = [1, 2, null, true, 0, "hi", false, -1];
+  const [strings, numbers, others] = partitionN(items, isStr, isNum);
+  expectType<string[]>(strings);
+  expectType<number[]>(numbers);
+  expectType<unknown[]>(others);
+}
+
+{
+  // partitionN with no type predicate
+  const items: unknown[] = [1, 2, null, true, 0, "hi", false, -1];
+  const [strings, numbers, others] = partitionN(
+    items,
+    (x) => String(typeof x) === "string",
+    (x) => String(typeof x) === "number",
+  );
+  expectType<unknown[]>(strings);
+  expectType<unknown[]>(numbers);
+  expectType<unknown[]>(others);
+}

test/more-itertools.test.ts

@@ -9,6 +9,7 @@ import {
   intersperse,
   pairwise,
   partition,
+  partitionN,
   roundrobin,
   take,
   uniqueEverseen,
@@ -20,6 +21,10 @@ const isEven = (x: number) => x % 2 === 0;
 const isEvenIndex = (_: unknown, index: number) => index % 2 === 0;
 const isPositive = (x: number) => x >= 0;
 
+function isString(value: unknown): value is string {
+  return typeof value === "string";
+}
+
 function isNum(value: unknown): value is number {
   return typeof value === "number";
 }
@@ -276,6 +281,45 @@ describe("partition", () => {
   });
 });
 
+describe("partitionN", () => {
+  it("partitionN without predicates throws error", () => {
+    expect(() => partitionN([])).toThrow("Must provide at least one predicate");
+  });
+
+  it("partitionN empty list", () => {
+    expect(partitionN([], isEven)).toEqual([[], []]);
+    expect(partitionN([], isEvenIndex)).toEqual([[], []]);
+    expect(partitionN([], isPositive)).toEqual([[], []]);
+    expect(partitionN([], isPositive, isEven)).toEqual([[], [], []]);
+    expect(partitionN([], isEven, isEvenIndex, isPositive)).toEqual([[], [], [], []]);
+  });
+
+  it("partitionN splits input list into N lists", () => {
+    const values = [1, -2, 3, 4, 5, 6, 8, 8, 0, -2, -3];
+    expect(partitionN(values, isEven)).toEqual([
+      [-2, 4, 6, 8, 8, 0, -2],
+      [1, 3, 5, -3],
+    ]);
+    expect(partitionN(values, isPositive)).toEqual([
+      [1, 3, 4, 5, 6, 8, 8, 0],
+      [-2, -2, -3],
+    ]);
+    expect(partitionN(values, isEven, isPositive)).toEqual([[-2, 4, 6, 8, 8, 0, -2], [1, 3, 5], [-3]]);
+    expect(partitionN(values, isPositive, isEven)).toEqual([[1, 3, 4, 5, 6, 8, 8, 0], [-2, -2], [-3]]);
+  });
+
+  it("partitionN retains rich type info", () => {
+    const values = ["hi", 3, null, "foo", -7];
+    const [nums, strings, rest] = partitionN(values, isNum, isString);
+    expect(nums).toEqual([3, -7]);
+    //     ^^^^ number[]
+    expect(strings).toEqual(["hi", "foo"]);
+    //     ^^^^ number[]
+    expect(rest).toEqual([null]);
+    //     ^^^ (string | null)[]
+  });
+});
+
 describe("roundrobin", () => {
   it("roundrobin on empty list", () => {
     expect(Array.from(roundrobin())).toEqual([]);