関数のオーバーロードをTypescriptでやってみた～SQLにおけるunnest関数の実装を通して～

こんにちわ。nap5です。

SQLにおけるunnest関数の実装を通して関数のオーバーロードをTypescriptでやってみましたので、紹介です。

関数のオーバーロードについては以下にまとまっています。
inferで頑張るのがしんどくなってきたときには重宝しそうです。
また、アロー関数にはオーバーロードの構文がデフォではないので、いわゆる関数で実現していくのが、よさそうです。

オーバーロード関数 (overload function) | TypeScript入門『サバイバルTypeScript』

unnestのイメージは以下を参考にしました。

配列の操作  |  BigQuery  |  Google Cloud

FLATTEN | Snowflake Documentation

定義側です

import { ArrayPath } from "dot-path-value";
import { omit } from "radash";
import type { Simplify } from "type-fest";

type ArrayElement<T> = T extends (infer U)[] ? U : never;

// オーバーロードのシグネチャ1: alias が指定された場合
export function unnest<T extends Record<string, unknown>, K extends ArrayPath<T> | keyof T, A extends string>(
  data: T[],
  arrayKey: K,
  alias: A
): Simplify<Omit<T, K> & { [P in A]: ArrayElement<T[K]> }>[];

// オーバーロードのシグネチャ2: alias が指定されない場合
export function unnest<T extends Record<string, unknown>, K extends ArrayPath<T> | keyof T>(
  data: T[],
  arrayKey: K
): Simplify<Omit<T, K> & { [P in K]: ArrayElement<T[K]> }>[];

// 関数の実装
export function unnest<T extends Record<string, unknown>, K extends ArrayPath<T> | keyof T, A extends string>(
  data: T[],
  arrayKey: K,
  alias?: A
) {
  const result: (Omit<T, K> & { [key: string]: ArrayElement<T[K]> })[] = [];

  data.forEach((item) => {
    const arrayData = item[arrayKey] as Array<ArrayElement<T[K]>>;
    arrayData.forEach((subItem) => {
      const clonedItem: Omit<T, K> & { [key: string]: ArrayElement<T[K]> } = {
        ...omit(item, [arrayKey]),
        [alias ?? arrayKey]: subItem,
      };
      result.push(clonedItem);
    });
  });

  return result;
};

// 関数オーバーロードを使用すると、関数の実装部分で any を使用していても、オーバーロードのシグネチャで指定された型情報が呼び出し側に適用される
export function unstable_unnest<T extends Record<string, unknown>, K extends ArrayPath<T>, A extends string>(
  data: T[],
  arrayKey: K,
  alias?: A
): any[] {
  const result: any[] = [];

  data.forEach((item) => {
    const arrayData = item[arrayKey] as Record<string, unknown>[];
    arrayData.forEach((subItem) => {
      const clonedItem = {
        ...omit(item, [arrayKey]),
        [alias ?? arrayKey]: subItem,
      };
      result.push(clonedItem);
    });
  });

  return result;
};

使用側です

import { describe, test, expect } from "vitest";
import { unnest } from "@/index";

describe("unnest", () => {
  test("Some data", () => {
    type Person = {
      id: number;
      name: string;
      roles: ("owner" | "member" | "admin")[];
    };

    const data: Person[] = [
      {
        id: 1,
        name: "John",
        roles: [
          "admin",
          "owner"
        ],
      },
      {
        id: 2,
        name: "Jane",
        roles: [
          "admin",
          "owner",
          "member"
        ],
      },
    ];
    const result = unnest(data, "roles", "role");
    expect(result).toStrictEqual([
      { id: 1, name: "John", role: "admin" },
      { id: 1, name: "John", role: "owner" },
      { id: 2, name: "Jane", role: "admin" },
      { id: 2, name: "Jane", role: "owner" },
      { id: 2, name: "Jane", role: "member" },
    ]);
  });


  test("Person data", () => {
    type Person = {
      id: number;
      name: string;
      hobbies: { id: number; name: string }[];
    };

    const data: Person[] = [
      {
        id: 1,
        name: "John",
        hobbies: [
          { id: 1, name: "reading" },
          { id: 2, name: "gaming" },
        ],
      },
      {
        id: 2,
        name: "Jane",
        hobbies: [
          { id: 3, name: "hiking" },
          { id: 4, name: "cooking" },
          { id: 5, name: "swimming" },
        ],
      },
    ];
    const result = unnest(data, "hobbies", "hobby");
    expect(result).toStrictEqual([
      { id: 1, name: "John", hobby: { id: 1, name: "reading" } },
      { id: 1, name: "John", hobby: { id: 2, name: "gaming" } },
      { id: 2, name: "Jane", hobby: { id: 3, name: "hiking" } },
      { id: 2, name: "Jane", hobby: { id: 4, name: "cooking" } },
      { id: 2, name: "Jane", hobby: { id: 5, name: "swimming" } },
    ]);
  });

  test("Classroom data", () => {
    type Student = {
      id: number;
      name: string;
    };

    type Classroom = {
      id: number;
      subject: string;
      students: Student[];
    };

    const data: Classroom[] = [
      {
        id: 1,
        subject: "Mathematics",
        students: [
          { id: 1, name: "Alice" },
          { id: 2, name: "Bob" },
        ],
      },
      {
        id: 2,
        subject: "History",
        students: [
          { id: 3, name: "Charlie" },
          { id: 4, name: "Daisy" },
          { id: 5, name: "Edward" },
        ],
      },
    ];

    const result = unnest(data, "students", "student");
    expect(result).toStrictEqual([
      { id: 1, subject: "Mathematics", student: { id: 1, name: "Alice" } },
      { id: 1, subject: "Mathematics", student: { id: 2, name: "Bob" } },
      { id: 2, subject: "History", student: { id: 3, name: "Charlie" } },
      { id: 2, subject: "History", student: { id: 4, name: "Daisy" } },
      { id: 2, subject: "History", student: { id: 5, name: "Edward" } },
    ]);
  });


  test("Classroom data non alias", () => {
    type Student = {
      id: number;
      name: string;
    };

    type Classroom = {
      id: number;
      subject: string;
      students: Student[];
    };

    const data: Classroom[] = [
      {
        id: 1,
        subject: "Mathematics",
        students: [
          { id: 1, name: "Alice" },
          { id: 2, name: "Bob" },
        ],
      },
      {
        id: 2,
        subject: "History",
        students: [
          { id: 3, name: "Charlie" },
          { id: 4, name: "Daisy" },
          { id: 5, name: "Edward" },
        ],
      },
    ];

    const result = unnest(data, "students");
    expect(result).toStrictEqual([
      { id: 1, subject: "Mathematics", students: { id: 1, name: "Alice" } },
      { id: 1, subject: "Mathematics", students: { id: 2, name: "Bob" } },
      { id: 2, subject: "History", students: { id: 3, name: "Charlie" } },
      { id: 2, subject: "History", students: { id: 4, name: "Daisy" } },
      { id: 2, subject: "History", students: { id: 5, name: "Edward" } },
    ]);
  });
});

demo code.

elastic-jang-9yg6wq

簡単ですが、以上です。