From 3d1dfb045378df107c60051fd4b4e96f80ac5d8c Mon Sep 17 00:00:00 2001
From: Ayaz Salikhov <mathbunnyru@users.noreply.github.com>
Date: Fri, 7 Jul 2023 17:44:42 +0400
Subject: [PATCH] Fix applying tags in local development (#1939)

---
 Makefile                |  4 ++--
 tagging/apply_tags.py   |  5 ++++-
 tagging/get_platform.py | 10 +++++++---
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/Makefile b/Makefile
index 8c7334a620..8a8ce75b85 100644
--- a/Makefile
+++ b/Makefile
@@ -69,8 +69,8 @@ linkcheck-docs: ## check broken links
 
 hook/%: ## run post-build hooks for an image
 	python3 -m tagging.write_tags_file --short-image-name "$(notdir $@)" --tags-dir /tmp/jupyter/tags/ --owner "$(OWNER)" && \
-	python3 -m tagging.apply_tags --short-image-name "$(notdir $@)" --tags-dir /tmp/jupyter/tags/ --platform "$(shell uname -m)" --owner "$(OWNER)" && \
-	python3 -m tagging.write_manifest --short-image-name "$(notdir $@)" --hist-line-dir /tmp/jupyter/hist_lines/ --manifest-dir /tmp/jupyter/manifests/ --owner "$(OWNER)"
+	python3 -m tagging.write_manifest --short-image-name "$(notdir $@)" --hist-line-dir /tmp/jupyter/hist_lines/ --manifest-dir /tmp/jupyter/manifests/ --owner "$(OWNER)" && \
+	python3 -m tagging.apply_tags --short-image-name "$(notdir $@)" --tags-dir /tmp/jupyter/tags/ --platform "$(shell uname -m)" --owner "$(OWNER)"
 hook-all: $(foreach I, $(ALL_IMAGES), hook/$(I)) ## run post-build hooks for all images
 
 
diff --git a/tagging/apply_tags.py b/tagging/apply_tags.py
index 04148216f0..64c391bdd2 100755
--- a/tagging/apply_tags.py
+++ b/tagging/apply_tags.py
@@ -7,6 +7,8 @@
 
 import plumbum
 
+from tagging.get_platform import unify_aarch64
+
 docker = plumbum.local["docker"]
 
 LOGGER = logging.getLogger(__name__)
@@ -55,7 +57,7 @@ def apply_tags(
         "--platform",
         required=True,
         type=str,
-        choices=["x86_64", "aarch64"],
+        choices=["x86_64", "aarch64", "arm64"],
         help="Image platform",
     )
     arg_parser.add_argument(
@@ -64,5 +66,6 @@ def apply_tags(
         help="Owner of the image",
     )
     args = arg_parser.parse_args()
+    args.platform = unify_aarch64(args.platform)
 
     apply_tags(args.short_image_name, args.owner, args.tags_dir, args.platform)
diff --git a/tagging/get_platform.py b/tagging/get_platform.py
index 139e3697f3..187a259262 100644
--- a/tagging/get_platform.py
+++ b/tagging/get_platform.py
@@ -5,10 +5,14 @@
 ALL_PLATFORMS = {"x86_64", "aarch64"}
 
 
-def get_platform() -> str:
-    machine = platform.machine()
+def unify_aarch64(platform: str) -> str:
     return {
         "aarch64": "aarch64",
         "arm64": "aarch64",  # To support local building on aarch64 Macs
         "x86_64": "x86_64",
-    }[machine]
+    }[platform]
+
+
+def get_platform() -> str:
+    machine = platform.machine()
+    return unify_aarch64(machine)