Enhancements to PYTHONPATH.

* Fix #177. * Fix #178.
wala · Apr 8, 2024 · 544499d · 544499d
1 parent 5b1e9fc
commit 544499d
Show file tree

Hide file tree

Showing 83 changed files with 1,125 additions and 20 deletions.
diff --git a/...m.wala.cast.python.jython3/source/com/ibm/wala/cast/python/parser/PythonModuleParser.java b/...m.wala.cast.python.jython3/source/com/ibm/wala/cast/python/parser/PythonModuleParser.java
@@ -10,6 +10,9 @@
  *****************************************************************************/
 package com.ibm.wala.cast.python.parser;
 
+import static com.google.common.io.Files.getNameWithoutExtension;
+import static com.ibm.wala.cast.python.ir.PythonLanguage.MODULE_INITIALIZATION_FILENAME;
+
 import com.ibm.wala.cast.python.ir.PythonCAstToIRTranslator;
 import com.ibm.wala.cast.python.util.Util;
 import com.ibm.wala.cast.tree.CAstEntity;
@@ -38,12 +41,17 @@
 import java.util.stream.Collectors;
 import org.antlr.runtime.ANTLRInputStream;
 import org.antlr.runtime.CharStream;
+import org.python.antlr.ast.Import;
 import org.python.antlr.ast.ImportFrom;
 
 public class PythonModuleParser extends PythonParser<ModuleEntry> {
 
   private static final Logger LOGGER = Logger.getLogger(PythonModuleParser.class.getName());
 
+  /** Name of the Python initialization file without the extension. */
+  private static final String MODULE_INITIALIZATION_ENTITY_NAME =
+      getNameWithoutExtension(MODULE_INITIALIZATION_FILENAME);
+
   private final Set<SourceModule> localModules = HashSetFactory.make();
 
   private final SourceModule fileName;
@@ -62,6 +70,85 @@ protected PythonParser<ModuleEntry>.CAstVisitor makeVisitor(
       WalkContext context, WalaPythonParser parser) {
     return new CAstVisitor(context, parser) {
 
+      @Override
+      public CAstNode visitImport(Import imp) throws Exception {
+        Optional<String> s =
+            imp.getInternalNames().stream()
+                .map(alias::getInternalName)
+                .reduce((a, b) -> a + "/" + b);
+
+        if (s.isPresent()) {
+          String moduleName = s.get().replace('.', '/');
+          LOGGER.finer("Module name from " + imp + " is: " + moduleName + ".");
+
+          // if it's a package.
+          if (moduleName.indexOf('/') != -1) {
+            if (!isLocalModule(moduleName)) moduleName += "/" + MODULE_INITIALIZATION_ENTITY_NAME;
+
+            LOGGER.finer("Module name from " + imp + " is: " + moduleName + ".");
+
+            if (isLocalModule(moduleName)) {
+              List<File> pythonPath = PythonModuleParser.this.getPythonPath();
+              LOGGER.info("PYTHONPATH is: " + pythonPath + ".");
+
+              // If there is a PYTHONPATH specified.
+              if (pythonPath != null && !pythonPath.isEmpty()) {
+                // Adjust the module name per the PYTHONPATH.
+                Optional<SourceModule> localModule = getLocalModule(moduleName);
+
+                for (File pathEntry : pythonPath) {
+                  Path modulePath =
+                      localModule
+                          .map(SourceModule::getURL)
+                          .map(URL::getFile)
+                          .map(Path::of)
+                          .orElseThrow(IllegalStateException::new);
+                  LOGGER.finer("Found module path: " + modulePath + ".");
+
+                  if (modulePath.startsWith(pathEntry.toPath())) {
+                    // Found it.
+                    Path scriptRelativePath = pathEntry.toPath().relativize(modulePath);
+                    LOGGER.finer("Relativized path is: " + scriptRelativePath + ".");
+
+                    // Remove the file extension if it exists.
+                    moduleName = scriptRelativePath.toString().replaceFirst("\\.py$", "");
+
+                    // Use the beginning segment initialization file.
+                    moduleName = moduleName.split("/")[0] + "/" + MODULE_INITIALIZATION_ENTITY_NAME;
+
+                    LOGGER.fine("Using module name: " + moduleName + ".");
+                    break;
+                  }
+                }
+              }
+
+              String yuck = moduleName;
+              return Ast.makeNode(
+                  CAstNode.BLOCK_STMT,
+                  imp.getInternalNames().stream()
+                      .map(alias::getInternalName)
+                      .map(
+                          n -> {
+                            n = n.split("\\.")[0];
+
+                            return Ast.makeNode(
+                                CAstNode.DECL_STMT,
+                                Ast.makeConstant(
+                                    new CAstSymbolImpl(n, PythonCAstToIRTranslator.Any)),
+                                Ast.makeNode(
+                                    CAstNode.PRIMITIVE,
+                                    Ast.makeConstant("import"),
+                                    Ast.makeConstant(yuck),
+                                    Ast.makeConstant(n)));
+                          })
+                      .collect(Collectors.toList()));
+            }
+          }
+        }
+
+        return super.visitImport(imp);
+      }
+
       @Override
       public CAstNode visitImportFrom(ImportFrom importFrom) throws Exception {
         Optional<String> s =
@@ -78,10 +165,7 @@ public CAstNode visitImportFrom(ImportFrom importFrom) throws Exception {
           String moduleName = s.get();
           LOGGER.finer("Module name from " + importFrom + " is: " + moduleName + ".");
 
-          if (!isLocalModule(moduleName)) {
-            LOGGER.finer("Module: " + moduleName + " isn't local.");
-            moduleName = s.get() + "/__init__";
-          } else LOGGER.finer("Module: " + moduleName + " is local.");
+          if (!isLocalModule(moduleName)) moduleName += "/" + MODULE_INITIALIZATION_ENTITY_NAME;
 
           LOGGER.finer("Module name from " + importFrom + " is: " + moduleName + ".");
 
@@ -107,7 +191,7 @@ public CAstNode visitImportFrom(ImportFrom importFrom) throws Exception {
                   Path scriptRelativePath = pathEntry.toPath().relativize(modulePath);
                   LOGGER.finer("Relativized path is: " + scriptRelativePath + ".");
 
-                  // Remove the file extension.
+                  // Remove the file extension if it exists.
                   moduleName = scriptRelativePath.toString().replaceFirst("\\.py$", "");
                   LOGGER.fine("Using module name: " + moduleName + ".");
                   break;