Fix whole module import deletion in feature tree (#6456)

* Fix whole module import deletion in feature tree
Fixes #6447

* Delete both

* Fix tests

* Lint

* Clean up for review

* Update src/lang/modifyAst.ts

Co-authored-by: graphite-app[bot] <96075541+graphite-app[bot]@users.noreply.github.com>

* Add extra error case

* Update e2e/playwright/point-click-assemblies.spec.ts

Co-authored-by: graphite-app[bot] <96075541+graphite-app[bot]@users.noreply.github.com>

---------

Co-authored-by: graphite-app[bot] <96075541+graphite-app[bot]@users.noreply.github.com>

e2e/playwright/point-click-assemblies.spec.ts

@@ -358,6 +358,24 @@ test.describe('Point-and-click assemblies tests', () => {
         await scene.expectPixelColor(bgColor, midPoint, tolerance)
         await scene.expectPixelColor(partColor, moreToTheRightPoint, tolerance)
       })
+
+      await test.step('Delete the part using the feature tree', async () => {
+        await toolbar.openPane('feature-tree')
+        const op = await toolbar.getFeatureTreeOperation('bracket', 0)
+        await op.click({ button: 'right' })
+        await page.getByTestId('context-menu-delete').click()
+        await scene.settled(cmdBar)
+        await toolbar.closePane('feature-tree')
+
+        // Expect empty editor and scene
+        await toolbar.openPane('code')
+        await editor.expectEditor.not.toContain('import')
+        await editor.expectEditor.not.toContain('bracket')
+        await editor.expectEditor.not.toContain('|> translate')
+        await editor.expectEditor.not.toContain('|> rotate')
+        await toolbar.closePane('code')
+        await scene.expectPixelColorNotToBe(partColor, midPoint, tolerance)
+      })
     }
   )
 
@@ -471,6 +489,46 @@ test.describe('Point-and-click assemblies tests', () => {
         await toolbar.closePane('code')
         await scene.expectPixelColor(partColor, partPoint, tolerance)
       })
+
+      await test.step('Delete first part using the feature tree', async () => {
+        await toolbar.openPane('feature-tree')
+        const op = await toolbar.getFeatureTreeOperation('cube', 0)
+        await op.click({ button: 'right' })
+        await page.getByTestId('context-menu-delete').click()
+        await scene.settled(cmdBar)
+        await toolbar.closePane('feature-tree')
+
+        // Expect only the import statement to be there
+        await toolbar.openPane('code')
+        await editor.expectEditor.not.toContain(`import "cube.step" as cube`)
+        await toolbar.closePane('code')
+        await editor.expectEditor.toContain(
+          `
+        import "${complexPlmFileName}" as cubeSw
+        cubeSw
+      `,
+          { shouldNormalise: true }
+        )
+        await toolbar.closePane('code')
+      })
+
+      await test.step('Delete second part using the feature tree', async () => {
+        await toolbar.openPane('feature-tree')
+        const op = await toolbar.getFeatureTreeOperation('cubeSw', 0)
+        await op.click({ button: 'right' })
+        await page.getByTestId('context-menu-delete').click()
+        await scene.settled(cmdBar)
+        await toolbar.closePane('feature-tree')
+
+        // Expect empty editor and scene
+        await toolbar.openPane('code')
+        await editor.expectEditor.not.toContain(
+          `import "${complexPlmFileName}" as cubeSw`
+        )
+        await editor.expectEditor.not.toContain('cubeSw')
+        await toolbar.closePane('code')
+        await scene.expectPixelColorNotToBe(partColor, midPoint, tolerance)
+      })
     }
   )
 

src/lang/modifyAst.ts

@@ -70,6 +70,7 @@ import type {
   CallExpression,
   CallExpressionKw,
   Expr,
+  ExpressionStatement,
   KclValue,
   Literal,
   PathToNode,
@@ -1334,6 +1335,49 @@ export async function deleteFromSelection(
     }
   }
 
+  // Module import and expression case, need to find and delete both
+  const statement = getNodeFromPath<ExpressionStatement>(
+    astClone,
+    selection.codeRef.pathToNode,
+    'ExpressionStatement'
+  )
+  if (!err(statement) && statement.node.type === 'ExpressionStatement') {
+    let expressionIndexToDelete: number | undefined
+    let importAliasToDelete: string | undefined
+    if (
+      statement.node.expression.type === 'Name' &&
+      statement.node.expression.name.type === 'Identifier'
+    ) {
+      expressionIndexToDelete = Number(selection.codeRef.pathToNode[1][0])
+      importAliasToDelete = statement.node.expression.name.name
+    } else if (
+      statement.node.expression.type === 'PipeExpression' &&
+      statement.node.expression.body[0].type === 'Name' &&
+      statement.node.expression.body[0].name.type === 'Identifier'
+    ) {
+      expressionIndexToDelete = Number(selection.codeRef.pathToNode[1][0])
+      importAliasToDelete = statement.node.expression.body[0].name.name
+    } else {
+      return new Error('Expected expression to be a Name or PipeExpression')
+    }
+
+    astClone.body.splice(expressionIndexToDelete, 1)
+    const importIndexToDelete = astClone.body.findIndex(
+      (n) =>
+        n.type === 'ImportStatement' &&
+        n.selector.type === 'None' &&
+        n.selector.alias?.type === 'Identifier' &&
+        n.selector.alias.name === importAliasToDelete
+    )
+    if (importIndexToDelete >= 0) {
+      astClone.body.splice(importIndexToDelete, 1)
+    } else {
+      return new Error("Couldn't find import to delete")
+    }
+
+    return astClone
+  }
+
   // Below is all AST-based deletion logic
   const varDec = getNodeFromPath<VariableDeclarator>(
     ast,