Python project dependency fixer: Pin all package versions in requirements.txt to a specified date.

pin_requirements.py

#!/usr/bin/env python3
"""
Pin all package versions in requirements.txt to a specified date.
This is for all those ML codelabs and github repos that forgot to pin their dependencies and now have broken installs.
Just run this script with the last commit date of the repo, and it will update your requirements.txt with the closest
versions available on PyPI as of that date.
Usage: python pin_requirements.py <date> [requirements.txt]
Date format: YYYY-MM-DD
"""

import sys
import re
from datetime import datetime
import requests
from pathlib import Path


def parse_requirement(line):
    """Parse a requirement line and extract package name and version if present."""
    line = line.strip()
    if not line or line.startswith('#'):
        return None

    # Handle various requirement formats
    match = re.match(r'^([a-zA-Z0-9\-_.]+)', line)
    if match:
        return match.group(1)
    return None


def get_releases_for_package(pkg_name):
    """Fetch all releases for a package from PyPI."""
    try:
        url = f"https://pypi.org/pypi/{pkg_name}/json"
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        data = response.json()
        releases = data.get('releases', {})
        return releases
    except requests.RequestException as e:
        print(f"Error fetching {pkg_name}: {e}", file=sys.stderr)
        return {}


def find_closest_version(releases, target_date):
    """
    Find the version closest to (but not after) the target date.
    Returns (version, release_date) or (None, None) if not found.
    """
    closest_version = None
    closest_date = None

    for version, release_info in releases.items():
        if not release_info:  # Skip empty release info
            continue

        # Get the upload date from the first file in the release
        upload_time = release_info[0].get('upload_time_iso_8601')
        if not upload_time:
            continue

        # Parse the ISO format date
        try:
            release_date = datetime.fromisoformat(upload_time.replace('Z', '+00:00')).date()
        except (ValueError, KeyError):
            continue

        # Find the closest version that's not after the target date
        if release_date <= target_date:
            if closest_date is None or release_date > closest_date:
                closest_date = release_date
                closest_version = version

    return closest_version, closest_date


def pin_requirements(requirements_file, target_date):
    """
    Read requirements.txt, pin all versions to the target date, and return new content.
    """
    # Parse the target date
    try:
        target_date = datetime.strptime(target_date, '%Y-%m-%d').date()
    except ValueError:
        raise ValueError(f"Invalid date format: {target_date}. Use YYYY-MM-DD")

    # Read the requirements file
    try:
        with open(requirements_file, 'r') as f:
            lines = f.readlines()
    except FileNotFoundError:
        raise FileNotFoundError(f"Requirements file not found: {requirements_file}")

    new_lines = []
    for line in lines:
        pkg_name = parse_requirement(line)

        if pkg_name is None:
            # Keep comments and blank lines as-is
            new_lines.append(line)
            continue

        print(f"Processing {pkg_name}...", file=sys.stderr)
        releases = get_releases_for_package(pkg_name)

        if not releases:
            print(f"  Warning: No releases found for {pkg_name}, keeping original", file=sys.stderr)
            new_lines.append(line)
            continue

        version, release_date = find_closest_version(releases, target_date)

        if version is None:
            print(f"  Warning: No release found for {pkg_name} before {target_date}, keeping original", file=sys.stderr)
            new_lines.append(line)
            continue

        print(f"  Found version {version} (released {release_date})", file=sys.stderr)
        new_lines.append(f"{pkg_name}=={version}\n")

    return ''.join(new_lines)


def main():
    if len(sys.argv) < 2:
        print("Usage: python pin_requirements.py <date> [requirements.txt]")
        print("Date format: YYYY-MM-DD")
        print("Example: python pin_requirements.py 2023-06-15 requirements.txt")
        sys.exit(1)

    target_date = sys.argv[1]
    requirements_file = sys.argv[2] if len(sys.argv) > 2 else 'requirements.txt'

    try:
        new_content = pin_requirements(requirements_file, target_date)

        # Create backup
        backup_file = f"{requirements_file}.bak"
        Path(requirements_file).rename(backup_file)
        print(f"Backed up to {backup_file}", file=sys.stderr)

        # Write new content
        with open(requirements_file, 'w') as f:
            f.write(new_content)
        print(f"Updated {requirements_file}", file=sys.stderr)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()