Fixing Extract method doesn't return assigned variable (microsoft#759)

heejaechang · Nov 25, 2020 · ff00054 · ff00054
1 parent 769afea
commit ff00054
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 31 deletions.
diff --git a/packages/pylance-internal/src/languageService/refactoringProvider.ts b/packages/pylance-internal/src/languageService/refactoringProvider.ts
@@ -24,7 +24,7 @@ import {
     convertOffsetToPosition,
     convertRangeToTextRange,
 } from 'pyright-internal/common/positionUtils';
-import { Position, Range, TextRange } from 'pyright-internal/common/textRange';
+import { comparePositions, Position, Range, TextRange } from 'pyright-internal/common/textRange';
 import { TextRangeCollection } from 'pyright-internal/common/textRangeCollection';
 import { FindReferencesTreeWalker, ReferencesResult } from 'pyright-internal/languageService/referencesProvider';
 import {
@@ -923,55 +923,68 @@ export class ExtractMethodProvider {
         parseResults: ParseResults,
         token: CancellationToken
     ): string[] {
-        const symbolsWrittenInSelection = new Map<string, string>();
+        const symbolsWrittenInSelection = findSymbolsInSelection(symbolReferences, token, selRange);
 
-        symbolReferences.forEach((refResults, symbolName) => {
-            refResults.declarations.forEach((decl, _) => {
-                if (symbolsWrittenInSelection.has(symbolName)) {
-                    return;
-                }
-
-                if (token.isCancellationRequested) {
-                    return;
-                }
-
-                const isInSelection = TextRange.contains(selRange, decl.node.start);
-                if (isInSelection) {
-                    symbolsWrittenInSelection.set(symbolName, symbolName);
-                }
-            });
-        });
+        const selectionEndPosition = convertOffsetToPosition(
+            TextRange.getEnd(selRange),
+            parseResults.tokenizerOutput.lines
+        );
 
+        // Search the symbols written inside the selection for reads after the selection,
+        // but also check for writes after the selection.
         const outputSymbol = new Map<string, string>();
         symbolReferences.forEach((refResults, symbolName) => {
             if (!symbolsWrittenInSelection.has(symbolName)) {
                 return;
             }
 
-            if (token.isCancellationRequested) {
-                return;
-            }
+            const readsAfterSelection = refResults.locations.filter(
+                (docRange) => comparePositions(docRange.range.start, selectionEndPosition) > 0
+            );
+
+            const declarationsAfterSelection = refResults.declarations.filter(
+                (decl) => decl.node.start > TextRange.getEnd(selRange)
+            );
+
+            readsAfterSelection.forEach((docRange) => {
+                if (token.isCancellationRequested) {
+                    return;
+                }
 
-            refResults.locations.forEach((docRange) => {
                 if (outputSymbol.has(symbolName)) {
                     return;
                 }
+
                 const readLocation = convertRangeToTextRange(docRange.range, parseResults.tokenizerOutput.lines);
                 if (readLocation === undefined) {
                     return;
                 }
 
-                const isReadAfterSelection = readLocation.start > TextRange.getEnd(selRange);
-                if (isReadAfterSelection) {
-                    // Note all declaration locations are also refResult locations so allow
-                    // declaration locations to equal reference locations
-                    const isReassignedAfterSelectionAndBeforeRead = refResults.declarations.some((decl) => {
-                        return decl.node.start > TextRange.getEnd(selRange) && decl.node.start <= readLocation.start;
-                    });
+                // Note all declaration locations are also refResult locations, so skip checking those reads
+                const isDeclaration = declarationsAfterSelection.find((decl) => readLocation.start === decl.node.start);
+                if (isDeclaration) {
+                    return;
+                }
 
-                    if (!isReassignedAfterSelectionAndBeforeRead) {
-                        outputSymbol.set(symbolName, symbolName);
+                const isWrittenBeforeReadAfterSelection = declarationsAfterSelection.some((decl) => {
+                    const writePosition = convertOffsetToPosition(decl.node.start, parseResults.tokenizerOutput.lines);
+                    const readPosition = convertOffsetToPosition(
+                        readLocation.start,
+                        parseResults.tokenizerOutput.lines
+                    );
+
+                    // We need to compare line numbers instead of offsets because in the case
+                    // (x = x) the read is before the write but the offset of the read is greater than the offset of the write
+                    let isWriteBeforeRead = writePosition.line < readPosition.line;
+                    if (writePosition.line === readPosition.line) {
+                        isWriteBeforeRead = writePosition.character > readPosition.character;
                     }
+
+                    return isWriteBeforeRead;
+                });
+
+                if (!isWrittenBeforeReadAfterSelection) {
+                    outputSymbol.set(symbolName, symbolName);
                 }
             });
         });
@@ -1159,6 +1172,32 @@ export class ExtractMethodProvider {
     }
 }
 
+function findSymbolsInSelection(
+    symbolReferences: Map<string, ReferencesResult>,
+    token: CancellationToken,
+    selRange: TextRange
+) {
+    const symbolsWrittenInSelection = new Map<string, string>();
+
+    symbolReferences.forEach((refResults, symbolName) => {
+        refResults.declarations.forEach((decl, _) => {
+            if (symbolsWrittenInSelection.has(symbolName)) {
+                return;
+            }
+
+            if (token.isCancellationRequested) {
+                return;
+            }
+
+            const isInSelection = TextRange.contains(selRange, decl.node.start);
+            if (isInSelection) {
+                symbolsWrittenInSelection.set(symbolName, symbolName);
+            }
+        });
+    });
+    return symbolsWrittenInSelection;
+}
+
 function adjustRangeForWhitespace(selRange: TextRange, fileContents: string): TextRange {
     let offset = selRange.start;
 

diff --git a/packages/pylance-internal/src/tests/fourslash/extractmethod.codeAction.fourslash.ts b/packages/pylance-internal/src/tests/fourslash/extractmethod.codeAction.fourslash.ts
@@ -663,3 +663,17 @@ more than just one line """
         `new_func()`,
     ],
 });
+
+// @filename: TestReadBeforeWriteAfterSelection.py
+////[|/*marker44*/x = 1|]
+////x = x
+// @ts-ignore
+await helper.verifyExtractMethod('marker44', {
+    ['file:///TestReadBeforeWriteAfterSelection.py']: [
+        `\n
+def new_func():
+    x = 1
+    return x\n\n`,
+        `x = new_func()`,
+    ],
+});