python_proto_importer/postprocess/

mod.rs

//! Post-processing modules for generated Python protobuf code.
//!
//! This module provides utilities to transform and enhance generated Python code
//! after the initial protoc/buf generation phase. The post-processing steps ensure
//! the generated code integrates seamlessly with your project structure and
//! follows Python best practices.
//!
//! # Available Post-processors
//!
//! - **Import Rewriting** ([`apply`]): Converts absolute imports to relative imports
//! - **Package Creation** ([`create_packages`]): Automatically creates `__init__.py` files
//! - **Type Checker Headers** ([`add_pyright_header`]): Adds suppression headers for type checkers
//! - **FileDescriptorSet Processing** ([`fds`]): Extracts metadata from protoc output
//! - **Import Analysis** ([`rel_imports`]): Scans and reports import conversion opportunities
//!
//! # Post-processing Pipeline
//!
//! ```no_run
//! use python_proto_importer::postprocess::{create_packages, add_pyright_header};
//! use std::path::Path;
//!
//! let output_dir = Path::new("generated");
//!
//! // 1. Create __init__.py files for Python package structure
//! let packages_created = create_packages(output_dir)?;
//! println!("Created {} __init__.py files", packages_created);
//!
//! // 2. Add type checker suppression headers
//! let headers_added = add_pyright_header(output_dir)?;
//! println!("Added headers to {} files", headers_added);
//!
//! # Ok::<(), anyhow::Error>(())
//! ```

use anyhow::{Context, Result};
use std::collections::BTreeSet;
use std::fs;
use std::path::{Path, PathBuf};
use walkdir::WalkDir;

pub mod apply;
pub mod fds;
pub mod rel_imports;

/// Add Pyright suppression headers to generated Python protobuf files.
///
/// This function adds type checker suppression headers to generated `_pb2.py` and `_pb2_grpc.py`
/// files. These headers help suppress false positive warnings from type checkers when working
/// with dynamically generated protobuf code that may reference experimental APIs.
///
/// # Arguments
///
/// * `root` - Root directory to recursively scan for protobuf Python files
///
/// # Returns
///
/// Returns the number of files that were modified with headers.
///
/// # Behavior
///
/// - Only modifies files ending with `_pb2.py` or `_pb2_grpc.py`
/// - Skips files that already have the suppression header
/// - Recursively processes all subdirectories
/// - Preserves existing file content, only prepending the header
///
/// # Header Content
///
/// Adds the following header to suppress common type checker issues:
/// ```python
/// # pyright: reportAttributeAccessIssue=false
/// # This file is generated by grpcio-tools and may reference grpc.experimental which lacks stubs in types-grpcio.
/// ```
pub fn add_pyright_header(root: &Path) -> Result<usize> {
    use std::io::Write;
    let mut modified = 0usize;
    for entry in WalkDir::new(root).into_iter().filter_map(Result::ok) {
        let p = entry.path();
        if p.is_file() && p.extension().and_then(|s| s.to_str()) == Some("py") {
            let name = p.file_name().and_then(|s| s.to_str()).unwrap_or("");
            if !name.ends_with("_pb2.py") && !name.ends_with("_pb2_grpc.py") {
                continue;
            }
            let content = fs::read_to_string(p).with_context(|| format!("read {}", p.display()))?;
            let header = "# pyright: reportAttributeAccessIssue=false\n# This file is generated by grpcio-tools and may reference grpc.experimental which lacks stubs in types-grpcio.\n";
            if content.starts_with(header) {
                continue;
            }
            let mut f = fs::OpenOptions::new()
                .write(true)
                .truncate(true)
                .open(p)
                .with_context(|| format!("open {} for write", p.display()))?;
            f.write_all(header.as_bytes())?;
            f.write_all(content.as_bytes())?;
            modified += 1;
        }
    }
    Ok(modified)
}

/// Create `__init__.py` files for all directories in the output tree.
///
/// This function ensures that all directories in the generated output have `__init__.py`
/// files, making them proper Python packages. This is essential for Python's import
/// system to recognize directories as packages and enables proper module imports.
///
/// # Arguments
///
/// * `root` - Root directory to recursively process for package creation
///
/// # Returns
///
/// Returns the number of `__init__.py` files that were created.
///
/// # Behavior
///
/// - Recursively scans all directories under the root path
/// - Creates empty `__init__.py` files in directories that don't have them
/// - Skips directories that already have `__init__.py` files
/// - Uses `BTreeSet` for consistent ordering of directory processing
///
/// # Package Types
///
/// This creates regular packages (with `__init__.py`). For namespace packages
/// (PEP 420), disable package creation in your configuration:
///
/// ```toml
/// [tool.python_proto_importer.postprocess]
/// create_package = false
/// ```
pub fn create_packages(root: &Path) -> Result<usize> {
    let mut dirs: BTreeSet<PathBuf> = BTreeSet::new();
    for entry in WalkDir::new(root).into_iter().filter_map(Result::ok) {
        let path = entry.path();
        if path.is_dir() {
            dirs.insert(path.to_path_buf());
        }
    }

    let mut created = 0usize;
    for dir in dirs {
        let init_py = dir.join("__init__.py");
        if !init_py.exists() {
            fs::write(&init_py, b"")
                .with_context(|| format!("failed to write {}", init_py.display()))?;
            created += 1;
        }
    }
    Ok(created)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::tempdir;

    #[test]
    fn add_pyright_header_to_pb2_files() {
        let dir = tempdir().unwrap();

        // Create test files
        let pb2_file = dir.path().join("service_pb2.py");
        let grpc_file = dir.path().join("service_pb2_grpc.py");
        let regular_file = dir.path().join("regular.py");

        fs::write(&pb2_file, "# Generated code\nclass Service:\n    pass\n").unwrap();
        fs::write(
            &grpc_file,
            "# Generated gRPC code\nclass ServiceServicer:\n    pass\n",
        )
        .unwrap();
        fs::write(&regular_file, "# Regular Python file\nprint('hello')\n").unwrap();

        let modified = add_pyright_header(dir.path()).unwrap();
        assert_eq!(modified, 2); // Only pb2 and grpc files should be modified

        // Verify headers were added
        let pb2_content = fs::read_to_string(&pb2_file).unwrap();
        assert!(pb2_content.starts_with("# pyright: reportAttributeAccessIssue=false"));
        assert!(pb2_content.contains("# Generated code"));

        let grpc_content = fs::read_to_string(&grpc_file).unwrap();
        assert!(grpc_content.starts_with("# pyright: reportAttributeAccessIssue=false"));

        // Regular file should be unchanged
        let regular_content = fs::read_to_string(&regular_file).unwrap();
        assert!(!regular_content.contains("pyright"));
    }

    #[test]
    fn add_pyright_header_skips_existing() {
        let dir = tempdir().unwrap();
        let pb2_file = dir.path().join("service_pb2.py");

        // File already has header
        let existing_content = "# pyright: reportAttributeAccessIssue=false\n# This file is generated by grpcio-tools and may reference grpc.experimental which lacks stubs in types-grpcio.\n# Generated code\n";
        fs::write(&pb2_file, existing_content).unwrap();

        let modified = add_pyright_header(dir.path()).unwrap();
        assert_eq!(modified, 0); // Should skip files that already have header

        let content = fs::read_to_string(&pb2_file).unwrap();
        assert_eq!(content, existing_content); // Content should be unchanged
    }

    #[test]
    fn add_pyright_header_nested_directories() {
        let dir = tempdir().unwrap();
        let nested_dir = dir.path().join("services");
        fs::create_dir_all(&nested_dir).unwrap();

        let pb2_file = nested_dir.join("api_pb2.py");
        fs::write(&pb2_file, "# Generated code\n").unwrap();

        let modified = add_pyright_header(dir.path()).unwrap();
        assert_eq!(modified, 1);

        let content = fs::read_to_string(&pb2_file).unwrap();
        assert!(content.starts_with("# pyright: reportAttributeAccessIssue=false"));
    }

    #[test]
    fn create_packages_all_directories() {
        let dir = tempdir().unwrap();

        // Create nested directory structure
        let nested_dirs = ["services", "services/api", "services/auth", "common"];

        for nested in &nested_dirs {
            fs::create_dir_all(dir.path().join(nested)).unwrap();
        }

        let created = create_packages(dir.path()).unwrap();
        // Should create __init__.py in root + 4 nested directories = 5 total
        assert_eq!(created, 5);

        // Verify all __init__.py files exist
        assert!(dir.path().join("__init__.py").exists());
        assert!(dir.path().join("services/__init__.py").exists());
        assert!(dir.path().join("services/api/__init__.py").exists());
        assert!(dir.path().join("services/auth/__init__.py").exists());
        assert!(dir.path().join("common/__init__.py").exists());
    }

    #[test]
    fn create_packages_skips_existing() {
        let dir = tempdir().unwrap();

        let nested_dir = dir.path().join("services");
        fs::create_dir_all(&nested_dir).unwrap();

        // Pre-create one __init__.py file
        fs::write(nested_dir.join("__init__.py"), "# Existing content").unwrap();

        let created = create_packages(dir.path()).unwrap();
        // Should only create __init__.py in root directory
        assert_eq!(created, 1);

        // Verify existing file content is preserved
        let content = fs::read_to_string(nested_dir.join("__init__.py")).unwrap();
        assert_eq!(content, "# Existing content");

        // New file should be empty
        let root_content = fs::read_to_string(dir.path().join("__init__.py")).unwrap();
        assert_eq!(root_content, "");
    }

    #[test]
    fn create_packages_empty_directory() {
        let dir = tempdir().unwrap();
        // Empty directory should still get __init__.py

        let created = create_packages(dir.path()).unwrap();
        assert_eq!(created, 1);
        assert!(dir.path().join("__init__.py").exists());
    }

    #[test]
    fn add_pyright_header_file_extension_filtering() {
        let dir = tempdir().unwrap();

        // Create files with various extensions
        fs::write(dir.path().join("service_pb2.py"), "# Python file").unwrap();
        fs::write(dir.path().join("service_pb2.pyi"), "# Stub file").unwrap();
        fs::write(dir.path().join("service_pb2.txt"), "# Text file").unwrap();
        fs::write(dir.path().join("service.py"), "# Regular Python").unwrap();

        let modified = add_pyright_header(dir.path()).unwrap();
        // Only .py files with correct naming should be modified
        assert_eq!(modified, 1);

        let py_content = fs::read_to_string(dir.path().join("service_pb2.py")).unwrap();
        assert!(py_content.starts_with("# pyright"));

        let pyi_content = fs::read_to_string(dir.path().join("service_pb2.pyi")).unwrap();
        assert!(!pyi_content.contains("pyright"));

        let regular_content = fs::read_to_string(dir.path().join("service.py")).unwrap();
        assert!(!regular_content.contains("pyright"));
    }
}