Add JsonSchemaExporterContext.ParentTypeInfo (#104550)

Author: eiriktsarpalis

src/libraries/System.Text.Json/ref/System.Text.Json.cs

@@ -914,6 +914,7 @@ public readonly partial struct JsonSchemaExporterContext
     {
         private readonly object _dummy;
         private readonly int _dummyPrimitive;
+        public System.Text.Json.Serialization.Metadata.JsonTypeInfo? BaseTypeInfo { get { throw null; } }
         public System.Text.Json.Serialization.Metadata.JsonPropertyInfo? PropertyInfo { get { throw null; } }
         public System.ReadOnlySpan<string> Path { get { throw null; } }
         public System.Text.Json.Serialization.Metadata.JsonTypeInfo TypeInfo { get { throw null; } }

src/libraries/System.Text.Json/src/System/Text/Json/Schema/JsonSchemaExporter.cs

@@ -68,7 +68,7 @@ private static JsonSchema MapJsonSchemaCore(
             JsonPropertyInfo? propertyInfo = null,
             JsonConverter? customConverter = null,
             JsonNumberHandling? customNumberHandling = null,
-            Type? parentPolymorphicType = null,
+            JsonTypeInfo? parentPolymorphicTypeInfo = null,
             bool parentPolymorphicTypeContainsTypesWithoutDiscriminator = false,
             bool parentPolymorphicTypeIsNonNullable = false,
             KeyValuePair<string, JsonSchema>? typeDiscriminator = null,
@@ -90,7 +90,7 @@ private static JsonSchema MapJsonSchemaCore(
                 return CompleteSchema(ref state, schema);
             }
 
-            if (parentPolymorphicType is null && typeInfo.PolymorphismOptions is { DerivedTypes.Count: > 0 } polyOptions)
+            if (parentPolymorphicTypeInfo is null && typeInfo.PolymorphismOptions is { DerivedTypes.Count: > 0 } polyOptions)
             {
                 // This is the base type of a polymorphic type hierarchy. The schema for this type
                 // will include an "anyOf" property with the schemas for all derived types.
@@ -133,7 +133,7 @@ private static JsonSchema MapJsonSchemaCore(
                     JsonSchema derivedSchema = MapJsonSchemaCore(
                         ref state,
                         derivedTypeInfo,
-                        parentPolymorphicType: typeInfo.Type,
+                        parentPolymorphicTypeInfo: typeInfo,
                         typeDiscriminator: derivedTypeDiscriminator,
                         parentPolymorphicTypeContainsTypesWithoutDiscriminator: containsTypesWithoutDiscriminator,
                         parentPolymorphicTypeIsNonNullable: propertyInfo is { IsGetNullable: false, IsSetNullable: false },
@@ -372,7 +372,7 @@ JsonSchema CompleteSchema(ref GenerationState state, JsonSchema schema)
                 if (state.ExporterOptions.TransformSchemaNode != null)
                 {
                     // Prime the schema for invocation by the JsonNode transformer.
-                    schema.ExporterContext = state.CreateContext(typeInfo, propertyInfo);
+                    schema.ExporterContext = state.CreateContext(typeInfo, propertyInfo, parentPolymorphicTypeInfo);
                 }
 
                 return schema;
@@ -454,9 +454,9 @@ public void PopGeneratedType()
                 _generationStack.RemoveAt(_generationStack.Count - 1);
             }
 
-            public JsonSchemaExporterContext CreateContext(JsonTypeInfo typeInfo, JsonPropertyInfo? propertyInfo)
+            public JsonSchemaExporterContext CreateContext(JsonTypeInfo typeInfo, JsonPropertyInfo? propertyInfo, JsonTypeInfo? baseTypeInfo)
             {
-                return new JsonSchemaExporterContext(typeInfo, propertyInfo, _currentPath.ToArray());
+                return new JsonSchemaExporterContext(typeInfo, propertyInfo, baseTypeInfo, _currentPath.ToArray());
             }
 
             private static string FormatJsonPointer(List<string> currentPathList, int depth)

src/libraries/System.Text.Json/src/System/Text/Json/Schema/JsonSchemaExporterContext.cs

@@ -12,10 +12,15 @@ public readonly struct JsonSchemaExporterContext
     {
         private readonly string[] _path;
 
-        internal JsonSchemaExporterContext(JsonTypeInfo typeInfo, JsonPropertyInfo? propertyInfo, string[] path)
+        internal JsonSchemaExporterContext(
+            JsonTypeInfo typeInfo,
+            JsonPropertyInfo? propertyInfo,
+            JsonTypeInfo? baseTypeInfo,
+            string[] path)
         {
             TypeInfo = typeInfo;
             PropertyInfo = propertyInfo;
+            BaseTypeInfo = baseTypeInfo;
             _path = path;
         }
 
@@ -29,6 +34,11 @@ internal JsonSchemaExporterContext(JsonTypeInfo typeInfo, JsonPropertyInfo? prop
         /// </summary>
         public JsonPropertyInfo? PropertyInfo { get; }
 
+        /// <summary>
+        /// Gets the <see cref="JsonTypeInfo"/> for polymorphic base type if the schema is being generated for a derived type.
+        /// </summary>
+        public JsonTypeInfo? BaseTypeInfo { get; }
+
         /// <summary>
         /// The path to the current node in the generated JSON schema.
         /// </summary>

src/libraries/System.Text.Json/tests/Common/JsonSchemaExporterTests.cs

@@ -112,6 +112,39 @@ public void JsonSerializerOptions_SmallMaxDepth_ThrowsInvalidOperationException(
             Assert.Contains("depth", ex.Message);
         }
 
+        [Theory]
+        [InlineData(typeof(int))]
+        [InlineData(typeof(string))]
+        [InlineData(typeof(SimplePoco))]
+        [InlineData(typeof(DiscriminatedUnion))]
+        public void JsonSchemaExporterContext_BaseTypeInfo_ReturnsExpectedValue(Type type)
+        {
+            bool isCallbackInvoked = false;
+            JsonSerializerOptions options = Serializer.DefaultOptions;
+            JsonSchemaExporterOptions exporterOptions = new()
+            {
+                TransformSchemaNode = (ctx, node) =>
+                {
+                    if (typeof(DiscriminatedUnion).IsAssignableFrom(ctx.TypeInfo.Type) &&
+                       typeof(DiscriminatedUnion) != ctx.TypeInfo.Type)
+                    {
+                        Assert.NotNull(ctx.BaseTypeInfo);
+                        Assert.Equal(typeof(DiscriminatedUnion), ctx.BaseTypeInfo.Type);
+                    }
+                    else
+                    {
+                        Assert.Null(ctx.BaseTypeInfo);
+                    }
+
+                    isCallbackInvoked = true;
+                    return node;
+                }
+            };
+
+            options.GetJsonSchemaAsNode(type, exporterOptions);
+            Assert.True(isCallbackInvoked);
+        }
+
         [Fact]
         public void ReferenceHandlePreserve_Enabled_ThrowsNotSupportedException()
         {