Implement zip_safe=False support

example/BUILD

@@ -1,5 +1,6 @@
 load("@rules_zapp//zapp:zapp.bzl",
      "zapp_binary",
+     "zapp_test",
 )
 
 load("@rules_python//python:defs.bzl", "py_library")
@@ -8,13 +9,13 @@ load("@my_deps//:requirements.bzl",
      py_requirement="requirement",
 )
 
-zapp_binary(
+zapp_test(
     name = "hello_script",
     main = "hello.py",
     # entry_point is inferred from main =
 )
 
-zapp_binary(
+zapp_test(
     name = "hello_deps",
     main = "hello.py",
     # deps also get zapped via their underlying wheels
@@ -31,10 +32,20 @@ py_library(
     ]
 )
 
-zapp_binary(
+zapp_test(
     name = "hello_lib_deps",
     main = "hello.py",
     deps = [
         ":lib_hello",
     ],
 )
+
+
+zapp_test(
+    name = "hello_unzipped",
+    zip_safe = False,
+    main = "hello.py",
+    deps = [
+        ":lib_hello",
+    ],
+)

example/WORKSPACE

@@ -20,6 +20,9 @@ git_repository(
 load("@bazel_skylib//:workspace.bzl", "bazel_skylib_workspace")
 bazel_skylib_workspace()
 
+####################################################################################################
+# rules_python
+####################################################################################################
 git_repository(
     name = "rules_python",
     remote = "https://github.com/bazelbuild/rules_python.git",
@@ -41,6 +44,6 @@ local_repository(
 load("@rules_python//python:pip.bzl", "pip_install")
 
 pip_install(
-   name = "my_deps",
+    name = "my_deps",
-   requirements = "//:requirements.txt",
+    requirements = "//:requirements.txt",
 )

example/hello.py

@@ -5,7 +5,7 @@ def main():
     for e in sys.path:
         print(" -", e)
 
-    print("hello, world!")
+    print(f"hello, world! This is {__file__}")
 
     try:
         import yaml
@@ -14,6 +14,8 @@ def main():
     except ImportError:
         print("Don't have YAML.")
 
+    exit(0)
+
 
 if __name__ == "__main__":
     main()

zapp/compiler/__main__.py

@@ -104,7 +104,7 @@ def load_wheel(opts, manifest, path):
 
     prefix = os.path.dirname(path)
 
-    sources = {k: v for k, v in manifest["sources"].items() if v.startswith(prefix)}
+    sources = {k: v for k, v in manifest["sources"].items() if v["source"].startswith(prefix)}
 
     return {
         # "record": record,
@@ -114,7 +114,6 @@ def load_wheel(opts, manifest, path):
     }
 
 
-
 def wheel_name(wheel):
     """Construct the "canonical" filename of the wheel."""
 
@@ -145,7 +144,7 @@ def zip_wheel(tmpdir, wheel):
 
     with zipfile.ZipFile(wheel_file, "w") as whl:
         for dest, src in wheel["sources"].items():
-            whl.write(src, dest)
+            whl.write(src["source"], dest)
 
     return wheel_file
 
@@ -158,15 +157,17 @@ def rezip_wheels(opts, manifest):
     Files sourced from unzipped wheels should be removed, and a single wheel reference inserted."""
 
     wheels = [
-        load_wheel(opts, manifest, os.path.dirname(p))
+        load_wheel(opts, manifest, os.path.dirname(s["source"]))
-        for p in manifest["sources"].values()
+        for s in manifest["sources"].values()
-        if p.endswith("/WHEEL")
+        if s["source"].endswith("/WHEEL")
     ]
 
     # Zip up the wheels and insert wheel records to the manifest
     for w in wheels:
         # Try to cheat and hit in the local cache first rather than building wheels every time
         wn = wheel_name(w)
+        # Expunge sources available in the wheel
+        manifest["sources"] = dsub(manifest["sources"], w["sources"])
 
         # We may have a double-path dependency.
         # If we DON'T, we have to zip
@@ -184,13 +185,10 @@ def rezip_wheels(opts, manifest):
             # Insert a new wheel source
             manifest["wheels"][wn] = {"hashes": [], "source": wf}
 
-        # Expunge sources available in the wheel
-        manifest["sources"] = dsub(manifest["sources"], w["sources"])
-
     return manifest
 
 
-def generate_dunder_inits(manifest):
+def generate_dunder_inits(opts, manifest):
     """Hack the manifest to insert __init__ files as needed."""
 
     sources = manifest["sources"]
@@ -207,20 +205,46 @@ def generate_dunder_inits(manifest):
 def insert_manifest_json(opts, manifest):
     """Insert the manifest.json file."""
 
-    manifest["sources"]["zapp/manifest.json"] = opts.manifest
+    tempf = os.path.join(opts.tmpdir, "manifest.json")
+
+    # Note ordering to enable somewhat self-referential manifest
+    manifest["sources"]["zapp/manifest.json"] = {"source": tempf, "hashes": []}
+
+    with open(tempf, "w") as fp:
+        fp.write(json.dumps(manifest))
 
     return manifest
 
 
-def enable_unzipping(manifest):
+def generate_dunder_main(opts, manifest):
+    """Insert the __main__.py to the manifest."""
+
+    if "__main__.py" in manifest["sources"]:
+        print("Error: __main__.py conflict.", file=sys.stderr)
+        exit(1)
+
+    tempf = os.path.join(opts.tmpdir, "__main__.py")
+    # Note ordering to enable somewhat self-referential manifest
+    manifest["sources"]["__main__.py"] = {"source": tempf, "hashes": []}
+    with open(tempf, "w") as fp:
+        fp.write(make_dunder_main(manifest))
+
+    return manifest
+
+
+def enable_unzipping(opts, manifest):
     """Inject unzipping behavior as needed."""
 
     if manifest["wheels"]:
-        manifest["prelude_points"].append("zapp.support.unpack:unpack_deps")
+        manifest["prelude_points"].extend([
+            "zapp.support.unpack:unpack_deps",
+            "zapp.support.unpack:install_deps",
+        ])
 
-    # FIXME:
+    if not manifest["zip_safe"]:
-    # if not manifest["zip_safe"]:
+        manifest["prelude_points"].extend([
-    # enable a similar injection for unzipping
+            "zapp.support.unpack:unpack_zapp",
+        ])
 
     return manifest
 
@@ -235,11 +259,13 @@ def main():
         setattr(opts, "tmpdir", d)
 
         manifest = rezip_wheels(opts, manifest)
-        manifest = insert_manifest_json(opts, manifest)
+        manifest = enable_unzipping(opts, manifest)
-        manifest = enable_unzipping(manifest)
         # Patch the manifest to insert needed __init__ files
+        manifest = generate_dunder_inits(opts, manifest)
+        manifest = generate_dunder_main(opts, manifest)
+        # Generate and insert the manifest
         # NOTE: This has to be the LAST thing we do
-        manifest = generate_dunder_inits(manifest)
+        manifest = insert_manifest_json(opts, manifest)
 
         if opts.debug:
             from pprint import pprint
@@ -257,22 +283,14 @@ def main():
             shebang = "#!" + manifest["shebang"] + "\n"
             zapp.write(shebang)
 
-        if "__main__.py" in manifest["sources"]:
-            print("Error: __main__.py conflict.", file=sys.stderr)
-            exit(1)
-
         # Now we're gonna build the zapp from the manifest
         with zipfile.ZipFile(opts.output, "a") as zapp:
-
-            # Append the __main__.py generated record
-            zapp.writestr("__main__.py", make_dunder_main(manifest))
-
             # Append user-specified sources
             for dest, src in sorted(manifest["sources"].items(), key=lambda x: x[0]):
                 if src is None:
                     zapp.writestr(dest, "")
                 else:
-                    zapp.write(src, dest)
+                    zapp.write(src["source"], dest)
 
             # Append user-specified libraries
             for whl, config in manifest["wheels"].items():

zapp/support/unpack.py

@@ -3,13 +3,26 @@
 import os
 import sys
 from pathlib import Path
-from zipfile import ZipFile
+from tempfile import mkdtemp
+from zipfile import ZipFile, is_zipfile
 
-from zapp.support.manifest import manifest
+from zapp.support.manifest import manifest, once
 
 
+@once
 def cache_root() -> Path:
-    return Path(os.path.join(os.path.expanduser("~"))) / ".cache" / "zapp"
+    """Find a root directory for cached behaviors."""
+
+    shared_cache = Path(os.path.join(os.path.expanduser("~"))) / ".cache" / "zapp"
+
+    # Happy path, read+write filesystem
+    if os.access(shared_cache, os.X_OK | os.W_OK):
+        return shared_cache
+
+    # Unhappy path, we need a tempdir.
+    # At least make one that's stable for the program's lifetime.
+    else:
+        return Path(os.getenv("ZAPP_TMPDIR") or mkdtemp(), "deps")
 
 
 def cache_wheel_root():
@@ -31,24 +44,79 @@ def cache_zapp_path(fingerprint):
 def unpack_deps():
     """Unpack deps, populating and updating the host's cache."""
 
-    # Create the cache dir as needed
+    if is_zipfile(sys.argv[0]):
-    cache_wheel_root().mkdir(parents=True, exist_ok=True)
+        # Create the cache dir as needed
+        cache_wheel_root().mkdir(parents=True, exist_ok=True)
 
-    # For each wheel, touch the existing cached wheel or unpack this one.
+        # For each wheel, touch the existing cached wheel or unpack this one.
-    with ZipFile(sys.argv[0], "r") as zf:
+        with ZipFile(sys.argv[0], "r") as zf:
-        for whl, config in manifest()["wheels"].items():
+            for whl, config in manifest()["wheels"].items():
-            cached_whl = cache_wheel_path(whl)
+                cached_whl = cache_wheel_path(whl)
-            if cached_whl.exists():
+                if cached_whl.exists():
-                cached_whl.touch()
+                    cached_whl.touch()
 
-            else:
+                else:
-                with open(cached_whl, "wb") as of:
+                    with open(cached_whl, "wb") as of:
-                    of.write(zf.read(".deps/" + whl))
+                        of.write(zf.read(".deps/" + whl))
 
-            sys.path.insert(0, str(cached_whl))
+    else:
+        pass
 
 
-def main():
+def install_deps():
-    """Inspect the manifest."""
+    """Validate that the PYTHONPATH has been configured correctly."""
 
-    unpack_deps()
+    # FIXME: Can we reference the requirements, not specific PATH entries?
+    for whl, config in manifest()["wheels"].items():
+        cached_whl = cache_wheel_path(whl)
+        if cached_whl.exists():
+            cached_whl.touch()
+
+        # Remove any references to the dep and shift the cached whl to the front
+        p = str(cached_whl.resolve())
+        try:
+            sys.path.remove(p)
+        except ValueError:
+            pass
+        sys.path.insert(0, p)
+
+
+def canonicalize_path():
+    """Fixup sys.path entries to use absolute paths. De-dupe in the same pass."""
+
+    # Note that this HAS to be mutative/in-place
+    shift = 0
+    for i in range(len(sys.path)):
+        idx = i - shift
+        el = str(Path.resolve(sys.path[idx]))
+        if el in sys.path:
+            shift += 1
+            sys.path.pop(idx)
+        else:
+            sys.path[idx] = el
+
+
+def unpack_zapp():
+    """Unzip a zapp (excluding the .deps/* tree) into a working directory.
+
+    Note that unlike PEX, these directories are per-run which prevents local mutable state somewhat.
+
+    """
+
+    # Short circuit
+    if is_zipfile(sys.argv[0]):
+        # Extract
+        tmpdir = mkdtemp()
+        with ZipFile(sys.argv[0], "r") as zf:
+            for src in manifest()["sources"].keys():
+                ofp = Path(tmpdir, "usr", src)
+                ofp.parent.mkdir(parents=True, exist_ok=True)
+                with open(ofp, "wb") as of:
+                    of.write(zf.read(src))
+
+        # Re-exec the current interpreter
+        args = [sys.executable, "--", os.path.join(tmpdir, "usr", "__main__.py")] + sys.argv[1:]
+        os.execvpe(args[0], args[1:], {"PYTHONPATH": "", "ZAPP_TMPDIR": tmpdir})
+
+    else:
+        pass

zapp/zapp.bzl

@@ -125,7 +125,7 @@ def _zapp_impl(ctx):
         output = manifest_file,
         content = json.encode({
             "shebang": ctx.attr.shebang,
-            "sources": sources_map,
+            "sources": {d: {"hashes": [], "source": s} for d, s in sources_map.items()},
             "zip_safe": ctx.attr.zip_safe,
             "prelude_points": ctx.attr.prelude_points,
             "entry_point": main_py_ref,
@@ -153,25 +153,32 @@ def _zapp_impl(ctx):
     )
 
     # .zapp file itself has no runfiles and no providers
-    return []
+    return struct(
+      runfiles = ctx.runfiles(
+          files = [ctx.outputs.executable],
+          transitive_files = None,
+          collect_default = True,
+      )
+  )
 
+_zapp_attrs = {
+    "src": attr.label(mandatory = True),
+    "main": attr.label(allow_single_file = True),
+    "wheels": attr.label_list(),
+    "entry_point": attr.string(),
+    "prelude_points": attr.string_list(),
+    "compiler": attr.label(
+        default = Label(DEFAULT_COMPILER),
+        executable = True,
+        cfg = "host",
+    ),
+    "shebang": attr.string(default = "/usr/bin/env python3"),
+    "zip_safe": attr.bool(default = True),
+    "root_import": attr.bool(default = False),
+}
 
-zapp = rule(
+_zapp = rule(
-    attrs = {
+    attrs = _zapp_attrs,
-        "src": attr.label(mandatory = True),
-        "main": attr.label(allow_single_file = True),
-        "wheels": attr.label_list(),
-        "entry_point": attr.string(),
-        "prelude_points": attr.string_list(),
-        "compiler": attr.label(
-            default = Label(DEFAULT_COMPILER),
-            executable = True,
-            cfg = "host",
-        ),
-        "shebang": attr.string(default = "/usr/bin/env python3"),
-        "zip_safe": attr.bool(default = True),
-        "root_import": attr.bool(default = False),
-    },
     executable = True,
     implementation = _zapp_impl,
 )
@@ -186,6 +193,7 @@ def zapp_binary(name,
                 test=False,
                 compiler=None,
                 zip_safe=True,
+                rule=_zapp,
                 **kwargs):
     """A self-contained, single-file Python program, with a .zapp file extension.
 
@@ -232,7 +240,7 @@ def zapp_binary(name,
         **kwargs
     )
 
-    zapp(
+    rule(
         name = name,
         src = name + ".lib",
         compiler = compiler,
@@ -244,8 +252,14 @@ def zapp_binary(name,
     )
 
 
+_zapp_test = rule(
+    attrs = _zapp_attrs,
+    test = True,
+    implementation = _zapp_impl,
+)
+
+
 def zapp_test(name, **kwargs):
     """Same as zapp_binary, just sets the test=True bit."""
 
-    kwargs.pop("test")
+    zapp_binary(name, rule=_zapp_test, **kwargs)
-    zapp_binary(name, test=True, **kwargs)