Skip to content

Commit

Permalink
attributes: allow usage of string literals as field idents
Browse files Browse the repository at this point in the history
Replaces the Field.name type with an enum that can be the current
Punctuated type or a LitStr
  • Loading branch information
valkum committed Mar 16, 2024
1 parent 908cc43 commit a6841ce
Show file tree
Hide file tree
Showing 3 changed files with 75 additions and 6 deletions.
36 changes: 34 additions & 2 deletions tracing-attributes/src/attr.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use std::collections::HashSet;
use syn::parse::discouraged::Speculative;
use syn::{punctuated::Punctuated, Expr, Ident, LitInt, LitStr, Path, Token};

use proc_macro2::TokenStream;
Expand Down Expand Up @@ -268,7 +269,7 @@ pub(crate) struct Fields(pub(crate) Punctuated<Field, Token![,]>);

#[derive(Clone, Debug)]
pub(crate) struct Field {
pub(crate) name: Punctuated<Ident, Token![.]>,
pub(crate) name: FieldName,
pub(crate) value: Option<Expr>,
pub(crate) kind: FieldKind,
}
Expand All @@ -280,6 +281,12 @@ pub(crate) enum FieldKind {
Value,
}

#[derive(Clone, Debug)]
pub(crate) enum FieldName {
Ident(Punctuated<Ident, Token![.]>),
Literal(LitStr),
}

impl Parse for Fields {
fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
let _ = input.parse::<kw::fields>();
Expand All @@ -306,7 +313,7 @@ impl Parse for Field {
input.parse::<Token![?]>()?;
kind = FieldKind::Debug;
};
let name = Punctuated::parse_separated_nonempty_with(input, Ident::parse_any)?;
let name = FieldName::parse(input)?;
let value = if input.peek(Token![=]) {
input.parse::<Token![=]>()?;
if input.peek(Token![%]) {
Expand Down Expand Up @@ -347,6 +354,31 @@ impl ToTokens for Field {
}
}

impl Parse for FieldName {
fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
let ahead = input.fork();
if let Ok(ident) = Punctuated::parse_separated_nonempty_with(&ahead, Ident::parse_any) {
input.advance_to(&ahead);
return Ok(Self::Ident(ident));
}
if let Ok(lit) = input.parse::<LitStr>() {
return Ok(Self::Literal(lit));
}
Err(ahead.error(
"expected \"field.name\" (string literal) or field.name (punctuated identifier)",
))
}
}

impl ToTokens for FieldName {
fn to_tokens(&self, tokens: &mut TokenStream) {
match self {
FieldName::Ident(ident) => ident.to_tokens(tokens),
FieldName::Literal(lit) => lit.to_tokens(tokens),
}
}
}

impl ToTokens for FieldKind {
fn to_tokens(&self, tokens: &mut TokenStream) {
match self {
Expand Down
16 changes: 12 additions & 4 deletions tracing-attributes/src/expand.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::iter;

use proc_macro2::TokenStream;
use quote::{quote, quote_spanned, ToTokens};
use syn::parse::Parse;
use syn::visit_mut::VisitMut;
use syn::{
punctuated::Punctuated, spanned::Spanned, Block, Expr, ExprAsync, ExprCall, FieldPat, FnArg,
Expand All @@ -10,7 +11,7 @@ use syn::{
};

use crate::{
attr::{Field, Fields, FormatMode, InstrumentArgs, Level},
attr::{Field, FieldName, Fields, FormatMode, InstrumentArgs, Level},
MaybeItemFn, MaybeItemFnRef,
};

Expand Down Expand Up @@ -189,9 +190,16 @@ fn gen_block<B: ToTokens>(
// If any parameters have the same name as a custom field, skip
// and allow them to be formatted by the custom field.
if let Some(ref fields) = args.fields {
fields.0.iter().all(|Field { ref name, .. }| {
let first = name.first();
first != name.last() || !first.iter().any(|name| name == &param)
fields.0.iter().all(|Field { ref name, .. }| match name {
FieldName::Ident(name) => {
let first = name.first();
first != name.last() || !first.iter().any(|name| name == &param)
}
FieldName::Literal(name) => {
// If the literal string would be a valid ident, apply the same overwrite logic.
let literal_ident = name.parse_with(syn::Ident::parse);
!literal_ident.iter().any(|name| name == param)
}
})
} else {
true
Expand Down
29 changes: 29 additions & 0 deletions tracing-attributes/tests/instrument.rs
Original file line number Diff line number Diff line change
Expand Up @@ -238,3 +238,32 @@ fn impl_trait_return_type() {

handle.assert_finished();
}

#[test]
fn keywords_in_fields() {
#[instrument(fields("d.type" = "test", "x" = ?x))]
fn my_fn(x: u64) {
tracing::event!(Level::TRACE, "r.type" = "test", "event name");
}

let span = expect::span().named("my_fn");

let (subscriber, handle) = collector::mock()
.new_span(
span.clone().with_fields(
expect::field("d.type")
.with_value(&"test")
.and(expect::field("x").with_value(&tracing::field::display(10))),
),
)
.enter(span.clone())
.event(expect::event().with_fields(expect::field("r.type").with_value(&"test")))
.exit(span.clone())
.drop_span(span)
.only()
.run_with_handle();

with_default(subscriber, || my_fn(10));

handle.assert_finished();
}

0 comments on commit a6841ce

Please sign in to comment.