feat: filter labels when label_type_in is provided (#1783)

Author: florianlega

src/kili/services/export/tools.py

@@ -165,6 +165,31 @@ def fetch_assets(
         kili.kili_api_gateway, download_media, fields, ProjectId(project_id), local_media_dir
     )
     assets_gen = kili.kili_api_gateway.list_assets(filters, fields, options)
+
+    if (label_type_in is not None) and (len(label_type_in) > 0):
+        if export_type == "latest":
+            assets_gen = filter(
+                lambda asset: asset["latestLabel"].get("labelType") in label_type_in, assets_gen
+            )
+        else:
+            assets_gen = filter(
+                lambda asset: any(
+                    label.get("labelType") in label_type_in for label in asset["labels"]
+                ),
+                assets_gen,
+            )
+            assets_gen = (
+                {
+                    **asset,
+                    "labels": [
+                        label
+                        for label in asset["labels"]
+                        if label.get("labelType") in label_type_in
+                    ],
+                }
+                for asset in assets_gen
+            )
+
     if download_media_function is not None:
         assets: List[Dict] = []
         # TODO: modify download_media function so it can take a generator of assets

tests/fakes/fake_data.py

@@ -65,6 +65,7 @@
     "latestLabel": {
         "jsonResponse": job_object_detection,
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
+        "labelType": "DEFAULT",
     },
     "externalId": "car_1",
     "content": "https://storage.googleapis.com/label-public-staging/car/car_1.jpg",
@@ -77,6 +78,7 @@
     "latestLabel": {
         "jsonResponse": job_object_detection_with_classification,
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
+        "labelType": "DEFAULT",
     },
     "externalId": "car_1",
     "content": "https://storage.googleapis.com/label-public-staging/car/car_1.jpg",
@@ -87,6 +89,7 @@
     "latestLabel": {
         "jsonResponse": {},
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
+        "labelType": "DEFAULT",
     },
     "externalId": "car_1",
     "content": "https://storage.googleapis.com/label-public-staging/car/car_1.jpg",
@@ -97,6 +100,7 @@
     "latestLabel": {
         "jsonResponse": job_object_detection,
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
+        "labelType": "DEFAULT",
     },
     "externalId": "car_2",
     "content": "https://storage.googleapis.com/label-public-staging/car/car_2.jpg",
@@ -107,6 +111,7 @@
     "latestLabel": {
         "jsonResponse": job_object_detection,
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
+        "labelType": "DEFAULT",
     },
     "externalId": "car_3",
     "content": "",
@@ -134,6 +139,7 @@
             "1": job_object_detection,
             **{str(i): {} for i in range(2, 28)},
         },
+        "labelType": "DEFAULT",
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
     },
     "externalId": "short_video",
@@ -148,6 +154,7 @@
             "1": job_object_detection,
             **{str(i): {} for i in range(2, 130)},
         },
+        "labelType": "DEFAULT",
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
     },
     "externalId": "video2",

tests/unit/adapters/kili_api_gateway/label/test_data/test_case_13.py

@@ -106,6 +106,7 @@
             "isLatestLabelForUser": True,
             "isSentBackToQueue": False,
             "jsonResponse": json_response,
+            "labelType": "DEFAULT",
         },
         "priority": 0,
         "skipped": False,

tests/unit/services/export/test_export.py

@@ -826,7 +826,7 @@ def test_export_with_asset_filter_kwargs(mocker):
             "inference_mark_gte": 0.7,
             "inference_mark_lte": 0.8,
         },
-        label_type_in=["PREDICTION"],
+        label_type_in=["PREDICTION", "DEFAULT", "REVIEW"],
     )
     expected_where = AssetFilters(
         project_id=ProjectId("fake_proj_id"),
@@ -837,7 +837,7 @@ def test_export_with_asset_filter_kwargs(mocker):
         honeypot_mark_lte=0.4,
         created_at_gte="2022-02-03",
         created_at_lte="2023-02-03",
-        label_type_in=["PREDICTION"],
+        label_type_in=["PREDICTION", "DEFAULT", "REVIEW"],
         inference_mark_gte=0.7,
         inference_mark_lte=0.8,
         issue_type="QUESTION",