A simple script to fetch spotify song audio features and use a regression forest to predict it's year

spotify.py

#!/usr/bin/env python3

import enum
import spotipy
import os
import numpy as np
from typing import List

from regtree.tree import RandomForest

from spotipy.oauth2 import SpotifyClientCredentials


def login():
    sp = spotipy.Spotify(
        auth_manager=SpotifyClientCredentials(
            client_id=os.environ.get("SPOTIFY_CLIENT_ID"),
            client_secret=os.environ.get("SPOTIFY_CLIENT_SECRET"),
        )
    )

    return sp


def search(sp, song_title) -> List:
    results = sp.search(q=song_title, type="track", limit=10)
    return results["tracks"]["items"]


def select_track(results):
    for i, track in enumerate(results):
        print(f"{i + 1} - {track['name']} - {track['artists'][0]['name']}")

    selected = int(input("Select track number: ")) - 1
    return results[selected]


# load forest from json file
def main():
    sp = login()

    song_title = input("Song title: ")
    results = search(sp, song_title)
    if len(results) <= 0:
        print("No songs found")
        exit(1)

    track = select_track(results)

    features = sp.audio_features(track["uri"])[0]

    attributes = np.array(
        [
            0,
            features["duration_ms"],
            features["danceability"],
            features["energy"],
            features["key"],
            features["loudness"],
            features["mode"],
            features["speechiness"],
            features["acousticness"],
            features["instrumentalness"],
            features["liveness"],
            features["valence"],
            features["tempo"],
            features["time_signature"]
        ]
    )
    print(f"Real year: {track['album']['release_date'][:4]}")

    print("Predictions\n=================")
    # load all forests from all json files from the data directory and run the prediction on all of them
    # and display the results
    for file in os.listdir("data"):
        if file.endswith(".json"):
            path =  os.path.join("data", file)
            with open(path, "r") as f:
                forest = RandomForest.from_json(f.read())
                print(f"Forest {os.path.splitext(file)[0]}: {round(forest.predict_median(attributes))}")

if __name__ == "__main__":
    main()