Customising generated code with marker attributes: Creating a source generator - Part 4

In the previous posts in this series I showed how to create an incremental source generator, how to unit and integration test it, and how to package it in a NuGet package. In this post I describe how to customise the source generator's behaviour by extending the marker attribute with additional properties.

Extending the source generator marker attribute

One of the first steps for any source generator is to identify which code in the project needs to partake in the source generation. The source generator might look for specific types or members, but another common approach is to use a marker attribute. This is the approach I described in the first post in this series.

The [EnumExtensions] attribute I described in the first post was a simple attribute with no other properties. That meant there was no way to customise the code generated by the source generator. That was one of the limitations I discussed at the end of the post.

A common way to provide this functionality is to add additional properties to the marker attribute. In this post, I'm going to show how to do this for a single setting—the name of the extension method class to generate.

By default, the name EnumExtensions is used for the extension method class. With this change, you'll be able to specify an alternative name by setting the ExtensionClassName property. For example the following:

[EnumExtensions(ExtensionClassName = "DirectionExtensions")]
public enum Direction
{
    Left,
    Right,
    Up,
    Down,
}

would generate a class called DirectionExtensions, that looks something like this:

//HintName: EnumExtensions.g.cs

namespace NetEscapades.EnumGenerators
{
    public static partial class DirectionExtensions // 👈 Note the custom name
    {
        public static string ToStringFast(this Direction value)
            => value switch
            {
                Direction.Left => nameof(Direction.Left),
                Direction.Right => nameof(Direction.Right),
                Direction.Up => nameof(Direction.Up),
                Direction.Down => nameof(Direction.Down),
                _ => value.ToString(),
            };
    }
}

For the remainder of the post, I'll walk through the changes needed to the original source generator to achieve this.

I'm not going to show the full code for the source generator here, just the incremental changes over the original in the first post. You can find the full code on GitHub.

1. Update the marker attribute

The first step is to update the marker attribute with the new property:

[System.AttributeUsage(System.AttributeTargets.Enum)]
public class EnumExtensionsAttribute : System.Attribute
{
    public string ExtensionClassName { get; set; } // 👈 New property
}

This marker attribute is automatically added to the compilation by the source generator, as described in the first post, so we're actually updating a string here rather than an attribute. If you want to add more customisation, like the ability to customise the generated code's namespace for example, then you can add extra properties to this attribute.

2. Allow setting separate extension class name for each enum

With this change, users can now set a different name for the extension class for each enum, so we need to record the extension name when we're extracting the details about the enum into an EnumToGenerate object:

public readonly struct EnumToGenerate
{
    public readonly string ExtensionName; // 👈 New field
    public readonly string Name;
    public readonly List<string> Values;

    public EnumToGenerate(string extensionName, string name, List<string> values)
    {
        Name = name;
        Values = values;
        ExtensionName = extensionName;
    }
}

Note that as we make the extension method partial, and each ToStringFast() method will be a different overload, it doesn't matter if a user specifies the same extension class name more than once.

3. Update code generation

We're working backwards somewhat here, so the following shows the updated code for the extension generator. There's nothing complicated here, it's just a bit fiddly working with the StringBuilder. The main difference from the previous iteration is that we generate a separate class for each enum (instead of one class with many methods), and that the class name comes from the EnumToGenerate:

public static string GenerateExtensionClass(List<EnumToGenerate> enumsToGenerate)
{
    var sb = new StringBuilder();
    sb.Append(@"
namespace NetEscapades.EnumGenerators
{");
    foreach(var enumToGenerate in enumsToGenerate)
    {
        sb.Append(@"
public static partial class ").Append(enumToGenerate.ExtensionName).Append(@"
{
    public static string ToStringFast(this ").Append(enumToGenerate.Name).Append(@" value)
        => value switch
        {");
        foreach (var member in enumToGenerate.Values)
        {
            sb.Append(@"
            ")
                .Append(enumToGenerate.Name).Append('.').Append(member)
                .Append(" => nameof(")
                .Append(enumToGenerate.Name).Append('.').Append(member).Append("),");
        }

        sb.Append(@"
            _ => value.ToString(),
        };
}
");
    }
    sb.Append('}');

    return sb.ToString();
}

All that's left is to update the source generator code itself to read the value of ExtensionClassName from the marker attribute.

4. Reading the property value from a marker attribute

So far we've only had to make small changes to support this new functionality, but we haven't done the hard part yet—reading the value from the compilation. When you set a property on an attribute, semantically you're setting a named constructor argument.

To find the value of the property ExtensionClassName, we first need to find the AttributeData for the [EnumExtensions] attribute. We can then check the NamedArguments for the specific property. The following shows a stripped down version of the code to extract the property value if it's provided:

static List<EnumToGenerate> GetTypesToGenerate(Compilation compilation, IEnumerable<EnumDeclarationSyntax> enums, CancellationToken ct)
{
    var enumsToGenerate = new List<EnumToGenerate>();
    // Get a reference to the [EnumExtensions] symbol
    INamedTypeSymbol? enumAttribute = compilation.GetTypeByMetadataName("NetEscapades.EnumGenerators.EnumExtensionsAttribute");

    // ... error checking and verification elided

    foreach (var enumDeclarationSyntax in enums)
    {
        // Get the semantic model of the enum symbol
        SemanticModel semanticModel = compilation.GetSemanticModel(enumDeclarationSyntax.SyntaxTree);
        INamedTypeSymbol enumSymbol = semanticModel.GetDeclaredSymbol(enumDeclarationSyntax);

        // Set the default extension name
        string extensionName = "EnumExtensions";

        // Loop through all of the attributes on the enum
        foreach (AttributeData attributeData in enumSymbol.GetAttributes())
        {
            if (!enumAttribute.Equals(attributeData.AttributeClass, SymbolEqualityComparer.Default))
            {
                // This isn't the [EnumExtensions] attribute
                continue;
            }

            // This is the attribute, check all of the named arguments
            foreach (KeyValuePair<string, TypedConstant> namedArgument in attributeData.NamedArguments)
            {
                // Is this the ExtensionClassName argument?
                if (namedArgument.Key == "ExtensionClassName"
                    && namedArgument.Value.Value?.ToString() is { } n)
                {
                    extensionName = n;
                }
            }

            break;
        }

        // ... Not shown: existing code to retrieve the enum name and members

        // Record the extension name
        enumsToGenerate.Add(new EnumToGenerate(extensionName, enumName, members));
    }

    return enumsToGenerate;
}

With these changes, you can add arbitrarily more customisation to your source generator by extending the marker attribute.

5. Supporting attribute constructors

In the example above, we're only checking the NamedArguments of the attribute, because the attribute doesn't have a constructor, so it's the only way to specify the ExtensionClassName property. But what if the marker attribute was defined differently, and did have a constructor? For example, what if we make the ExtensionClassName required, and add a new optional property, ExtensionNamespaceName:

[System.AttributeUsage(System.AttributeTargets.Enum)]
public class EnumExtensionsAttribute : System.Attribute
{
    public EnumExtensionsAttribute(string extensionClassName)
    {
        ExtensionClassName = extensionClassName;
    }

    public string ExtensionClassName { get; }
    public string ExtensionNamespaceName { get; set; }
}

Then the code in the previous section won't work. And if you have multiple properties, and multiple constructors, then things become more complicated again. The following code shows the general approach to extract these values inside the source generator. Specifically, you need to read both the ConstructorArguments and the NamedArguments of the AttributeData, and infer the values set correctly:

INamedTypeSymbol enumSymbol = semanticModel.GetDeclaredSymbol(enumDeclarationSyntax);

// Placeholder variables for the specififed ExtensionClassName and ExtensionNamespaceName
string className = null;
string namespaceName = null;

// Loop through all of the attributes on the enum until we find the [EnumExtensions] attribute
foreach (AttributeData attributeData in enumSymbol.GetAttributes())
{
    if (!enumAttribute.Equals(attributeData.AttributeClass, SymbolEqualityComparer.Default))
    {
        // This isn't the [EnumExtensions] attribute
        continue;
    }

    // This is the right attribute, check the constructor arguments
    if (!attribute.ConstructorArguments.IsEmpty)
    {
        ImmutableArray<TypedConstant> args = attribute.ConstructorArguments;

        // make sure we don't have any errors
        foreach (TypedConstant arg in args)
        {
            if (arg.Kind == TypedConstantKind.Error)
            {
                // have an error, so don't try and do any generation
                return;
            }
        }

        // Use the position of the argument to infer which value is set
        switch (args.Length)
        {
            case 1:
                className = (string)args[0].Value;
                break;
        }
    }


    // now check for named arguments
    if (!attribute.NamedArguments.IsEmpty)
    {
        foreach (KeyValuePair<string, TypedConstant> arg in attribute.NamedArguments)
        {
            TypedConstant typedConstant = arg.Value;
            if (typedConstant.Kind == TypedConstantKind.Error)
            {
                // have an error, so don't try and do any generation
                return;
            }
            else
            {
                // Use the constructor argument or property name to infer which value is set
                switch (arg.Key)
                {
                    case "extensionClassName":
                        className = (string)typedConstant.Value;
                        break;
                    case "ExtensionNamespaceName":
                        namespaceName = (string)typedConstant.Value;
                        break;
                }
            }
        }
    }

    break;
}

This is obviously more complex, but may well be necessary to provide a better user experience for the consumer of your source generator.

Summary

In this post I described how you can provide customisation options to consumers of a source generator by adding properties to a marker attribute. This requires a few gymnastics to parse the provided values, especially if you use required constructor arguments in your attribute, as well as named properties. Overall though this is generally a good way to expand the capabilities of your source generator.