Change the DropLowProbabilityItemPipeline logic for nested items.

zyte_common_items/pipelines.py

@@ -115,36 +115,42 @@ def get_threshold_for_item(self, item, spider):
     def get_item_name(self, item):
         return item.__class__.__name__
 
-    def process_item(self, item, spider):
-        from scrapy.exceptions import DropItem
-
-        if isinstance(item, dict):
-            # support for custom attrs
-            for item_type, item_instance in item.items():
-                if item_type is not CustomAttributes:
-                    real_item = item_instance
-                    break
-            else:
-                return item
-        else:
-            real_item = item
-
-        item_name = self.get_item_name(real_item)
-        item_proba = real_item.get_probability()
-        threshold = self.get_threshold_for_item(real_item, spider)
-
+    def _process_probability(self, item, threshold):
+        item_name = self.get_item_name(item)
+        item_proba = item.get_probability()
         self.stats.inc_value("drop_low_probability_item/processed")
         self.stats.inc_value(f"drop_low_probability_item/processed/{item_name}")
-
         if item_proba is None or item_proba >= threshold:
             self.stats.inc_value("drop_low_probability_item/kept")
             self.stats.inc_value(f"drop_low_probability_item/kept/{item_name}")
-            return item
+            return True
+        else:
+            self.stats.inc_value("drop_low_probability_item/dropped")
+            self.stats.inc_value(f"drop_low_probability_item/dropped/{item_name}")
+            return False
 
-        self.stats.inc_value("drop_low_probability_item/dropped")
-        self.stats.inc_value(f"drop_low_probability_item/dropped/{item_name}")
+    def process_item(self, item, spider):
+        from scrapy.exceptions import DropItem
 
-        raise DropItem(
-            f"This item is dropped since the probability ({item_proba}) "
-            f"is below the threshold ({threshold}):\n{item!r}"
-        )
+        if isinstance(item, dict):
+            # for nested items remove sub-items that have low probability
+            # instead of dropping the whole result
+            new_item = {}
+            for item_type, sub_item in item.items():
+                if item_type is CustomAttributes:
+                    continue
+                threshold = self.get_threshold_for_item(sub_item, spider)
+                if self._process_probability(sub_item, threshold):
+                    new_item[item_type] = sub_item
+            if not new_item:
+                # everything has been removed
+                raise DropItem
+            return new_item
+        else:
+            threshold = self.get_threshold_for_item(item, spider)
+            if self._process_probability(item, threshold):
+                return item
+            raise DropItem(
+                f"This item is dropped since the probability ({item.get_probability()}) "
+                f"is below the threshold ({threshold}):\n{item!r}"
+            )