Python: Allow absolute imports in directories with scripts

Fixes the import logic to account for absolute imports.

We do this by classifying which files and folders may serve as the
entry point for execution, based on a few simple heuristics. If the
file `module.py` is in the same folder as a file `main.py` that may be
executed directly, then we allow `module` to be a valid name for
`module.py` so that `import module` will work as expected.
This commit is contained in:
Taus Brock-Nannestad
2021-03-23 18:20:53 +01:00
parent 4289e358bf
commit 17d1768259
3 changed files with 38 additions and 1 deletions

View File

@@ -72,6 +72,33 @@ class File extends Container {
* are specified to be extracted.
*/
string getContents() { file_contents(this, result) }
/** Holds if this file is likely to get executed directly, and thus act as an entry point for execution. */
predicate maybeExecutedDirectly() {
// Only consider files in the source code, and not things like the standard library
exists(this.getRelativePath()) and
(
// The file doesn't have the extension `.py` but still contains Python statements
not this.getExtension() = "py" and
exists(Stmt s | s.getLocation().getFile() = this)
or
// The file contains the usual `if __name__ == '__main__':` construction
exists(If i, Name name, StrConst main, Cmpop op |
i.getScope().(Module).getFile() = this and
op instanceof Eq and
i.getTest().(Compare).compares(name, op, main) and
name.getId() = "__name__" and
main.getText() = "__main__"
)
or
// The file contains a `#!` line referencing the python interpreter
exists(Comment c |
c.getLocation().getFile() = this and
c.getLocation().getStartLine() = 1 and
c.getText().regexpMatch("^#! */.*python(2|3)?[ \\\\t]*$")
)
)
}
}
private predicate occupied_line(File f, int n) {
@@ -121,6 +148,9 @@ class Folder extends Container {
this.getBaseName().regexpMatch("[^\\d\\W]\\w*") and
result = this.getParent().getImportRoot(n)
}
/** Holds if execution may start in a file in this directory. */
predicate mayContainEntryPoint() { any(File f | f.getParent() = this).maybeExecutedDirectly() }
}
/**

View File

@@ -204,8 +204,13 @@ private string moduleNameFromBase(Container file) {
string moduleNameFromFile(Container file) {
exists(string basename |
basename = moduleNameFromBase(file) and
legalShortName(basename) and
legalShortName(basename)
|
result = moduleNameFromFile(file.getParent()) + "." + basename
or
// If execution can start in the folder containing this module, then we will assume `file` can
// be imported as an absolute import, and hence return `basename` as a possible name.
file.getParent().(Folder).mayContainEntryPoint() and result = basename
)
or
isPotentialSourcePackage(file) and