From 1c79dd5dceaf74fcf40f200f3fe046ec0e84ee42 Mon Sep 17 00:00:00 2001
From: Doridian <git@doridian.net>
Date: Tue, 5 Mar 2024 16:08:11 -0800
Subject: [PATCH] Allow adding more pump entity classes

---
 lua/entities/rd_pump/init.lua   | 2 +-
 lua/entities/rd_pump/shared.lua | 6 +++++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/lua/entities/rd_pump/init.lua b/lua/entities/rd_pump/init.lua
index 7fa79b1..7b8d7b8 100644
--- a/lua/entities/rd_pump/init.lua
+++ b/lua/entities/rd_pump/init.lua
@@ -94,7 +94,7 @@ function ENT:LinkToPump(ply, ent)
 		return false, "You are not allowed to control this pump!"
 	end
 
-	if ent:GetClass() ~= self:GetClass() then
+	if not (IsValid(ent) and ent.IsRDPump and ent:IsRDPump()) then
 		return false, "The other entity is not a pump"
 	end
 
diff --git a/lua/entities/rd_pump/shared.lua b/lua/entities/rd_pump/shared.lua
index 4dbeb58..b7e6634 100644
--- a/lua/entities/rd_pump/shared.lua
+++ b/lua/entities/rd_pump/shared.lua
@@ -12,4 +12,8 @@ ENT.IsPump = true
 
 list.Set("LSEntOverlayText", "rd_pump", {
 	HasOOO = true
-})
\ No newline at end of file
+})
+
+function ENT:IsRDPump()
+	return true
+end