diff --git a/Sources/AnyLanguageModelMacros/GenerableMacro.swift b/Sources/AnyLanguageModelMacros/GenerableMacro.swift index 31ff2f8..f22c640 100644 --- a/Sources/AnyLanguageModelMacros/GenerableMacro.swift +++ b/Sources/AnyLanguageModelMacros/GenerableMacro.swift @@ -168,26 +168,127 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { return GuideInfo(description: nil, guides: [], pattern: nil) } - private static func isDictionaryType(_ type: String) -> Bool { - let trimmed = type.trimmingCharacters(in: .whitespacesAndNewlines) - return trimmed.hasPrefix("[") && trimmed.contains(":") && trimmed.hasSuffix("]") + private static func topLevelColonIndex(in text: String) -> String.Index? { + var totalDepth = 0 + + for index in text.indices { + switch text[index] { + case "[": + totalDepth += 1 + case "]": + totalDepth -= 1 + case "<": + totalDepth += 1 + case ">": + totalDepth -= 1 + case "(": + totalDepth += 1 + case ")": + totalDepth -= 1 + case ":" where totalDepth == 0: + return index + default: + break + } + + if totalDepth < 0 { + return nil + } + } + + return nil } private static func extractDictionaryTypes(_ type: String) -> (key: String, value: String)? { let trimmed = type.trimmingCharacters(in: .whitespacesAndNewlines) - guard trimmed.hasPrefix("[") && trimmed.hasSuffix("]") && trimmed.contains(":") else { + guard trimmed.hasPrefix("[") && trimmed.hasSuffix("]") else { return nil } let inner = String(trimmed.dropFirst().dropLast()) - let parts = inner.split(separator: ":", maxSplits: 1).map { - $0.trimmingCharacters(in: .whitespacesAndNewlines) + guard let colonIndex = topLevelColonIndex(in: inner) else { + return nil + } + + let key = inner[.. Bool { + extractDictionaryTypes(type) != nil + } + + private static func baseTypeName(_ type: String) -> String { + let trimmed = type.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.hasSuffix("?") { + return String(trimmed.dropLast()).trimmingCharacters(in: .whitespacesAndNewlines) + } + return trimmed + } + + private static func arrayElementType(from type: String) -> String? { + let trimmed = type.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.hasPrefix("[") && trimmed.hasSuffix("]") { + let inner = String(trimmed.dropFirst().dropLast()) + guard topLevelColonIndex(in: inner) == nil else { + return nil + } + return inner.trimmingCharacters(in: .whitespacesAndNewlines) } + if trimmed.hasPrefix("Array<") && trimmed.hasSuffix(">") { + return String(trimmed.dropFirst("Array<".count).dropLast()) + .trimmingCharacters(in: .whitespacesAndNewlines) + } + return nil + } - guard parts.count == 2 else { return nil } + private static let primitiveTypes: Set = [ + "String", + "Int", + "Double", + "Float", + "Bool", + "Decimal", + ] + + private static func partiallyGeneratedTypeName(for type: String) -> String { + partiallyGeneratedTypeName(for: type, preserveOptional: false) + } - return (key: parts[0], value: parts[1]) + private static func partiallyGeneratedTypeName(for type: String, preserveOptional: Bool) -> String { + let trimmed = type.trimmingCharacters(in: .whitespacesAndNewlines) + if preserveOptional { + var normalized = trimmed + var optionalCount = 0 + while normalized.hasSuffix("?") { + normalized = String(normalized.dropLast()) + optionalCount += 1 + } + if optionalCount > 1 { + return "\(partiallyGeneratedTypeName(for: normalized, preserveOptional: false))?" + } + if optionalCount == 1 { + return "\(partiallyGeneratedTypeName(for: normalized, preserveOptional: true))?" + } + } + + let baseType = baseTypeName(trimmed) + if primitiveTypes.contains(baseType) || isDictionaryType(baseType) { + return baseType + } + if let elementType = arrayElementType(from: baseType) { + let elementPartial = partiallyGeneratedTypeName(for: elementType, preserveOptional: true) + return "[\(elementPartial)]" + } + return "\(baseType).PartiallyGenerated" } private static func getDefaultValue(for type: String) -> String { @@ -383,19 +484,20 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { propertyName: String, propertyType: String ) -> String { - switch propertyType { - case "String", "String?": + let baseType = baseTypeName(propertyType) + + switch baseType { + case "String": return "self.\(propertyName) = try? properties[\"\(propertyName)\"]?.value(String.self)" - case "Int", "Int?": + case "Int": return "self.\(propertyName) = try? properties[\"\(propertyName)\"]?.value(Int.self)" - case "Double", "Double?": + case "Double": return "self.\(propertyName) = try? properties[\"\(propertyName)\"]?.value(Double.self)" - case "Float", "Float?": + case "Float": return "self.\(propertyName) = try? properties[\"\(propertyName)\"]?.value(Float.self)" - case "Bool", "Bool?": + case "Bool": return "self.\(propertyName) = try? properties[\"\(propertyName)\"]?.value(Bool.self)" default: - let baseType = propertyType.replacingOccurrences(of: "?", with: "") if isDictionaryType(baseType) { return """ if let value = properties[\"\(propertyName)\"] { @@ -404,10 +506,21 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { self.\(propertyName) = nil } """ + } else if let elementType = arrayElementType(from: baseType) { + let elementPartial = partiallyGeneratedTypeName(for: elementType, preserveOptional: true) + let arrayPartial = "[\(elementPartial)]" + return """ + if let value = properties[\"\(propertyName)\"] { + self.\(propertyName) = try? \(arrayPartial)(value) + } else { + self.\(propertyName) = nil + } + """ } else { + let partialType = partiallyGeneratedTypeName(for: baseType) return """ if let value = properties[\"\(propertyName)\"] { - self.\(propertyName) = try? \(propertyType)(value) + self.\(propertyName) = try? \(partialType)(value) } else { self.\(propertyName) = nil } @@ -676,12 +789,8 @@ public struct GenerableMacro: MemberMacro, ExtensionMacro { properties: [PropertyInfo] ) -> DeclSyntax { let optionalProperties = properties.map { prop in - let propertyType = prop.type - if propertyType.hasSuffix("?") { - return "public let \(prop.name): \(propertyType)" - } else { - return "public let \(prop.name): \(propertyType)?" - } + let partialType = partiallyGeneratedTypeName(for: prop.type) + return "public let \(prop.name): \(partialType)?" }.joined(separator: "\n ") let propertyExtractions = properties.map { prop in diff --git a/Tests/AnyLanguageModelTests/GenerableMacroTests.swift b/Tests/AnyLanguageModelTests/GenerableMacroTests.swift index 3607b6c..7134821 100644 --- a/Tests/AnyLanguageModelTests/GenerableMacroTests.swift +++ b/Tests/AnyLanguageModelTests/GenerableMacroTests.swift @@ -34,6 +34,69 @@ struct TestArguments { var age: Int } +@Generable +private struct ArrayItem { + @Guide(description: "A name") + var name: String +} + +@Generable +private struct ArrayContainer { + @Guide(description: "Items", .count(2)) + var items: [ArrayItem] +} + +@Generable +private struct PrimitiveContainer { + @Guide(description: "A title") + var title: String + + @Guide(description: "A count") + var count: Int +} + +@Generable +private struct PrimitiveArrayContainer { + @Guide(description: "Names", .count(2)) + var names: [String] +} + +@Generable +private struct OptionalArrayContainer { + @Guide(description: "Optional names", .count(2)) + var names: [String]? +} + +@Generable +private struct NestedArrayContainer { + @Guide(description: "Nested items", .count(2)) + var items: [[ArrayItem]] +} + +@Generable +private struct OptionalPrimitiveContainer { + @Guide(description: "Optional title") + var title: String? + + @Guide(description: "Optional count") + var count: Int? + + @Guide(description: "Optional flag") + var flag: Bool? +} + +@Generable +private struct OptionalItemContainer { + @Guide(description: "Optional item") + var item: ArrayItem? +} + +@Generable +private struct OptionalItemsContainer { + @Guide(description: "Optional items", .count(2)) + var items: [ArrayItem]? +} + @Suite("Generable Macro") struct GenerableMacroTests { @Test("@Guide description with multiline string") @@ -140,6 +203,188 @@ struct GenerableMacroTests { #expect(args.age == 25) #expect(args.asPartiallyGenerated().id == generationID) } + + @Test("Array properties use partially generated element types") + func arrayPropertiesUsePartiallyGeneratedElements() throws { + let content = GeneratedContent( + properties: [ + "items": GeneratedContent( + kind: .array([ + GeneratedContent(properties: ["name": "Alpha"]), + GeneratedContent(properties: ["name": "Beta"]), + ]) + ) + ] + ) + + let container = try ArrayContainer(content) + let partial = container.asPartiallyGenerated() + #expect(partial.items?.count == 2) + #expect(partial.items?.first?.name == "Alpha") + } + + @Test("Primitive properties use concrete partial types") + func primitivePropertiesRemainUnchanged() throws { + let content = GeneratedContent( + properties: [ + "title": "Hello", + "count": 3, + ] + ) + + let container = try PrimitiveContainer(content) + let partial = container.asPartiallyGenerated() + #expect(partial.title == "Hello") + #expect(partial.count == 3) + } + + @Test("Array primitives use concrete element types") + func arrayPrimitivesRemainConcrete() throws { + let content = GeneratedContent( + properties: [ + "names": GeneratedContent( + kind: .array([ + GeneratedContent("Alpha"), + GeneratedContent("Beta"), + ]) + ) + ] + ) + + let container = try PrimitiveArrayContainer(content) + let partial = container.asPartiallyGenerated() + #expect(partial.names?.count == 2) + #expect(partial.names?.first == "Alpha") + } + + @Test("Optional primitive arrays remain concrete") + func optionalPrimitiveArraysRemainConcrete() throws { + let content = GeneratedContent( + properties: [ + "names": GeneratedContent( + kind: .array([ + GeneratedContent("Alpha"), + GeneratedContent("Beta"), + ]) + ) + ] + ) + + let container = try OptionalArrayContainer(content) + let partial = container.asPartiallyGenerated() + #expect(partial.names?.count == 2) + #expect(partial.names?.first == "Alpha") + } + + @Test("Nested arrays of generable types are handled") + func nestedArraysGenerateNestedPartialTypes() throws { + let content = GeneratedContent( + properties: [ + "items": GeneratedContent( + kind: .array([ + GeneratedContent( + kind: .array([ + GeneratedContent(properties: ["name": "Alpha"]), + GeneratedContent(properties: ["name": "Beta"]), + ]) + ), + GeneratedContent( + kind: .array([ + GeneratedContent(properties: ["name": "Gamma"]) + ]) + ), + ]) + ) + ] + ) + + let container = try NestedArrayContainer(content) + let partial = container.asPartiallyGenerated() + #expect(partial.items?.count == 2) + #expect(partial.items?.first?.count == 2) + #expect(partial.items?.first?.first?.name == "Alpha") + #expect(partial.items?.last?.first?.name == "Gamma") + } + + @Test("Optional primitive properties are handled") + func optionalPrimitivePropertiesHandled() throws { + let content = GeneratedContent( + properties: [ + "title": "Hello", + "count": 3, + "flag": true, + ] + ) + + let container = try OptionalPrimitiveContainer(content) + let partial = container.asPartiallyGenerated() + #expect(partial.title == "Hello") + #expect(partial.count == 3) + #expect(partial.flag == true) + } + + @Test("Optional generable properties are handled") + func optionalGenerableItemBecomesPartial() throws { + let content = GeneratedContent( + properties: [ + "item": GeneratedContent(properties: ["name": "Alpha"]) + ] + ) + + let container = try OptionalItemContainer(content) + let partial = container.asPartiallyGenerated() + #expect(partial.item?.name == "Alpha") + } + + @Test("Optional arrays of generable types are handled") + func optionalGenerableArraysTransformToPartialArrays() throws { + let content = GeneratedContent( + properties: [ + "items": GeneratedContent( + kind: .array([ + GeneratedContent(properties: ["name": "Alpha"]), + GeneratedContent(properties: ["name": "Beta"]), + ]) + ) + ] + ) + + let container = try OptionalItemsContainer(content) + let partial = container.asPartiallyGenerated() + #expect(partial.items?.count == 2) + #expect(partial.items?.first?.name == "Alpha") + } + + @Test("Missing optional properties become nil in partials") + func missingOptionalProperties() throws { + let content = GeneratedContent(properties: [:]) + + let primitive = try OptionalPrimitiveContainer(content).asPartiallyGenerated() + #expect(primitive.title == nil) + #expect(primitive.count == nil) + #expect(primitive.flag == nil) + + let item = try OptionalItemContainer(content).asPartiallyGenerated() + #expect(item.item == nil) + + let items = try OptionalItemsContainer(content).asPartiallyGenerated() + #expect(items.items == nil) + + let names = try OptionalArrayContainer(content).asPartiallyGenerated() + #expect(names.names == nil) + } + + @Test("Schema generation includes optional properties") + func schemaIncludesOptionalProperties() throws { + let schema = OptionalPrimitiveContainer.generationSchema + let encoder = JSONEncoder() + let jsonData = try encoder.encode(schema) + let jsonString = String(data: jsonData, encoding: .utf8) ?? "" + + #expect(jsonString.contains("\"title\"")) + #expect(jsonString.contains("\"count\"")) + #expect(jsonString.contains("\"flag\"")) + } } // MARK: - #Playground Usage