use anyhow::Result;
use regex::Regex;
use std::fs;
use std::path::Path;
#[allow(dead_code)]
pub fn rewrite_file_for_relative_imports(path: &Path) -> Result<usize> {
let content = fs::read_to_string(path)?;
let import_re = Regex::new(r"(?m)^import\s+([A-Za-z0-9_\.]+_pb2(?:_grpc)?)\b").unwrap();
let from_re =
Regex::new(r"(?m)^from\s+([A-Za-z0-9_\.]+)\s+import\s+([A-Za-z0-9_]+_pb2(?:_grpc)?)\b")
.unwrap();
let mut hits = 0usize;
hits += import_re.find_iter(&content).count();
hits += from_re.find_iter(&content).count();
Ok(hits)
}
#[allow(dead_code)]
pub fn scan_and_report(root: &Path) -> Result<(usize, usize)> {
let mut files = 0usize;
let mut lines = 0usize;
for entry in walkdir::WalkDir::new(root)
.into_iter()
.filter_map(Result::ok)
{
let p = entry.path();
if p.is_file() && p.extension().and_then(|e| e.to_str()) == Some("py") {
files += 1;
lines += rewrite_file_for_relative_imports(p)?;
}
}
Ok((files, lines))
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
use tempfile::tempdir;
#[test]
fn rewrite_file_for_relative_imports_basic() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("test.py");
let content = r#"
import service_pb2
from api.v1 import user_pb2
import other_module
from package import regular_module
import grpc_service_pb2_grpc
"#;
fs::write(&file_path, content).unwrap();
let hits = rewrite_file_for_relative_imports(&file_path).unwrap();
assert_eq!(hits, 3);
}
#[test]
fn rewrite_file_for_relative_imports_no_matches() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("test.py");
let content = r#"
import os
from typing import List
import requests
from dataclasses import dataclass
"#;
fs::write(&file_path, content).unwrap();
let hits = rewrite_file_for_relative_imports(&file_path).unwrap();
assert_eq!(hits, 0);
}
#[test]
fn rewrite_file_for_relative_imports_complex_patterns() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("test.py");
let content = r#"
# Import statements
import api.v1.service_pb2
import api.v2.user_pb2_grpc
from package.subpackage import module_pb2
from api import payment_pb2_grpc
from . import local_pb2 # Should not match (already relative)
# Mixed content
def function():
pass
import another_service_pb2
"#;
fs::write(&file_path, content).unwrap();
let hits = rewrite_file_for_relative_imports(&file_path).unwrap();
assert_eq!(hits, 6);
}
#[test]
fn rewrite_file_for_relative_imports_multiline() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("test.py");
let content = "import service_pb2\nfrom api import user_pb2\nimport normal_module";
fs::write(&file_path, content).unwrap();
let hits = rewrite_file_for_relative_imports(&file_path).unwrap();
assert_eq!(hits, 2); }
#[test]
fn scan_and_report_basic() {
let dir = tempdir().unwrap();
let file1 = dir.path().join("service.py");
let file2 = dir.path().join("api.py");
let file3 = dir.path().join("utils.txt");
fs::write(&file1, "import service_pb2\nfrom api import user_pb2").unwrap();
fs::write(&file2, "import payment_pb2_grpc").unwrap();
fs::write(&file3, "import service_pb2").unwrap();
let (files, lines) = scan_and_report(dir.path()).unwrap();
assert_eq!(files, 2); assert_eq!(lines, 3); }
#[test]
fn scan_and_report_nested_directories() {
let dir = tempdir().unwrap();
let nested_dir = dir.path().join("services");
fs::create_dir_all(&nested_dir).unwrap();
let file1 = dir.path().join("main.py");
let file2 = nested_dir.join("api.py");
fs::write(&file1, "import main_service_pb2").unwrap();
fs::write(&file2, "from proto import api_pb2\nimport grpc_pb2_grpc").unwrap();
let (files, lines) = scan_and_report(dir.path()).unwrap();
assert_eq!(files, 2);
assert_eq!(lines, 3); }
#[test]
fn scan_and_report_empty_directory() {
let dir = tempdir().unwrap();
let (files, lines) = scan_and_report(dir.path()).unwrap();
assert_eq!(files, 0);
assert_eq!(lines, 0);
}
#[test]
fn scan_and_report_no_proto_imports() {
let dir = tempdir().unwrap();
let file = dir.path().join("normal.py");
fs::write(&file, "import os\nfrom typing import List").unwrap();
let (files, lines) = scan_and_report(dir.path()).unwrap();
assert_eq!(files, 1); assert_eq!(lines, 0); }
#[test]
fn rewrite_file_nonexistent_file() {
let nonexistent = std::path::Path::new("/nonexistent/file.py");
let result = rewrite_file_for_relative_imports(nonexistent);
assert!(result.is_err());
}
}