diff --git a/artiq/frontend/artiq_coremgmt.py b/artiq/frontend/artiq_coremgmt.py index e5ccd7215..d2163a5bc 100755 --- a/artiq/frontend/artiq_coremgmt.py +++ b/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) diff --git a/artiq/frontend/artiq_flash.py b/artiq/frontend/artiq_flash.py index b1914a4b5..b9563a339 100755 --- a/artiq/frontend/artiq_flash.py +++ b/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) diff --git a/artiq/frontend/fetch_bin.py b/artiq/frontend/fetch_bin.py new file mode 100644 index 000000000..890c3413c --- /dev/null +++ b/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