use tree_sitter::Node;
#[derive(Debug, Default)]
pub struct YardTags {
pub params: Vec<ParamTag>,
pub returns: Option<String>,
pub type_annotation: Option<String>,
}
#[derive(Debug)]
pub struct ParamTag {
pub name: String,
pub type_str: String,
#[allow(dead_code)] pub description: Option<String>,
}
pub fn extract_yard_comment(node: Node, content: &[u8]) -> Option<String> {
if let Some(comment) = try_extract_comment(node, content) {
return Some(comment);
}
if let Some(parent) = node.parent() {
let parent_kind = parent.kind();
if parent_kind == "visibility_modifier"
|| parent_kind == "body_statement"
|| parent_kind.contains("_statement")
{
return try_extract_comment(parent, content);
}
}
None
}
fn try_extract_comment(node: Node, content: &[u8]) -> Option<String> {
let node_start_line = node.start_position().row;
let mut prev_sibling = node.prev_sibling();
let mut comment_lines = Vec::new();
let mut expected_line = node_start_line;
while let Some(sibling) = prev_sibling {
if sibling.kind() == "comment" {
let comment_text = sibling.utf8_text(content).ok()?;
let comment_end_line = sibling.end_position().row;
let line_distance = expected_line.saturating_sub(comment_end_line);
if line_distance <= 2 {
if comment_text.trim().starts_with('#') {
comment_lines.push(comment_text.to_string());
expected_line = sibling.start_position().row;
prev_sibling = sibling.prev_sibling();
continue;
}
}
break;
} else if !sibling.kind().contains("whitespace") {
break;
}
prev_sibling = sibling.prev_sibling();
}
if comment_lines.is_empty() {
return None;
}
comment_lines.reverse();
Some(comment_lines.join("\n"))
}
pub fn parse_yard_tags(yard_comment: &str) -> YardTags {
let mut tags = YardTags::default();
for line in yard_comment.lines() {
let line = line.trim().trim_start_matches('#').trim();
if line.starts_with("@param") {
if let Some(param) = parse_param_tag(line) {
tags.params.push(param);
}
}
else if line.starts_with("@return") {
if let Some(type_str) = extract_type_from_tag(line) {
tags.returns = Some(type_str);
}
}
else if line.starts_with("@type")
&& let Some(type_str) = extract_type_from_tag(line)
{
tags.type_annotation = Some(type_str);
}
}
tags
}
fn parse_param_tag(line: &str) -> Option<ParamTag> {
let after_keyword = line.trim_start_matches("@param").trim();
let (type_str, end_index) = extract_balanced_brackets_with_index(after_keyword)?;
let after_type = &after_keyword[end_index + 1..].trim();
let name = extract_param_name(after_type)?;
let description = after_type
.split_once(&name)
.and_then(|(_, rest)| rest.trim().strip_prefix('-'))
.map(|s| s.trim().to_string());
Some(ParamTag {
name,
type_str,
description,
})
}
fn extract_param_name(text: &str) -> Option<String> {
let text = text.trim();
let text_to_parse = if let Some(rest) = text.strip_prefix('*') {
rest.trim()
} else {
text
};
let text_to_parse = if let Some(rest) = text_to_parse.strip_prefix('*') {
rest.trim()
} else {
text_to_parse
};
let text_to_parse = if let Some(rest) = text_to_parse.strip_prefix('&') {
rest.trim()
} else {
text_to_parse
};
let name = text_to_parse
.split(|c: char| c.is_whitespace() || c == '-' || c == ':')
.next()?
.trim();
if name.is_empty() {
None
} else {
Some(name.to_string())
}
}
fn extract_type_from_tag(line: &str) -> Option<String> {
extract_balanced_brackets(line)
}
fn extract_balanced_brackets_with_index(s: &str) -> Option<(String, usize)> {
let start = s.find('[')?;
let mut depth = 0;
let mut end = start;
for (i, ch) in s[start..].char_indices() {
match ch {
'[' => depth += 1,
']' => {
depth -= 1;
if depth == 0 {
end = start + i;
break;
}
}
_ => {}
}
}
if depth != 0 {
return None; }
Some((s[start + 1..end].to_string(), end))
}
fn extract_balanced_brackets(s: &str) -> Option<String> {
extract_balanced_brackets_with_index(s).map(|(content, _)| content)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_extract_balanced_brackets() {
assert_eq!(
extract_balanced_brackets("@param [String] name"),
Some("String".to_string())
);
assert_eq!(
extract_balanced_brackets("@param [User, Admin] user"),
Some("User, Admin".to_string())
);
assert_eq!(
extract_balanced_brackets("@param [Array<String>] items"),
Some("Array<String>".to_string())
);
assert_eq!(
extract_balanced_brackets("@return [Hash{String => Integer}]"),
Some("Hash{String => Integer}".to_string())
);
}
#[test]
fn test_extract_param_name() {
assert_eq!(extract_param_name("name"), Some("name".to_string()));
assert_eq!(
extract_param_name("count - the count"),
Some("count".to_string())
);
assert_eq!(extract_param_name("*args"), Some("args".to_string()));
assert_eq!(
extract_param_name("**kwargs - keyword args"),
Some("kwargs".to_string())
);
assert_eq!(extract_param_name("&block"), Some("block".to_string()));
assert_eq!(extract_param_name("name:"), Some("name".to_string()));
assert_eq!(
extract_param_name("value: 10 - default value"),
Some("value".to_string())
);
}
#[test]
fn test_parse_param_tag() {
let tag = parse_param_tag("@param [String] name - description").unwrap();
assert_eq!(tag.name, "name");
assert_eq!(tag.type_str, "String");
assert_eq!(tag.description, Some("description".to_string()));
let tag = parse_param_tag("@param [Integer] count").unwrap();
assert_eq!(tag.name, "count");
assert_eq!(tag.type_str, "Integer");
}
#[test]
fn test_parse_yard_tags() {
let yard = r"# @param [String] name
# @param [Integer] age
# @return [User]";
let tags = parse_yard_tags(yard);
assert_eq!(tags.params.len(), 2);
assert_eq!(tags.params[0].name, "name");
assert_eq!(tags.params[1].name, "age");
assert_eq!(tags.returns, Some("User".to_string()));
}
#[test]
fn test_parse_yard_type_tag() {
let yard = r"# @type [String]";
let tags = parse_yard_tags(yard);
assert_eq!(tags.type_annotation, Some("String".to_string()));
}
#[test]
fn test_parse_yard_union_types() {
let yard = r"# @param [String, Integer] value
# @return [Boolean]";
let tags = parse_yard_tags(yard);
assert_eq!(tags.params[0].type_str, "String, Integer");
assert_eq!(tags.returns, Some("Boolean".to_string()));
}
#[test]
fn test_parse_yard_array_types() {
let yard = r"# @param [Array<User>] users
# @return [Hash{String => Integer}]";
let tags = parse_yard_tags(yard);
assert_eq!(tags.params[0].type_str, "Array<User>");
assert_eq!(tags.returns, Some("Hash{String => Integer}".to_string()));
}
#[test]
fn test_parse_yard_nullable_types() {
let yard = r"# @param [String, nil] value
# @return [User, nil]";
let tags = parse_yard_tags(yard);
assert_eq!(tags.params[0].type_str, "String, nil");
assert_eq!(tags.returns, Some("User, nil".to_string()));
}
}