ImproveEnumFieldsCommand.cs

using System;
using System.CodeDom;
using System.Collections.Generic;
using System.ComponentModel;
using System.Xml.Serialization;

namespace NHibernate.Tool.HbmXsd
{
	/// <summary>
	/// Responsible for improving the enum field names in code generated by an
	/// <see cref="XsdCodeGenerator" />.
	/// </summary>
	public class ImproveEnumFieldsCommand
	{
		private readonly Dictionary<string, Dictionary<string, string>> changedEnums =
			new Dictionary<string, Dictionary<string, string>>();

		private readonly CodeNamespace code;

		public ImproveEnumFieldsCommand(CodeNamespace code)
		{
			if (code == null)
				throw new ArgumentNullException("code");

			this.code = code;
		}

		/// <summary>Changes enum field names to use camel casing.</summary>
		public void Execute()
		{
			ChangeDeclaredEnumFields();
			UpdateReferences();
		}

		private void ChangeDeclaredEnumFields()
		{
			foreach (CodeTypeDeclaration type in code.Types)
				if (type.IsEnum)
				{
					changedEnums[type.Name] = new Dictionary<string, string>();

					foreach (CodeTypeMember member in type.Members)
					{
						string xmlEnumValue = GetXmlEnumValue(member);
						string newEnumName = GetNewEnumName(type.Name, member.Name, xmlEnumValue);
						changedEnums[type.Name][member.Name] = newEnumName;
						member.Name = newEnumName;
					}
				}
		}

		private void UpdateReferences()
		{
			foreach (CodeTypeDeclaration type in code.Types)
				if (type.IsClass)
					foreach (CodeTypeMember member in type.Members)
					{
						CodeMemberField field = member as CodeMemberField;
						CodeConstructor constructor = member as CodeConstructor;

						if (field != null)
							UpdateFieldEnumValueReferences(field);

						else if (constructor != null)
							UpdateMethodEnumValueReferences(constructor);
					}
		}

		private static string GetXmlEnumValue(CodeTypeMember member)
		{
			if (member.CustomAttributes.Count == 1)
			{
				CodeAttributeArgument argument = member.CustomAttributes[0].Arguments[0];
				return ((CodePrimitiveExpression) argument.Value).Value.ToString();
			}
			else
			{
				AddXmlEnumAttribute(member);
				return null;
			}
		}

		protected virtual string GetNewEnumName(string typeName, string memberName, string xmlEnumValue)
		{
			if (xmlEnumValue != null)
				return StringTools.CamelCase(xmlEnumValue);
			else
				return StringTools.CamelCase(memberName);
		}

		private static void AddXmlEnumAttribute(CodeTypeMember member)
		{
			CodeTypeReference attributeType = new CodeTypeReference(typeof (XmlEnumAttribute));
			CodePrimitiveExpression argumentValue = new CodePrimitiveExpression(member.Name);
			CodeAttributeArgument argument = new CodeAttributeArgument(argumentValue);
			CodeAttributeDeclaration attribute = new CodeAttributeDeclaration(attributeType, argument);

			member.CustomAttributes.Add(attribute);
		}

		private void UpdateFieldEnumValueReferences(CodeTypeMember field)
		{
			foreach (CodeAttributeDeclaration attribute in field.CustomAttributes)
				if (attribute.Name == typeof (DefaultValueAttribute).FullName)
				{
					CodeFieldReferenceExpression reference = attribute.Arguments[0].Value
						as CodeFieldReferenceExpression;

					if (reference != null)
						UpdateReference(reference);
				}
		}

		private void UpdateMethodEnumValueReferences(CodeMemberMethod method)
		{
			foreach (CodeStatement statement in method.Statements)
			{
				CodeAssignStatement assignment = (CodeAssignStatement) statement;
				CodeFieldReferenceExpression right = assignment.Right as CodeFieldReferenceExpression;

				if (right != null)
					UpdateReference(right);
			}
		}

		private void UpdateReference(CodeFieldReferenceExpression reference)
		{
			string type = ((CodeTypeReferenceExpression) reference.TargetObject).Type.BaseType;

			if (changedEnums[type].ContainsKey(reference.FieldName))
				reference.FieldName = changedEnums[type][reference.FieldName];
		}
	}
}