Condense live charts to one function

tests/scrobbles_tests/test_aggregators.py

@@ -5,12 +5,7 @@ import time_machine
 from django.contrib.auth import get_user_model
 from django.urls import reverse
 from django.utils import timezone
-from music.aggregators import (
-    scrobble_counts,
-    top_artists,
-    top_tracks,
-    week_of_scrobbles,
-)
+from music.aggregators import live_charts, scrobble_counts, week_of_scrobbles
 from profiles.models import UserProfile
 from scrobbles.models import Scrobble
 
@@ -55,7 +50,7 @@ def test_week_of_scrobbles_data(client, mopidy_track_request_data):
 def test_top_tracks_by_day(client, mopidy_track_request_data):
     build_scrobbles(client, mopidy_track_request_data, 7, 1)
     user = get_user_model().objects.first()
-    tops = top_tracks(user)
+    tops = live_charts(user)
     assert tops[0].title == "Same in the End"
 
 
@@ -63,7 +58,7 @@ def test_top_tracks_by_day(client, mopidy_track_request_data):
 def test_top_tracks_by_week(client, mopidy_track_request_data):
     build_scrobbles(client, mopidy_track_request_data, 7, 1)
     user = get_user_model().objects.first()
-    tops = top_tracks(user, filter='week')
+    tops = live_charts(user, chart_period='week')
     assert tops[0].title == "Same in the End"
 
 
@@ -71,7 +66,7 @@ def test_top_tracks_by_week(client, mopidy_track_request_data):
 def test_top_tracks_by_month(client, mopidy_track_request_data):
     build_scrobbles(client, mopidy_track_request_data, 7, 1)
     user = get_user_model().objects.first()
-    tops = top_tracks(user, filter='month')
+    tops = live_charts(user, chart_period='month')
     assert tops[0].title == "Same in the End"
 
 
@@ -79,7 +74,7 @@ def test_top_tracks_by_month(client, mopidy_track_request_data):
 def test_top_tracks_by_year(client, mopidy_track_request_data):
     build_scrobbles(client, mopidy_track_request_data, 7, 1)
     user = get_user_model().objects.first()
-    tops = top_tracks(user, filter='year')
+    tops = live_charts(user, chart_period='year')
     assert tops[0].title == "Same in the End"
 
 
@@ -87,7 +82,7 @@ def test_top_tracks_by_year(client, mopidy_track_request_data):
 def test_top__artists_by_week(client, mopidy_track_request_data):
     build_scrobbles(client, mopidy_track_request_data, 7, 1)
     user = get_user_model().objects.first()
-    tops = top_artists(user, filter='week')
+    tops = live_charts(user, chart_period='week', media_type="Artist")
     assert tops[0].name == "Sublime"
 
 
@@ -95,7 +90,7 @@ def test_top__artists_by_week(client, mopidy_track_request_data):
 def test_top__artists_by_month(client, mopidy_track_request_data):
     build_scrobbles(client, mopidy_track_request_data, 7, 1)
     user = get_user_model().objects.first()
-    tops = top_artists(user, filter='month')
+    tops = live_charts(user, chart_period='month', media_type="Artist")
     assert tops[0].name == "Sublime"
 
 
@@ -103,5 +98,5 @@ def test_top__artists_by_month(client, mopidy_track_request_data):
 def test_top__artists_by_year(client, mopidy_track_request_data):
     build_scrobbles(client, mopidy_track_request_data, 7, 1)
     user = get_user_model().objects.first()
-    tops = top_artists(user, filter='year')
+    tops = live_charts(user, chart_period='year', media_type="Artist")
     assert tops[0].name == "Sublime"

vrobbler/apps/music/aggregators.py

@@ -1,22 +1,14 @@
 from datetime import datetime, timedelta
-from typing import List, Optional
+from typing import List
 
-import pytz
-from django.db.models import Count, Q, Sum
+from django.apps import apps
+from django.db.models import Count, Q, QuerySet
 from django.utils import timezone
 from music.models import Artist, Track
 from scrobbles.models import Scrobble
 from videos.models import Video
 from vrobbler.apps.profiles.utils import now_user_timezone
 
-NOW = timezone.now()
-START_OF_TODAY = datetime.combine(NOW.date(), datetime.min.time(), NOW.tzinfo)
-STARTING_DAY_OF_CURRENT_WEEK = NOW.date() - timedelta(
-    days=NOW.today().isoweekday() % 7
-)
-STARTING_DAY_OF_CURRENT_MONTH = NOW.date().replace(day=1)
-STARTING_DAY_OF_CURRENT_YEAR = NOW.date().replace(month=1, day=1)
-
 
 def scrobble_counts(user=None):
 
@@ -92,73 +84,65 @@ def week_of_scrobbles(
     return scrobble_day_dict
 
 
-def top_tracks(
-    user: "User", filter: str = "today", limit: int = 30
-) -> List["Track"]:
-
+def live_charts(
+    user: "User",
+    media_type: str = "Track",
+    chart_period: str = "all",
+    limit: int = 15,
+) -> QuerySet:
     now = timezone.now()
+    tzinfo = now.tzinfo
+    now = now.date()
     if user.is_authenticated:
         now = now_user_timezone(user.profile)
+        tzinfo = now.tzinfo
 
-    start_of_today = datetime.combine(
-        now.date(), datetime.min.time(), now.tzinfo
-    )
-    starting_day_of_current_week = now.date() - timedelta(
-        days=now.today().isoweekday() % 7
-    )
-    starting_day_of_current_month = now.date().replace(day=1)
-    starting_day_of_current_year = now.date().replace(month=1, day=1)
+    start_of_today = datetime.combine(now, datetime.min.time(), tzinfo)
+    start_day_of_week = now - timedelta(days=now.today().isoweekday() % 7)
+    start_day_of_month = now.replace(day=1)
+    start_day_of_year = now.replace(month=1, day=1)
 
-    time_filter = Q()
-    if filter == "today":
-        time_filter = Q(scrobble__timestamp__gte=start_of_today)
-    if filter == "week":
-        time_filter = Q(scrobble__timestamp__gte=starting_day_of_current_week)
-    if filter == "month":
-        time_filter = Q(scrobble__timestamp__gte=starting_day_of_current_month)
-    if filter == "year":
-        time_filter = Q(scrobble__timestamp__gte=starting_day_of_current_year)
-
-    return (
-        Track.objects.filter(time_filter)
-        .annotate(
-            num_scrobbles=Count(
-                "scrobble",
-                filter=Q(scrobble__played_to_completion=True),
-                distinct=True,
-            )
-        )
-        .order_by("-num_scrobbles")[:limit]
-    )
+    media_model = apps.get_model(app_label='music', model_name=media_type)
 
+    period_queries = {
+        'today': {'scrobble__timestamp__gte': start_of_today},
+        'week': {'scrobble__timestamp__gte': start_day_of_week},
+        'month': {'scrobble__timestamp__gte': start_day_of_month},
+        'year': {'scrobble__timestamp__gte': start_day_of_year},
+        'all': {},
+    }
 
-def top_artists(
-    user: "User", filter: str = "today", limit: int = 15
-) -> List["Artist"]:
-    time_filter = Q(track__scrobble__timestamp__gte=START_OF_TODAY)
-    if filter == "week":
-        time_filter = Q(
-            track__scrobble__timestamp__gte=STARTING_DAY_OF_CURRENT_WEEK
-        )
-    if filter == "month":
-        time_filter = Q(
-            track__scrobble__timestamp__gte=STARTING_DAY_OF_CURRENT_MONTH
-        )
-    if filter == "year":
-        time_filter = Q(
-            track__scrobble__timestamp__gte=STARTING_DAY_OF_CURRENT_YEAR
+    time_filter = Q()
+    completion_filter = Q(
+        scrobble__user=user, scrobble__played_to_completion=True
+    )
+    user_filter = Q(scrobble__user=user)
+    count_field = "scrobble"
+
+    if media_type == "Artist":
+        for period, query_dict in period_queries.items():
+            period_queries[period] = {
+                "track__" + k: v for k, v in query_dict.items()
+            }
+        completion_filter = Q(
+            track__scrobble__user=user,
+            track__scrobble__played_to_completion=True,
         )
+        count_field = "track__scrobble"
+        user_filter = Q(track__scrobble__user=user)
+
+    time_filter = Q(**period_queries[chart_period])
 
     return (
-        Artist.objects.filter(time_filter)
+        media_model.objects.filter(user_filter, time_filter)
         .annotate(
             num_scrobbles=Count(
-                "track__scrobble",
-                filter=Q(track__scrobble__played_to_completion=True),
+                count_field,
+                filter=completion_filter,
                 distinct=True,
             )
         )
-        .order_by("-num_scrobbles")
+        .order_by("-num_scrobbles")[:limit]
     )