diff --git a/CHANGELOG.md b/CHANGELOG.md index 9e12f7a00c421e9848787ec8f951b326e419072b..f49588c9b66a9e496248030709ab4c9f3879e91c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,6 +42,7 @@ release. ### Added - Apollo Metric drivers, tests, and data [#533](https://github.com/DOI-USGS/ale/pull/533) - Rosetta Virtis drivers, tests, and data [#520](https://github.com/DOI-USGS/ale/pull/520) +- Added compress and decompress ISD functions and added --compress flag to isd_generate[#604](https://github.com/DOI-USGS/ale/issues/604) ### Changed - Changed how push frame sensor drivers compute the `ephemeris_time` property [#595](https://github.com/DOI-USGS/ale/pull/595) diff --git a/ale/isd_generate.py b/ale/isd_generate.py index 6b9695b2ab3c4178922292894de897ed9d8be5bd..67d39fedead51c90e0d0e95bbc685ca3bcc90e36 100755 --- a/ale/isd_generate.py +++ b/ale/isd_generate.py @@ -19,6 +19,9 @@ from pathlib import Path, PurePath import sys import ale +import brotli +import json +from ale.drivers import AleJsonEncoder logger = logging.getLogger(__name__) @@ -57,6 +60,13 @@ def main(): action="store_true", help="Display information as program runs." ) + parser.add_argument( + "-c", "--compress", + action="store_true", + help="Output a compressed isd json file with .br file extension. " + "Ale uses the brotli compression algorithm. " + "To decompress an isd file run: python -c \"import ale.isd_generate as isdg; isdg.decompress_json('/path/to/isd.br')\"" + ) parser.add_argument( "-i", "--only_isis_spice", action="store_true", @@ -109,7 +119,7 @@ def main(): if len(args.input) == 1: try: - file_to_isd(args.input[0], args.out, kernels=k, log_level=log_level, only_isis_spice=args.only_isis_spice, only_naif_spice=args.only_naif_spice, local=args.local) + file_to_isd(args.input[0], args.out, kernels=k, log_level=log_level, compress=args.compress, only_isis_spice=args.only_isis_spice, only_naif_spice=args.only_naif_spice, local=args.local) except Exception as err: # Seriously, this just throws a generic Exception? sys.exit(f"File {args.input[0]}: {err}") @@ -143,6 +153,7 @@ def file_to_isd( out: os.PathLike = None, kernels: list = None, log_level=logging.WARNING, + compress=False, only_isis_spice=False, only_naif_spice=False, local=False, @@ -185,11 +196,61 @@ def file_to_isd( else: usgscsm_str = ale.loads(file, props=props, verbose=log_level>logging.INFO, only_isis_spice=only_isis_spice, only_naif_spice=only_naif_spice) - logger.info(f"Writing: {isd_file}") - isd_file.write_text(usgscsm_str) + if compress: + logger.info(f"Writing: {os.path.splitext(isd_file)[0] + '.br'}") + compress_json(usgscsm_str, os.path.splitext(isd_file)[0] + '.br') + else: + logger.info(f"Writing: {isd_file}") + isd_file.write_text(usgscsm_str) return +def compress_json(json_data, output_file): + """ + Compresses inputted JSON data using brotli compression algorithm. + + Parameters + ---------- + json_data : str + JSON data + + output_file : str + The output compressed file path with .br extension. + + """ + binary_json = json.dumps(json_data).encode('utf-8') + + if not os.path.splitext(output_file)[1] == '.br': + raise ValueError("Output file {} is not a valid .br file extension".format(output_file.split(".")[1])) + + with open(output_file, 'wb') as f: + f.write(brotli.compress(binary_json)) + + +def decompress_json(compressed_json_file): + """ + Decompresses inputted .br file. + + Parameters + ---------- + compressed_json_file : str + .br file path + + Returns + ------- + str + Decompressed .json file path + """ + if not os.path.splitext(compressed_json_file)[1] == '.br': + raise ValueError("Inputted file {} is not a valid .br file extension".format(compressed_json_file)) + with open(compressed_json_file, 'rb') as f: + data = f.read() + with open(compressed_json_file, 'wb') as f: + f.write(brotli.decompress(data)) + + os.rename(compressed_json_file, os.path.splitext(compressed_json_file)[0] + '.json') + + return os.path.splitext(compressed_json_file)[0] + '.json' if __name__ == "__main__": try: diff --git a/environment.yml b/environment.yml index 062d3e934e7af97e7d0fc7db2b9e786a02b791af..1703afda5b1bce4a661fe9f2999f0ae7d99d383d 100644 --- a/environment.yml +++ b/environment.yml @@ -22,3 +22,4 @@ dependencies: - pytest-cov - networkx - breathe + - brotli diff --git a/tests/pytests/test_compression.py b/tests/pytests/test_compression.py new file mode 100644 index 0000000000000000000000000000000000000000..3bc6cdcde169c4994e0cf805f7cf721571cebfb6 --- /dev/null +++ b/tests/pytests/test_compression.py @@ -0,0 +1,28 @@ +import pytest +import json +import os + +import ale + +import ale.isd_generate as isdg + +from conftest import get_image_label, get_isd, compare_dicts + + +def test_compress_decompress(): + label = get_image_label("EN1072174528M") + isd_str = get_isd("messmdis_isis") + + compressed_file = os.path.splitext(label)[0] + '.br' + + isdg.compress_json(isd_str, compressed_file) + + decompressed_file = isdg.decompress_json(compressed_file) + + with open(decompressed_file, 'r') as fp: + isis_dict = json.load(fp) + + comparison = compare_dicts(isis_dict, isd_str) + assert comparison == [] + +