AssignedRolesEventArgs

Author: marchellcx

LabExtended/Events/Round/AssignedRolesEventArgs.cs

@@ -0,0 +1,24 @@
+using LabExtended.API;
+
+using PlayerRoles;
+
+namespace LabExtended.Events.Round
+{
+    /// <summary>
+    /// Gets called after the server starts assigns roles when the round starts.
+    /// </summary>
+    public class AssignedRolesEventArgs : EventArgs
+    {
+        /// <summary>
+        /// A dictionary of players and their decided roles.
+        /// </summary>
+        public IReadOnlyDictionary<ExPlayer, RoleTypeId> Roles { get; }
+
+        /// <summary>
+        /// Creates a new <see cref="AssignedRolesEventArgs"/> instance.
+        /// </summary>
+        /// <param name="roles">The assigned roles.</param>
+        public AssignedRolesEventArgs(Dictionary<ExPlayer, RoleTypeId> roles)
+            => Roles = roles;
+    }
+}

LabExtended/Events/Round/AssigningRolesEventArgs.cs

@@ -12,13 +12,73 @@ public class AssigningRolesEventArgs : BooleanEventArgs
         /// <summary>
         /// A dictionary of players and their decided roles.
         /// </summary>
-        public Dictionary<ExPlayer?, RoleTypeId> Roles { get; }
+        public Dictionary<ExPlayer, RoleTypeId> Roles { get; }
 
         /// <summary>
         /// Creates a new <see cref="AssigningRolesEventArgs"/> instance.
         /// </summary>
-        /// <param name="roles"></param>
-        public AssigningRolesEventArgs(Dictionary<ExPlayer?, RoleTypeId> roles)
+        /// <param name="roles">The assigned roles.</param>
+        public AssigningRolesEventArgs(Dictionary<ExPlayer, RoleTypeId> roles)
             => Roles = roles;
+
+        /// <summary>
+        /// Sets the role for the specified player if the player exists in the collection.
+        /// </summary>
+        /// <param name="player">The player whose role is to be set. Cannot be null.</param>
+        /// <param name="role">The role to assign to the specified player.</param>
+        /// <exception cref="ArgumentNullException">Thrown if <paramref name="player"/> is null.</exception>
+        public void Set(ExPlayer player, RoleTypeId role)
+        {
+            if (player is null)
+                throw new ArgumentNullException(nameof(player));
+
+            if (!Roles.ContainsKey(player))
+                return;
+
+            Roles[player] = role;
+        }
+
+        /// <summary>
+        /// Sets the role for all players by applying the specified selector function to each player.
+        /// </summary>
+        /// <param name="selector">A function that takes an ExPlayer and returns the RoleTypeId to assign to that player.</param>
+        /// <exception cref="ArgumentNullException">Thrown if selector is null.</exception>
+        public void SetAll(Func<ExPlayer, RoleTypeId> selector)
+        {
+            if (selector is null)
+                throw new ArgumentNullException(nameof(selector));
+
+            foreach (var player in ExPlayer.Players)
+            {
+                if (!Roles.ContainsKey(player))
+                    continue;
+
+                Roles[player] = selector(player);
+            }
+        }
+
+        /// <summary>
+        /// Sets the specified role for all players that match the given predicate.
+        /// </summary>
+        /// <param name="predicate">A predicate used to determine which players should have their role set. The method applies the role to each
+        /// player for whom this predicate returns <see langword="true"/>.</param>
+        /// <param name="role">The role to assign to each player that matches the predicate.</param>
+        /// <exception cref="ArgumentNullException">Thrown if <paramref name="predicate"/> is <see langword="null"/>.</exception>
+        public void SetWhere(Predicate<ExPlayer> predicate, RoleTypeId role)
+        {
+            if (predicate is null)
+                throw new ArgumentNullException(nameof(predicate));
+
+            foreach (var player in ExPlayer.Players)
+            {
+                if (!Roles.ContainsKey(player))
+                    continue;
+
+                if (!predicate(player))
+                    continue;
+
+                Roles[player] = role;
+            }
+        }
     }
 }

LabExtended/Patches/Events/Round/AssigningRolesPatch.cs

@@ -5,17 +5,17 @@
 using PlayerRoles.RoleAssign;
 using PlayerRoles;
 
-using LabExtended.Core.Pooling.Pools;
-
 using LabExtended.API;
+
 using LabExtended.Events;
+
 using LabExtended.Utilities;
 using LabExtended.Utilities.RoleSelection;
 
 namespace LabExtended.Patches.Events.Round;
 
 /// <summary>
-/// Implements the <see cref="ExRoundEvents.AssigningRoles"/> event.
+/// Implements the <see cref="ExRoundEvents.AssigningRoles"/> and <see cref="ExRoundEvents.AssignedRoles"/> events.
 /// </summary>
 public static class AssigningRolesPatch
 {
@@ -29,7 +29,7 @@ public static class AssigningRolesPatch
     /// </summary>
     public static readonly List<ExPlayer> Players = new();
 
-    public static FastEvent<Action> OnPlayersSpawned { get; } =
+    private static FastEvent<Action> OnPlayersSpawned { get; } =
         FastEvents.DefineEvent<Action>(typeof(RoleAssigner), nameof(RoleAssigner.OnPlayersSpawned));
 
     [HarmonyPatch(typeof(RoleAssigner), nameof(RoleAssigner.OnRoundStarted))]
@@ -62,14 +62,23 @@ private static bool Prefix()
 
         foreach (var pair in Roles)
         {
-            pair.Key.Role.RoundStartRole = pair.Value;
-            pair.Key.Role.Set(pair.Value, RoleChangeReason.RoundStart, RoleSpawnFlags.All);
-
-            if (pair.Key.Role.IsAlive)
-                RoleAssigner.AlreadySpawnedPlayers.Add(pair.Key.UserId);
+            if (pair.Value != RoleTypeId.None)
+            {
+                pair.Key.Role.RoundStartRole = pair.Value;
+                pair.Key.Role.Set(pair.Value, RoleChangeReason.RoundStart, RoleSpawnFlags.All);
+
+                if (pair.Key.Role.IsAlive)
+                    RoleAssigner.AlreadySpawnedPlayers.Add(pair.Key.UserId);
+            }
         }
 
-        OnPlayersSpawned.InvokeEvent(null, Array.Empty<object>());
+        OnPlayersSpawned.InvokeEvent(null!, Array.Empty<object>());
+
+        ExRoundEvents.OnAssignedRoles(new(Roles));
+
+        Roles.Clear();
+        Players.Clear();
+
         return false;
     }
 }