Skip to content

Commit 101df66

Browse files
porglezompporglezomp
committed
Add support for parameters passed to the @Traced macro
First, allow overriding the regular parameters of the withSpan call: setting the operationName, the context, and the kind. If it's not specified, it's not passed to the function, which means the function remains the source of truth for default arguments. Also allow overriding the span name, which controls the variable binding in the closure body. This is primarily useful for avoiding shadowing an outer "span" variable.
1 parent 50446c1 commit 101df66

File tree

2 files changed

+276
-3
lines changed

2 files changed

+276
-3
lines changed

‎Sources/TracingMacrosImplementation/TracedMacro.swift

Lines changed: 74 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,28 @@ public struct TracedMacro: BodyMacro {
3030
}
3131

3232
// Construct a withSpan call matching the invocation of the @Traced macro
33+
let (operationName, context, kind, spanName) = try extractArguments(from: node)
3334

34-
let operationName = StringLiteralExprSyntax(content: function.name.text)
35-
let withSpanCall: ExprSyntax = "withSpan(\(operationName))"
35+
var withSpanCall = FunctionCallExprSyntax("withSpan()" as ExprSyntax)!
36+
withSpanCall.arguments.append(LabeledExprSyntax(
37+
expression: operationName ?? ExprSyntax(StringLiteralExprSyntax(content: function.name.text))))
38+
func appendComma() {
39+
withSpanCall.arguments[withSpanCall.arguments.index(before: withSpanCall.arguments.endIndex)].trailingComma = .commaToken()
40+
}
41+
if let context {
42+
appendComma()
43+
withSpanCall.arguments.append(LabeledExprSyntax(label: "context", expression: context))
44+
}
45+
if let kind {
46+
appendComma()
47+
withSpanCall.arguments.append(LabeledExprSyntax(label: "ofKind", expression: kind))
48+
}
49+
50+
// Introduce a span identifier in scope
51+
var spanIdentifier: TokenSyntax = "span"
52+
if let spanName {
53+
spanIdentifier = .identifier(spanName)
54+
}
3655

3756
// We want to explicitly specify the closure effect specifiers in order
3857
// to avoid warnings about unused try/await expressions.
@@ -47,7 +66,7 @@ public struct TracedMacro: BodyMacro {
4766
throwsClause?.throwsSpecifier = .keyword(.throws)
4867
}
4968
var withSpanExpr: ExprSyntax = """
50-
\(withSpanCall) { span \(asyncClause)\(throwsClause)\(returnClause)in \(body.statements) }
69+
\(withSpanCall) { \(spanIdentifier) \(asyncClause)\(throwsClause)\(returnClause)in \(body.statements) }
5170
"""
5271

5372
// Apply a try / await as necessary to adapt the withSpan expression
@@ -62,6 +81,58 @@ public struct TracedMacro: BodyMacro {
6281

6382
return ["\(withSpanExpr)"]
6483
}
84+
85+
static func extractArguments(
86+
from node: AttributeSyntax
87+
) throws -> (
88+
operationName: ExprSyntax?,
89+
context: ExprSyntax?,
90+
kind: ExprSyntax?,
91+
spanName: String?
92+
) {
93+
// If there are no arguments, we don't have to do any of these bindings
94+
guard let arguments = node.arguments?.as(LabeledExprListSyntax.self) else {
95+
return (nil, nil, nil, nil)
96+
}
97+
98+
func getArgument(label: String) -> ExprSyntax? {
99+
arguments.first(where: { $0.label?.identifier?.name == label })?.expression
100+
}
101+
102+
// The operation name is the first argument if it's unlabeled
103+
var operationName: ExprSyntax?
104+
if let firstArgument = arguments.first, firstArgument.label == nil {
105+
operationName = firstArgument.expression
106+
}
107+
108+
let context = getArgument(label: "context")
109+
let kind = getArgument(label: "ofKind")
110+
var spanName: String?
111+
let spanNameExpr = getArgument(label: "span")
112+
if let spanNameExpr {
113+
guard let stringLiteral = spanNameExpr.as(StringLiteralExprSyntax.self),
114+
stringLiteral.segments.count == 1,
115+
let segment = stringLiteral.segments.first,
116+
let segmentText = segment.as(StringSegmentSyntax.self)
117+
else {
118+
throw MacroExpansionErrorMessage("span name must be a simple string literal")
119+
}
120+
let text = segmentText.content.text
121+
let isValidIdentifier = DeclReferenceExprSyntax("\(raw: text)" as ExprSyntax)?.hasError == false
122+
let isValidWildcard = text == "_"
123+
guard isValidIdentifier || isValidWildcard else {
124+
throw MacroExpansionErrorMessage("'\(text)' is not a valid parameter name")
125+
}
126+
spanName = text
127+
}
128+
return (
129+
operationName: operationName,
130+
context: context,
131+
kind: kind,
132+
spanName: spanName,
133+
)
134+
}
135+
65136
}
66137
#endif
67138

‎Tests/TracingMacrosTests/TracedTests.swift

Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,203 @@ final class TracedMacroTests: XCTestCase {
214214
macros: ["Traced": TracedMacro.self]
215215
)
216216
}
217+
218+
func test_tracedMacro_specifyOperationName() {
219+
assertMacroExpansion(
220+
"""
221+
@Traced("example but with a custom operationName")
222+
func example(param: Int) {
223+
span.attributes["param"] = param
224+
}
225+
""",
226+
expandedSource: """
227+
func example(param: Int) {
228+
withSpan("example but with a custom operationName") { span in
229+
span.attributes["param"] = param
230+
}
231+
}
232+
""",
233+
macros: ["Traced": TracedMacro.self]
234+
)
235+
236+
assertMacroExpansion(
237+
"""
238+
let globalName = "example"
239+
240+
@Traced(globalName)
241+
func example(param: Int) {
242+
span.attributes["param"] = param
243+
}
244+
""",
245+
expandedSource: """
246+
let globalName = "example"
247+
func example(param: Int) {
248+
withSpan(globalName) { span in
249+
span.attributes["param"] = param
250+
}
251+
}
252+
""",
253+
macros: ["Traced": TracedMacro.self]
254+
)
255+
}
256+
257+
func test_tracedMacro_specifyContext() {
258+
assertMacroExpansion(
259+
"""
260+
@Traced(context: .topLevel)
261+
func example() {
262+
print("Hello")
263+
}
264+
""",
265+
expandedSource: """
266+
func example() {
267+
withSpan("example", context: .topLevel) { span in
268+
print("Hello")
269+
}
270+
}
271+
""",
272+
macros: ["Traced": TracedMacro.self]
273+
)
274+
}
275+
276+
func test_tracedMacro_specifyKind() {
277+
assertMacroExpansion(
278+
"""
279+
@Traced(ofKind: .client)
280+
func example() {
281+
print("Hello")
282+
}
283+
""",
284+
expandedSource: """
285+
func example() {
286+
withSpan("example", ofKind: .client) { span in
287+
print("Hello")
288+
}
289+
}
290+
""",
291+
macros: ["Traced": TracedMacro.self]
292+
)
293+
}
294+
295+
func test_tracedMacro_specifySpanBindingName() {
296+
assertMacroExpansion(
297+
"""
298+
@Traced(span: "customSpan")
299+
func example(span: String) throws {
300+
customSpan.attributes["span"] = span
301+
}
302+
""",
303+
expandedSource: """
304+
func example(span: String) throws {
305+
try withSpan("example") { customSpan throws in
306+
customSpan.attributes["span"] = span
307+
}
308+
}
309+
""",
310+
macros: ["Traced": TracedMacro.self]
311+
)
312+
313+
assertMacroExpansion(
314+
"""
315+
@Traced(span: "_")
316+
func example(span: String) {
317+
print(span)
318+
}
319+
""",
320+
expandedSource: """
321+
func example(span: String) {
322+
withSpan("example") { _ in
323+
print(span)
324+
}
325+
}
326+
""",
327+
macros: ["Traced": TracedMacro.self]
328+
)
329+
}
330+
331+
func test_tracedMacro_specifySpanBindingName_invalid() {
332+
assertMacroExpansion(
333+
"""
334+
@Traced(span: 1)
335+
func example(span: String) throws {
336+
customSpan.attributes["span"] = span
337+
}
338+
""",
339+
expandedSource: """
340+
func example(span: String) throws {
341+
customSpan.attributes["span"] = span
342+
}
343+
""",
344+
diagnostics: [
345+
.init(message: "span name must be a simple string literal", line: 1, column: 1),
346+
],
347+
macros: ["Traced": TracedMacro.self]
348+
)
349+
350+
assertMacroExpansion(
351+
"""
352+
@Traced(span: "invalid name")
353+
func example(span: String) throws {
354+
customSpan.attributes["span"] = span
355+
}
356+
357+
@Traced(span: "123")
358+
func example2(span: String) throws {
359+
customSpan.attributes["span"] = span
360+
}
361+
""",
362+
expandedSource: """
363+
func example(span: String) throws {
364+
customSpan.attributes["span"] = span
365+
}
366+
func example2(span: String) throws {
367+
customSpan.attributes["span"] = span
368+
}
369+
""",
370+
diagnostics: [
371+
.init(message: "'invalid name' is not a valid parameter name", line: 1, column: 1),
372+
.init(message: "'123' is not a valid parameter name", line: 6, column: 1),
373+
],
374+
macros: ["Traced": TracedMacro.self]
375+
)
376+
377+
assertMacroExpansion(
378+
"""
379+
@Traced(span: "Hello \\(1)")
380+
func example(span: String) throws {
381+
customSpan.attributes["span"] = span
382+
}
383+
""",
384+
expandedSource: """
385+
func example(span: String) throws {
386+
customSpan.attributes["span"] = span
387+
}
388+
""",
389+
diagnostics: [
390+
.init(message: "span name must be a simple string literal", line: 1, column: 1),
391+
],
392+
macros: ["Traced": TracedMacro.self]
393+
)
394+
}
395+
396+
func test_tracedMacro_multipleMacroParameters() {
397+
assertMacroExpansion(
398+
"""
399+
@Traced("custom span name", context: .topLevel, ofKind: .client, span: "customSpan")
400+
func example(span: Int) {
401+
customSpan.attributes["span"] = span + 1
402+
}
403+
""",
404+
expandedSource: """
405+
func example(span: Int) {
406+
withSpan("custom span name", context: .topLevel, ofKind: .client) { customSpan in
407+
customSpan.attributes["span"] = span + 1
408+
}
409+
}
410+
""",
411+
macros: ["Traced": TracedMacro.self]
412+
)
413+
}
217414
}
218415

219416
// MARK: Compile tests
@@ -261,4 +458,9 @@ func example(param: Int) {
261458
span.attributes["param"] = param
262459
}
263460

461+
@Traced("custom span name", context: .topLevel, ofKind: .client, span: "customSpan")
462+
func exampleWithParams(span: Int) {
463+
customSpan.attributes["span"] = span + 1
464+
}
465+
264466
#endif

0 commit comments

Comments
 (0)