Fixed the invalidation of the internal entity cache before all entity deletion listeners were called.

Fixed a KeyError being silenced in EntityDictionary.on_automatically_removed overrides.
Slightly improved the retrieval speed of the cached entity instances.

Author: jordanbriere

addons/source-python/packages/source-python/entities/_base.py

@@ -101,9 +101,10 @@ def __call__(cls, index, caching=True):
         """
         # Let's first lookup for a cached instance
         if caching:
-            obj = cls._cache.get(index, None)
-            if obj is not None:
-                return obj
+            try:
+                return cls._cache[index]
+            except KeyError:
+                pass
 
         # Nothing in cache, let's create a new instance
         obj = super().__call__(index)
@@ -129,6 +130,17 @@ def cache(cls):
         """
         return cls._cache
 
+    @staticmethod
+    def _invalidate_cache(base_entity):
+        """Invalidates the cache for the given entity."""
+        try:
+            index = base_entity.index
+        except ValueError:
+            return
+
+        for cls in _entity_classes:
+            cls.cache.pop(index, None)
+
 
 class Entity(BaseEntity, metaclass=_EntityCaching):
     """Class used to interact directly with entities.
@@ -1225,10 +1237,6 @@ def _on_entity_deleted(base_entity):
     except ValueError:
         return
 
-    # Cleanup the cache
-    for cls in _entity_classes:
-        cls.cache.pop(index, None)
-
     with suppress(KeyError):
         # Loop through all delays...
         for delay in _entity_delays[index]:

addons/source-python/packages/source-python/entities/dictionary.py

@@ -88,7 +88,7 @@ def _on_entity_deleted(self, base_entity):
         except ValueError:
             return
 
-        with suppress(KeyError):
+        if index in self:
             # Call the deletion callback for the index...
             self.on_automatically_removed(index)
 

src/core/sp_main.cpp

@@ -631,6 +631,10 @@ void CSourcePython::OnEntitySpawned( CBaseEntity *pEntity )
 void CSourcePython::OnEntityDeleted( CBaseEntity *pEntity )
 {
 	CALL_LISTENERS(OnEntityDeleted, ptr((CBaseEntityWrapper*) pEntity));
+
+	// Invalidate the internal entity cache once all callbacks have been called.
+	static object oCacheInvalidator = import("entities").attr("_base").attr("_EntityCaching").attr("_invalidate_cache");
+	oCacheInvalidator(ptr((CBaseEntityWrapper*) pEntity));
 }
 
 void CSourcePython::OnDataLoaded( MDLCacheDataType_t type, MDLHandle_t handle )