Modified how DataViewTypes are registered

src/Microsoft.ML.Data/Data/DataViewTypeManager.cs

@@ -53,6 +53,8 @@ public static class DataViewTypeManager
         /// </summary>
         internal static DataViewType GetDataViewType(Type type, IEnumerable<Attribute> typeAttributes = null)
         {
+            //Filter attributes as we only care about TypeAttributes
+            typeAttributes = typeAttributes == null ? typeAttributes : typeAttributes.Where(attr => attr.GetType().IsSubclassOf(typeof(TypeAttribute)));
             lock (_lock)
             {
                 // Compute the ID of type with extra attributes.
@@ -73,6 +75,8 @@ internal static DataViewType GetDataViewType(Type type, IEnumerable<Attribute> t
         /// </summary>
         internal static bool Knows(Type type, IEnumerable<Attribute> typeAttributes = null)
         {
+            //Filter attributes as we only care about TypeAttributes
+            typeAttributes = typeAttributes == null ? typeAttributes : typeAttributes.Where(attr => attr.GetType().IsSubclassOf(typeof(TypeAttribute)));
             lock (_lock)
             {
                 // Compute the ID of type with extra attributes.
@@ -113,6 +117,16 @@ internal static bool Knows(DataViewType dataViewType)
         /// <param name="typeAttributes">The <see cref="Attribute"/>s attached to <paramref name="type"/>.</param>
         public static void Register(DataViewType dataViewType, Type type, IEnumerable<Attribute> typeAttributes = null)
         {
+            if (typeAttributes != null)
+            {
+                foreach (var attr in typeAttributes)
+                {
+                    if (!attr.GetType().IsSubclassOf(typeof(TypeAttribute)))
+                    {
+                        throw Contracts.ExceptParam(nameof(type), $"Type {type} has an attribute that is not of TypeAttribute.");
+                    }
+                }
+            }
             lock (_lock)
             {
                 if (_bannedRawTypes.Contains(type))

src/Microsoft.ML.Data/Data/SchemaDefinition.cs

@@ -18,7 +18,7 @@ namespace Microsoft.ML.Data
     /// Can be applied only for member of following types: <see cref="byte"/>, <see cref="ushort"/>, <see cref="uint"/>, <see cref="ulong"/>
     /// </remarks>
     [AttributeUsage(AttributeTargets.Field | AttributeTargets.Property, AllowMultiple = false, Inherited = true)]
-    public sealed class KeyTypeAttribute : Attribute
+    public sealed class KeyTypeAttribute : TypeAttribute
     {
         /// <summary>
         /// Marks member as <see cref="KeyDataViewType"/>.
@@ -55,7 +55,7 @@ public KeyTypeAttribute(ulong count)
     /// the dimensionality of the resulting array.
     /// </summary>
     [AttributeUsage(AttributeTargets.Field | AttributeTargets.Property, AllowMultiple = false, Inherited = true)]
-    public sealed class VectorTypeAttribute : Attribute
+    public sealed class VectorTypeAttribute : TypeAttribute
     {
         /// <summary>
         /// The length of the vectors from this vector valued field.

src/Microsoft.ML.DataView/DataViewType.cs

@@ -473,7 +473,7 @@ public override bool Equals(DataViewType other)
     /// in <see cref="IDataView"/> would be the associated <see cref="DataViewType"/>.
     /// </summary>
     [AttributeUsage(AttributeTargets.Field | AttributeTargets.Property, AllowMultiple = false, Inherited = true)]
-    public abstract class DataViewTypeAttribute : Attribute, IEquatable<DataViewTypeAttribute>
+    public abstract class DataViewTypeAttribute : TypeAttribute, IEquatable<DataViewTypeAttribute>
     {
         /// <summary>
         /// A function implicitly invoked by ML.NET when processing a custom type. It binds a DataViewType to a custom type plus its attributes.

src/Microsoft.ML.DataView/TypeAttribute.cs

@@ -0,0 +1,14 @@
+using System;
+using System.Collections.Generic;
+using System.Text;
+
+namespace Microsoft.ML.Data
+{
+    /*
+     * This class is should be used as the base type for all Attributes which define a Type.
+     * The DataViewManager relies upon this type to determine if the Attribute is relevant.
+     */
+    public abstract class TypeAttribute : Attribute
+    {
+    }
+}

test/Microsoft.ML.Core.Tests/UnitTests/TestCustomTypeRegister.cs

@@ -271,11 +271,26 @@ public void TestTypeManager()
             DataViewTypeManager.Register(e, typeof(AlienBody), new[] { f });
             Assert.True(DataViewTypeManager.Knows(e));
             Assert.True(DataViewTypeManager.Knows(typeof(AlienBody), new[] { f }));
-            Assert.True(DataViewTypeManager.Knows(typeof(AlienBody), new[] { f }));
             // "e" is associated with typeof(AlienBody) with "f," so the call below should return true.
             Assert.Equal(e, DataViewTypeManager.GetDataViewType(typeof(AlienBody), new[] { f }));
             // "a" is associated with typeof(AlienBody) without any attribute, so the call below should return false.
             Assert.NotEqual(a, DataViewTypeManager.GetDataViewType(typeof(AlienBody), new[] { f }));
         }
+
+        [Fact]
+        public void RegisterTypeWithAdditionalAttributes()
+        {
+            var a = new DataViewAlienBodyType(7788);
+            var b = new AlienTypeAttributeAttribute(8877);
+            var c = new ColumnNameAttribute("SomeName");
+
+            Assert.Throws<ArgumentOutOfRangeException>(() => DataViewTypeManager.Register(a, typeof(AlienBody), new Attribute[] { b, c }));
+
+            DataViewTypeManager.Register(a, typeof(AlienBody), new Attribute[] { b });
+            Assert.True(DataViewTypeManager.Knows(a));
+            Assert.True(DataViewTypeManager.Knows(typeof(AlienBody), new Attribute[] { b, c }));
+            // "a" is associated with typeof(AlienBody) with "b," so the call below should return true.
+            Assert.Equal(a, DataViewTypeManager.GetDataViewType(typeof(AlienBody), new Attribute[] { b, c }));
+        }
     }
 }