use std::error::Error;
#[cfg(feature = "walk")]
use std::fmt::{Display, Formatter};
use crate::CellStream;
#[cfg(feature = "walk")]
use crate::fold::{self, Fold};
#[cfg(feature = "walk")]
use crate::nodes::{Expr, ExprAssignInstallMethod, ExprLambda, ExprLiteral, TypeAnnotation};
#[cfg(feature = "walk")]
use crate::{CellStreamError, NativeParseError, NativeParser, Parser, Tokenize, parse_quote_m2};
pub trait ProcMacro<Args = ()> {
type Error: Error;
fn expand(&mut self, input: CellStream, args: &Args) -> Result<CellStream, Self::Error>;
}
#[cfg(feature = "walk")]
#[derive(Debug, Default)]
pub struct LiteralReturnTypes;
#[cfg(feature = "walk")]
#[derive(Debug)]
pub enum LiteralReturnTypesError {
Parse(NativeParseError),
Tokenize(CellStreamError),
}
#[cfg(feature = "walk")]
impl Display for LiteralReturnTypesError {
fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Parse(error) => Display::fmt(error, formatter),
Self::Tokenize(error) => Display::fmt(error, formatter),
}
}
}
#[cfg(feature = "walk")]
impl Error for LiteralReturnTypesError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Self::Parse(error) => Some(error),
Self::Tokenize(error) => Some(error),
}
}
}
#[cfg(feature = "walk")]
impl ProcMacro for LiteralReturnTypes {
type Error = LiteralReturnTypesError;
fn expand(&mut self, input: CellStream, _args: &()) -> Result<CellStream, Self::Error> {
let mut parser = NativeParser;
let syntax = parser
.parse_cells(input)
.map_err(LiteralReturnTypesError::Parse)?;
self.fold_source_file(syntax)
.to_cell_stream()
.map_err(LiteralReturnTypesError::Tokenize)
}
}
#[cfg(feature = "walk")]
impl Fold for LiteralReturnTypes {
fn fold_expr_assign_install_method(
&mut self,
node: ExprAssignInstallMethod,
) -> ExprAssignInstallMethod {
let mut node = fold::fold_expr_assign_install_method(self, node);
if node.typical_value.is_none() {
node.typical_value = literal_return_annotation(&node.function);
}
node
}
}
#[cfg(feature = "walk")]
fn literal_return_annotation(function: &Expr) -> Option<TypeAnnotation> {
let Expr::Lambda(lambda) = function else {
return None;
};
let body = match lambda {
ExprLambda::Variadic(lambda) => lambda.body.as_ref(),
ExprLambda::List(lambda) => lambda.body.as_ref(),
};
let Expr::Literal(literal) = body else {
return None;
};
let typical_value = match literal {
ExprLiteral::Integer(_) => parse_quote_m2!(ZZ),
ExprLiteral::Float(_) => parse_quote_m2!(RR),
ExprLiteral::String(_) | ExprLiteral::RawString(_) => parse_quote_m2!(String),
};
Some(TypeAnnotation::new(typical_value, parse_quote_m2!(=>)))
}