Use type checker to identify msg function and lit-localize module

Author: aomarks

src/outputters/transform.ts

@@ -52,7 +52,7 @@ export function transformOutput(
     }
     opts.outDir = pathLib.join(outRoot, '/', locale);
     program.emit(undefined, undefined, undefined, undefined, {
-      before: [litLocalizeTransform(translations)],
+      before: [litLocalizeTransform(translations, program)],
     });
   }
 }
@@ -61,10 +61,11 @@ export function transformOutput(
  * Return a TypeScript TransformerFactory for the lit-localize transformer.
  */
 export function litLocalizeTransform(
-  translations: Map<string, Message> | undefined
+  translations: Map<string, Message> | undefined,
+  program: ts.Program
 ): ts.TransformerFactory<ts.SourceFile> {
   return (context) => {
-    const transformer = new Transformer(context, translations);
+    const transformer = new Transformer(context, translations, program);
     return (file) => ts.visitNode(file, transformer.boundVisitNode);
   };
 }
@@ -75,21 +76,24 @@ export function litLocalizeTransform(
 class Transformer {
   private context: ts.TransformationContext;
   private translations: Map<string, Message> | undefined;
+  private typeChecker: ts.TypeChecker;
   boundVisitNode = this.visitNode.bind(this);
 
   constructor(
     context: ts.TransformationContext,
-    translations: Map<string, Message> | undefined
+    translations: Map<string, Message> | undefined,
+    program: ts.Program
   ) {
     this.context = context;
     this.translations = translations;
+    this.typeChecker = program.getTypeChecker();
   }
 
   /**
    * Top-level delegating visitor for all nodes.
    */
   visitNode(node: ts.Node): ts.VisitResult<ts.Node> {
-    if (isMsgCall(node)) {
+    if (isMsgCall(node, this.typeChecker)) {
       return this.replaceMsgCall(node);
     }
     if (isLitExpression(node)) {
@@ -101,8 +105,8 @@ class Transformer {
         )
       );
     }
-    if (ts.isImportDeclaration(node)) {
-      return this.removeMsgImport(node);
+    if (this.isLitLocalizeImport(node)) {
+      return undefined;
     }
     return ts.visitEachChild(node, this.boundVisitNode, this.context);
   }
@@ -328,29 +332,31 @@ class Transformer {
   }
 
   /**
-   * Remove import declarations for the lit-localize `msg` function, because we
-   * are transforming away all calls to that function.
+   * Return whether the given node is an import for the lit-localize module.
    */
-  removeMsgImport(
-    imprt: ts.ImportDeclaration
-  ): ts.ImportDeclaration | undefined {
-    const clause = imprt.importClause;
-    if (clause === undefined) {
-      return imprt;
+  isLitLocalizeImport(node: ts.Node): node is ts.ImportDeclaration {
+    if (!ts.isImportDeclaration(node)) {
+      return false;
     }
-    const bindings = clause.namedBindings;
-    if (bindings === undefined || !ts.isNamedImports(bindings)) {
-      return imprt;
+    const moduleSymbol = this.typeChecker.getSymbolAtLocation(
+      node.moduleSpecifier
+    );
+    if (!moduleSymbol) {
+      return false;
     }
-    // TODO(aomarks) This is too crude. We should do better to identify only our
-    // `msg` function.
-    if (
-      bindings.elements.length === 1 &&
-      bindings.elements[0].name.text === 'msg'
-    ) {
-      return undefined;
+    // TODO(aomarks) Is there a better way to reliably identify the lit-localize
+    // module that doesn't require this cast? We could export a const with a
+    // known name and then look through `exports`, but it doesn't seem good to
+    // polute the module like that.
+    const file = (moduleSymbol.valueDeclaration as unknown) as {
+      identifiers: Map<string, unknown>;
+    };
+    for (const id of file.identifiers.keys()) {
+      if (id === '_LIT_LOCALIZE_MSG_') {
+        return true;
+      }
     }
-    return imprt;
+    return false;
   }
 }
 

src/program-analysis.ts

@@ -18,11 +18,11 @@ import {createDiagnostic} from './typescript';
  * Extract translation messages from all files in a TypeScript program.
  */
 export function extractMessagesFromProgram(
-  node: ts.Program
+  program: ts.Program
 ): {messages: ProgramMessage[]; errors: ts.Diagnostic[]} {
   const messages: ProgramMessage[] = [];
   const errors: ts.Diagnostic[] = [];
-  for (const sourcefile of node.getSourceFiles()) {
+  for (const sourcefile of program.getSourceFiles()) {
     extractMessagesFromNode(sourcefile, sourcefile, messages, errors, []);
   }
   const deduped = dedupeMessages(messages);
@@ -493,14 +493,23 @@ export function isLitExpression(
 /**
  * Return whether this is a call to the lit-localize `msg` function.
  */
-export function isMsgCall(node: ts.Node): node is ts.CallExpression {
-  // TODO(aomarks) This is too crude. We should do better to identify only our
-  // `msg` function.
-  return (
-    ts.isCallExpression(node) &&
-    ts.isIdentifier(node.expression) &&
-    node.expression.escapedText === 'msg'
-  );
+export function isMsgCall(
+  node: ts.Node,
+  typeChecker?: ts.TypeChecker
+): node is ts.CallExpression {
+  if (!ts.isCallExpression(node)) {
+    return false;
+  }
+  if (typeChecker === undefined) {
+    // TODO(aomarks) Remove this branch once migration to static lit-localize
+    // library is done.
+    return (
+      ts.isIdentifier(node.expression) && node.expression.escapedText === 'msg'
+    );
+  }
+  const type = typeChecker.getTypeAtLocation(node.expression);
+  const props = typeChecker.getPropertiesOfType(type);
+  return props.some((prop) => prop.escapedName === '_LIT_LOCALIZE_MSG_');
 }
 
 /**

src/tests/compile-ts-fragment.ts

@@ -42,7 +42,7 @@ export function compileTsFragment(
   inputCode: string,
   options: ts.CompilerOptions,
   cache: CompilerHostCache,
-  transformers?: ts.CustomTransformers
+  transformers?: (program: ts.Program) => ts.CustomTransformers
 ): CompileResult {
   const dummyTsFilename = '__DUMMY__.ts';
   const dummyJsFilename = '__DUMMY__.js';
@@ -124,7 +124,7 @@ export function compileTsFragment(
     undefined,
     undefined,
     undefined,
-    transformers
+    transformers ? transformers(program) : undefined
   );
   return {
     code: outputCode,

src/tests/transform.unit.test.ts

@@ -29,19 +29,22 @@ function checkTransform(
   t: ExecutionContext,
   inputTs: string,
   expectedJs: string,
-  messages: Message[]
+  messages: Message[],
+  autoImport = true
 ) {
-  // Rather than fuss with imports in all the test cases, this little hack
-  // automatically imports for `msg` and `html` (assuming those strings aren't
-  // used with any other meanings).
-  if (inputTs.includes('msg')) {
-    inputTs = IMPORT_MSG + inputTs;
-    // Note we don't expect to see the `msg` import in the output JS, since it
-    // should be un-used after litLocalizeTransformation.
-  }
-  if (inputTs.includes('html')) {
-    inputTs = IMPORT_LIT_HTML + inputTs;
-    expectedJs = IMPORT_LIT_HTML + expectedJs;
+  if (autoImport) {
+    // Rather than fuss with imports in all the test cases, this little hack
+    // automatically imports for `msg` and `html` (assuming those strings aren't
+    // used with any other meanings).
+    if (inputTs.includes('msg')) {
+      inputTs = IMPORT_MSG + inputTs;
+      // Note we don't expect to see the `msg` import in the output JS, since it
+      // should be un-used after litLocalizeTransformation.
+    }
+    if (inputTs.includes('html')) {
+      inputTs = IMPORT_LIT_HTML + inputTs;
+      expectedJs = IMPORT_LIT_HTML + expectedJs;
+    }
   }
   const options = ts.getDefaultCompilerOptions();
   options.target = ts.ScriptTarget.ES2015;
@@ -50,9 +53,9 @@ function checkTransform(
   // Don't automatically load typings from nodes_modules/@types, we're not using
   // them here, so it's a waste of time.
   options.typeRoots = [];
-  const result = compileTsFragment(inputTs, options, cache, {
-    before: [litLocalizeTransform(makeMessageIdMap(messages))],
-  });
+  const result = compileTsFragment(inputTs, options, cache, (program) => ({
+    before: [litLocalizeTransform(makeMessageIdMap(messages), program)],
+  }));
 
   let formattedExpected = prettier.format(expectedJs, {parser: 'typescript'});
   let formattedActual;
@@ -280,3 +283,41 @@ test('msg(fn(string), msg(string)) translated', (t) => {
     ]
   );
 });
+
+test('import * as litLocalize', (t) => {
+  checkTransform(
+    t,
+    `
+    import * as litLocalize from './lib_client/index.js';
+    litLocalize.msg("foo", "Hello World");
+  `,
+    '"Hello World";',
+    [],
+    false
+  );
+});
+
+test('import {msg as foo}', (t) => {
+  checkTransform(
+    t,
+    `
+    import {msg as foo} from './lib_client/index.js';
+    foo("foo", "Hello World");
+  `,
+    '"Hello World";',
+    [],
+    false
+  );
+});
+
+test('exclude different msg function', (t) => {
+  checkTransform(
+    t,
+    `function msg(id: string, template: string) { return template; }
+    msg("foo", "Hello World");`,
+    `function msg(id, template) { return template; }
+    msg("foo", "Hello World");`,
+    [],
+    false
+  );
+});

src_client/index.ts

@@ -166,23 +166,23 @@ export function localeReady(): Promise<void> {
  * @param args In the case that `template` is a function, it is invoked with
  * the 3rd and onwards arguments to `msg`.
  */
-export function msg(id: string, template: string): string;
+export function _msg(id: string, template: string): string;
 
-export function msg(id: string, template: TemplateResult): TemplateResult;
+export function _msg(id: string, template: TemplateResult): TemplateResult;
 
-export function msg<F extends (...args: any[]) => string>(
+export function _msg<F extends (...args: any[]) => string>(
   id: string,
   fn: F,
   ...params: Parameters<F>
 ): string;
 
-export function msg<F extends (...args: any[]) => TemplateResult>(
+export function _msg<F extends (...args: any[]) => TemplateResult>(
   id: string,
   fn: F,
   ...params: Parameters<F>
 ): TemplateResult;
 
-export function msg(
+export function _msg(
   id: string,
   template: TemplateLike,
   ...params: any[]
@@ -198,3 +198,5 @@ export function msg(
   }
   return template;
 }
+
+export const msg: typeof _msg & {_LIT_LOCALIZE_MSG_?: never} = _msg;