Condense live charts to one function

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,39 +84,61 @@ 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)
+
+    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': {},
+    }
 
     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)
+    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 (
-        Track.objects.filter(time_filter)
+        media_model.objects.filter(user_filter, time_filter)
         .annotate(
             num_scrobbles=Count(
-                "scrobble",
-                filter=Q(scrobble__played_to_completion=True),
+                count_field,
+                filter=completion_filter,
                 distinct=True,
             )
         )
@@ -132,35 +146,5 @@ def top_tracks(
     )
 
 
-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
-        )
-
-    return (
-        Artist.objects.filter(time_filter)
-        .annotate(
-            num_scrobbles=Count(
-                "track__scrobble",
-                filter=Q(track__scrobble__played_to_completion=True),
-                distinct=True,
-            )
-        )
-        .order_by("-num_scrobbles")
-    )
-
-
 def artist_scrobble_count(artist_id: int, filter: str = "today") -> int:
     return Scrobble.objects.filter(track__artist=artist_id).count()