Improve type registry types (#268)

Co-authored-by: aikoven <dan.lytkin@gmail.com>
stephenh · Apr 4, 2021 · 30b5f52 · 30b5f52
1 parent d929082
commit 30b5f52
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/integration/type-registry/typeRegistry.ts b/integration/type-registry/typeRegistry.ts
@@ -1,15 +1,18 @@
 /* eslint-disable */
 import { Writer, Reader } from 'protobufjs/minimal';
 
-export interface MessageType<Message> {
+export interface MessageType<Message extends UnknownMessage = UnknownMessage> {
+  $type: Message['$type'];
   encode(message: Message, writer?: Writer): Writer;
   decode(input: Reader | Uint8Array, length?: number): Message;
   fromJSON(object: any): Message;
   toJSON(message: Message): unknown;
   fromPartial(object: DeepPartial<Message>): Message;
 }
 
-export const messageTypeRegistry = new Map<string, MessageType<unknown>>();
+export type UnknownMessage = { $type: string };
+
+export const messageTypeRegistry = new Map<string, MessageType>();
 
 type Builtin = Date | Function | Uint8Array | string | number | undefined;
 export type DeepPartial<T> = T extends Builtin

diff --git a/src/generate-type-registry.ts b/src/generate-type-registry.ts
@@ -10,7 +10,11 @@ export function generateTypeRegistry(ctx: Context): Code {
   chunks.push(generateMessageType(ctx));
 
   chunks.push(code`
-    export const messageTypeRegistry = new Map<string, MessageType<unknown>>();
+    export type UnknownMessage = {$type: string};
+  `);
+
+  chunks.push(code`
+    export const messageTypeRegistry = new Map<string, MessageType>();
   `);
 
   chunks.push(code`${ctx.utils.DeepPartial.ifUsed}`);
@@ -21,7 +25,9 @@ export function generateTypeRegistry(ctx: Context): Code {
 function generateMessageType(ctx: Context): Code {
   const chunks: Code[] = [];
 
-  chunks.push(code`export interface MessageType<Message> {`);
+  chunks.push(code`export interface MessageType<Message extends UnknownMessage = UnknownMessage> {`);
+
+  chunks.push(code`$type: Message['$type'];`);
 
   if (ctx.options.outputEncodeMethods) {
     chunks.push(code`encode(message: Message, writer?: ${Writer}): ${Writer};`);