From d64a49a5ea6e17020096d0699370c6a87c038b5f Mon Sep 17 00:00:00 2001
From: Iajret Creature <122297233+Steals-The-PRs@users.noreply.github.com>
Date: Sat, 2 Dec 2023 18:20:33 +0300
Subject: [PATCH] [MIRROR] Fixes Callback Trait Addition/Removal Oversight [MDB
 IGNORE] (#25361) (#907)

* Fixes Callback Trait Addition/Removal Oversight (#80009)

## About The Pull Request

This pissed me off working yesterday, traits are supposed to work on the
`/datum` level, not just the `/atom` level. Let's fix that, while also
cleaning up the code a bunch to make it look better and have error
logging for unwanted cases.
## Changelog
Nothing that concerns players.

* Fixes Callback Trait Addition/Removal Oversight

---------

Co-authored-by: SkyratBot <59378654+SkyratBot@users.noreply.github.com>
Co-authored-by: san7890 <the@san7890.com>
---
 code/__HELPERS/traits.dm | 38 ++++++++++++++++++++------------------
 1 file changed, 20 insertions(+), 18 deletions(-)

diff --git a/code/__HELPERS/traits.dm b/code/__HELPERS/traits.dm
index ba99b2e1e7f..fcb0593bf23 100644
--- a/code/__HELPERS/traits.dm
+++ b/code/__HELPERS/traits.dm
@@ -2,32 +2,34 @@
 #define TRAIT_CALLBACK_REMOVE(target, trait, source) CALLBACK(GLOBAL_PROC, GLOBAL_PROC_REF(___TraitRemove), ##target, ##trait, ##source)
 
 ///DO NOT USE ___TraitAdd OR ___TraitRemove as a replacement for ADD_TRAIT / REMOVE_TRAIT defines. To be used explicitly for callback.
-/proc/___TraitAdd(target,trait,source)
+/proc/___TraitAdd(target, trait, source)
 	if(!target || !trait || !source)
 		return
+
 	if(islist(target))
-		for(var/i in target)
-			if(!isatom(i))
-				continue
-			var/atom/the_atom = i
-			ADD_TRAIT(the_atom,trait,source)
-	else if(isatom(target))
-		var/atom/the_atom2 = target
-		ADD_TRAIT(the_atom2,trait,source)
+		for(var/datum/listed_target in target)
+			ADD_TRAIT(listed_target, trait, source)
+		return
+
+	ASSERT(isdatum(target), "Invalid target used in TRAIT_CALLBACK_ADD! Expected a datum reference, got [target] instead.")
+
+	var/datum/datum_target = target
+	ADD_TRAIT(datum_target, trait, source)
 
 ///DO NOT USE ___TraitAdd OR ___TraitRemove as a replacement for ADD_TRAIT / REMOVE_TRAIT defines. To be used explicitly for callback.
-/proc/___TraitRemove(target,trait,source)
+/proc/___TraitRemove(target, trait, source)
 	if(!target || !trait || !source)
 		return
+
 	if(islist(target))
-		for(var/i in target)
-			if(!isatom(i))
-				continue
-			var/atom/the_atom = i
-			REMOVE_TRAIT(the_atom,trait,source)
-	else if(isatom(target))
-		var/atom/the_atom2 = target
-		REMOVE_TRAIT(the_atom2,trait,source)
+		for(var/datum/listed_target in target)
+			REMOVE_TRAIT(listed_target, trait, source)
+		return
+
+	ASSERT(isdatum(target), "Invalid target used in TRAIT_CALLBACK_REMOVE! Expected a datum reference, got [target] instead.")
+
+	var/datum/datum_target = target
+	REMOVE_TRAIT(datum_target, trait, source)
 
 
 /// Proc that handles adding multiple traits to a target via a list. Must have a common source and target.