Stripping Unnecessary Python Packages from AWS Lambda Layers

#!/usr/bin/env python3
"""prune_lambda_layer.py

Safely strip non-runtime artifacts from an AWS Lambda geospatial layer.

Usage:
    python prune_lambda_layer.py ./layer/python/lib/python3.11/site-packages

The script:
  1. Removes test/, docs/, examples/, benchmarks/ trees from every package.
  2. Deletes all .pyc / .pyo bytecode files and __pycache__/ directories.
  3. Removes .dist-info/ for packages that do NOT use pkg_resources at runtime.
  4. Strips DWARF debug symbols from every .so compiled extension.
  5. Prints a before/after size summary so you can verify savings.
"""
import os
import sys
import shutil
import subprocess
from pathlib import Path

# Packages whose .dist-info directories must survive — they call
# importlib.metadata or pkg_resources to locate proj.db or GDAL data files.
KEEP_DIST_INFO = {
    "pyproj",
    "rasterio",
    "fiona",
    "shapely",
    "gdal",
    "numpy",
    "certifi",     # rasterio uses certifi.where() via pkg_resources
}

# Directory names anywhere in the tree that are safe to delete entirely.
PRUNE_DIRS = {"tests", "test", "docs", "doc", "examples", "example", "benchmarks"}

def dir_size_mb(path: Path) -> float:
    """Return total size of a directory tree in megabytes."""
    total = sum(f.stat().st_size for f in path.rglob("*") if f.is_file())
    return total / (1024 * 1024)

def prune(root: Path) -> None:
    """Walk root and remove non-runtime artifacts in place."""
    # Collect dirs first to avoid mutating the tree mid-walk
    dirs_to_remove: list[Path] = []
    for item in sorted(root.rglob("*"), reverse=True):
        if not item.exists():
            # Already deleted by an ancestor removal
            continue

        if item.is_dir():
            # Test/docs/examples trees
            if item.name.lower() in PRUNE_DIRS:
                dirs_to_remove.append(item)
                continue

            # __pycache__ directories
            if item.name == "__pycache__":
                dirs_to_remove.append(item)
                continue

            # .dist-info: remove unless the package needs it at runtime
            if item.name.endswith(".dist-info"):
                pkg_name = item.name.split("-")[0].lower().replace("-", "_")
                if pkg_name not in KEEP_DIST_INFO:
                    dirs_to_remove.append(item)
                continue

            # .data directories created by wheel unpacking (rarely needed)
            if item.name.endswith(".data"):
                dirs_to_remove.append(item)
                continue

        elif item.is_file():
            # Compiled bytecode — Python regenerates .pyc on first import
            if item.suffix in {".pyc", ".pyo"}:
                item.unlink(missing_ok=True)
                continue

            # C source files left by certain sdist builds
            if item.suffix == ".c" and item.parent.name != "include":
                item.unlink(missing_ok=True)
                continue

            # Strip debug symbols from shared libraries.
            # --strip-unneeded removes debug sections and non-exported symbols
            # while preserving the dynamic symbol table required for dlopen().
            if item.suffix == ".so":
                result = subprocess.run(
                    ["strip", "--strip-unneeded", str(item)],
                    capture_output=True,
                )
                if result.returncode != 0:
                    # strip fails on already-stripped or non-ELF files; skip safely
                    print(f"[SKIP] strip failed on {item.name}: {result.stderr.decode().strip()}")

    for d in dirs_to_remove:
        if d.exists():
            shutil.rmtree(d, ignore_errors=True)

def main() -> None:
    layer_dir = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(
        "python/lib/python3.11/site-packages"
    )
    if not layer_dir.is_dir():
        print(f"ERROR: {layer_dir} does not exist or is not a directory.")
        sys.exit(1)

    before_mb = dir_size_mb(layer_dir)
    print(f"Before pruning : {before_mb:.1f} MB  ({layer_dir})")

    prune(layer_dir)

    after_mb = dir_size_mb(layer_dir)
    saved_mb = before_mb - after_mb
    print(f"After pruning  : {after_mb:.1f} MB")
    print(f"Saved          : {saved_mb:.1f} MB  ({saved_mb / before_mb * 100:.0f}%)")

    if after_mb > 250:
        print("WARNING: layer still exceeds 250 MB — consider removing unused packages.")
    else:
        print("OK: layer is within the 250 MB AWS Lambda unzipped limit.")

if __name__ == "__main__":
    main()

# Build and prune entirely inside the Lambda-compatible container
docker run --rm \
  -v "$(pwd)/layer:/layer" \
  -v "$(pwd)/prune_lambda_layer.py:/prune_lambda_layer.py:ro" \
  public.ecr.aws/sam/build-python3.11 \
  bash -c "
    # Install dependencies into the layer directory
    pip install --quiet \
      --target /layer/python/lib/python3.11/site-packages \
      rasterio==1.4.3 shapely==2.0.6 pyproj==3.7.0 fiona==1.10.1

    # Prune non-runtime artifacts
    python3 /prune_lambda_layer.py /layer/python/lib/python3.11/site-packages
  "

# Package the pruned layer
cd layer && zip -r9 ../geospatial-layer.zip python/ && cd ..
echo "Zipped size: $(du -sh geospatial-layer.zip | cut -f1)"

# Smoke test: mount the pruned layer directory and verify imports
docker run --rm \
  -v "$(pwd)/layer:/opt:ro" \
  -e GDAL_DATA=/opt/python/lib/python3.11/site-packages/rasterio/gdal_data \
  -e PROJ_LIB=/opt/python/lib/python3.11/site-packages/pyproj/proj_dir/share/proj \
  -e LD_LIBRARY_PATH=/opt/python/lib/python3.11/site-packages/rasterio.libs \
  public.ecr.aws/sam/build-python3.11 \
  python3 -c "
import rasterio
import shapely
import fiona
import pyproj
from pyproj import CRS

# Verify PROJ database is reachable (requires intact .dist-info for pyproj)
crs = CRS.from_epsg(4326)
print(f'rasterio {rasterio.__version__} OK')
print(f'shapely  {shapely.__version__} OK')
print(f'fiona    {fiona.__version__} OK')
print(f'pyproj   {pyproj.__version__} OK  CRS name: {crs.name}')
print('ALL IMPORTS SUCCESSFUL')
"

rasterio 1.4.3 OK
shapely  2.0.6 OK
fiona    1.10.1 OK
pyproj   3.7.0 OK  CRS name: WGS 84
ALL IMPORTS SUCCESSFUL

# Unzipped size of everything in the layer directory
du -sh layer/
# Should report ≤ 250M

# Zipped size for upload (50 MB limit for direct upload; use S3 above that)
ls -lh geospatial-layer.zip