cleanup_directory.py

#!/usr/bin/env python3
"""
Post-process markdown files such as:
- remove '_wandb' and everything after it from filenames. _wandb is added to
    avoid conflicts during generation.
- Delete empty directories.
- Converts .md files to .mdx files if specified. By default, this conversion is enabled.
- Lower cases all characters in filenames.
"""

import os
from pathlib import Path
import argparse
import yaml
import json


def extract_frontmatter(file_path):
    """
    Extract frontmatter from a markdown file.

    Args:
        file_path: Path to the markdown file

    Returns:
        Dictionary containing frontmatter data or empty dict if no frontmatter
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()

        # Check if file starts with frontmatter delimiter
        if content.startswith('---'):
            # Find the closing delimiter
            end_index = content.find('---', 3)
            if end_index != -1:
                frontmatter_text = content[3:end_index].strip()
                return yaml.safe_load(frontmatter_text) or {}
    except Exception as e:
        print(f"Error reading frontmatter from {file_path}: {e}")

    return {}


def clean_filename(filename):
    """
    Remove '_wandb' and everything after it from a filename.

    Args:
        filename: The original filename

    Returns:
        The cleaned filename

    Examples:
        >>> clean_filename("histogram_wandb.plot.md")
        'histogram.md'
        >>> clean_filename("Api_wandb.apis.public.md")
        'Api.md'
        >>> clean_filename("regular_file.md")
        'regular_file.md'
    """
    if '_wandb' in filename:
        # Split at '_wandb' and take only the first part
        base_name = filename.split('_wandb')[0]
        # Get the file extension (should be .md)
        _, ext = os.path.splitext(filename)
        return base_name + ext
    return filename


def clean_title(title):
    """
    Remove '_wandb' and everything after it from a title string.

    Args:
        title: The original title

    Returns:
        The cleaned title

    Examples:
        >>> clean_title("histogram_wandb.plot")
        'histogram'
        >>> clean_title("Api_wandb.apis.public")
        'Api'
    """
    if '_wandb' in title:
        return title.split('_wandb')[0]
    return title


def update_frontmatter_title(file_path, dry_run=False):
    """
    Update the title in the frontmatter of a markdown file.

    Args:
        file_path: Path to the markdown file
        dry_run: If True, only print what would be changed

    Returns:
        True if the file was updated, False otherwise
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()

        # Check if file has frontmatter
        if not content.startswith('---'):
            return False

        # Find the end of frontmatter
        end_index = content.find('---', 3)
        if end_index == -1:
            return False

        # Extract frontmatter
        frontmatter_text = content[3:end_index].strip()
        remaining_content = content[end_index:]

        # Parse frontmatter
        frontmatter = yaml.safe_load(frontmatter_text) or {}

        # Check if there's a title to clean
        if 'title' in frontmatter and '_wandb' in frontmatter['title']:
            old_title = frontmatter['title']
            new_title = clean_title(old_title)
            frontmatter['title'] = new_title

            if dry_run:
                print(f"  Would update title: '{old_title}' -> '{new_title}'")
                return False
            else:
                # Reconstruct the file with updated frontmatter
                new_frontmatter = yaml.dump(frontmatter, default_flow_style=False, sort_keys=False)
                new_content = f"---\n{new_frontmatter}---{remaining_content[3:]}"

                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write(new_content)

                print(f"  Updated title: '{old_title}' -> '{new_title}'")
                return True
    except Exception as e:
        print(f"  Error updating title in {file_path}: {e}")

    return False


def get_unique_filename(base_path):
    """
    Get a unique filename by appending numbers if necessary.

    Args:
        base_path: Path object for the desired filename

    Returns:
        Path object with a unique filename
    """
    if not base_path.exists():
        return base_path

    # Split the filename and extension
    stem = base_path.stem
    suffix = base_path.suffix
    parent = base_path.parent

    # Try appending numbers until we find a unique name
    counter = 1
    while True:
        new_path = parent / f"{stem}_{counter}{suffix}"
        if not new_path.exists():
            return new_path
        counter += 1

def lowercase_filename(filename):
    """
    Lowercase all characters in a filename.

    Args:
        filename: The original filename

    Returns:
        The lowercased filename
    """
    return filename.lower()


def rename_markdown_files(directory, dry_run=False, convert_to_mdx=False):
    """
    Rename all markdown files in a directory tree by removing '_wandb' and everything after.
    Also updates the title in the frontmatter to remove '_wandb' and everything after.
    Handles conflicts by checking frontmatter and either deleting duplicates or appending numbers.
    Optionally converts .md extensions to .mdx in the same pass.

    Args:
        directory: Path to the directory to process
        dry_run: If True, only print what would be renamed without actually renaming
        convert_to_mdx: If True, convert .md extensions to .mdx after other renaming

    Returns:
        List of tuples (old_path, new_path) for renamed files
    """
    directory = Path(directory)
    renamed_files = []
    deleted_files = []
    titles_updated = 0

    # Find all markdown files recursively
    for md_file in directory.rglob('*.md'):
        old_name = md_file.name
        new_name = lowercase_filename(clean_filename(old_name))

        # Update the title in frontmatter first (before renaming the file)
        if update_frontmatter_title(md_file, dry_run):
            titles_updated += 1

        # Determine if we need to rename the file
        old_path = md_file
        new_path = md_file.parent / new_name

        # Apply .mdx conversion if requested
        if convert_to_mdx:
            new_path = new_path.with_suffix('.mdx')

        # Only rename if the path actually changed (filename or extension)
        if old_path != new_path:
            if dry_run:
                if new_path.exists():
                    # Check frontmatter to determine action
                    old_frontmatter = extract_frontmatter(old_path)
                    new_frontmatter = extract_frontmatter(new_path)

                    old_namespace = old_frontmatter.get('namespace', '')
                    new_namespace = new_frontmatter.get('namespace', '')
                    old_obj_type = old_frontmatter.get('python_object_type', '')
                    new_obj_type = new_frontmatter.get('python_object_type', '')

                    if old_namespace == new_namespace and old_obj_type == new_obj_type:
                        print(f"Would delete duplicate: {old_path} (same as {new_path})")
                        deleted_files.append(str(old_path))
                    else:
                        unique_path = get_unique_filename(new_path)
                        print(f"Would rename with number: {old_path} -> {unique_path}")
                        renamed_files.append((str(old_path), str(unique_path)))
                else:
                    print(f"Would rename: {old_path} -> {new_path}")
                    renamed_files.append((str(old_path), str(new_path)))
            else:
                # Check if target file already exists
                if new_path.exists():
                    # Extract and compare frontmatter
                    old_frontmatter = extract_frontmatter(old_path)
                    new_frontmatter = extract_frontmatter(new_path)

                    old_namespace = old_frontmatter.get('namespace', '')
                    new_namespace = new_frontmatter.get('namespace', '')
                    old_obj_type = old_frontmatter.get('python_object_type', '')
                    new_obj_type = new_frontmatter.get('python_object_type', '')

                    if old_namespace == new_namespace and old_obj_type == new_obj_type:
                        # Same namespace and type - delete the file being renamed
                        print(f"Deleting duplicate: {old_path} (same as {new_path})")
                        old_path.unlink()
                        deleted_files.append(str(old_path))
                    else:
                        # Different namespace or type - append number to filename
                        unique_path = get_unique_filename(new_path)
                        old_path.rename(unique_path)
                        print(f"Renamed with number: {old_path} -> {unique_path}")
                        renamed_files.append((str(old_path), str(unique_path)))
                else:
                    old_path.rename(new_path)
                    print(f"Renamed: {old_path} -> {new_path}")
                    renamed_files.append((str(old_path), str(new_path)))

    if deleted_files:
        print(f"\nTotal files {'that would be' if dry_run else ''} deleted as duplicates: {len(deleted_files)}")

    if titles_updated > 0:
        print(f"Total titles {'that would be' if dry_run else ''} updated: {titles_updated}")

    return renamed_files


def add_public_apis_admonition(directory='python/public-api'):
    """
    Add public API admonition to markdown files in the public-apis directory.

    Args:
        directory: Path to the public-apis directory
    """
    directory = Path(directory)
    if not directory.exists():
        print(f"Directory {directory} does not exist")
        return

    admonition = "{{< readfile file=\"/_includes/public-api-use.md\" >}}\n\n"
    files_updated = 0

    # Process all markdown files in the directory
    for md_file in directory.rglob('*.md'):
        try:
            with open(md_file, 'r', encoding='utf-8') as f:
                content = f.read()

            # Check if file has frontmatter
            if content.startswith('---'):
                # Find the end of frontmatter
                end_index = content.find('---', 3)
                if end_index != -1:
                    # Insert admonition after frontmatter
                    end_of_frontmatter = content.find('\n', end_index) + 1
                    new_content = content[:end_of_frontmatter] + admonition + content[end_of_frontmatter:]
                else:
                    # No closing frontmatter, add at beginning
                    new_content = admonition + content
            else:
                # No frontmatter, add at beginning
                new_content = admonition + content

            # Write back the updated content
            with open(md_file, 'w', encoding='utf-8') as f:
                f.write(new_content)

            files_updated += 1
            print(f"Added admonition to: {md_file}")

        except Exception as e:
            print(f"Error processing {md_file}: {e}")

    print(f"Total files updated: {files_updated}")


def delete_empty_directories(root_directory):
    """Delete empty directories in the root directory."""
    for dirpath, dirnames, filenames in os.walk(root_directory, topdown=False):
        if not dirnames and not filenames:
            print(f"Deleting empty directory: {dirpath}")
            os.rmdir(dirpath)


def create_mdx_file_list(renamed_files):
    """
    Extract and create list of only the .mdx files from the renamed_files list.

    Args:
        renamed_files: List of tuples (old_path, new_path) for renamed files

    Returns:
        List of relative paths for .mdx files
    """
    mdx_files = []

    for old_path, new_path in renamed_files:
        # Check if the new file has .mdx extension
        if new_path.endswith('.mdx'):
            mdx_files.append(new_path)

    return sorted(mdx_files)


def check_json_output_directory(json_output):
    """
    Ensure the JSON output directory exists, create it if it doesn't.

    Args:
        json_output: Path to the JSON output directory
    """
    if json_output and not os.path.exists(json_output):
        print(f"Creating JSON output directory: {json_output}")
        os.makedirs(json_output, exist_ok=True)

def main(args):

    if not os.path.exists(args.directory):
        print(f"Error: Directory '{args.directory}' does not exist")
        return 1

    print(f"Processing directory: {args.directory}")
    print(f"Dry run: {args.dry_run}")
    print(f"Convert to MDX: {args.convert_to_mdx}")
    print()

    # Rename markdown files (and optionally convert to .mdx)
    renamed_files = rename_markdown_files(args.directory, args.dry_run, args.convert_to_mdx)

    print()
    print(f"Total files {'that would be' if args.dry_run else ''} renamed: {len(renamed_files)}")

    # Delete empty directories (unless skipped or in dry-run mode)
    if not args.skip_empty_cleanup and not args.dry_run:
        print("\nCleaning up empty directories...")
        delete_empty_directories(args.directory)

    # Add public API admonition
    print("\nAdding public API admonitions...")
    add_public_apis_admonition(directory=os.path.join(args.directory, 'public-api'))

    # Check and create JSON output directory if needed
    # Output to JSON and text files
    output_dir = args.json_output if args.json_output else '.'
    check_json_output_directory(output_dir)

    # Extract and output .mdx files as JSON
    print("\nCreating MDX file list...")
    mdx_files = create_mdx_file_list(renamed_files)

    filename = os.path.join(output_dir, 'mdx_file_list.json')
    with open(filename, 'w', encoding='utf-8') as f:
        json.dump(mdx_files, f, indent=2)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Remove "_wandb" and everything after from markdown filenames and clean up empty directories'
    )
    parser.add_argument(
        '--directory',
        default='python',
        help='Directory to process (will process all subdirectories)'
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Show what would be renamed without actually renaming'
    )
    parser.add_argument(
        '--skip-empty-cleanup',
        action='store_true',
        help='Skip deletion of empty directories'
    )
    parser.add_argument(
        '--convert-to-mdx',
        action='store_true',
        default=True,
        help='Convert .md extensions to .mdx'
    )
    parser.add_argument(
        '--json-output',
        default=None,
        help='Output the list of .mdx files to a JSON file in the specified directory'
    )
    args = parser.parse_args()
    main(args)