feat: allow strings for enums, bytes, and Long in .d.ts

.gitignore

@@ -5,7 +5,7 @@ npm-debug.log
 .coverage
 .nyc_output
 docs/
-google/
+protos/google/
 out/
 system-test/secrets.js
 system-test/*key.json

test/fixtures/dts-update/google/dts.proto

@@ -0,0 +1,33 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+syntax = "proto3";
+
+package google;
+
+service DtsUpdateTest {
+  rpc Test(TestMessage) returns (TestMessage) {
+  }
+}
+
+enum TestEnum {
+  TEST_ENUM_UNSPECIFIED = 0;
+  SOME_VALUE = 1;
+}
+
+message TestMessage {
+  int64 long_field = 1;
+  bytes bytes_field = 2; 
+  TestEnum enum_field = 3;
+}

test/fixtures/dts-update/test_proto_list.json

@@ -0,0 +1 @@
+["./google/dts.proto"]

test/unit/compileProtos.ts

@@ -31,8 +31,12 @@ const testDir = path.join(process.cwd(), '.compileProtos-test');
 const resultDir = path.join(testDir, 'protos');
 const cwd = process.cwd();
 
+const expectedJsonResultFile = path.join(resultDir, 'protos.json');
+const expectedJSResultFile = path.join(resultDir, 'protos.js');
+const expectedTSResultFile = path.join(resultDir, 'protos.d.ts');
+
 describe('compileProtos tool', () => {
-  before(async () => {
+  beforeEach(async () => {
     if (fs.existsSync(testDir)) {
       await rmrf(testDir);
     }
@@ -42,17 +46,14 @@ describe('compileProtos tool', () => {
     process.chdir(testDir);
   });
 
-  after(() => {
+  afterEach(() => {
     process.chdir(cwd);
   });
 
   it('compiles protos to JSON, JS, TS', async () => {
     await compileProtos.main([
       path.join(__dirname, '..', '..', 'test', 'fixtures', 'protoLists'),
     ]);
-    const expectedJsonResultFile = path.join(resultDir, 'protos.json');
-    const expectedJSResultFile = path.join(resultDir, 'protos.js');
-    const expectedTSResultFile = path.join(resultDir, 'protos.d.ts');
     assert(fs.existsSync(expectedJsonResultFile));
     assert(fs.existsSync(expectedJSResultFile));
     assert(fs.existsSync(expectedTSResultFile));
@@ -90,10 +91,29 @@ describe('compileProtos tool', () => {
         'empty'
       ),
     ]);
-    const expectedResultFile = path.join(resultDir, 'protos.json');
-    assert(fs.existsSync(expectedResultFile));
+    assert(fs.existsSync(expectedJsonResultFile));
 
-    const json = await readFile(expectedResultFile);
+    const json = await readFile(expectedJsonResultFile);
     assert.strictEqual(json.toString(), '{}');
   });
+
+  it('fixes types in the TS file', async () => {
+    await compileProtos.main([
+      path.join(__dirname, '..', '..', 'test', 'fixtures', 'dts-update'),
+    ]);
+    assert(fs.existsSync(expectedTSResultFile));
+    const ts = await readFile(expectedTSResultFile);
+
+    assert(ts.toString().includes('import * as Long'));
+    assert(
+      ts.toString().includes('http://www.apache.org/licenses/LICENSE-2.0')
+    );
+    assert(ts.toString().includes('longField?: (number|Long|string|null);'));
+    assert(ts.toString().includes('bytesField?: (Uint8Array|string|null);'));
+    assert(ts.toString().includes('enumField?: (google.TestEnum|keyof typeof google.TestEnum|null);'));
+    assert(ts.toString().includes('public longField: (number|Long|string);'));
+    assert(ts.toString().includes('public bytesField: (Uint8Array|string);'));
+    assert(ts.toString().includes('public enumField: (google.TestEnum|keyof typeof google.TestEnum);'));
+  });
+
 });

tools/compileProtos.ts

@@ -80,6 +80,95 @@ function normalizePath(filePath: string): string {
   return path.join(...filePath.split('/'));
 }
 
+function getAllEnums(dts: string): Set<string> {
+  const result = new Set<string>();
+  const lines = dts.split('\n');
+  const nestedIds = [];
+  let currentEnum = undefined;
+  for (const line of lines) {
+    const match = line.match(/^\s*(?:export )?(namespace|class|interface|enum) (\w+) .*{/);
+    if (match) {
+      const [, keyword, id] = match;
+      nestedIds.push(id);
+      if (keyword === 'enum') {
+        currentEnum = nestedIds.join('.');
+        result.add(currentEnum);
+      }
+      continue;
+    }
+    if (line.match(/^\s*}/)) {
+      nestedIds.pop();
+      currentEnum = undefined;
+      continue;
+    }
+  }
+
+  return result;
+}
+
+function updateDtsTypes(dts: string, enums: Set<string>): string {
+  const lines = dts.split('\n');
+  const result: string[] = [];
+
+  for (const line of lines) {
+    let typeName: string|undefined = undefined;
+    // Enums can be used in interfaces and in classes.
+    // For simplicity, we'll check these two cases independently.
+    // encoding?: (google.cloud.speech.v1p1beta1.RecognitionConfig.AudioEncoding|null);
+    const interfaceMatch = line.match(/\w+\?: \(([\w.]+)\|null\);/);
+    if (interfaceMatch) {
+      typeName = interfaceMatch[1];
+    }
+    // public encoding: google.cloud.speech.v1p1beta1.RecognitionConfig.AudioEncoding;
+    const classMatch = line.match(/public \w+: ([\w.]+);/);
+    if (classMatch) {
+      typeName = classMatch[1];
+    }
+
+    if (line.match(/\(number\|Long(?:\|null)?\)/)) {
+      typeName = 'Long';
+    }
+
+    let replaced = line;
+    if (typeName && enums.has(typeName)) {
+      // enum: E => E|keyof typeof E  to allow all string values
+      replaced = replaced.replace(typeName, `${typeName}|keyof typeof ${typeName}`);
+    } else if (typeName === 'Uint8Array') {
+      // bytes: Uint8Array => Uint8Array|string to allow base64-encoded strings
+      replaced = replaced.replace(typeName, `${typeName}|string`);
+    } else if (typeName === 'Long') {
+      // Longs can be passed as strings :(
+      // number|Long => number|Long|string
+      replaced = replaced.replace('number|Long', 'number|Long|string');
+    }
+
+    // add brackets if we have added a |
+    replaced = replaced.replace(/: ([\w.]+\|[ \w.|]+);/, ': ($1);');
+
+    result.push(replaced);
+  }
+
+  return result.join('\n');
+}
+
+function fixDtsFile(dts: string): string {
+  // 1. fix for pbts output: the corresponding protobufjs PR
+  // https://github.com/protobufjs/protobuf.js/pull/1166
+  // is merged but not yet released.
+  if (!dts.match(/import \* as Long/)) {
+    dts = 'import * as Long from "long";\n' + dts;
+  }
+
+  // 2. add Apache license to the generated .d.ts file
+  dts = apacheLicense + dts;
+
+  // 3. major hack: update types to allow passing strings 
+  // where enums, longs, or bytes are expected
+  const enums = getAllEnums(dts);
+  dts = updateDtsTypes(dts, enums);
+  return dts;
+}
+
 /**
  * Returns a combined list of proto files listed in all JSON files given.
  *
@@ -153,14 +242,7 @@ async function compileProtos(protos: string[]): Promise<void> {
   await pbtsMain(pbjsArgs4ts);
 
   let tsResult = (await readFile(tsOutput)).toString();
-  // fix for pbts output: the corresponding protobufjs PR
-  // https://github.com/protobufjs/protobuf.js/pull/1166
-  // is merged but not yet released.
-  if (!tsResult.match(/import \* as Long/)) {
-    tsResult = 'import * as Long from "long";\n' + tsResult;
-  }
-  // add Apache license to the generated .d.ts file
-  tsResult = apacheLicense + tsResult;
+  tsResult = fixDtsFile(tsResult);
   await writeFile(tsOutput, tsResult);
 }