Allow ANDROID_SDK override

Summary: Add override parameter to `redex.py`. Slightly refactor the code to find a tool.

Reviewed By: justinjhendrick

Differential Revision: D22803372

fbshipit-source-id: e69676e29e2bf2da88d3ca1d15d37a01170315e4

pyredex/utils.py

@@ -23,6 +23,7 @@
 from pyredex.logger import log
 
 
+android_sdk_path = None
 temp_dirs = []
 
 
@@ -60,7 +61,13 @@ def with_temp_cleanup(fn, always_clean=False):
             remove_temp_dirs()
 
 
-def find_android_build_tools():
+def set_android_sdk_path(new_path):
+    global android_sdk_path
+    android_sdk_path = new_path
+    log('Setting SDK path to "%s"' % new_path)
+
+
+def find_android_build_tools_env():
     VERSION_REGEXP = r"\d+\.\d+\.\d+$"
     android_home = os.environ["ANDROID_SDK"]
     build_tools = join(android_home, "build-tools")
@@ -71,11 +78,26 @@ def find_android_build_tools():
     return join(build_tools, version)
 
 
+def find_android_build_tools():
+    global android_sdk_path
+    if android_sdk_path:
+        return android_sdk_path
+    android_sdk_path = find_android_build_tools_env()
+    return android_sdk_path
+
+
+def find_android_build_tool(tool):
+    try:
+        candidate = join(find_android_build_tools(), tool)
+        if os.path.exists(candidate):
+            return candidate
+    except BaseException:
+        pass
+    return shutil.which(tool)
+
+
 def find_apksigner():
-    candidate = join(find_android_build_tools(), "apksigner")
-    if os.path.exists(candidate):
-        return candidate
-    return shutil.which("apksigner")
+    return find_android_build_tool("apksigner")
 
 
 def remove_signature_files(extracted_apk_dir):

redex.py

@@ -39,11 +39,12 @@
     argparse_yes_no_flag,
     dex_glob,
     ensure_libs_dir,
-    find_android_build_tools,
+    find_android_build_tool,
     get_file_ext,
     make_temp_dir,
     move_dexen_to_directories,
     remove_comments,
+    set_android_sdk_path,
     sign_apk,
     with_temp_cleanup,
 )
@@ -500,33 +501,32 @@ def run():
 
 def zipalign(unaligned_apk_path, output_apk_path, ignore_zipalign, page_align):
     # Align zip and optionally perform good compression.
-    try:
-        zipalign = [join(find_android_build_tools(), "zipalign")]
-    except Exception:
-        # We couldn't find zipalign via ANDROID_SDK.  Try PATH.
-        zipalign = ["zipalign"]
-    args = ["4", unaligned_apk_path, output_apk_path]
+    zipalign = [
+        find_android_build_tool("zipalign"),
+        "4",
+        unaligned_apk_path,
+        output_apk_path,
+    ]
     if page_align:
-        args = ["-p"] + args
-    success = False
+        zipalign.insert(1, "-p")
     try:
-        p = subprocess.Popen(zipalign + args, stderr=subprocess.PIPE)
-        err = p.communicate()[1]
-        if p.returncode != 0:
-            error = err.decode(sys.getfilesystemencoding())
-            print("Failed to execute zipalign, stderr: {}".format(error))
-        else:
-            success = True
+        p = subprocess.Popen(zipalign, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+        out, _ = p.communicate()
+        if p.returncode == 0:
+            os.remove(unaligned_apk_path)
+            return
+        out_str = out.decode(sys.getfilesystemencoding())
+        raise RuntimeError("Failed to execute zipalign, output: {}".format(out_str))
     except OSError as e:
         if e.errno == errno.ENOENT:
             print("Couldn't find zipalign. See README.md to resolve this.")
-        else:
-            print("Failed to execute zipalign, strerror: {}".format(e.strerror))
-    finally:
-        if not success:
-            if not ignore_zipalign:
-                raise Exception("Zipalign failed to run")
-            shutil.copy(unaligned_apk_path, output_apk_path)
+        if not ignore_zipalign:
+            raise e
+        shutil.copy(unaligned_apk_path, output_apk_path)
+    except BaseException:
+        if not ignore_zipalign:
+            raise
+        shutil.copy(unaligned_apk_path, output_apk_path)
     os.remove(unaligned_apk_path)
 
 
@@ -798,6 +798,8 @@ def arg_parser(binary=None, config=None, keystore=None, keyalias=None, keypass=N
         help="Do not be verbose, and override TRACE.",
     )
 
+    parser.add_argument("--android-sdk-path", type=str, help="Path to Android SDK")
+
     return parser
 
 
@@ -832,6 +834,9 @@ def __init__(
 def prepare_redex(args):
     debug_mode = args.unpack_only or args.debug
 
+    if args.android_sdk_path:
+        set_android_sdk_path(args.android_sdk_path)
+
     # avoid accidentally mixing up file formats since we now support
     # both apk files and Android bundle files
     if not args.unpack_only: