From dcc778520a2ca4f42cb5d6f4c139d99df27a0d17 Mon Sep 17 00:00:00 2001
From: Rasmus Lerchedahl Petersen <yoff@github.com>
Date: Thu, 26 Oct 2023 14:56:14 +0200
Subject: [PATCH] Python: refactor code

Also add explanatory comment.

Co-authored-by: Taus <tausbn@github.com>
---
 python/ql/lib/semmle/python/Module.qll | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/python/ql/lib/semmle/python/Module.qll b/python/ql/lib/semmle/python/Module.qll
index 9e0e32b96d35..fa756fc96557 100644
--- a/python/ql/lib/semmle/python/Module.qll
+++ b/python/ql/lib/semmle/python/Module.qll
@@ -231,17 +231,25 @@ private predicate isRegularPackage(Folder f, string name) {
   exists(f.getFile("__init__.py"))
 }
 
+/** Gets the name of a module imported in package `c`. */
+private string moduleImportedInPackage(Container c) {
+  legalShortName(result) and
+  // it has to be imported in this folder
+  result =
+    any(ImportExpr i | i.getLocation().getFile().getParent() = c)
+        .getName()
+        // strip everything after the first `.`
+        .regexpReplaceAll("\\..*", "") and
+  result != ""
+}
+
 /** Holds if the file `f` could be resolved to a module named `name`. */
 private predicate isPotentialModuleFile(File file, string name) {
   legalShortName(name) and
   name = file.getStem() and
   file.getExtension() = ["py", "pyc", "so", "pyd"] and
   // it has to be imported in this folder
-  name =
-    any(ImportExpr i | i.getLocation().getFile().getParent() = file.getParent())
-        .getName()
-        .regexpReplaceAll("\\..*", "") and
-  name != ""
+  name = moduleImportedInPackage(file.getParent())
 }
 
 /**
@@ -256,11 +264,7 @@ private predicate isNameSpacePackage(Folder f, string name) {
   not isRegularPackage(f, name) and
   // it has to be imported in a file
   // either in this folder or next to this folder
-  name =
-    any(ImportExpr i | i.getLocation().getFile().getParent() in [f, f.getParent()])
-        .getName()
-        .regexpReplaceAll("\\..*", "") and
-  name != "" and
+  name = moduleImportedInPackage([f, f.getParent()]) and
   // no sibling regular package
   // and no sibling module
   not exists(Folder sibling | sibling.getParent() = f.getParent() |