protoc-gen-reg - generate a reg.ts file for each proto file, with a registry for all types in the file

protoc-gen-reg.ts

#!/usr/bin/env -S npx tsx

// See here how to run this plugin: https://github.com/bufbuild/protobuf-es/tree/main/packages/protoplugin-example

import { createEcmaScriptPlugin, runNodeJs } from "@bufbuild/protoplugin";
import type { Schema } from "@bufbuild/protoplugin/ecmascript";
import type {
  DescEnum,
  DescExtension,
  DescFile,
  DescMessage,
  DescService,
} from "@bufbuild/protobuf";

const protocGenReg = createEcmaScriptPlugin({
  name: "protoc-gen-reg",
  version: `v1`,
  generateTs,
});

runNodeJs(protocGenReg);

function generateTs(schema: Schema) {
  for (const file of schema.files) {
    const f = schema.generateFile(file.name + "_reg.ts");
    f.preamble(file);
    const createRegistry = f.import("createRegistry", "@bufbuild/protobuf");
    f.print("export const registry = ", createRegistry, "(");
    for (const descType of getAllTypes(file)) {
      if (descType.kind == "message") {
        f.print("  ", descType, ", ");
      }
    }
    f.print(");");
  }
}

function* getAllTypes(
  desc: DescFile | DescMessage,
): Iterable<DescMessage | DescEnum | DescExtension | DescService> {
  switch (desc.kind) {
    case "file":
      for (const message of desc.messages) {
        yield message;
        yield* getAllTypes(message);
      }
      yield* desc.enums;
      yield* desc.services;
      yield* desc.extensions;
      break;
    case "message":
      for (const message of desc.nestedMessages) {
        yield message;
        yield* getAllTypes(message);
      }
      yield* desc.nestedEnums;
      yield* desc.nestedExtensions;
      break;
  }
}