Merge pull request github#5614 from tausbn/python-allow-absolute-imports-from-source-directory

Python: Allow absolute imports from source directory

Author: tausbn

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 isPossibleEntryPoint() {
+    // 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().matches("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) {

python/ql/src/semmle/python/Module.qll

@@ -205,11 +205,31 @@ private string moduleNameFromBase(Container file) {
   file instanceof File and result = file.getStem()
 }
 
+/**
+ * Holds if `file` may be transitively imported from a file that may serve as the entry point of
+ * the execution.
+ */
+private predicate transitively_imported_from_entry_point(File file) {
+  file.getExtension().matches("%py%") and
+  exists(File importer |
+    importer.getParent() = file.getParent() and
+    exists(ImportExpr i | i.getLocation().getFile() = importer and i.getName() = file.getStem())
+  |
+    importer.isPossibleEntryPoint() or transitively_imported_from_entry_point(importer)
+  )
+}
+
 string moduleNameFromFile(Container file) {
   exists(string basename |
     basename = moduleNameFromBase(file) and
-    legalShortName(basename) and
+    legalShortName(basename)
+  |
     result = moduleNameFromFile(file.getParent()) + "." + basename
+    or
+    // If `file` is a transitive import of a file that's executed directly, we allow references
+    // to it by its `basename`.
+    transitively_imported_from_entry_point(file) and
+    result = basename
   )
   or
   isPotentialSourcePackage(file) and

python/ql/test/3/library-tests/modules/entry_point/hash_bang/main.py

@@ -0,0 +1,7 @@
+#! /usr/bin/python3
+print(__file__)
+import module
+import package
+import namespace_package
+import namespace_package.namespace_package_main
+print(module.message)

python/ql/test/3/library-tests/modules/entry_point/hash_bang/module.py

@@ -0,0 +1,2 @@
+print(__file__.split("entry_point")[1])
+message = "Hello world!"

python/ql/test/3/library-tests/modules/entry_point/hash_bang/namespace_package/namespace_package_main.py

@@ -0,0 +1,2 @@
+print(__file__.split("entry_point")[1])
+import namespace_package.namespace_package_module

python/ql/test/3/library-tests/modules/entry_point/hash_bang/namespace_package/namespace_package_module.py

@@ -0,0 +1 @@
+print(__file__.split("entry_point")[1])

python/ql/test/3/library-tests/modules/entry_point/hash_bang/package/__init__.py

@@ -0,0 +1,2 @@
+print(__file__.split("entry_point")[1])
+from . import package_main

python/ql/test/3/library-tests/modules/entry_point/hash_bang/package/package_main.py

@@ -0,0 +1,2 @@
+print(__file__.split("entry_point")[1])
+from . import package_module

python/ql/test/3/library-tests/modules/entry_point/hash_bang/package/package_module.py

@@ -0,0 +1 @@
+print(__file__.split("entry_point")[1])

python/ql/test/3/library-tests/modules/entry_point/modules.expected

@@ -0,0 +1,14 @@
+| module | hash_bang/module.py:0:0:0:0 | Module module |
+| module | name_main/module.py:0:0:0:0 | Module module |
+| package | hash_bang/package:0:0:0:0 | Package package |
+| package | name_main/package:0:0:0:0 | Package package |
+| package | no_py_extension/package:0:0:0:0 | Package package |
+| package.__init__ | hash_bang/package/__init__.py:0:0:0:0 | Module package.__init__ |
+| package.__init__ | name_main/package/__init__.py:0:0:0:0 | Module package.__init__ |
+| package.__init__ | no_py_extension/package/__init__.py:0:0:0:0 | Module package.__init__ |
+| package.package_main | hash_bang/package/package_main.py:0:0:0:0 | Module package.package_main |
+| package.package_main | name_main/package/package_main.py:0:0:0:0 | Module package.package_main |
+| package.package_main | no_py_extension/package/package_main.py:0:0:0:0 | Module package.package_main |
+| package.package_module | hash_bang/package/package_module.py:0:0:0:0 | Module package.package_module |
+| package.package_module | name_main/package/package_module.py:0:0:0:0 | Module package.package_module |
+| package.package_module | no_py_extension/package/package_module.py:0:0:0:0 | Module package.package_module |