feat: Support sorting by multiple values.

homebound-team · Nov 25, 2024 · c047496 · c047496
1 parent aea5b45
commit c047496
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 4 deletions.
diff --git a/src/array/sortBy.test.ts b/src/array/sortBy.test.ts
@@ -5,6 +5,21 @@ describe("sortBy", () => {
     expect([1, 2, 3].sortBy((n: number) => -n)).toEqual([3, 2, 1]);
   });
 
+  it("sorts by two values", () => {
+    const foos: { foo: number; bar: number }[] = [
+      { foo: 3, bar: 2 },
+      { foo: 1, bar: 0 },
+      { foo: 3, bar: 1 },
+      { foo: 2, bar: 0 },
+    ];
+    expect(foos.sortBy((f) => [f.foo, f.bar])).toEqual([
+      { foo: 1, bar: 0 },
+      { foo: 2, bar: 0 },
+      { foo: 3, bar: 1 },
+      { foo: 3, bar: 2 },
+    ]);
+  });
+
   it("works for bigint fields", () => {
     // Using `any` here as an easy case for the compiler not being able to catch the error
     const foos: { foo: bigint }[] = [{ foo: 3n }, { foo: -1n }, { foo: 2n }];

diff --git a/src/array/sortBy.ts b/src/array/sortBy.ts
@@ -3,7 +3,7 @@ import { Comparable, compare, KeysOfType } from "../utils";
 declare global {
   interface Array<T> {
     /** Array is returned in ascending order. */
-    sortBy<K extends Comparable>(fn: (el: T) => K): T[];
+    sortBy<K extends Comparable>(fn: (el: T) => K | K[]): T[];
     /** Array is returned in ascending order. */
     sortByKey<K extends KeysOfType<T, Comparable>>(key: K): T[];
   }
@@ -16,8 +16,22 @@ declare global {
   }
 }
 
-Array.prototype.sortBy = function <T, K extends Comparable>(this: T[], fn: (el: T) => K): T[] {
-  return [...this].sort((a, b) => compare(fn(a), fn(b)));
+Array.prototype.sortBy = function <T, K extends Comparable>(this: T[], fn: (el: T) => K | K[]): T[] {
+  return [...this].sort((a, b) => {
+    const av = fn(a);
+    const bv = fn(b);
+    if (Array.isArray(av) && Array.isArray(bv)) {
+      for (let i = 0; i < av.length; i++) {
+        const comparison = compare(av[i], bv[i]);
+        if (comparison !== 0) return comparison;
+      }
+      return 0;
+    } else if (Array.isArray(av) || Array.isArray(bv)) {
+      throw new Error("Cannot compare array to non-array");
+    } else {
+      return compare(av, bv);
+    }
+  });
 };
 
 Array.prototype.sortByKey = function <T, K extends KeysOfType<T, Comparable>>(this: T[], key: K): T[] {

diff --git a/src/utils.ts b/src/utils.ts
@@ -39,7 +39,7 @@ export function isDefined<T extends any>(param: T | undefined | null): param is
   return param !== null && param !== undefined;
 }
 
-export type Comparable = string | number | bigint | Date | Temporal.PlainDate | Temporal.ZonedDateTime;
+export type Comparable = string | number | bigint | Date | Temporal.PlainDate | Temporal.ZonedDateTime | undefined;
 
 export function compare<T extends Comparable>(a: T, b: T): number {
   if (!isDefined(a) || !isDefined(b)) {