From 1e28ceff6dd10d78d6261bc9b018f4743e692b10 Mon Sep 17 00:00:00 2001
From: Quietust <quietust@gmail.com>
Date: Wed, 22 Aug 2012 16:54:00 -0500
Subject: [PATCH] Add getProfessionColor and getCasteProfessionColor to Units
 module

---
 library/include/modules/Units.h |  3 +++
 library/modules/Units.cpp       | 37 +++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/library/include/modules/Units.h b/library/include/modules/Units.h
index 7bf4f101b..9003dc3af 100644
--- a/library/include/modules/Units.h
+++ b/library/include/modules/Units.h
@@ -226,6 +226,9 @@ DFHACK_EXPORT bool getNoblePositions(std::vector<NoblePosition> *pvec, df::unit
 
 DFHACK_EXPORT std::string getProfessionName(df::unit *unit, bool ignore_noble = false, bool plural = false);
 DFHACK_EXPORT std::string getCasteProfessionName(int race, int caste, df::profession pid, bool plural = false);
+
+DFHACK_EXPORT int8_t getProfessionColor(df::unit *unit, bool ignore_noble = false);
+DFHACK_EXPORT int8_t getCasteProfessionColor(int race, int caste, df::profession pid);
 }
 }
 #endif
diff --git a/library/modules/Units.cpp b/library/modules/Units.cpp
index ee383cc07..874dabc3d 100644
--- a/library/modules/Units.cpp
+++ b/library/modules/Units.cpp
@@ -948,3 +948,40 @@ std::string DFHack::Units::getCasteProfessionName(int race, int casteid, df::pro
 
     return Translation::capitalize(prof, true);
 }
+
+int8_t DFHack::Units::getProfessionColor(df::unit *unit, bool ignore_noble)
+{
+    std::vector<NoblePosition> np;
+
+    if (!ignore_noble && getNoblePositions(&np, unit))
+    {
+        if (np[0].position->flags.is_set(entity_position_flags::COLOR))
+            return np[0].position->color[0] + np[0].position->color[2] * 8;
+    }
+
+    return getCasteProfessionColor(unit->race, unit->caste, unit->profession);
+}
+
+int8_t DFHack::Units::getCasteProfessionColor(int race, int casteid, df::profession pid)
+{
+    // make sure it's an actual profession
+    if (pid < 0 || !is_valid_enum_item(pid))
+        return 3;
+
+    // If it's not a Peasant, it's hardcoded
+    if (pid != profession::STANDARD)
+        return ENUM_ATTR(profession, color, pid);
+
+    if (auto creature = df::creature_raw::find(race))
+    {
+        if (auto caste = vector_get(creature->caste, casteid))
+        {
+            if (caste->flags.is_set(caste_raw_flags::CASTE_COLOR))
+                return caste->caste_color[0] + caste->caste_color[2] * 8;
+        }
+        return creature->color[0] + creature->color[2] * 8;
+    }
+
+    // default to dwarven peasant color
+    return 3;
+}