Python: Allow absolute imports in directories with scripts

tausbn · tausbn · commit 17d17682597b · 2021-03-23T18:32:17.000+01:00
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.
diff --git a/python/ql/src/semmle/python/Files.qll b/python/ql/src/semmle/python/Files.qll
@@ -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() }
 }
 
 /**
diff --git a/python/ql/src/semmle/python/Module.qll b/python/ql/src/semmle/python/Module.qll
@@ -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
diff --git a/python/ql/test/3/library-tests/modules/entry_point/modules.expected b/python/ql/test/3/library-tests/modules/entry_point/modules.expected
@@ -1,3 +1,5 @@
+| main | code/main.py:0:0:0:0 | Script main |
+| module | code/module.py:0:0:0:0 | Module module |
 | package | code/package:0:0:0:0 | Package package |
 | package.__init__ | code/package/__init__.py:0:0:0:0 | Module package.__init__ |
 | package.package_main | code/package/package_main.py:0:0:0:0 | Module package.package_main |

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+\| main \| code/main.py:0:0:0:0 \| Script main \|`
	`2`	`+\| module \| code/module.py:0:0:0:0 \| Module module \|`
`1`	`3`	`\| package \| code/package:0:0:0:0 \| Package package \|`
`2`	`4`	`\| package.__init__ \| code/package/__init__.py:0:0:0:0 \| Module package.__init__ \|`
`3`	`5`	`\| package.package_main \| code/package/package_main.py:0:0:0:0 \| Module package.package_main \|`