import assert from 'assert' ; import { CancellationToken } from 'vscode-languageserver' ; import { findNodeByOffset } from '../analyzer/parseTreeUtils' ; import { Program } from '../analyzer/program' ; import { createMapFromItems } from '../common/collectionUtils' ; import { ConfigOptions } from '../common/configOptions' ; import { TextRange } from '../common/textRange' ; import { DocumentSymbolCollector } from '../languageService/documentSymbolCollector' ; import { NameNode } from '../parser/parseNodes' ; import { Range } from './harness/fourslash/fourSlashTypes' ; import { parseAndGetTestState } from './harness/fourslash/testState' ; test ( 'folder reference' , ( ) => { const code = `
// @filename: common/__init__.py
//// from [|io2|] import tools as tools
//// from [|io2|].tools import pathUtils as pathUtils

// @filename: io2/empty.py
//// # empty

// @filename: io2/tools/__init__.py
//// def combine(a, b):
////     pass

// @filename: io2/tools/pathUtils.py
//// def getFilename(path):
////     pass

// @filename: test1.py
//// from common import *
////
//// tools.combine(1, 1)
//// pathUtils.getFilename("c")

// @filename: test2.py
//// from .[|io2|] import tools as t
////
//// t.combine(1, 1)

// @filename: test3.py
//// from .[|io2|].tools import pathUtils as p
////
//// p.getFilename("c")

// @filename: test4.py
//// from common import tools, pathUtils
////
//// tools.combine(1, 1)
//// pathUtils.getFilename("c")

// @filename: test5.py
//// from [|io2|] import tools as tools
//// from [|io2|].tools  import pathUtils as pathUtils
////
//// tools.combine(1, 1)
//// pathUtils.getFilename("c")
    ` ; const state = parseAndGetTestState ( code ) . state ; const ranges = state . getRangesByText ( ) . get ( 'io2' ) ! ; for ( const range of ranges ) { verifyReferencesAtPosition ( state . program , state . configOptions , 'io2' , range . fileName , range . pos , ranges ) ; } } ) ; test ( '__init__ wildcard import' , ( ) => { const code = `
// @filename: common/__init__.py
//// from io2 import [|tools|] as [|tools|]
//// from io2.[|tools|] import pathUtils as pathUtils

// @filename: io2/empty.py
//// # empty

// @filename: io2/tools/__init__.py
//// def combine(a, b):
////     pass

// @filename: io2/tools/pathUtils.py
//// def getFilename(path):
////     pass

// @filename: test1.py
//// from common import *
////
//// [|tools|].combine(1, 1)
//// pathUtils.getFilename("c")

// @filename: test2.py
//// from .io2 import [|tools|] as t
////
//// t.combine(1, 1)

// @filename: test3.py
//// from .io2.[|tools|] import pathUtils as p
////
//// p.getFilename("c")

// @filename: test4.py
//// from common import [|tools|], pathUtils
////
//// [|tools|].combine(1, 1)
//// pathUtils.getFilename("c")

// @filename: test5.py
//// from io2 import [|tools|] as [|tools|]
//// from io2.[|tools|]  import pathUtils as pathUtils
////
//// [|tools|].combine(1, 1)
//// pathUtils.getFilename("c")
    ` ; const state = parseAndGetTestState ( code ) . state ; const ranges = state . getRangesByText ( ) . get ( 'tools' ) ! ; for ( const range of ranges ) { verifyReferencesAtPosition ( state . program , state . configOptions , 'tools' , range . fileName , range . pos , ranges ) ; } } ) ; test ( 'submodule wildcard import' , ( ) => { const code = `
// @filename: common/__init__.py
//// from io2 import tools as tools
//// from io2.tools import [|pathUtils|] as [|pathUtils|]

// @filename: io2/empty.py
//// # empty

// @filename: io2/tools/__init__.py
//// def combine(a, b):
////     pass

// @filename: io2/tools/pathUtils.py
//// def getFilename(path):
////     pass

// @filename: test1.py
//// from common import *
////
//// tools.combine(1, 1)
//// [|pathUtils|].getFilename("c")

// @filename: test2.py
//// from .io2 import tools as t
////
//// t.combine(1, 1)

// @filename: test3.py
//// from .io2.tools import [|pathUtils|] as p
////
//// p.getFilename("c")

// @filename: test4.py
//// from common import tools, [|pathUtils|]
////
//// tools.combine(1, 1)
//// [|pathUtils|].getFilename("c")

// @filename: test5.py
//// from io2 import tools as tools
//// from io2.tools  import [|pathUtils|] as [|pathUtils|]
////
//// tools.combine(1, 1)
//// [|pathUtils|].getFilename("c")
    ` ; const state = parseAndGetTestState ( code ) . state ; const ranges = state . getRangesByText ( ) . get ( 'pathUtils' ) ! ; for ( const range of ranges ) { verifyReferencesAtPosition ( state . program , state . configOptions , 'pathUtils' , range . fileName , range . pos , ranges ) ; } } ) ; test ( 'use localName import alias' , ( ) => { const code = `
// @filename: common/__init__.py
//// from io2 import tools as [|/*marker1*/tools|]
//// from io2.tools import pathUtils as pathUtils

// @filename: io2/empty.py
//// # empty

// @filename: io2/tools/__init__.py
//// def combine(a, b):
////     pass

// @filename: io2/tools/pathUtils.py
//// def getFilename(path):
////     pass

// @filename: test1.py
//// from common import *
////
//// [|/*marker2*/tools|].combine(1, 1)
//// pathUtils.getFilename("c")

// @filename: test2.py
//// from .io2 import tools as t
////
//// t.combine(1, 1)

// @filename: test3.py
//// from .io2.tools import pathUtils as p
////
//// p.getFilename("c")

// @filename: test4.py
//// from common import [|/*marker3*/tools|], pathUtils
////
//// [|/*marker4*/tools|].combine(1, 1)
//// pathUtils.getFilename("c")

// @filename: test5.py
//// from io2 import tools as [|/*marker5*/tools|]
//// from io2.tools  import pathUtils as pathUtils
////
//// [|/*marker6*/tools|].combine(1, 1)
//// pathUtils.getFilename("c")
    ` ; const state = parseAndGetTestState ( code ) . state ; const references = state . getRangesByText ( ) . get ( 'tools' ) ! . map ( ( r ) => ( { path : r . fileName , range : state . convertPositionRange ( r ) } ) ) ; state . verifyFindAllReferences ( { marker1 : { references } , marker2 : { references } , marker3 : { references } , marker4 : { references } , marker5 : { references } , marker6 : { references } , } ) ; } ) ; test ( 'use localName import module' , ( ) => { const code = `
// @filename: common/__init__.py
//// from io2 import [|/*marker1*/tools|] as [|tools|]
//// from io2.[|/*marker2*/tools|] import pathUtils as pathUtils

// @filename: io2/empty.py
//// # empty

// @filename: io2/tools/__init__.py
//// def combine(a, b):
////     pass

// @filename: io2/tools/pathUtils.py
//// def getFilename(path):
////     pass

// @filename: test1.py
//// from common import *
////
//// [|tools|].combine(1, 1)
//// pathUtils.getFilename("c")

// @filename: test2.py
//// from .io2 import [|/*marker3*/tools|] as t
////
//// t.combine(1, 1)

// @filename: test3.py
//// from .io2.[|/*marker4*/tools|] import pathUtils as p
////
//// p.getFilename("c")

// @filename: test4.py
//// from common import [|tools|], pathUtils
////
//// [|tools|].combine(1, 1)
//// pathUtils.getFilename("c")

// @filename: test5.py
//// from io2 import [|/*marker5*/tools|] as [|tools|]
//// from io2.[|/*marker6*/tools|]  import pathUtils as pathUtils
////
//// [|tools|].combine(1, 1)
//// pathUtils.getFilename("c")
    ` ; const state = parseAndGetTestState ( code ) . state ; const references = state . getRangesByText ( ) . get ( 'tools' ) ! . map ( ( r ) => ( { path : r . fileName , range : state . convertPositionRange ( r ) } ) ) ; state . verifyFindAllReferences ( { marker1 : { references } , marker2 : { references } , marker3 : { references } , marker4 : { references } , marker5 : { references } , marker6 : { references } , } ) ; } ) ; test ( 'import dotted name' , ( ) => { const code = `
// @filename: nest1/__init__.py
//// # empty

// @filename: nest1/nest2/__init__.py
//// # empty

// @filename: nest1/nest2/module.py
//// def foo():
////     pass

// @filename: test1.py
//// import [|nest1|].[|nest2|].[|module|]
////
//// [|nest1|].[|nest2|].[|module|]

// @filename: nest1/test2.py
//// import [|nest1|].[|nest2|].[|module|]
////
//// [|nest1|].[|nest2|].[|module|]
    ` ; const state = parseAndGetTestState ( code ) . state ; function verify ( name ) { const ranges = state . getRangesByText ( ) . get ( name ) ! ; for ( const range of ranges ) { verifyReferencesAtPosition ( state . program , state . configOptions , name , range . fileName , range . pos , ranges ) ; } } verify ( 'nest1' ) ; verify ( 'nest2' ) ; verify ( 'module' ) ; } ) ; test ( 'import alias' , ( ) => { const code = `
// @filename: nest/__init__.py
//// # empty

// @filename: nest/module2.py
//// # empty

// @filename: module1.py
//// # empty

// @filename: test1.py
//// import [|/*marker1*/module1|] as [|module1|]

// @filename: test2.py
//// import nest.[|/*marker2*/module2|] as [|module2|]
    ` ; const state = parseAndGetTestState ( code ) . state ; const marker1 = state . getMarkerByName ( 'marker1' ) ; const ranges1 = state . getRangesByText ( ) . get ( 'module1' ) ! ; verifyReferencesAtPosition ( state . program , state . configOptions , 'module1' , marker1 . fileName , marker1 . position , ranges1 ) ; const marker2 = state . getMarkerByName ( 'marker2' ) ; const ranges2 = state . getRangesByText ( ) . get ( 'module2' ) ! ; verifyReferencesAtPosition ( state . program , state . configOptions , 'module2' , marker2 . fileName , marker2 . position , ranges2 ) ; } ) ; test ( 'string in __all__' , ( ) => { const code = `
// @filename: test1.py
//// class [|/*marker1*/A|]:
////     pass
////
//// a: "[|A|]" = "A"
////
//// __all__ = [ "[|A|]" ]
    ` ; const state = parseAndGetTestState ( code ) . state ; const marker1 = state . getMarkerByName ( 'marker1' ) ; const ranges1 = state . getRangesByText ( ) . get ( 'A' ) ! ; verifyReferencesAtPosition ( state . program , state . configOptions , 'A' , marker1 . fileName , marker1 . position , ranges1 ) ; } ) ; function verifyReferencesAtPosition ( program , configOption , symbolName , fileName , position , ranges ) { const sourceFile = program . getBoundSourceFile ( fileName ) ; assert ( sourceFile ) ; const node = findNodeByOffset ( sourceFile . getParseResults ( ) ! . parseTree , position ) ; const decls = DocumentSymbolCollector . getDeclarationsForNode ( node as NameNode , program . evaluator ! , true , CancellationToken . None , program . test_createSourceMapper ( configOption . findExecEnvironment ( fileName ) ) ) ; const rangesByFile = createMapFromItems ( ranges , ( r ) => r . fileName ) ; for ( const rangeFileName of rangesByFile . keys ( ) ) { const collector = new DocumentSymbolCollector ( symbolName , decls , program . evaluator ! , CancellationToken . None , program . getBoundSourceFile ( rangeFileName ) ! . getParseResults ( ) ! . parseTree , true ) ; const results = collector . collect ( ) ; const rangesOnFile = rangesByFile . get ( rangeFileName ) ! ; assert . strictEqual ( results . length , rangesOnFile . length , `${ rangeFileName }@${ symbolName }` ) ; for ( const result of results ) { assert ( rangesOnFile . some ( ( r ) => r . pos === result . range . start && r . end === TextRange . getEnd ( result . range ) ) ) ; } } }