add test and fix for erroneously frozen package paths in pre-initialized python modules

(cherry picked from commit 353f2bb)

Author: timfel

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/PythonContext.java

@@ -37,6 +37,7 @@
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Supplier;
 
+import org.graalvm.nativeimage.ImageInfo;
 import org.graalvm.options.OptionValues;
 
 import com.oracle.graal.python.PythonLanguage;
@@ -45,7 +46,10 @@
 import com.oracle.graal.python.builtins.objects.cext.PythonNativeClass;
 import com.oracle.graal.python.builtins.objects.common.HashingStorage;
 import com.oracle.graal.python.builtins.objects.dict.PDict;
+import com.oracle.graal.python.builtins.objects.list.PList;
 import com.oracle.graal.python.builtins.objects.module.PythonModule;
+import com.oracle.graal.python.builtins.objects.str.PString;
+import com.oracle.graal.python.nodes.SpecialAttributeNames;
 import com.oracle.graal.python.runtime.AsyncHandler.AsyncAction;
 import com.oracle.graal.python.runtime.exception.PException;
 import com.oracle.truffle.api.Assumption;
@@ -234,17 +238,50 @@ public boolean isInitialized() {
 
     public void initialize() {
         core.initialize(this);
-        setupRuntimeInformation();
+        setupRuntimeInformation(false);
         core.postInitialize();
     }
 
     public void patch(Env newEnv) {
         setEnv(newEnv);
-        setupRuntimeInformation();
+        setupRuntimeInformation(true);
         core.postInitialize();
     }
 
-    private void setupRuntimeInformation() {
+    /**
+     * During pre-initialization, we're also loading code from the Python standard library. Since
+     * some of those modules may be packages, they will have their __path__ attribute set to the
+     * absolute path of the package on the build system. We use this function to patch the paths
+     * during build time and after starting up from a pre-initialized context so they point to the
+     * run-time package paths.
+     */
+    private void patchPackagePaths(String from, String to) {
+        for (Object v : sysModules.getDictStorage().values()) {
+            if (v instanceof PythonModule) {
+                Object path = ((PythonModule) v).getAttribute(SpecialAttributeNames.__PATH__);
+                if (path instanceof PList) {
+                    Object[] paths = ((PList) path).getSequenceStorage().getCopyOfInternalArray();
+                    for (int i = 0; i < paths.length; i++) {
+                        Object pathElement = paths[i];
+                        String strPath;
+                        if (pathElement instanceof PString) {
+                            strPath = ((PString) pathElement).getValue();
+                        } else if (pathElement instanceof String) {
+                            strPath = (String) pathElement;
+                        } else {
+                            continue;
+                        }
+                        if (strPath.startsWith(from)) {
+                            paths[i] = strPath.replace(from, to);
+                        }
+                    }
+                    ((PythonModule) v).setAttribute(SpecialAttributeNames.__PATH__, core.factory().createList(paths));
+                }
+            }
+        }
+    }
+
+    private void setupRuntimeInformation(boolean isPatching) {
         PythonModule sysModule = core.lookupBuiltinModule("sys");
         sysModules = (PDict) sysModule.getAttribute("modules");
 
@@ -256,6 +293,18 @@ private void setupRuntimeInformation() {
 
         sysModules.setItem(__MAIN__, mainModule);
 
+        final String stdLibPlaceholder = "!stdLibHome!";
+        final String stdLibHome = PythonCore.getStdlibHome(getEnv());
+        if (ImageInfo.inImageBuildtimeCode()) {
+            // Patch any pre-loaded packages' paths if we're running
+            // pre-initialization
+            patchPackagePaths(stdLibHome, stdLibPlaceholder);
+        } else if (isPatching && ImageInfo.inImageRuntimeCode()) {
+            // Patch any pre-loaded packages' paths to the new stdlib home if
+            // we're patching a pre-initialized context
+            patchPackagePaths(stdLibPlaceholder, stdLibHome);
+        }
+
         currentException = null;
         isInitialized = true;
     }

mx.graalpython/mx_graalpython.py

@@ -394,6 +394,19 @@ def graalpython_gate_runner(args, tasks):
             ])
             if success not in out.data:
                 mx.abort('Output from generated SVM image "' + svm_image + '" did not match success pattern:\n' + success)
+            # Test that stdlib paths are not cached on e.g. encodings module
+            out = mx.OutputCapture()
+            mx.run([svm_image, "-S", "--python.StdLibHome=/foobar", "-c", "import encodings; print(encodings.__path__)"], out=mx.TeeOutputCapture(out))
+            if "/foobar" not in out.data:
+                mx.abort('Output from generated SVM image "' + svm_image + '" did not have patched std lib path "/foobar", got:\n' + success)
+            # Finally, test that we can start even if the graalvm was moved
+            out = mx.OutputCapture()
+            graalvm_home = svm_image.replace(os.path.sep.join(["", "bin", "graalpython"]), "")
+            new_graalvm_home = graalvm_home + "_new"
+            shutil.move(graalvm_home, new_graalvm_home)
+            launcher = os.path.join(new_graalvm_home, "bin", "graalpython")
+            mx.log(launcher)
+            mx.run([launcher, "-S", "-c", "print(b'abc'.decode('ascii'))"]) # Should not fail
 
 
 mx_gate.add_gate_runner(SUITE, graalpython_gate_runner)