Merge --base-docker-image and --docker-image flag

src/xpk/core/docker_image.py

@@ -18,6 +18,7 @@
 import os
 import random
 import string
+import subprocess
 
 from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
@@ -169,50 +170,47 @@ def setup_docker_image(args) -> tuple[int, str]:
       0 if successful and 1 otherwise.
       Name of the docker image to use.
   """
-  use_base_docker_image = use_base_docker_image_or_docker_image(args)
-
-  docker_image = args.base_docker_image
-  if use_base_docker_image:
-    validate_docker_image_code = validate_docker_image(docker_image, args)
-    if validate_docker_image_code != 0:
-      xpk_exit(validate_docker_image_code)
-    build_docker_image_code, docker_image = build_docker_image_from_base_image(
-        args
-    )
-    if build_docker_image_code != 0:
-      xpk_exit(build_docker_image_code)
-  else:
-    docker_image = args.docker_image
-    validate_docker_image_code = validate_docker_image(args.docker_image, args)
-    if validate_docker_image_code != 0:
-      xpk_exit(validate_docker_image_code)
 
-  return 0, docker_image
+  docker_image = args.docker_image
 
+  if (
+      docker_image
+      and args.base_docker_image
+      and args.base_docker_image is not DEFAULT_DOCKER_IMAGE
+  ):
+    xpk_print(
+        '`--base-docker-image` and --docker-image can not be used together.'
+        ' Please see `--help` command for more details.'
+    )
+    xpk_exit(1)
 
-def use_base_docker_image_or_docker_image(args) -> bool:
-  """Checks for correct docker image arguments.
+  if not docker_image and args.base_docker_image:
+    docker_image = args.base_docker_image  # fallback for legacy users
 
-  Args:
-    args: user provided arguments for running the command.
+  if not docker_image:
+    xpk_print(
+        f'No docker image specified, using default: {DEFAULT_DOCKER_IMAGE}'
+    )
+    docker_image = DEFAULT_DOCKER_IMAGE
 
-  Returns:
-    True if intended to use base docker image, False to use docker image.
-  """
-  use_base_docker_image = True
-  # Check if (base_docker_image and script_dir) or (docker_image) is set.
-  if args.docker_image is not None:
-    if args.script_dir is not DEFAULT_SCRIPT_DIR:
-      xpk_print(
-          '`--script-dir` and --docker-image can not be used together. Please'
-          ' see `--help` command for more details.'
-      )
-      xpk_exit(1)
-    if args.base_docker_image is not DEFAULT_DOCKER_IMAGE:
-      xpk_print(
-          '`--base-docker-image` and --docker-image can not be used together.'
-          ' Please see `--help` command for more details.'
-      )
-      xpk_exit(1)
-    use_base_docker_image = False
-  return use_base_docker_image
+  result = subprocess.run(
+      ['docker', 'images', '-q', docker_image],
+      capture_output=True,
+      text=True,
+      check=True,
+  )
+  is_local = bool(result.stdout.strip())
+
+  if is_local or (args.script_dir and args.script_dir != DEFAULT_SCRIPT_DIR):
+    validate_code = validate_docker_image(docker_image, args)
+    if validate_code != 0:
+      xpk_exit(validate_code)
+    build_code, docker_image = build_docker_image_from_base_image(args)
+    if build_code != 0:
+      xpk_exit(build_code)
+  else:
+    validate_code = validate_docker_image(docker_image, args)
+    if validate_code != 0:
+      xpk_exit(validate_code)
+
+  return 0, docker_image