From c1cb029bc28130020522b230b78464faacfde581 Mon Sep 17 00:00:00 2001
From: Sergio Pedri <sergio0694@live.com>
Date: Mon, 14 Feb 2022 11:28:47 +0100
Subject: [PATCH] Fix [ObservableRecipient] with inherited base attributes

---
 .../ObservableRecipientGenerator.cs           |  4 ++--
 .../Extensions/INamedTypeSymbolExtensions.cs  | 21 +++++++++++++++++++
 .../Test_ObservableRecipientAttribute.cs      | 21 +++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/CommunityToolkit.Mvvm.SourceGenerators/ComponentModel/ObservableRecipientGenerator.cs b/CommunityToolkit.Mvvm.SourceGenerators/ComponentModel/ObservableRecipientGenerator.cs
index da4b1d69..0363523b 100644
--- a/CommunityToolkit.Mvvm.SourceGenerators/ComponentModel/ObservableRecipientGenerator.cs
+++ b/CommunityToolkit.Mvvm.SourceGenerators/ComponentModel/ObservableRecipientGenerator.cs
@@ -82,8 +82,8 @@ public sealed class ObservableRecipientGenerator : TransitiveMembersGenerator<Ob
         // In order to use [ObservableRecipient], the target type needs to inherit from ObservableObject,
         // or be annotated with [ObservableObject] or [INotifyPropertyChanged] (with additional helpers).
         if (!typeSymbol.InheritsFrom("global::CommunityToolkit.Mvvm.ComponentModel.ObservableObject") &&
-            !typeSymbol.GetAttributes().Any(static a => a.AttributeClass?.HasFullyQualifiedName("global::CommunityToolkit.Mvvm.ComponentModel.ObservableObjectAttribute") == true) &&
-            !typeSymbol.GetAttributes().Any(static a =>
+            !typeSymbol.HasOrInheritsAttribute(static a => a.AttributeClass?.HasFullyQualifiedName("global::CommunityToolkit.Mvvm.ComponentModel.ObservableObjectAttribute") == true) &&
+            !typeSymbol.HasOrInheritsAttribute(static a =>
                 a.AttributeClass?.HasFullyQualifiedName("global::CommunityToolkit.Mvvm.ComponentModel.INotifyPropertyChangedAttribute") == true &&
                 !a.HasNamedArgument("IncludeAdditionalHelperMethods", false)))
         {
diff --git a/CommunityToolkit.Mvvm.SourceGenerators/Extensions/INamedTypeSymbolExtensions.cs b/CommunityToolkit.Mvvm.SourceGenerators/Extensions/INamedTypeSymbolExtensions.cs
index ffc97b59..0abb152b 100644
--- a/CommunityToolkit.Mvvm.SourceGenerators/Extensions/INamedTypeSymbolExtensions.cs
+++ b/CommunityToolkit.Mvvm.SourceGenerators/Extensions/INamedTypeSymbolExtensions.cs
@@ -2,6 +2,8 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System;
+using System.Linq;
 using System.Text;
 using Microsoft.CodeAnalysis;
 
@@ -62,4 +64,23 @@ internal static class INamedTypeSymbolExtensions
 
         return false;
     }
+
+    /// <summary>
+    /// Checks whether or not a given <see cref="INamedTypeSymbol"/> has or inherits a specified attribute.
+    /// </summary>
+    /// <param name="typeSymbol">The target <see cref="INamedTypeSymbol"/> instance to check.</param>
+    /// <param name="predicate">The predicate used to match available attributes.</param>
+    /// <returns>Whether or not <paramref name="typeSymbol"/> has an attribute matching <paramref name="predicate"/>.</returns>
+    public static bool HasOrInheritsAttribute(this INamedTypeSymbol typeSymbol, Func<AttributeData, bool> predicate)
+    {
+        for (INamedTypeSymbol? currentType = typeSymbol; currentType is not null; currentType = currentType.BaseType)
+        {
+            if (currentType.GetAttributes().Any(predicate))
+            {
+                return true;
+            }
+        }
+
+        return false;
+    }
 }
diff --git a/tests/CommunityToolkit.Mvvm.UnitTests/Test_ObservableRecipientAttribute.cs b/tests/CommunityToolkit.Mvvm.UnitTests/Test_ObservableRecipientAttribute.cs
index 222061c8..602fee32 100644
--- a/tests/CommunityToolkit.Mvvm.UnitTests/Test_ObservableRecipientAttribute.cs
+++ b/tests/CommunityToolkit.Mvvm.UnitTests/Test_ObservableRecipientAttribute.cs
@@ -260,4 +260,25 @@ public partial class Test_ObservableRecipientAttribute
             OnDeactivatedResult = true;
         }
     }
+
+    // See https://github.com/CommunityToolkit/dotnet/issues/109
+    [TestMethod]
+    public void Test_ObservableRecipientAttribute_WorksWithBaseClassWithObservableObjectAttribute()
+    {
+        ViewModelWithOnlyObservableRecipientAttribute model = new();
+
+        // This test method really only needs the two classes below to compile at all
+        Assert.IsTrue(model is INotifyPropertyChanged); // From [ObservableObject]
+        Assert.IsFalse(model.IsActive); // From [ObservableRecipient]
+    }
+
+    [ObservableObject]
+    public partial class BaseViewModelWithObservableObjectAttribute
+    {
+    }
+
+    [ObservableRecipient]
+    public partial class ViewModelWithOnlyObservableRecipientAttribute : BaseViewModelWithObservableObjectAttribute
+    {
+    }
 }