feat(protoc-gen-ng): add support for custom WKT

smnbbrv · Sep 28, 2021 · 8a64bde · 8a64bde
1 parent e95366c
commit 8a64bde
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 6 deletions.
diff --git a/packages/protoc-gen-ng/src/config.ts b/packages/protoc-gen-ng/src/config.ts
@@ -127,6 +127,10 @@ export class Config {
    * By default false
    */
   public embedWellKnownTypes: boolean;
+  /**
+   * Custom well-known type package names
+   */
+  public customWellKnownTypes: {[key: string]: string};
   /**
    * Per-generated-file-type config
    */
@@ -147,6 +151,7 @@ export class Config {
   constructor(config: Partial<Config> = {}) {
     this.debug = config.debug ?? false;
     this.embedWellKnownTypes = config.embedWellKnownTypes ?? false;
+    this.customWellKnownTypes = config.customWellKnownTypes ?? {};
     this.files = new ConfigFiles(config.files ?? {});
   }
 

diff --git a/packages/protoc-gen-ng/src/input/proto.ts b/packages/protoc-gen-ng/src/input/proto.ts
@@ -145,14 +145,25 @@ export class Proto {
     const root = Array(this.name.split('/').length - 1).fill('..').join('/');
 
     return this.resolved.allDependencies.map(pp => {
-      const isWKT = pp.pb_package === 'google.protobuf';
+      const wktDependency = this.getWktDependency(pp.pb_package);
+      const isWKT = wktDependency != null;
       const genwkt = Services.Config.embedWellKnownTypes;
-      const path = (genwkt || !genwkt && !isWKT) ? `${root || '.'}/${pp.getGeneratedFileBaseName()}` : '@ngx-grpc/well-known-types';
+      const path = (genwkt || !genwkt && !isWKT) ? `${root || '.'}/${pp.getGeneratedFileBaseName()}` : wktDependency;
 
       return `import * as ${this.getDependencyPackageName(pp)} from '${path}';`;
     }).join('\n');
   }
 
+  getWktDependency(pbPackage: string) {
+    if (pbPackage === 'google.protobuf') {
+      return '@ngx-grpc/well-known-types';
+    }
+    if (!!Services.Config.customWellKnownTypes && !!Services.Config.customWellKnownTypes[pbPackage]) {
+      return Services.Config.customWellKnownTypes[pbPackage];
+    }
+    return null;
+  }
+
   getGeneratedFileBaseName() {
     return `${dasherize(this.name.replace(/\.proto$/, ''))}.pb`;
   }

diff --git a/packages/protoc-gen-ng/src/main.ts b/packages/protoc-gen-ng/src/main.ts
@@ -52,7 +52,8 @@ async function main() {
     const genwkt = Services.Config.embedWellKnownTypes;
 
     protos
-      .filter(p => genwkt || !genwkt && p.pb_package !== 'google.protobuf')
+      .filter(p => genwkt || !genwkt && (p.pb_package !== 'google.protobuf' &&
+        (!Services.Config.customWellKnownTypes || !Services.Config.customWellKnownTypes[p.pb_package])))
       .forEach(proto => {
         Services.Logger.debug(`Start processing proto ${proto.name}`);
 

diff --git a/packages/protoc-gen-ng/src/output/types/message.ts b/packages/protoc-gen-ng/src/output/types/message.ts
@@ -25,6 +25,7 @@ import { TimestampWKT } from './well-known-types/timestamp.wkt';
 import { EnumValueWKT, EnumWKT, FieldWKT, OptionWKT, TypeWKT } from './well-known-types/type.wkt';
 import { BoolValueWKT, BytesValueWKT, DoubleValueWKT, FloatValueWKT, Int32ValueWKT, Int64ValueWKT, StringValueWKT, UInt32ValueWKT, UInt64ValueWKT } from './well-known-types/wrappers.wkt';
 import { WKT } from './wkt';
+import { Services } from '../../services';
 
 export class Message {
 
@@ -33,12 +34,15 @@ export class Message {
   private oneOfs: OneOf[];
 
   private wkt?: WKT;
+  private isWkt = false;
 
   constructor(
     private proto: Proto,
     private message: ProtoMessage,
   ) {
-    if (this.proto.pb_package === 'google.protobuf') {
+    if (this.proto.pb_package === 'google.protobuf' ||
+      (!!Services.Config.customWellKnownTypes && !!Services.Config.customWellKnownTypes[this.proto.pb_package])) {
+      this.isWkt = true;
       switch (this.message.name) {
         case 'Any': this.wkt = new AnyWKT(); break;
         case 'Api': this.wkt = new ApiWKT(); break;
@@ -120,8 +124,16 @@ export class Message {
     );
 
     const messageId = (this.proto.pb_package ? this.proto.pb_package + '.' : '') + this.message.name;
-    const wktUrl = 'https://developers.google.com/protocol-buffers/docs/reference/google.protobuf';
-    const constructorComment = this.wkt ? `Well known type, more at ${wktUrl}` : `Message implementation for ${messageId}`;
+    let constructorComment = `Message implementation for ${messageId}`;
+    if (this.isWkt){
+      if (this.proto.pb_package  === 'google.protobuf') {
+        constructorComment = `Well known type, more at https://developers.google.com/protocol-buffers/docs/reference/google.protobuf`;
+      } else{
+
+        constructorComment = `Custom well known type`;
+      }
+    }
+
 
     printer.addLine(`
     /**