diff --git a/tests/genrule_sandbox_test.py b/tests/genrule_sandbox_test.py
new file mode 100755
index 0000000..39a60d9
--- /dev/null
+++ b/tests/genrule_sandbox_test.py
@@ -0,0 +1,174 @@
+#!/usr/bin/env python3
+
+# Copyright (C) 2023 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import collections
+import json
+import os.path
+import subprocess
+import tempfile
+
+SRC_ROOT_DIR = os.path.abspath(__file__ + "/../../../..")
+
+
+def _module_graph_path(out_dir):
+  return os.path.join(SRC_ROOT_DIR, out_dir, "soong", "module-actions.json")
+
+
+def _build_with_soong(targets, target_product, out_dir, extra_env={}):
+  env = {
+      "TARGET_PRODUCT": target_product,
+      "TARGET_BUILD_VARIANT": "userdebug",
+  }
+  env.update(os.environ)
+  env.update(extra_env)
+  args = [
+      "build/soong/soong_ui.bash",
+      "--make-mode",
+      "--skip-soong-tests",
+  ]
+  args.extend(targets)
+  try:
+    out = subprocess.check_output(
+        args,
+        cwd=SRC_ROOT_DIR,
+        env=env,
+    )
+  except subprocess.CalledProcessError as e:
+    print(e)
+    print(e.stdout)
+    print(e.stderr)
+    exit(1)
+
+
+def _find_outputs_for_modules(modules, out_dir, target_product):
+  module_path = os.path.join(
+      SRC_ROOT_DIR, out_dir, "soong", "module-actions.json"
+  )
+
+  if not os.path.exists(module_path):
+    _build_with_soong(["json-module-graph"], target_product, out_dir)
+
+  action_graph = json.load(open(_module_graph_path(out_dir)))
+
+  module_to_outs = collections.defaultdict(set)
+  for mod in action_graph:
+    name = mod["Name"]
+    if name in modules:
+      for act in mod["Module"]["Actions"]:
+        if "}generate " in act["Desc"]:
+          module_to_outs[name].update(act["Outputs"])
+  return module_to_outs
+
+
+def _store_outputs_to_tmp(output_files):
+  try:
+    tempdir = tempfile.TemporaryDirectory()
+    for f in output_files:
+      out = subprocess.check_output(
+          ["cp", "--parents", f, tempdir.name],
+          cwd=SRC_ROOT_DIR,
+      )
+    return tempdir
+  except subprocess.CalledProcessError as e:
+    print(e)
+    print(e.stdout)
+    print(e.stderr)
+
+
+def _diff_outs(file1, file2, show_diff):
+  base_args = ["diff"]
+  if not show_diff:
+    base_args.append("--brief")
+  try:
+    args = base_args + [file1, file2]
+    output = subprocess.check_output(
+        args,
+        cwd=SRC_ROOT_DIR,
+    )
+  except subprocess.CalledProcessError as e:
+    if e.returncode == 1:
+      if show_diff:
+        return output
+      return True
+  return None
+
+
+def _compare_outputs(module_to_outs, tempdir, show_diff):
+  different_modules = collections.defaultdict(list)
+  for module, outs in module_to_outs.items():
+    for out in outs:
+      output = None
+      diff = _diff_outs(os.path.join(tempdir.name, out), out, show_diff)
+      if diff:
+        different_modules[module].append(diff)
+
+  tempdir.cleanup()
+  return different_modules
+
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      "--target_product",
+      "-t",
+      default="aosp_cf_arm64_phone",
+      help="optional, target product, always runs as eng",
+  )
+  parser.add_argument(
+      "modules",
+      nargs="+",
+      help="modules to compare builds with genrule sandboxing enabled/not",
+  )
+  parser.add_argument(
+      "--show-diff",
+      "-d",
+      action="store_true",
+      required=False,
+      help="whether to display differing files",
+  )
+  args = parser.parse_args()
+
+  out_dir = os.environ.get("OUT_DIR", "out")
+  target_product = args.target_product
+  modules = set(args.modules)
+
+  module_to_outs = _find_outputs_for_modules(modules, out_dir, target_product)
+  all_outs = set()
+  for outs in module_to_outs.values():
+    all_outs.update(outs)
+  print("build without sandboxing")
+  _build_with_soong(list(all_outs), target_product, out_dir)
+  tempdir = _store_outputs_to_tmp(all_outs)
+  print("build with sandboxing")
+  _build_with_soong(
+      list(all_outs),
+      target_product,
+      out_dir,
+      extra_env={"GENRULE_SANDBOXING": "true"},
+  )
+  diffs = _compare_outputs(module_to_outs, tempdir, args.show_diff)
+  if len(diffs) == 0:
+    print("All modules are correct")
+  elif args.show_diff:
+    for m, d in diffs.items():
+      print(f"Module {m} has diffs {d}")
+  else:
+    print(f"Modules {list(diffs.keys())} have diffs")
+
+
+if __name__ == "__main__":
+  main()
diff --git a/tests/genrule_sandbox_test.sh b/tests/genrule_sandbox_test.sh
deleted file mode 100755
index 21b476b..0000000
--- a/tests/genrule_sandbox_test.sh
+++ /dev/null
@@ -1,111 +0,0 @@
-#!/bin/bash
-
-# Copyright (C) 2023 The Android Open Source Project
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-set -e
-
-# Build the given genrule modules with GENRULE_SANDBOXING enabled and disabled,
-# then compare the output of the modules and report result.
-
-function die() { format=$1; shift; printf >&2 "$format\n" $@; exit 1; }
-
-function usage() {
-  die "usage: ${0##*/} <-t lunch_target> [module]..."
-}
-
-if [ ! -e "build/make/core/Makefile" ]; then
-  die "$0 must be run from the top of the Android source tree."
-fi
-
-declare TARGET=
-while getopts "t:" opt; do
-  case $opt in
-    t)
-      TARGET=$OPTARG ;;
-    *) usage ;;
-  esac
-done
-
-shift $((OPTIND-1))
-MODULES="$@"
-
-source build/envsetup.sh
-
-if [[ -n $TARGET ]]; then
-  lunch $TARGET
-fi
-
-if [[ -z ${OUT_DIR+x} ]]; then
-  OUT_DIR="out"
-fi
-
-OUTPUT_DIR="$(mktemp -d tmp.XXXXXX)"
-PASS=true
-
-function cleanup {
-  if [ $PASS = true ]; then
-    rm -rf "${OUTPUT_DIR}"
-  fi
-}
-trap cleanup EXIT
-
-declare -A GEN_PATH_MAP
-
-function find_gen_paths() {
-  for module in $MODULES; do
-    module_path=$(pathmod "$module")
-    package_path=${module_path#$ANDROID_BUILD_TOP}
-    gen_path=$OUT_DIR/soong/.intermediates$package_path/$module
-    GEN_PATH_MAP[$module]=$gen_path
-  done
-}
-
-function store_outputs() {
-  local dir=$1; shift
-
-  for module in $MODULES; do
-    dest_dir=$dir/${module}
-    mkdir -p $dest_dir
-    gen_path=${GEN_PATH_MAP[$module]}
-    cp -r $gen_path $dest_dir
-  done
-}
-
-function cmp_outputs() {
-  local dir1=$1; shift
-  local dir2=$1; shift
-
-  for module in $MODULES; do
-    if ! diff -rq --exclude=genrule.sbox.textproto $dir1/$module $dir2/$module; then
-      PASS=false
-      echo "$module differ"
-    fi
-  done
-  if [ $PASS = true ]; then
-    echo "Test passed"
-  fi
-}
-
-if [ ! -f "$ANDROID_PRODUCT_OUT/module-info.json" ]; then
-  refreshmod
-fi
-
-find_gen_paths
-m --skip-soong-tests GENRULE_SANDBOXING=true "${MODULES[@]}"
-store_outputs "$OUTPUT_DIR/sandbox"
-m --skip-soong-tests GENRULE_SANDBOXING=false "${MODULES[@]}"
-store_outputs "$OUTPUT_DIR/non_sandbox"
-
-cmp_outputs "$OUTPUT_DIR/non_sandbox" "$OUTPUT_DIR/sandbox"