fix: bugs with addImports/removeImports/replaceImports

Author: jedwards1211

src/Astx.ts

@@ -22,7 +22,7 @@ import { SimpleReplacementInterface } from './util/SimpleReplacementCollector'
 import forEachNode from './util/forEachNode'
 import addImports from './util/addImports'
 import findImports from './util/findImports'
-import removeImports from './util/removeImports'
+import removeImports, { removeFoundImport } from './util/removeImports'
 import createReplacementConverter from './convertReplacement'
 import compileReplacement, { CompiledReplacement } from './compileReplacement'
 
@@ -780,7 +780,7 @@ export default class Astx extends ExtendableProxy implements Iterable<Astx> {
           )
         }
         const found = findImports(this, pattern)
-        return new ImportReplacer(this, found, pattern)
+        return new ImportReplacer(this, found, decl)
       },
       arg0,
       ...rest
@@ -791,8 +791,8 @@ export default class Astx extends ExtendableProxy implements Iterable<Astx> {
 class ImportReplacer {
   constructor(
     public astx: Astx,
-    public match: Astx,
-    public findPattern: readonly NodePath<Node, any>[]
+    public found: Astx,
+    public decl: ImportDeclaration
   ) {}
 
   with(
@@ -824,8 +824,8 @@ class ImportReplacer {
     const { parsePatternToNodes } = backend
 
     const doReplace = (rawReplacement: any): boolean => {
-      if (!this.match.matched) return false
-      const match = this.match.match
+      if (!this.found.matched) return false
+      const match = this.found.match
       const path =
         match.path.parentPath?.node?.type === 'ExpressionStatement'
           ? match.path.parentPath
@@ -847,7 +847,7 @@ class ImportReplacer {
         Array.isArray(generated) ? generated[0] : generated
       )
 
-      removeImports(this.astx, this.findPattern)
+      removeFoundImport(this.astx, this.decl as any, this.found)
       this.astx.addImports(converted)
       return true
     }
@@ -857,7 +857,7 @@ class ImportReplacer {
         // Always replace in reverse so that if there are matches inside of
         // matches, the inner matches get replaced first (since they come
         // later in the code)
-        return doReplace((arg0 as any)(this.match, parsePatternToNodes))
+        return doReplace((arg0 as any)(this.found, parsePatternToNodes))
       } else if (typeof arg0 === 'string') {
         return doReplace(parsePatternToNodes(arg0))
       } else if (isNode(arg0) || isNodeArray(arg0)) {

src/util/addImports.ts

@@ -23,6 +23,7 @@ export default function addImports(
   }
 
   function addDeclaration(decl: t.ImportDeclaration) {
+    astx.context.simpleReplacements?.bail()
     const before =
       astx.find(
         (a) =>
@@ -100,6 +101,7 @@ export default function addImports(
         if (existingImportKind.matched) {
           const existingDecl: t.ImportDeclaration =
             existingImportKind.node as any
+          astx.context.simpleReplacements?.bail()
           addSpecifierToDeclaration(existingDecl, t.cloneNode(specifier))
         } else {
           addDeclaration(t.cloneNode({ ...decl, specifiers: [specifier] }))

src/util/removeImport.ts

@@ -2,6 +2,7 @@ import { ImportDeclaration } from '@babel/types'
 import Astx from '../Astx'
 import { Node, NodePath } from '../types'
 import * as t from '@babel/types'
+import { stripImportKind } from './imports'
 
 export default function removeImports(
   astx: Astx,
@@ -26,22 +27,56 @@ export default function removeImports(
           specifiers: [specifier],
         }).matched
         if (!found) continue
+
         result = true
-        found.find(specifier).remove()
-        for (const path of found) {
-          if (!(path.node as ImportDeclaration)?.specifiers.length) {
-            path.remove()
-          }
-        }
+
+        removeFoundImport(astx, decl, found)
       }
     } else {
       const found = astx.findImports(
         t.importDeclaration([], decl.source)
       ).matched
       if (!found) continue
       result = true
-      found.remove()
+      removeFoundImport(astx, decl, found)
     }
   }
+  if (result) astx.context.simpleReplacements?.bail()
   return result
 }
+
+export function removeFoundImport(
+  astx: Astx,
+  decl: ImportDeclaration,
+  found: Astx
+) {
+  if (decl.specifiers?.length) {
+    if (decl.specifiers.length > 1) {
+      throw new Error(`decl must have a single specifier`)
+    }
+    const [specifier] = decl.specifiers
+    found.find(specifier).remove()
+    if (
+      specifier.type === 'ImportSpecifier' &&
+      specifier.importKind === 'type'
+    ) {
+      found
+        .find(
+          (a) =>
+            a.node.type === 'ImportDeclaration' && a.node.importKind === 'type'
+        )
+        .find(stripImportKind(specifier))
+        .remove()
+    }
+    for (const path of found) {
+      if (
+        path.node.type === 'ImportDeclaration' &&
+        !path.node.specifiers?.length
+      ) {
+        path.remove()
+      }
+    }
+  } else {
+    found.remove()
+  }
+}

src/util/removeImports.ts

@@ -0,0 +1,94 @@
+import { TransformOptions } from '../../src'
+import { astxTestcase } from '../astxTestcase'
+import dedent from 'dedent-js'
+
+astxTestcase({
+  file: __filename,
+  parsers: ['babel/tsx'],
+  input: dedent`
+    import * as React from 'react'
+    import { withRouter } from 'react-router-dom'
+    import type { LocationShape, RouterHistory } from 'react-router-dom'
+    export type Props = {
+      history: RouterHistory
+      to: string | LocationShape
+      replace?: boolean
+      children: (props: ChildProps) => React.ReactElement
+    }
+    export type ChildProps = {
+      bind: {
+        onClick: (e: Event) => any
+      }
+    }
+    class WorkaroundLink extends React.Component<Props> {
+      _handleClick = () => {
+        const { history, to, replace } = this.props
+        if (replace) history.replace(to)
+        else history.push(to)
+      }
+      render(): React.ReactElement {
+        const { children } = this.props
+        return children({
+          bind: {
+            onClick: this._handleClick,
+          },
+        })
+      }
+    }
+    export default withRouter(WorkaroundLink as any)
+  `,
+  astx: ({ astx }: TransformOptions) => {
+    if (
+      !astx.find(
+        (a) =>
+          a.node.type === 'ImportDeclaration' &&
+          a.node.source.value === 'react-router-dom'
+      ).matched
+    ) {
+      return null
+    }
+
+    const stringOrLocationShapeMatch =
+      astx.find`type T = /**/ string | LocationShape`().matched
+    if (stringOrLocationShapeMatch) {
+      astx.replaceImport`import { type LocationShape } from 'react-router-dom'`()
+        .with`import type { LocationDescriptor } from 'history'`()
+      astx.addImports`import type { LocationState } from '${'../react-router/LocationState'}'`()
+      stringOrLocationShapeMatch.replace`LocationDescriptor<LocationState>`()
+    }
+  },
+  expected: dedent`
+    import * as React from 'react'
+    import { withRouter } from 'react-router-dom'
+    import type { RouterHistory } from 'react-router-dom'
+    import type { LocationDescriptor } from 'history'
+    import type { LocationState } from '../react-router/LocationState'
+    export type Props = {
+      history: RouterHistory
+      to: LocationDescriptor<LocationState>
+      replace?: boolean
+      children: (props: ChildProps) => React.ReactElement
+    }
+    export type ChildProps = {
+      bind: {
+        onClick: (e: Event) => any
+      }
+    }
+    class WorkaroundLink extends React.Component<Props> {
+      _handleClick = () => {
+        const { history, to, replace } = this.props
+        if (replace) history.replace(to)
+        else history.push(to)
+      }
+      render(): React.ReactElement {
+        const { children } = this.props
+        return children({
+          bind: {
+            onClick: this._handleClick,
+          },
+        })
+      }
+    }
+    export default withRouter(WorkaroundLink as any)    
+  `,
+})