pipeline/download_imdb_data.py

import os
import gzip
import shutil
from urllib.request import urlretrieve
from tqdm import tqdm


def download_large_file(url: str, output_file: str):
    if not os.path.exists(output_file):
        urlretrieve(url, output_file)


def unzip_file(input_file):
    output_file = os.path.splitext(input_file)[0]
    if not os.path.exists(output_file):
        with gzip.open(input_file, "rb") as f_in:
            # Input file has the format xxx.tsv.gz
            with open(output_file, "wb") as f_out:
                shutil.copyfileobj(f_in, f_out)


if __name__ == "__main__":
    imdb_url = "https://datasets.imdbws.com"
    filenames = [
        "name.basics.tsv.gz",
        "title.basics.tsv.gz",
        "title.ratings.tsv.gz",
        "title.principals.tsv.gz",
    ]
    for filename in tqdm(filenames):
        url = f"{imdb_url}/{filename}"
        output_file = os.path.join("data", filename)
        download_large_file(url, output_file)
        unzip_file(output_file)