use ruff_python_ast::relocate::relocate_expr;
use ruff_python_ast::{Expr, ExprStringLiteral, ModExpression, StringLiteral};
use ruff_text_size::Ranged;
use crate::{ParseError, Parsed, parse_expression, parse_string_annotation};
type AnnotationParseResult = Result<ParsedAnnotation, ParseError>;
#[derive(Debug)]
pub struct ParsedAnnotation {
parsed: Parsed<ModExpression>,
kind: AnnotationKind,
}
impl ParsedAnnotation {
pub fn parsed(&self) -> &Parsed<ModExpression> {
&self.parsed
}
pub fn expression(&self) -> &Expr {
self.parsed.expr()
}
pub fn kind(&self) -> AnnotationKind {
self.kind
}
}
#[derive(Copy, Clone, Debug)]
pub enum AnnotationKind {
Simple,
Complex,
}
impl AnnotationKind {
pub const fn is_simple(self) -> bool {
matches!(self, AnnotationKind::Simple)
}
}
pub fn parse_type_annotation(
string_expr: &ExprStringLiteral,
source: &str,
) -> AnnotationParseResult {
if let Some(string_literal) = string_expr.as_single_part_string() {
if &source[string_literal.content_range()] == string_literal.as_str() {
parse_simple_type_annotation(string_literal, source)
} else {
parse_complex_type_annotation(string_expr)
}
} else {
parse_complex_type_annotation(string_expr)
}
}
fn parse_simple_type_annotation(
string_literal: &StringLiteral,
source: &str,
) -> AnnotationParseResult {
Ok(ParsedAnnotation {
parsed: parse_string_annotation(source, string_literal)?,
kind: AnnotationKind::Simple,
})
}
fn parse_complex_type_annotation(string_expr: &ExprStringLiteral) -> AnnotationParseResult {
let mut parsed = parse_expression(string_expr.value.to_str())?;
relocate_expr(parsed.expr_mut(), string_expr.range());
Ok(ParsedAnnotation {
parsed,
kind: AnnotationKind::Complex,
})
}