124 lines
4.1 KiB
Python
124 lines
4.1 KiB
Python
import logging
|
|
import time
|
|
|
|
from django.core.management.base import BaseCommand
|
|
|
|
from music.listenbrainz import get_similar_artists
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Fetch similar artists from ListenBrainz for artists with a musicbrainz_id"
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument(
|
|
"--commit",
|
|
action="store_true",
|
|
help="Commit changes to the database",
|
|
)
|
|
parser.add_argument(
|
|
"--batch-size",
|
|
type=int,
|
|
default=100,
|
|
help="Number of artists to process per batch (default: 100)",
|
|
)
|
|
parser.add_argument(
|
|
"--artist-id",
|
|
type=str,
|
|
default="",
|
|
help="Only process the artist with this musicbrainz_id",
|
|
)
|
|
parser.add_argument(
|
|
"--overwrite",
|
|
action="store_true",
|
|
help="Re-fetch similar artists even if already populated",
|
|
)
|
|
|
|
def handle(self, *args, **options):
|
|
from music.models import Artist
|
|
|
|
commit = options["commit"]
|
|
batch_size = options["batch_size"]
|
|
artist_id = options["artist_id"]
|
|
overwrite = options["overwrite"]
|
|
|
|
if not commit:
|
|
self.stdout.write(
|
|
"Dry run — no changes will be saved. Use --commit to apply."
|
|
)
|
|
|
|
qs = Artist.objects.exclude(musicbrainz_id__isnull=True).exclude(
|
|
musicbrainz_id=""
|
|
)
|
|
|
|
if artist_id:
|
|
qs = qs.filter(musicbrainz_id=artist_id)
|
|
self.stdout.write(f"Filtering to artist with musicbrainz_id: {artist_id}")
|
|
elif not overwrite:
|
|
qs = qs.filter(similar_artists__isnull=True)
|
|
else:
|
|
qs = qs.all()
|
|
|
|
total = qs.count()
|
|
if total == 0:
|
|
self.stdout.write("No artists to process.")
|
|
return
|
|
|
|
self.stdout.write(
|
|
f"Found {total} artists with musicbrainz_id"
|
|
+ (" (overwrite mode)" if overwrite else "")
|
|
+ (" (--artist-id filter)" if artist_id else "")
|
|
)
|
|
|
|
if not commit:
|
|
self.stdout.write(
|
|
"\nSkipping API lookups in dry-run mode. Use --commit to run against ListenBrainz."
|
|
)
|
|
return
|
|
|
|
found_count = 0
|
|
empty_count = 0
|
|
error_count = 0
|
|
artist_ids = list(qs.values_list("pk", flat=True))
|
|
i = 0
|
|
|
|
for batch_num, offset in enumerate(
|
|
range(0, len(artist_ids), batch_size)
|
|
):
|
|
batch_pks = artist_ids[offset : offset + batch_size]
|
|
for artist in Artist.objects.filter(pk__in=batch_pks).iterator():
|
|
i += 1
|
|
self.stdout.write(
|
|
f" [{i}/{total}] Fetching similar artists for {artist}...",
|
|
ending="",
|
|
)
|
|
try:
|
|
similar = get_similar_artists(artist.musicbrainz_id)
|
|
if similar:
|
|
artist.similar_artists = similar
|
|
artist.save(update_fields=["similar_artists"])
|
|
self.stdout.write(f" {len(similar)} similar artists found")
|
|
found_count += 1
|
|
else:
|
|
artist.similar_artists = []
|
|
artist.save(update_fields=["similar_artists"])
|
|
self.stdout.write(" none found")
|
|
empty_count += 1
|
|
except Exception as e:
|
|
self.stdout.write(f" error: {e}")
|
|
error_count += 1
|
|
|
|
self.stdout.write(
|
|
f" Batch {batch_num + 1}: {offset + len(batch_pks)}/{total} processed, "
|
|
f"found: {found_count}, empty: {empty_count}, errors: {error_count}"
|
|
)
|
|
time.sleep(1)
|
|
|
|
self.stdout.write(
|
|
f"\nResults (commit={commit}):\n"
|
|
f" Similar artists found: {found_count}\n"
|
|
f" No similar artists: {empty_count}\n"
|
|
f" Errors: {error_count}"
|
|
)
|