From b321ab8da2185fb1cb0f233c8a7b022b268e451c Mon Sep 17 00:00:00 2001
From: Georg Krause <mail@georg-krause.net>
Date: Sat, 20 Feb 2021 12:54:11 +0100
Subject: [PATCH] Add number of tracks and discs of an album to API

---
 api/funkwhale_api/music/models.py      |  9 +++++++-
 api/funkwhale_api/music/serializers.py |  5 +++++
 api/tests/music/test_serializers.py    | 29 ++++++++++++++++++++++++++
 changes/changelog.d/1238.enhancement   |  1 +
 4 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 changes/changelog.d/1238.enhancement

diff --git a/api/funkwhale_api/music/models.py b/api/funkwhale_api/music/models.py
index 3e39e91e2..4325a1697 100644
--- a/api/funkwhale_api/music/models.py
+++ b/api/funkwhale_api/music/models.py
@@ -20,6 +20,7 @@ from django.db.models.signals import post_save, pre_save
 from django.dispatch import receiver
 from django.urls import reverse
 from django.utils import timezone
+from django.db.models import Prefetch, Count
 
 from funkwhale_api import musicbrainz
 from funkwhale_api.common import fields
@@ -420,7 +421,13 @@ def import_album(v):
 class TrackQuerySet(common_models.LocalFromFidQuerySet, models.QuerySet):
     def for_nested_serialization(self):
         return self.prefetch_related(
-            "artist", "album__artist", "album__attachment_cover"
+            "artist",
+            Prefetch(
+                "album",
+                queryset=Album.objects.select_related(
+                    "artist", "attachment_cover"
+                ).annotate(_prefetched_tracks_count=Count("tracks")),
+            ),
         )
 
     def annotate_playable_by_actor(self, actor):
diff --git a/api/funkwhale_api/music/serializers.py b/api/funkwhale_api/music/serializers.py
index bed611c33..cc78358cb 100644
--- a/api/funkwhale_api/music/serializers.py
+++ b/api/funkwhale_api/music/serializers.py
@@ -227,6 +227,10 @@ class AlbumSerializer(OptionalDescriptionMixin, serializers.Serializer):
 class TrackAlbumSerializer(serializers.ModelSerializer):
     artist = serializers.SerializerMethodField()
     cover = cover_field
+    tracks_count = serializers.SerializerMethodField()
+
+    def get_tracks_count(self, o):
+        return getattr(o, "_prefetched_tracks_count", len(o.tracks.all()))
 
     class Meta:
         model = models.Album
@@ -240,6 +244,7 @@ class TrackAlbumSerializer(serializers.ModelSerializer):
             "cover",
             "creation_date",
             "is_local",
+            "tracks_count",
         )
 
     def get_artist(self, o):
diff --git a/api/tests/music/test_serializers.py b/api/tests/music/test_serializers.py
index b1d561cb1..c0bb2298f 100644
--- a/api/tests/music/test_serializers.py
+++ b/api/tests/music/test_serializers.py
@@ -196,6 +196,35 @@ def test_album_serializer(factories, to_api_date):
     assert serializer.data == expected
 
 
+def test_track_album_serializer(factories, to_api_date):
+    actor = factories["federation.Actor"]()
+    track1 = factories["music.Track"](
+        position=2, album__attributed_to=actor, album__with_cover=True
+    )
+    factories["music.Track"](position=1, album=track1.album)
+    album = track1.album
+    expected = {
+        "id": album.id,
+        "fid": album.fid,
+        "mbid": str(album.mbid),
+        "title": album.title,
+        "artist": serializers.serialize_artist_simple(album.artist),
+        "creation_date": to_api_date(album.creation_date),
+        "is_playable": False,
+        "cover": common_serializers.AttachmentSerializer(album.attachment_cover).data,
+        "release_date": to_api_date(album.release_date),
+        "tracks_count": 2,
+        "is_local": album.is_local,
+        "tags": [],
+        "attributed_to": federation_serializers.APIActorSerializer(actor).data,
+    }
+    serializer = serializers.AlbumSerializer(
+        album.__class__.objects.with_tracks_count().get(pk=album.pk)
+    )
+
+    assert serializer.data == expected
+
+
 def test_track_serializer(factories, to_api_date):
     actor = factories["federation.Actor"]()
     upload = factories["music.Upload"](
diff --git a/changes/changelog.d/1238.enhancement b/changes/changelog.d/1238.enhancement
new file mode 100644
index 000000000..d826949d4
--- /dev/null
+++ b/changes/changelog.d/1238.enhancement
@@ -0,0 +1 @@
+Add number of tracks and discs of an album to API (#1238)