feat: auto-import assert function

src/assertify-statement.ts

@@ -1,14 +1,19 @@
 import { NodePath } from '@babel/traverse';
 import * as BabelTypes from '@babel/types';
 
-export default (t: typeof BabelTypes) => (
+import { InternalConfig } from './config';
+import generateAssertIdentifier from './generate-assert-identifier';
+
+export default (t: typeof BabelTypes, config: InternalConfig) => (
   statementPath: NodePath<BabelTypes.Statement>,
 ) => {
   const statement = statementPath.node;
   if (t.isExpressionStatement(statement)) {
     const origExpr = statement.expression;
 
-    const assertIdentifier = t.identifier('assert');
+    const assertIdentifier = generateAssertIdentifier(t, config)(
+      statementPath.scope,
+    );
     assertIdentifier.loc = {
       start: origExpr.loc.start,
       end: origExpr.loc.start,

src/config.ts

@@ -1,12 +1,49 @@
-export interface Config {
+/**
+ * Config after sanitizing in extractConfigFromState.
+ */
+export interface InternalConfig {
   powerAssert: boolean;
+  autoImport: string; // empty => no auto import
+}
+
+/**
+ * Config as it can be provided by the user.
+ * Comments indicate transformations made to conform to InternalConfig format.
+ */
+export interface Config {
+  powerAssert?: boolean;
+  autoImport?: boolean | string; // false => '', true => 'power-assert'
 }
 
-export const defaultConfig: Config = {
+export const defaultConfig = {
   powerAssert: true,
+  autoImport: true,
 };
 
-export const extractConfigFromState = (state: any): Config => ({
-  ...defaultConfig,
-  ...state.opts,
-});
+// tslint:disable no-parameter-reassignment
+export const extractConfigFromState = ({
+  opts: { powerAssert, autoImport },
+}: {
+  opts: Config;
+}): InternalConfig => {
+  // powerAssert
+  if (powerAssert === undefined) {
+    ({ powerAssert } = defaultConfig);
+  }
+
+  // autoImport
+  if (autoImport === undefined) {
+    ({ autoImport } = defaultConfig);
+  }
+  if (autoImport === false) {
+    autoImport = '';
+  }
+  if (autoImport === true) {
+    autoImport = 'power-assert';
+  }
+
+  return {
+    powerAssert,
+    autoImport,
+  };
+};

src/generate-assert-identifier.ts

@@ -0,0 +1,65 @@
+import { Scope } from '@babel/traverse';
+import * as BabelTypes from '@babel/types';
+
+import { InternalConfig } from './config';
+
+const ASSERT_IDENTIFIER_NAME = 'assert';
+
+const findExistingImportFromSource = (
+  scope: Scope,
+  t: typeof BabelTypes,
+  source: string,
+) => {
+  const program = scope.getProgramParent().path.node as BabelTypes.Program;
+
+  // try to find existing default import from source
+  for (const stmt of program.body) {
+    if (t.isImportDeclaration(stmt) && stmt.source.value === source) {
+      const defaultSpecifier = stmt.specifiers.find(specifier =>
+        t.isImportDefaultSpecifier(specifier),
+      ) as BabelTypes.ImportDefaultSpecifier | undefined;
+      if (defaultSpecifier) {
+        const { local } = defaultSpecifier;
+        const { name } = local;
+        // make sure the import is not shadowed in the scope
+        if (scope.bindingIdentifierEquals(name, local)) {
+          return name;
+        }
+      }
+    }
+  }
+  return;
+};
+
+const addImport = (scope: Scope, t: typeof BabelTypes, source: string) => {
+  const program = scope.getProgramParent().path;
+
+  // generate default import from source
+  const id = scope.generateUidIdentifier(ASSERT_IDENTIFIER_NAME);
+  (program as any).unshiftContainer(
+    'body',
+    t.importDeclaration(
+      [t.importDefaultSpecifier(id)],
+      t.stringLiteral(source),
+    ),
+  );
+  // recrawl scope so the added import is found for the next assertion
+  (scope as any).crawl();
+  return id;
+};
+
+export default (
+  t: typeof BabelTypes,
+  { autoImport: importSource }: InternalConfig,
+) => (scope: Scope) => {
+  if (importSource) {
+    const name = findExistingImportFromSource(scope, t, importSource);
+    if (name) {
+      return t.identifier(name);
+    }
+
+    return addImport(scope, t, importSource);
+  }
+
+  return t.identifier(ASSERT_IDENTIFIER_NAME);
+};

src/index.ts

@@ -9,7 +9,6 @@ import { extractConfigFromState } from './config';
 const assertionBlockLabels = ['expect', 'then'];
 
 const plugin = (babel: { types: typeof BabelTypes }): PluginObj => {
-  const assertify = assertifyStatement(babel.types);
   const espowerVisitor = createEspowerVisitor(babel, {
     embedAst: true,
     patterns: ['assert(value)'],
@@ -19,6 +18,7 @@ const plugin = (babel: { types: typeof BabelTypes }): PluginObj => {
     visitor: {
       LabeledStatement(path, state) {
         const config = extractConfigFromState(state);
+        const assertify = assertifyStatement(babel.types, config);
 
         if (assertionBlockLabels.includes(path.node.label.name)) {
           const bodyPath = path.get('body') as NodePath<BabelTypes.Statement>;

test/__snapshots__/auto-import.test.ts.snap

@@ -0,0 +1,53 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`does not attempt to use a shadowed existing default import 1`] = `
+"import _assert from \\"power-assert\\";
+import fancyAssert from 'power-assert';
+{
+  let fancyAssert;
+
+  expect: _assert(1 === 1);
+}"
+`;
+
+exports[`does not attempt to use an existing named import 1`] = `
+"import _assert from \\"power-assert\\";
+import { fancyAssert } from 'power-assert';
+
+expect: _assert(1 === 1);"
+`;
+
+exports[`does not break preset-env module transform and generates code runnable in node 1`] = `"false == true"`;
+
+exports[`does not clash with an existing "_assert" import 1`] = `
+"import _assert2 from \\"power-assert\\";
+import _assert from 'fancy-assert';
+
+expect: _assert2(1 === 1);"
+`;
+
+exports[`imports from a custom source 1`] = `
+"import _assert from \\"fancy-assert\\";
+
+expect: _assert(1 === 1);"
+`;
+
+exports[`imports from power-assert by default 1`] = `
+"import _assert from \\"power-assert\\";
+
+expect: _assert(1 === 1);"
+`;
+
+exports[`reuses the same import for multiple assertions 1`] = `
+"import _assert from \\"power-assert\\";
+
+expect: _assert(1 === 1);
+
+expect: _assert(2 === 2);"
+`;
+
+exports[`uses an existing default import 1`] = `
+"import fancyAssert from 'power-assert';
+
+expect: fancyAssert(1 === 1);"
+`;

test/auto-import.test.ts

@@ -0,0 +1,80 @@
+import { transform } from '@babel/core';
+
+import plugin from '../src';
+import { Config } from '../src/config';
+
+test('imports from power-assert by default', () => {
+  const { code } = transform(`expect: 1 === 1;`, {
+    plugins: [[plugin, { powerAssert: false } as Config]],
+  });
+  expect(code).toMatchSnapshot();
+});
+
+test('does not clash with an existing "_assert" import', () => {
+  const { code } = transform(
+    `import _assert from 'fancy-assert';
+    expect: 1 === 1;`,
+    { plugins: [[plugin, { powerAssert: false } as Config]] },
+  );
+  expect(code).toMatchSnapshot();
+});
+
+test('imports from a custom source', () => {
+  const { code } = transform(`expect: 1 === 1;`, {
+    plugins: [
+      [plugin, { powerAssert: false, autoImport: 'fancy-assert' } as Config],
+    ],
+  });
+  expect(code).toMatchSnapshot();
+});
+
+test('uses an existing default import', () => {
+  const { code } = transform(
+    `import fancyAssert from 'power-assert';
+    expect: 1 === 1;`,
+    { plugins: [[plugin, { powerAssert: false } as Config]] },
+  );
+  expect(code).toMatchSnapshot();
+});
+
+test('does not attempt to use an existing named import', () => {
+  const { code } = transform(
+    `import { fancyAssert } from 'power-assert';
+    expect: 1 === 1;`,
+    { plugins: [[plugin, { powerAssert: false } as Config]] },
+  );
+  expect(code).toMatchSnapshot();
+});
+
+test('does not attempt to use a shadowed existing default import', () => {
+  const { code } = transform(
+    `import fancyAssert from 'power-assert';
+    {
+      let fancyAssert;
+      expect: 1 === 1;
+    }`,
+    { plugins: [[plugin, { powerAssert: false } as Config]] },
+  );
+  expect(code).toMatchSnapshot();
+});
+
+test('reuses the same import for multiple assertions', () => {
+  const { code } = transform(
+    `expect: 1 === 1;
+    expect: 2 === 2;`,
+    {
+      plugins: [[plugin, { powerAssert: false } as Config]],
+    },
+  );
+  expect(code).toMatchSnapshot();
+});
+
+test('does not break preset-env module transform and generates code runnable in node', () => {
+  const { code } = transform(`expect: 1 === 2;`, {
+    plugins: [[plugin, { powerAssert: false } as Config]],
+    presets: ['@babel/preset-env'],
+  });
+  expect(() =>
+    new Function('require', code as string)(require),
+  ).toThrowErrorMatchingSnapshot();
+});

test/power-assert.test.ts

@@ -1,13 +1,14 @@
 import { transform } from '@babel/core';
 
 import plugin from '../src';
+import { Config } from '../src/config';
 
 test('prints a nice error for an "expected"-labeled expression statement', () => {
   const { code } = transform(
     `const assert = require('power-assert');
     expect: 1 === 2;`,
     {
-      plugins: [plugin],
+      plugins: [[plugin, { autoImport: false } as Config]],
       filename: 'test.js',
     },
   );
@@ -22,7 +23,7 @@ test('passes a truthy expression', () => {
     `const assert = require('power-assert');
     expect: 2 === 2;`,
     {
-      plugins: [plugin],
+      plugins: [[plugin, { autoImport: false } as Config]],
       filename: 'test.js',
     },
   );
@@ -35,7 +36,7 @@ test('leaves unrelated assert statements untouched', () => {
     `const assert = require('power-assert');
     assert(1 === 2);`,
     {
-      plugins: [plugin],
+      plugins: [[plugin, { autoImport: false } as Config]],
       filename: 'test.js',
     },
   );
@@ -49,7 +50,7 @@ test('still works if babel-plugin-espower is used for other assertions in the fi
     assert(x >= 0);
     expect: x > 0;`,
     {
-      plugins: [plugin, 'espower'],
+      plugins: [[plugin, { autoImport: false } as Config], 'espower'],
       filename: 'test.js',
     },
   );
@@ -74,7 +75,7 @@ test('supports non-standard JSX syntax', () => {
     `const assert = require('power-assert');
     expect: (<div></div>).prop === 'expected';`,
     {
-      plugins: [plugin],
+      plugins: [[plugin, { autoImport: false } as Config]],
       presets: ['@babel/preset-react'],
       filename: 'test.js',
     },