From 73a88ab359a87e9e82018baa118b17ef9beeb110 Mon Sep 17 00:00:00 2001 From: Rbrtcs1 Date: Fri, 27 Sep 2024 08:49:47 -0400 Subject: [PATCH] fix for censorship label bug (#454) * fix for censorship label bug * fix: don't iterate on `None` from gen_metadata key `dict.get(...)` returns the value of a key if it is set. The default parameter (the second parameter) only returns a default value **if the key does not exist**. Therefore, if the key `gen_metadata` is set in `kwargs`, and is set to `None`, then `kwargs.get("gen_metadata", [])` will return `None` and **not** an empty list (`[]`). --------- Co-authored-by: tazlin --- horde/classes/stable/processing_generation.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/horde/classes/stable/processing_generation.py b/horde/classes/stable/processing_generation.py index 94d9360a..13fe77d6 100644 --- a/horde/classes/stable/processing_generation.py +++ b/horde/classes/stable/processing_generation.py @@ -79,11 +79,15 @@ def log_aborted_generation(self): def set_generation(self, generation, things_per_sec, **kwargs): state = kwargs.get("state", "ok") - for metadata in kwargs.get("gen_metadata", []): - if metadata.get('value') == 'csam': - state = 'csam' + gen_metadata = kwargs.get("gen_metadata") or [] + for metadata in gen_metadata: + if metadata.get("type") != "censorship": + # this metadata isnt about censorship + continue + if metadata.get("value") == "csam": + state = "csam" else: - state = 'censored' + state = "censored" if state in ["censored", "csam"]: self.censored = True db.session.commit()