fix class naming

metabrainz · Jan 28, 2025 · 4e25412 · 4e25412
1 parent 5156d0b
commit 4e25412
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/listenbrainz_spark/stats/incremental/listener/entity.py b/listenbrainz_spark/stats/incremental/listener/entity.py
@@ -20,7 +20,7 @@ def get_entity_id(self):
         raise NotImplementedError()
 
 
-class EntityStatsMessageCreator(UserStatsMessageCreator):
+class EntityListenerStatsMessageCreator(UserStatsMessageCreator):
 
     def items_per_message(self):
         return 10000
diff --git a/listenbrainz_spark/stats/listener/entity.py b/listenbrainz_spark/stats/listener/entity.py
@@ -3,7 +3,7 @@
 
 from listenbrainz_spark.stats.incremental.incremental_stats_engine import IncrementalStatsEngine
 from listenbrainz_spark.stats.incremental.listener.artist import ArtistEntityListenerStatsQuery
-from listenbrainz_spark.stats.incremental.listener.entity import EntityListenerStatsQueryProvider, EntityStatsMessageCreator
+from listenbrainz_spark.stats.incremental.listener.entity import EntityListenerStatsQueryProvider, EntityListenerStatsMessageCreator
 from listenbrainz_spark.stats.incremental.listener.release_group import ReleaseGroupEntityListenerStatsQuery
 from listenbrainz_spark.stats.incremental.range_selector import StatsRangeListenRangeSelector
 
@@ -23,6 +23,6 @@ def get_listener_stats(entity: str, stats_range: str, database: str = None) -> I
     selector = StatsRangeListenRangeSelector(stats_range)
     entity_cls = incremental_entity_obj_map[entity]
     entity_obj = entity_cls(selector, NUMBER_OF_TOP_LISTENERS)
-    message_creator = EntityStatsMessageCreator(entity, "entity_listener", selector, database)
+    message_creator = EntityListenerStatsMessageCreator(entity, "entity_listener", selector, database)
     engine = IncrementalStatsEngine(entity_obj, message_creator)
     return engine.run()
diff --git a/listenbrainz_spark/stats/user/entity.py b/listenbrainz_spark/stats/user/entity.py
@@ -4,7 +4,8 @@
 from listenbrainz_spark.stats.incremental.incremental_stats_engine import IncrementalStatsEngine
 from listenbrainz_spark.stats.incremental.range_selector import StatsRangeListenRangeSelector
 from listenbrainz_spark.stats.incremental.user.artist import ArtistUserEntity
-from listenbrainz_spark.stats.incremental.user.entity import UserEntityStatsQueryProvider, UserStatsMessageCreator
+from listenbrainz_spark.stats.incremental.user.entity import UserEntityStatsQueryProvider, \
+    UserEntityStatsMessageCreator
 from listenbrainz_spark.stats.incremental.user.recording import RecordingUserEntity
 from listenbrainz_spark.stats.incremental.user.release import ReleaseUserEntity
 from listenbrainz_spark.stats.incremental.user.release_group import ReleaseGroupUserEntity
@@ -26,6 +27,6 @@ def get_entity_stats(entity: str, stats_range: str, database: str = None) -> Ite
     logger.debug(f"Calculating user_{entity}_{stats_range}...")
     selector = StatsRangeListenRangeSelector(stats_range)
     entity_obj = incremental_entity_map[entity](selector, NUMBER_OF_TOP_ENTITIES)
-    message_creator = UserStatsMessageCreator(entity, "user_entity", selector, database)
+    message_creator = UserEntityStatsMessageCreator(entity, "user_entity", selector, database)
     engine = IncrementalStatsEngine(entity_obj, message_creator)
     return engine.run()