use std::collections::HashMap;
use std::path::PathBuf;
#[derive(Debug, Clone)]
pub struct PublicApiConfig {
pub naming_convention_weight: f32,
pub docstring_weight: f32,
pub type_annotation_weight: f32,
pub symmetric_pair_weight: f32,
pub module_export_weight: f32,
pub public_api_threshold: f32,
pub custom_public_prefixes: Vec<String>,
pub custom_symmetric_pairs: Vec<(String, String)>,
}
impl Default for PublicApiConfig {
fn default() -> Self {
Self {
naming_convention_weight: 0.3,
docstring_weight: 0.25,
type_annotation_weight: 0.15,
symmetric_pair_weight: 0.2,
module_export_weight: 0.1,
public_api_threshold: 0.6,
custom_public_prefixes: vec![],
custom_symmetric_pairs: vec![],
}
}
}
#[derive(Debug, Clone)]
pub struct PublicApiScore {
pub is_public: bool,
pub confidence: f32,
pub heuristic_scores: HashMap<String, f32>,
pub reasoning: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct FileContext {
pub file_path: PathBuf,
pub language: Language,
pub module_all: Option<Vec<String>>,
pub functions: HashMap<String, FunctionDef>,
pub used_functions: Vec<String>,
pub init_exports: Vec<String>,
}
impl FileContext {
pub fn new(file_path: PathBuf, language: Language) -> Self {
Self {
file_path,
language,
module_all: None,
functions: HashMap::new(),
used_functions: vec![],
init_exports: vec![],
}
}
pub fn language(&self) -> Language {
self.language
}
pub fn is_module_level(&self, function: &FunctionDef) -> bool {
!function.is_method
}
pub fn is_class_method(&self, function: &FunctionDef) -> bool {
function.is_method
}
pub fn is_in_module_all(&self, name: &str) -> bool {
self.module_all
.as_ref()
.map(|all| all.contains(&name.to_string()))
.unwrap_or(false)
}
pub fn is_exported_in_init(&self, name: &str) -> bool {
self.init_exports.iter().any(|export| export == name)
}
pub fn find_function(&self, name: &str) -> Option<&FunctionDef> {
self.functions.get(name)
}
pub fn is_function_used(&self, function: &FunctionDef) -> bool {
self.used_functions
.iter()
.any(|used| used == &function.name)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Language {
Python,
Rust,
JavaScript,
TypeScript,
}
#[derive(Debug, Clone)]
pub struct FunctionDef {
pub name: String,
pub docstring: Option<String>,
pub parameters: Vec<Parameter>,
pub return_type: Option<String>,
pub decorators: Vec<String>,
pub is_method: bool,
pub class_name: Option<String>,
pub line: usize,
pub visibility: Option<String>,
pub is_trait_impl: bool,
}
impl FunctionDef {
pub fn has_visibility_keyword(&self, keyword: &str) -> bool {
self.visibility
.as_ref()
.map(|v| v.contains(keyword))
.unwrap_or(false)
}
pub fn is_trait_implementation(&self) -> bool {
self.is_trait_impl
}
}
#[derive(Debug, Clone)]
pub struct Parameter {
pub name: String,
pub type_annotation: Option<String>,
pub default_value: Option<String>,
}
pub trait ApiHeuristic: Send + Sync {
fn name(&self) -> &str;
fn evaluate(&self, function: &FunctionDef, context: &FileContext) -> f32;
fn explain(&self, function: &FunctionDef) -> String;
}
pub struct PublicApiDetector {
config: PublicApiConfig,
}
impl PublicApiDetector {
pub fn new(config: PublicApiConfig) -> Self {
Self { config }
}
pub fn is_public_api(&self, function: &FunctionDef, context: &FileContext) -> PublicApiScore {
let mut heuristic_scores = HashMap::new();
let mut reasoning = Vec::new();
let mut weighted_score = 0.0f32;
let naming_heuristic = NamingConventionHeuristic;
let naming_score = naming_heuristic.evaluate(function, context);
heuristic_scores.insert(naming_heuristic.name().to_string(), naming_score);
weighted_score += naming_score * self.config.naming_convention_weight;
reasoning.push(naming_heuristic.explain(function));
let docstring_heuristic = DocstringHeuristic;
let docstring_score = docstring_heuristic.evaluate(function, context);
heuristic_scores.insert(docstring_heuristic.name().to_string(), docstring_score);
weighted_score += docstring_score * self.config.docstring_weight;
if docstring_score > 0.0 {
reasoning.push(docstring_heuristic.explain(function));
}
let type_annotation_heuristic = TypeAnnotationHeuristic;
let type_score = type_annotation_heuristic.evaluate(function, context);
heuristic_scores.insert(type_annotation_heuristic.name().to_string(), type_score);
weighted_score += type_score * self.config.type_annotation_weight;
if type_score > 0.5 {
reasoning.push(type_annotation_heuristic.explain(function));
}
let symmetric_heuristic = SymmetricPairHeuristic::new(&self.config);
let symmetric_score = symmetric_heuristic.evaluate(function, context);
heuristic_scores.insert(symmetric_heuristic.name().to_string(), symmetric_score);
weighted_score += symmetric_score * self.config.symmetric_pair_weight;
if symmetric_score > 0.0 {
reasoning.push(symmetric_heuristic.explain(function));
}
let export_heuristic = ModuleExportHeuristic;
let export_score = export_heuristic.evaluate(function, context);
heuristic_scores.insert(export_heuristic.name().to_string(), export_score);
weighted_score += export_score * self.config.module_export_weight;
if export_score > 0.0 {
reasoning.push(export_heuristic.explain(function));
}
if context.language() == Language::Rust {
let rust_heuristic = RustVisibilityHeuristic;
let rust_score = rust_heuristic.evaluate(function, context);
heuristic_scores.insert(rust_heuristic.name().to_string(), rust_score);
weighted_score = weighted_score.max(rust_score);
if rust_score > 0.0 {
reasoning.push(rust_heuristic.explain(function));
}
}
let is_public = weighted_score >= self.config.public_api_threshold;
PublicApiScore {
is_public,
confidence: weighted_score.clamp(0.0, 1.0),
heuristic_scores,
reasoning,
}
}
pub fn find_symmetric_pair<'a>(
&self,
function: &FunctionDef,
context: &'a FileContext,
) -> Option<&'a FunctionDef> {
let heuristic = SymmetricPairHeuristic::new(&self.config);
heuristic.find_pair(function, context)
}
}
struct NamingConventionHeuristic;
impl ApiHeuristic for NamingConventionHeuristic {
fn name(&self) -> &str {
"naming_convention"
}
fn evaluate(&self, function: &FunctionDef, context: &FileContext) -> f32 {
let name = &function.name;
if name.starts_with("__") && name.ends_with("__") {
return 0.5;
}
if name.starts_with('_') {
return 0.0;
}
if context.is_module_level(function) {
return 1.0;
}
if context.is_class_method(function) {
return 0.8;
}
0.5 }
fn explain(&self, function: &FunctionDef) -> String {
if function.name.starts_with('_') {
"Function has leading underscore (private convention)".to_string()
} else if function.name.starts_with("__") && function.name.ends_with("__") {
"Function is a dunder method (special method)".to_string()
} else {
"Function has no underscore prefix (public convention)".to_string()
}
}
}
struct DocstringHeuristic;
impl ApiHeuristic for DocstringHeuristic {
fn name(&self) -> &str {
"docstring_quality"
}
fn evaluate(&self, function: &FunctionDef, _context: &FileContext) -> f32 {
let docstring = match &function.docstring {
Some(doc) => doc,
None => return 0.0,
};
let length = docstring.len();
if length < 20 {
return 0.2;
}
if length < 50 {
return 0.5;
}
if length < 100 {
return 0.8;
}
if self.is_structured_docstring(docstring) {
return 1.0;
}
0.9
}
fn explain(&self, function: &FunctionDef) -> String {
match &function.docstring {
Some(doc) if doc.len() >= 50 => "Function has comprehensive docstring".to_string(),
Some(_) => "Function has basic docstring".to_string(),
None => "Function has no docstring".to_string(),
}
}
}
impl DocstringHeuristic {
fn is_structured_docstring(&self, doc: &str) -> bool {
let markers = [
"Args:",
"Returns:",
"Raises:",
"Yields:",
"Parameters:",
":param",
":return",
];
markers.iter().any(|marker| doc.contains(marker))
}
}
struct TypeAnnotationHeuristic;
impl ApiHeuristic for TypeAnnotationHeuristic {
fn name(&self) -> &str {
"type_annotations"
}
fn evaluate(&self, function: &FunctionDef, _context: &FileContext) -> f32 {
let param_annotations = function
.parameters
.iter()
.filter(|p| p.type_annotation.is_some())
.count();
let total_params = function.parameters.len();
if total_params == 0 {
return 0.5;
}
let annotation_ratio = param_annotations as f32 / total_params as f32;
let has_return_type = function.return_type.is_some();
if annotation_ratio >= 1.0 && has_return_type {
return 1.0;
}
if has_return_type {
return 0.5 + (annotation_ratio * 0.3);
}
annotation_ratio * 0.7
}
fn explain(&self, function: &FunctionDef) -> String {
let annotated = function
.parameters
.iter()
.filter(|p| p.type_annotation.is_some())
.count();
let total = function.parameters.len();
if annotated == total && function.return_type.is_some() {
"Function has full type annotations".to_string()
} else if annotated > 0 || function.return_type.is_some() {
format!(
"Function has partial type annotations ({}/{})",
annotated, total
)
} else {
"Function has no type annotations".to_string()
}
}
}
struct SymmetricPairHeuristic {
pairs: Vec<(String, String)>,
}
impl SymmetricPairHeuristic {
fn new(config: &PublicApiConfig) -> Self {
let mut pairs = vec![
("load".to_string(), "save".to_string()),
("get".to_string(), "set".to_string()),
("open".to_string(), "close".to_string()),
("create".to_string(), "destroy".to_string()),
("start".to_string(), "stop".to_string()),
("acquire".to_string(), "release".to_string()),
("add".to_string(), "remove".to_string()),
("push".to_string(), "pop".to_string()),
("read".to_string(), "write".to_string()),
];
pairs.extend(config.custom_symmetric_pairs.clone());
Self { pairs }
}
fn find_pair<'a>(
&self,
function: &FunctionDef,
context: &'a FileContext,
) -> Option<&'a FunctionDef> {
let func_name = &function.name;
let components: Vec<&str> = func_name.split('_').collect();
for (first, second) in &self.pairs {
let has_first = components.iter().any(|&c| c == first);
let has_second = components.iter().any(|&c| c == second);
if has_first || has_second {
let pair_name = if has_first {
components
.iter()
.map(|&c| if c == first { second.as_str() } else { c })
.collect::<Vec<_>>()
.join("_")
} else {
components
.iter()
.map(|&c| if c == second { first.as_str() } else { c })
.collect::<Vec<_>>()
.join("_")
};
if let Some(pair_func) = context.find_function(&pair_name) {
return Some(pair_func);
}
}
}
None
}
}
impl ApiHeuristic for SymmetricPairHeuristic {
fn name(&self) -> &str {
"symmetric_pair"
}
fn evaluate(&self, function: &FunctionDef, context: &FileContext) -> f32 {
if let Some(pair_func) = self.find_pair(function, context) {
if context.is_function_used(pair_func) {
return 1.0;
}
return 0.7;
}
0.0
}
fn explain(&self, function: &FunctionDef) -> String {
format!(
"Function '{}' may be part of a symmetric API pair",
function.name
)
}
}
struct ModuleExportHeuristic;
impl ApiHeuristic for ModuleExportHeuristic {
fn name(&self) -> &str {
"module_export"
}
fn evaluate(&self, function: &FunctionDef, context: &FileContext) -> f32 {
if context.is_in_module_all(&function.name) {
return 1.0;
}
if context.is_exported_in_init(&function.name) {
return 1.0;
}
0.0
}
fn explain(&self, function: &FunctionDef) -> String {
format!("Function '{}' is explicitly exported", function.name)
}
}
struct RustVisibilityHeuristic;
impl ApiHeuristic for RustVisibilityHeuristic {
fn name(&self) -> &str {
"rust_visibility"
}
fn evaluate(&self, function: &FunctionDef, context: &FileContext) -> f32 {
if context.language() != Language::Rust {
return 0.0;
}
if function.is_trait_implementation() {
return 1.0;
}
if function.has_visibility_keyword("pub(crate)") {
return 0.5;
} else if function.has_visibility_keyword("pub(super)") {
return 0.3;
} else if function.has_visibility_keyword("pub") {
return 1.0;
}
0.0
}
fn explain(&self, function: &FunctionDef) -> String {
if function.is_trait_implementation() {
"Function implements trait method (required by trait)".to_string()
} else if function.has_visibility_keyword("pub") {
"Function has `pub` visibility (Rust public API)".to_string()
} else {
"Function has no `pub` keyword (Rust private)".to_string()
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn create_test_function(name: &str) -> FunctionDef {
FunctionDef {
name: name.to_string(),
docstring: None,
parameters: vec![],
return_type: None,
decorators: vec![],
is_method: false,
class_name: None,
line: 10,
visibility: None,
is_trait_impl: false,
}
}
fn create_python_context() -> FileContext {
FileContext::new(PathBuf::from("test.py"), Language::Python)
}
#[test]
fn test_naming_convention_public() {
let function = create_test_function("create_bots");
let context = create_python_context();
let heuristic = NamingConventionHeuristic;
let score = heuristic.evaluate(&function, &context);
assert!(score >= 0.8, "Public function should score high");
}
#[test]
fn test_naming_convention_private() {
let function = create_test_function("_internal_helper");
let context = create_python_context();
let heuristic = NamingConventionHeuristic;
let score = heuristic.evaluate(&function, &context);
assert_eq!(score, 0.0, "Private function should score 0");
}
#[test]
fn test_docstring_structured() {
let mut function = create_test_function("process_data");
function.docstring = Some(
r#"
Process input data and return results.
Args:
data: Input data to process
Returns:
Processed data
"#
.to_string(),
);
let context = create_python_context();
let heuristic = DocstringHeuristic;
let score = heuristic.evaluate(&function, &context);
assert!(score >= 0.9, "Structured docstring should score very high");
}
#[test]
fn test_type_annotations_full() {
let mut function = create_test_function("calculate");
function.parameters = vec![
Parameter {
name: "x".to_string(),
type_annotation: Some("int".to_string()),
default_value: None,
},
Parameter {
name: "y".to_string(),
type_annotation: Some("int".to_string()),
default_value: None,
},
];
function.return_type = Some("int".to_string());
let context = create_python_context();
let heuristic = TypeAnnotationHeuristic;
let score = heuristic.evaluate(&function, &context);
assert!(score >= 0.9, "Fully annotated function should score high");
}
#[test]
fn test_symmetric_pair_detection() {
let save_func = create_test_function("save_chat_history");
let load_func = create_test_function("load_chat_history");
let mut context = create_python_context();
context
.functions
.insert("load_chat_history".to_string(), load_func.clone());
context.used_functions.push("load_chat_history".to_string());
let config = PublicApiConfig::default();
let heuristic = SymmetricPairHeuristic::new(&config);
let score = heuristic.evaluate(&save_func, &context);
assert!(
score >= 0.8,
"Function with used symmetric pair should score high"
);
}
#[test]
fn test_module_all_export() {
let function = create_test_function("exported_func");
let mut context = create_python_context();
context.module_all = Some(vec!["exported_func".to_string()]);
let heuristic = ModuleExportHeuristic;
let score = heuristic.evaluate(&function, &context);
assert_eq!(score, 1.0, "Function in __all__ should score 1.0");
}
#[test]
fn test_rust_pub_keyword() {
let mut function = create_test_function("analyze_code");
function.visibility = Some("pub".to_string());
let context = FileContext::new(PathBuf::from("lib.rs"), Language::Rust);
let heuristic = RustVisibilityHeuristic;
let score = heuristic.evaluate(&function, &context);
assert_eq!(score, 1.0, "pub function should score 1.0");
}
#[test]
fn test_rust_trait_implementation() {
let mut function = create_test_function("clone");
function.is_trait_impl = true;
let context = FileContext::new(PathBuf::from("lib.rs"), Language::Rust);
let heuristic = RustVisibilityHeuristic;
let score = heuristic.evaluate(&function, &context);
assert_eq!(
score, 1.0,
"Trait implementation should score 1.0 (never dead)"
);
}
#[test]
fn test_public_api_detector_integration() {
let mut function = create_test_function("create_bots_from_list");
function.docstring =
Some("Create bots from a list of bot configuration files.".to_string());
function.parameters = vec![Parameter {
name: "bot_files".to_string(),
type_annotation: Some("list".to_string()),
default_value: Some("None".to_string()),
}];
let context = create_python_context();
let detector = PublicApiDetector::new(PublicApiConfig::default());
let score = detector.is_public_api(&function, &context);
assert!(score.is_public, "Function should be detected as public API");
assert!(
score.confidence >= 0.6,
"Confidence should exceed threshold"
);
}
#[test]
fn test_underscore_prefix_override() {
let mut function = create_test_function("_internal_complex_algorithm");
function.docstring = Some(
r#"
Performs complex internal processing.
Args:
data: List of integers to process
Returns:
Dictionary containing processed results
"#
.to_string(),
);
function.parameters = vec![Parameter {
name: "data".to_string(),
type_annotation: Some("List[int]".to_string()),
default_value: None,
}];
function.return_type = Some("Dict[str, Any]".to_string());
let context = create_python_context();
let heuristic = NamingConventionHeuristic;
let score = heuristic.evaluate(&function, &context);
assert_eq!(
score, 0.0,
"Underscore prefix should score 0.0 regardless of docs"
);
}
}