artiq_flash/coremgmt flash -> fetch_bin

artiq/frontend/artiq_coremgmt.py

@@ -13,6 +13,7 @@ from artiq.master.databases import DeviceDB
 from artiq.coredevice.comm_kernel import CommKernel
 from artiq.coredevice.comm_mgmt import CommMgmt
 from artiq.frontend.bit2bin import bit2bin
+from artiq.frontend.fetch_bin import fetch_bin
 
 
 def get_argparser():
@@ -171,46 +172,9 @@ def main():
             boot = os.path.join(args.directory, "boot.bin")
             bins = [boot]
         else:
-            def artifact_path(this_binary_dir, *path_filename):
-                if args.srcbuild:
-                    # source tree - use path elements to locate file
-                    return os.path.join(this_binary_dir, *path_filename)
-                else:
-                    # flat tree
-                    # all files in the same directory, discard path elements
-                    *_, filename = path_filename
-                    return os.path.join(this_binary_dir, filename)
-
-            def convert_gateware(bit_filename):
-                bin_handle, bin_filename = tempfile.mkstemp(
-                    prefix="artiq_",
-                    suffix="_" + os.path.basename(bit_filename))
-                with open(bit_filename, "rb") as bit_file, \
-                        open(bin_handle, "wb") as bin_file:
-                    bit2bin(bit_file, bin_file)
-                atexit.register(lambda: os.unlink(bin_filename))
-                return bin_filename
-
-            gateware = convert_gateware(artifact_path(
-                args.directory, "gateware", "top.bit"))
-            bootloader = artifact_path(
-                args.directory, "software", "bootloader", "bootloader.bin")
-
-            firmwares = []
-            for firmware in "satman", "runtime":
-                filename = artifact_path(
-                    args.directory, "software", firmware, firmware + ".fbi")
-                if os.path.exists(filename):
-                    firmwares.append(filename)
-            if not firmwares:
-                raise FileNotFoundError("no firmware found")
-            if len(firmwares) > 1:
-                raise ValueError("more than one firmware file, "
-                                 "please clean up your build directory. "
-                                 "Found firmware files: {}".format(
-                                    " ".join(firmwares)))
-            firmware = firmwares[0]
-
+            gateware = fetch_bin(args.directory, "gateware", args.srcbuild)
+            bootloader = fetch_bin(args.directory, "bootloader", args.srcbuild)
+            firmware = fetch_bin(args.directory, ["runtime", "satman"], args.srcbuild)
             bins = [gateware, bootloader, firmware]
 
         mgmt.flash(bins)

artiq/frontend/artiq_flash.py

@@ -15,6 +15,7 @@ from sipyco import common_args
 from artiq import __version__ as artiq_version
 from artiq.remoting import SSHClient, LocalClient
 from artiq.frontend.bit2bin import bit2bin
+from artiq.frontend.fetch_bin import fetch_bin
 
 
 def get_argparser():
@@ -302,46 +303,19 @@ def main():
 
     programmer = config["programmer"](client, preinit_script=args.preinit_command)
 
-    def artifact_path(this_binary_dir, *path_filename):
-        if args.srcbuild:
-            # source tree - use path elements to locate file
-            return os.path.join(this_binary_dir, *path_filename)
-        else:
-            # flat tree - all files in the same directory, discard path elements
-            *_, filename = path_filename
-            return os.path.join(this_binary_dir, filename)
-
-    def convert_gateware(bit_filename):
-        bin_handle, bin_filename = tempfile.mkstemp(
-            prefix="artiq_", suffix="_" + os.path.basename(bit_filename))
-        with open(bit_filename, "rb") as bit_file, open(bin_handle, "wb") as bin_file:
-            bit2bin(bit_file, bin_file)
-        atexit.register(lambda: os.unlink(bin_filename))
-        return bin_filename
-
     for action in args.action:
         if action == "gateware":
-            gateware_bin = convert_gateware(
-                artifact_path(binary_dir, "gateware", "top.bit"))
+            gateware_bin = fetch_bin(binary_dir, "gateware", args.srcbuild)
             programmer.write_binary(*config["gateware"], gateware_bin)
         elif action == "bootloader":
-            bootloader_bin = artifact_path(binary_dir, "software", "bootloader", "bootloader.bin")
+            bootloader_bin = fetch_bin(binary_dir, "bootloader", args.srcbuild)
             programmer.write_binary(*config["bootloader"], bootloader_bin)
         elif action == "storage":
             storage_img = args.storage
             programmer.write_binary(*config["storage"], storage_img)
         elif action == "firmware":
-            firmware_fbis = []
-            for firmware in "satman", "runtime":
-                filename = artifact_path(binary_dir, "software", firmware, firmware + ".fbi")
-                if os.path.exists(filename):
-                    firmware_fbis.append(filename)
-            if not firmware_fbis:
-                raise FileNotFoundError("no firmware found")
-            if len(firmware_fbis) > 1:
-                raise ValueError("more than one firmware file, please clean up your build directory. "
-                   "Found firmware files: {}".format(" ".join(firmware_fbis)))
-            programmer.write_binary(*config["firmware"], firmware_fbis[0])
+            firmware_fbi = fetch_bin(binary_dir, ["satman", "runtime"], args.srcbuild)
+            programmer.write_binary(*config["firmware"], firmware_fbi)
         elif action == "load":
             gateware_bit = artifact_path(binary_dir, "gateware", "top.bit")
             programmer.load(gateware_bit, 0)

artiq/frontend/fetch_bin.py

@@ -0,0 +1,59 @@
+import atexit
+import os
+import tempfile
+
+from artiq.frontend.bit2bin import bit2bin
+
+
+def fetch_bin(binary_dir, component, srcbuild=False):
+    def artifact_path(this_binary_dir, *path_filename):
+        if srcbuild:
+            # source tree - use path elements to locate file
+            return os.path.join(this_binary_dir, *path_filename)
+        else:
+            # flat tree - all files in the same directory, discard path elements
+            *_, filename = path_filename
+            return os.path.join(this_binary_dir, filename)
+
+    def convert_gateware(bit_filename):
+        bin_handle, bin_filename = tempfile.mkstemp(
+            prefix="artiq_", suffix="_" + os.path.basename(bit_filename))
+        with open(bit_filename, "rb") as bit_file, open(bin_handle, "wb") as bin_file:
+            bit2bin(bit_file, bin_file)
+        atexit.register(lambda: os.unlink(bin_filename))
+        return bin_filename
+
+    if type(component) == list:
+        bins = []
+        for option in component:
+            try:
+                bins.append(fetch_bin(binary_dir, option, srcbuild))
+            except FileNotFoundError:
+                pass
+
+        if bins is None:
+            raise FileNotFoundError("multiple components not found: {}".format(
+                                        " ".join(component)))
+        
+        if len(bins) > 1:
+            raise ValueError("more than one file, "
+                             "please clean up your build directory. "
+                             "Found files: {}".format(
+                             " ".join(bins)))
+
+        return bins[0]
+
+    path = artifact_path(binary_dir, *{
+        "gateware": ["gateware", "top.bit"],
+        "bootloader": ["software", "bootloader", "bootloader.bin"],
+        "runtime": ["software", "runtime", "runtime.fbi"],
+        "satman": ["software", "satman", "satman.fbi"],
+    }[component])
+
+    if not os.path.exists(path):
+        raise FileNotFoundError("{} not found".format(component))
+
+    if component == "gateware":
+        path = convert_gateware(path)
+
+    return path