Skip packages with bad status during the export to JPEG

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

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",

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)

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()