feat(tool): Update the way we handle zod and enums

TypeContractor.Tests/TypeScript/ApiClientWriterTests.cs

@@ -238,6 +238,32 @@ [new EndpointParameter("year", typeof(int), null, false, false, true, false, fal
 			.And.Contain("url.searchParams.append('reportType', reportType.toString());");
 	}
 
+	[Fact]
+	public void Handles_Enum_Return_With_Zod_Schema()
+	{
+		// Arrange
+		var outputTypes = new List<OutputType>
+		{
+			_converter.Convert(typeof(ReportType))
+		};
+
+		var apiClient = new ApiClient("TestClient", "TestController", "test", null);
+		apiClient.AddEndpoint(new ApiClientEndpoint("getReportType", "report-type", EndpointMethod.GET, typeof(ReportType), typeof(ReportType), false, [], null));
+
+		// Act
+		var result = Sut.Write(apiClient, outputTypes, _converter, true, _templateFn, Casing.Pascal);
+
+		// Assert
+		var file = File.ReadAllText(result).Trim();
+		file.Should()
+			.NotBeEmpty()
+			.And.Contain("import { z } from 'zod';")
+			.And.Contain("import { ReportType, ReportTypeEnum } from '~/type-contractor/tests/type-script/report-type';")
+			.And.Contain("public async getReportType(cancellationToken: AbortSignal = null): Promise<ReportType> {")
+			.And.Contain("return await response.parseJson<ReportType>(ReportTypeEnum);")
+			.And.NotContain("ReportTypeSchema");
+	}
+
 	[Fact]
 	public void Unpacks_Complex_Object_To_Query()
 	{

TypeContractor.Tests/TypeScript/TypeScriptWriterTests.cs

@@ -296,9 +296,9 @@ public void Writes_Zod_Schema_With_Enum_Reference()
 		file.Should()
 			.NotBeEmpty()
 			.And.Contain("import { z } from 'zod';")
-			.And.Contain("import { ObsoleteEnum } from './ObsoleteEnum';")
+			.And.Contain("import { ObsoleteEnum, ObsoleteEnumEnum } from './ObsoleteEnum';")
 			.And.Contain("export const TypeWithEnumSchema = z.object({")
-			.And.Contain("  status: z.enum(ObsoleteEnum),")
+			.And.Contain("  status: ObsoleteEnumEnum,")
 			.And.Contain("});");
 	}
 
@@ -316,9 +316,9 @@ public void Writes_Zod_Schema_With_Nullable_Enum()
 		file.Should()
 			.NotBeEmpty()
 			.And.Contain("import { z } from 'zod';")
-			.And.Contain("import { ObsoleteEnum } from './ObsoleteEnum';")
+			.And.Contain("import { ObsoleteEnum, ObsoleteEnumEnum } from './ObsoleteEnum';")
 			.And.Contain("export const TypeWithNullableEnumSchema = z.object({")
-			.And.Contain("  status: z.enum(ObsoleteEnum).nullable(),")
+			.And.Contain("  status: ObsoleteEnumEnum.nullable(),")
 			.And.Contain("});");
 	}
 

TypeContractor/Templates/EndpointTemplateDto.cs

@@ -8,6 +8,7 @@ public record EndpointTemplateDto(
 	string HttpMethodUppercase,
 	string ReturnType,
 	string? UnwrappedReturnType,
+	string? UnwrappedReturnSchema,
 	bool EnumerableReturnType,
 	string? MappedReturnType,
 	bool? EnumerableMappedReturnType,

TypeContractor/Templates/aurelia.hbs

@@ -66,7 +66,7 @@ export class {{Name}} {
     const response = await this.http.{{HttpMethod}}(`${url.pathname}${url.search}`.slice(1), {{#if RequiresBody}}{{#if BodyParameter}}json({{BodyParameter}}){{else}}null{{/if}}, {{/if}}{ signal: cancellationToken });
     {{#if ../BuildZodSchema}}
     {{#if UnwrappedReturnType}}
-    return await response.parseJson<{{UnwrappedReturnType}}{{#if EnumerableReturnType}}[]{{/if}}>({{UnwrappedReturnType}}Schema{{#if EnumerableReturnType}}.array(){{/if}});
+    return await response.parseJson<{{UnwrappedReturnType}}{{#if EnumerableReturnType}}[]{{/if}}>({{UnwrappedReturnSchema}}{{#if EnumerableReturnType}}.array(){{/if}});
     {{else}}
     {{#if MappedReturnType}}
     return await response.parseJson(z.{{MappedReturnType}}(){{#if EnumerableMappedReturnType}}.array(){{/if}});

TypeContractor/TypeScript/ApiClientWriter.cs

@@ -24,7 +24,7 @@
 	[GeneratedRegex(@"([^\$])\{([A-Za-z0-9]+)\}")]
 	private static partial Regex RouteParameterRegex();
 
 	public string Write(ApiClient apiClient, IEnumerable<OutputType> allTypes, TypeScriptConverter converter, bool buildZodSchema, HandlebarsTemplate<object, ApiClientTemplateDto> template, Casing casing)
 	{
 		var _builder = new StringBuilder();
 		ArgumentNullException.ThrowIfNull(apiClient);
@@ -112,6 +112,11 @@
 			var targetType = buildZodSchema && endpoint.ReturnType is not null
 				? converter.GetDestinationType(endpoint.ReturnType, endpoint.ReturnType.CustomAttributes, false, TypeChecks.IsNullable(endpoint.ReturnType))
 				: null;
+			var unwrappedReturnSchema = endpoint.UnwrappedReturnType is null
+				? null
+				: endpoint.UnwrappedReturnType.IsEnum
+					? $"{endpoint.UnwrappedReturnType.Name}Enum"
+					: $"{endpoint.UnwrappedReturnType.Name}Schema";
 
 			endpoints.Add(new EndpointTemplateDto(
 				endpoint.Name,
@@ -121,6 +126,7 @@
 				method,
 				returnType,
 				endpoint.UnwrappedReturnType?.Name,
+				unwrappedReturnSchema,
 				endpoint.EnumerableReturnType,
 				targetType?.TypeName,
 				targetType?.IsArray,
@@ -171,19 +177,14 @@
 	private List<string> BuildImports(IEnumerable<ApiClientEndpoint> endpoints, IEnumerable<OutputType> allTypes, TypeScriptConverter converter, bool buildZodSchema)
 	{
 		var imports = new List<string>();
-		var needZodLibrary = false;
 
 		foreach (var endpoint in endpoints)
 		{
 			var returnType = endpoint.ReturnType is null
 				? null
 				: converter.GetDestinationType(endpoint.ReturnType, endpoint.ReturnType.CustomAttributes, false, TypeChecks.IsNullable(endpoint.ReturnType));
 
-			if (returnType is not null && returnType.IsBuiltin)
-			{
-				needZodLibrary = true;
-			}
-			else if (returnType is not null && !returnType.IsBuiltin)
+			if (returnType is not null && !returnType.IsBuiltin)
 			{
 				var importTypes = new List<string> { returnType.ImportType };
 				if (buildZodSchema)
@@ -227,7 +228,7 @@
 			}
 		}
 
-		if (buildZodSchema && needZodLibrary)
+		if (buildZodSchema)
 			imports.Insert(0, ZodSchemaWriter.LibraryImport);
 
 		return imports.Distinct().ToList();

TypeContractor/TypeScript/ZodSchemaWriter.cs

@@ -39,10 +39,6 @@ public static void Write(OutputType type, IEnumerable<OutputType> allTypes, Stri
 			if (TypeChecks.IsNullable(sourceType))
 				sourceType = TypeChecks.GetGenericType(sourceType);
 
-			// We don't currently import any schema for enums
-			if (sourceType.IsEnum)
-				return null;
-
 			var suffix = sourceType.IsEnum ? "Enum" : "Schema";
 			var name = GetImportType(import.InnerSourceType, sourceType);
 			return $"{name}{suffix}";
@@ -53,6 +49,10 @@ public static void Write(OutputType type, IEnumerable<OutputType> allTypes, Stri
 			if (returnType.IsBuiltin)
 				return null;
 
+			var sourceType = returnType.InnerType ?? returnType.SourceType;
+			if (sourceType?.IsEnum == true)
+				return $"{returnType.ImportType}Enum";
+
 			return $"{returnType.ImportType}Schema";
 		}
 
@@ -69,12 +69,12 @@ private static string GetImportType(Type? innerSourceType, Type sourceType)
 		private static string? GetZodOutputType(OutputProperty property, IEnumerable<OutputType> allTypes)
 		{
 			if (!property.IsBuiltin && property.SourceType.IsEnum)
-				return $"z.enum({property.SourceType.Name})";
+				return $"{property.SourceType.Name.Split('`').First()}Enum";
 			else if (!property.IsBuiltin && property.IsNullable && property.SourceType.IsGenericType)
 			{
 				var sourceType = TypeChecks.GetGenericType(property.SourceType);
 				if (sourceType.IsEnum)
-					return $"z.enum({sourceType.Name}).nullable()";
+					return $"{sourceType.Name.Split('`').First()}Enum.nullable()";
 			}
 
 			string? output;
@@ -87,19 +87,23 @@ private static string GetImportType(Type? innerSourceType, Type sourceType)
 			}
 			else if (!property.IsBuiltin && !property.IsNullable)
 			{
-				var name = property.InnerSourceType?.Name ?? property.SourceType.Name;
+				var sourceType = property.InnerSourceType ?? property.SourceType;
+				var name = sourceType.Name;
 				name = name.Split('`').First();
-				output = $"{name}Schema";
+				var suffix = sourceType.IsEnum ? "Enum" : "Schema";
+				output = $"{name}{suffix}";
 			}
 			else if (property.IsBuiltin)
 			{
 				output = GetZodOutputType(property.SourceType, allTypes);
 			}
 			else if (!property.IsBuiltin && property.IsArray && property.InnerSourceType is not null)
 			{
-				var name = property.InnerSourceType.Name;
+				var sourceType = property.InnerSourceType;
+				var name = sourceType.Name;
 				name = name.Split('`').First();
-				output = $"{name}Schema";
+				var suffix = sourceType.IsEnum ? "Enum" : "Schema";
+				output = $"{name}{suffix}";
 			}
 			else
 			{
@@ -142,7 +146,8 @@ private static string GetImportType(Type? innerSourceType, Type sourceType)
 			else if (allTypes.Any(x => IsOfType(sourceType, x.ContractedType.Type)))
 			{
 				var targetType = allTypes.First(x => IsOfType(sourceType, x.ContractedType.Type));
-				output = $"{targetType.Name}Schema";
+				var suffix = targetType.IsEnum ? "Enum" : "Schema";
+				output = $"{targetType.Name}{suffix}";
 			}
 			else
 				output = "z.any()";