Support multiple Power traits

OpenRA.Mods.Common/Effects/PowerdownIndicator.cs

@@ -19,10 +19,12 @@ namespace OpenRA.Mods.Common.Effects
 	{
 		readonly Actor a;
 		readonly Animation anim;
+		readonly CanPowerDown canPowerDown;
 
 		public PowerdownIndicator(Actor a)
 		{
 			this.a = a;
+			canPowerDown = a.Trait<CanPowerDown>();
 
 			anim = new Animation(a.World, "poweroff");
 			anim.PlayRepeating("offline");
@@ -30,7 +32,7 @@ namespace OpenRA.Mods.Common.Effects
 
 		public void Tick(World world)
 		{
-			if (!a.IsInWorld || a.IsDead || !a.Trait<CanPowerDown>().Disabled)
+			if (!a.IsInWorld || a.IsDead || !canPowerDown.Disabled)
 				world.AddFrameEndTask(w => w.Remove(this));
 
 			anim.Tick();

OpenRA.Mods.Common/Scripting/Properties/PowerProperties.cs

@@ -9,6 +9,8 @@
 #endregion
 
 using System;
+using System.Collections.Generic;
+using System.Linq;
 using Eluant;
 using OpenRA.Mods.Common.Power;
 using OpenRA.Scripting;
@@ -56,18 +58,18 @@ namespace OpenRA.Mods.Common.Scripting
 	[ScriptPropertyGroup("Power")]
 	public class ActorPowerProperties : ScriptActorProperties, Requires<PowerInfo>
 	{
-		readonly PowerInfo pi;
+		readonly IEnumerable<Power.Power> power;
 
 		public ActorPowerProperties(ScriptContext context, Actor self)
 			: base(context, self)
 		{
-			pi = self.Info.Traits.GetOrDefault<PowerInfo>();
+			power = self.TraitsImplementing<Power.Power>();
 		}
 
 		[Desc("Returns the power drained/provided by this actor.")]
 		public int Power
 		{
-			get { return pi.Amount; }
+			get { return power.Sum(p => p.GetEnabledPower()); }
 		}
 	}
 }

OpenRA.Mods.Common/Traits/Power/CanPowerDown.cs

@@ -16,18 +16,21 @@ namespace OpenRA.Mods.Common.Power
 	[Desc("The player can disable the power individually on this actor.")]
 	public class CanPowerDownInfo : UpgradableTraitInfo, ITraitInfo, Requires<PowerInfo>
 	{
+		[Desc("Restore power when this trait is disabled.")]
+		public readonly bool CancelWhenDisabled = false;
+
 		public object Create(ActorInitializer init) { return new CanPowerDown(init.self, this); }
 	}
 
-	public class CanPowerDown : UpgradableTrait<CanPowerDownInfo>, IPowerModifier, IResolveOrder, IDisable
+	public class CanPowerDown : UpgradableTrait<CanPowerDownInfo>, IPowerModifier, IResolveOrder, IDisable, INotifyOwnerChanged
 	{
 		[Sync] bool disabled = false;
-		readonly Power power;
+		PowerManager power;
 
 		public CanPowerDown(Actor self, CanPowerDownInfo info)
 			: base(info)
 		{
-			power = self.Trait<Power>();
+			power = self.Owner.PlayerActor.Trait<PowerManager>();
 		}
 
 		public bool Disabled { get { return disabled; } }
@@ -38,7 +41,7 @@ namespace OpenRA.Mods.Common.Power
 			{
 				disabled = !disabled;
 				Sound.PlayNotification(self.World.Map.Rules, self.Owner, "Sounds", disabled ? "EnablePower" : "DisablePower", self.Owner.Country.Race);
-				power.PlayerPower.UpdateActor(self);
+				power.UpdateActor(self);
 
 				if (disabled)
 					self.World.AddFrameEndTask(w => w.Add(new PowerdownIndicator(self)));
@@ -49,5 +52,19 @@ namespace OpenRA.Mods.Common.Power
 		{
 			return !IsTraitDisabled && disabled ? 0 : 100;
 		}
+
+		public void OnOwnerChanged(Actor self, Player oldOwner, Player newOwner)
+		{
+			power = newOwner.PlayerActor.Trait<PowerManager>();
+		}
+
+		protected override void UpgradeDisabled(Actor self)
+		{
+			if (!disabled || !Info.CancelWhenDisabled)
+				return;
+			disabled = false;
+			Sound.PlayNotification(self.World.Map.Rules, self.Owner, "Sounds", "EnablePower", self.Owner.Country.Race);
+			power.UpdateActor(self);
+		}
 	}
 }

OpenRA.Mods.Common/Traits/Power/ScalePowerWithHealth.cs

@@ -18,14 +18,14 @@ namespace OpenRA.Mods.Common.Power
 		public object Create(ActorInitializer init) { return new ScalePowerWithHealth(init.self); }
 	}
 
-	public class ScalePowerWithHealth : IPowerModifier, INotifyDamage
+	public class ScalePowerWithHealth : IPowerModifier, INotifyDamage, INotifyOwnerChanged
 	{
-		readonly Power power;
 		readonly Health health;
+		PowerManager power;
 
 		public ScalePowerWithHealth(Actor self)
 		{
-			power = self.Trait<Power>();
+			power = self.Owner.PlayerActor.Trait<PowerManager>();
 			health = self.Trait<Health>();
 		}
 
@@ -34,9 +34,10 @@ namespace OpenRA.Mods.Common.Power
 			return 100 * health.HP / health.MaxHP;
 		}
 
-		public void Damaged(Actor self, AttackInfo e)
+		public void Damaged(Actor self, AttackInfo e) { power.UpdateActor(self); }
+		public void OnOwnerChanged(Actor self, Player oldOwner, Player newOwner)
 		{
-			power.PlayerPower.UpdateActor(self);
+			power = newOwner.PlayerActor.Trait<PowerManager>();
 		}
 	}
 }