diff --git a/covjsonkit/api.py b/covjsonkit/api.py index 76346dd..008921f 100644 --- a/covjsonkit/api.py +++ b/covjsonkit/api.py @@ -82,7 +82,8 @@ def encode(self, type, domaintype): elif domaintype == "trajectory": domaintype = "path" feature = self._feature_factory(domaintype.lower(), "encoder") - return feature(self.conf, domaintype) + coveragejson = feature(self.conf, domaintype) + return coveragejson def decode(self, covjson): if "domainType" not in covjson: diff --git a/covjsonkit/config.py b/covjsonkit/config.py index 9fe2abd..0009509 100644 --- a/covjsonkit/config.py +++ b/covjsonkit/config.py @@ -1,5 +1,8 @@ +from typing import Literal + from conflator import ConfigModel class CovjsonKitConfig(ConfigModel): param_db: str = "ecmwf" + compression: Literal["zstd", "lz4", "binpack", None] = None diff --git a/covjsonkit/encoder/encoder.py b/covjsonkit/encoder/encoder.py index 81e4f39..f2ed368 100644 --- a/covjsonkit/encoder/encoder.py +++ b/covjsonkit/encoder/encoder.py @@ -182,7 +182,6 @@ def calculate_index_bounds(level_len, num_len, para_len, step_len, l, i, j, k): return start_index, end_index def append_composite_coords(dates, tree_values, lat, coords): - # for date in dates: for value in tree_values: coords[dates]["composite"].append([lat, value]) diff --git a/covjsonkit/utils.py b/covjsonkit/utils.py index 90c177f..64ff843 100644 --- a/covjsonkit/utils.py +++ b/covjsonkit/utils.py @@ -1,3 +1,7 @@ +import os +from pathlib import Path + + def merge_coverage_collections(collection1, collection2): """ Merges two coverage collections into one. @@ -87,3 +91,76 @@ def coverage_to_coveragecollection(coverage: dict) -> dict: collection[key] = coverage[key] return collection + + +def compress(in_path, out_path=None, compression=None, level=3): + """Compress a file using the specified compression algorithm. + Args: + in_path (str): Path to the input file. + out_path (str, optional): Path to the output file. If None, the output + file will be created in the same directory as the input file with an + appropriate suffix. Defaults to None. + compression (str, optional): Type of compression to use. Options are 'zstd', 'LZ4'. + Defaults to None. + level (int, optional): Compression level. Defaults to 3. + Returns: + str: Path to the compressed file. + """ + if compression is None: + print("No compression specified, please specify type of compression (zstd, LZ4, binpack)") + else: + if compression == "zstd": + import zstandard as zstd + + in_path = Path(in_path) + in_file_size = os.path.getsize(in_path) + + if out_path is None: + out_path = in_path + output_file = in_path.with_suffix(out_path.suffix + ".zst") + + cctx = zstd.ZstdCompressor(level=level) + + with open(in_path, "rb") as f_in, open(output_file, "wb") as f_out: + f_out.write(cctx.compress(f_in.read())) + + output_file_size = os.path.getsize(output_file) + + print(f"Compressed {in_path} → {output_file}") + print(f"Input file size: {in_file_size} bytes") + print(f"Output file size: {output_file_size} bytes") + print(f"Compression ratio: {in_file_size / output_file_size:.2f}") + + return output_file + + elif compression == "LZ4": + import lz4.frame + + input_file = Path(in_path) + in_file_size = os.path.getsize(input_file) + + if out_path is None: + output_file = input_file.with_suffix(input_file.suffix + ".lz4") + else: + output_file = Path(out_path) + + with open(input_file, "rb") as f_in: + data = f_in.read() + + compressed = lz4.frame.compress(data, compression_level=level) + + with open(output_file, "wb") as f_out: + f_out.write(compressed) + + output_file_size = os.path.getsize(output_file) + + print(f"Compressed {input_file} → {output_file}") + print(f"Input file size: {in_file_size} bytes") + print(f"Output file size: {output_file_size} bytes") + print(f"Compression ratio: {in_file_size / output_file_size:.2f}") + + return output_file + # import binpacking + + # compressed = binpacking.pack(data) + # return compressed diff --git a/requirements.txt b/requirements.txt index 6c4b775..fb272a7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,5 @@ covjson-pydantic conflator scipy pre-commit +zstandard +lz4