diff --git a/CHANGELOG.md b/CHANGELOG.md index 676ba22..8705ffd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Add `--debug` flag to see stack of exceptions, [PR-7](https://github.com/panda-official/DriftCLI/pull/7) +### Fixed + +- Skip packages with bad status during the export to JPEG, [PR-8](https://github.com/panda-official/DriftCLI/pull/8) + ## [0.4.0] - 2023-03-15 ### Added diff --git a/drift_cli/export_impl/raw.py b/drift_cli/export_impl/raw.py index befd412..a235c49 100644 --- a/drift_cli/export_impl/raw.py +++ b/drift_cli/export_impl/raw.py @@ -5,6 +5,7 @@ import numpy as np from drift_client import DriftClient +from drift_protocol.common import StatusCode from drift_protocol.meta import MetaInfo from rich.progress import Progress from wavelet_buffer.img import codecs @@ -37,8 +38,16 @@ async def _export_jpeg( **kwargs, ): async for package, task in read_topic(pool, client, topic, progress, sem, **kwargs): + if package.status_code != StatusCode.GOOD: + progress.console.print( + f"Can't extract picture from {topic}/{package.package_id}.dp: {StatusCode.Name(package.status_code)}" + ) + continue + meta = package.meta + if meta.type != MetaInfo.IMAGE: + progress.update( task, description=f"[SKIPPED] Topic {topic} is not an image", diff --git a/drift_cli/utils/helpers.py b/drift_cli/utils/helpers.py index 60d07a7..6fd007f 100644 --- a/drift_cli/utils/helpers.py +++ b/drift_cli/utils/helpers.py @@ -137,8 +137,6 @@ async def _read_package(pkg): if len(stats) > 1: speed = sum(s[0] for s in stats) / (stats[-1][1] - stats[0][1]) - yield drift_pkg, task - count += 1 progress.update( task, @@ -149,6 +147,7 @@ async def _read_package(pkg): refresh=True, ) + yield drift_pkg, task last_time = timestamp progress.update(task, total=1, completed=True) diff --git a/tests/export_test.py b/tests/export_test.py index 79aaf91..a1f9cfe 100644 --- a/tests/export_test.py +++ b/tests/export_test.py @@ -10,6 +10,7 @@ from drift_protocol.common import ( DriftPackage, DataPayload, + StatusCode, ) from drift_protocol.meta import TimeSeriesInfo, MetaInfo, ImageInfo from google.protobuf.any_pb2 import Any # pylint: disable=no-name-in-module @@ -80,7 +81,7 @@ def _make_images(): pkg.status = 0 payload = DataPayload() - payload.data = buffer.serialize(compression_level=16) + payload.data = buffer.serialize(compression_level=0) msg = Any() msg.Pack(payload) @@ -132,8 +133,8 @@ def test__export_raw_data(runner, client, conf, export_path, topics, timeseries) result = runner( f"-c {conf} -p 2 export raw test {export_path} --start 2022-01-01 --stop 2022-01-02" ) - assert f"Topic '{topics[0]}' (copied 2 packages (403 B)" in result.output - assert f"Topic '{topics[1]}' (copied 2 packages (403 B)" in result.output + assert f"Topic '{topics[0]}' (copied 2 packages (943 B)" in result.output + assert f"Topic '{topics[1]}' (copied 2 packages (943 B)" in result.output assert result.exit_code == 0 assert (export_path / topics[0] / "1.dp").exists() @@ -150,8 +151,8 @@ def test__export_raw_data_as_csv(runner, client, conf, export_path, topics, time f"-c {conf} -p 2 export raw test {export_path} --start 2022-01-01 --stop 2022-01-02 --csv" ) - assert f"Topic '{topics[0]}' (copied 2 packages (403 B)" in result.output - assert f"Topic '{topics[1]}' (copied 2 packages (403 B)" in result.output + assert f"Topic '{topics[0]}' (copied 2 packages (943 B)" in result.output + assert f"Topic '{topics[1]}' (copied 2 packages (943 B)" in result.output assert result.exit_code == 0 @@ -212,13 +213,36 @@ def test__export_raw_data_topics_jpeg( f"--jpeg" ) - assert f"Topic '{topics[0]}' (copied 2 packages (1 KB)" in result.output + assert f"Topic '{topics[0]}' (copied 2 packages (241 KB)" in result.output assert result.exit_code == 0 - img = WaveletImage([100, 100], 3, 1, WaveletType.NONE) + img = WaveletImage([100, 100], 3, 1, WaveletType.DB1) img.import_from_file( str(export_path / topics[0] / "1.jpeg"), denoise.Null(), codecs.RgbJpeg() ) img.import_from_file( str(export_path / topics[1] / "2.jpeg"), denoise.Null(), codecs.RgbJpeg() ) + + +@pytest.mark.usefixtures("set_alias") +def test__export_raw_jpeg_skip_bad_packages(runner, client, conf, export_path, topics): + """Should skip bad packages and continue""" + bad_package = DriftPackage() + bad_package.id = 1 + bad_package.status = StatusCode.BAD + + client.get_item.side_effect = ( + [DriftDataPackage(bad_package.SerializeToString())] * 2 * len(topics) + ) + result = runner( + f"-c {conf} -p 1 export raw test {export_path} --start 2022-01-01 --stop 2022-01-02 " + f"--jpeg" + ) + + assert "Can't extract picture from topic1/1.dp" in result.output + assert "Can't extract picture from topic2/1.dp" in result.output + assert result.exit_code == 0 + + assert not (export_path / topics[0] / "1.jpeg").exists() + assert not (export_path / topics[1] / "1.jpeg").exists()