refactor: improve flatten function

- uses grok llm, so could be further issues yet

src/util/flatten.ts

@@ -13,43 +13,69 @@ type Flattened<T> = T extends Array<infer U>
 		? ExcludeFunctionsAndPrivate<T>
 		: T
 
-function flatten<T>(obj: T): Flattened<T> {
-	// Handle arrays
-	if (Array.isArray(obj)) {
-		return obj.map((item) => flatten(item)) as Flattened<T>
-	}
-
+function flatten<T>(obj: T, ...props: Record<string, boolean | string>[]): Flattened<T> {
 	// Handle null or non-objects
 	if (!obj || typeof obj !== 'object') {
 		return obj as Flattened<T>
 	}
 
-	// Handle Date objects
-	if (obj instanceof Date) {
-		return obj.toISOString() as Flattened<T>
+	// Handle arrays
+	if (Array.isArray(obj)) {
+		return obj.map((item) => flatten(item)) as Flattened<T>
 	}
 
-	const result: Record<string, unknown> = {}
+	// Get all non-private properties (exclude keys starting with '_')
+	const objProps = Object.keys(obj)
+		.filter((key) => !key.startsWith('_'))
+		.map((key) => ({ [key]: true }))
 
-	for (const key in obj) {
-		// biome-ignore lint/suspicious/noExplicitAny: Safe
-		if (key.startsWith('_') || typeof (obj as any)[key] === 'function') {
-			continue
-		}
+	// Merge provided props with object props safely
+	const mergedProps: Record<string, boolean | string> = Object.assign(
+		{},
+		...objProps, // Spread array of objects
+		...props, // Spread additional props
+	)
+
+	const out: Record<string, unknown> = {}
+
+	for (const [prop, newProp] of Object.entries(mergedProps)) {
+		if (!newProp) continue
+		const outputKey = newProp === true ? prop : newProp
 
 		// biome-ignore lint/suspicious/noExplicitAny: Safe
-		const value = (obj as any)[key]
-
-		if (value && typeof value === 'object' && 'toJSON' in value) {
-			result[key] = flatten(value)
-		} else if (Array.isArray(value)) {
-			result[key] = flatten(value)
-		} else {
-			result[key] = value
+		const element = (obj as any)[prop]
+		const elemIsObj = element && typeof element === 'object'
+		// biome-ignore lint/suspicious/noShadowRestrictedNames: Safe
+		const valueOf = elemIsObj && typeof element.valueOf === 'function' ? element.valueOf() : null
+		const hasToJSON = elemIsObj && typeof element.toJSON === 'function'
+
+		// Handle Date objects explicitly
+		if (element instanceof Date) {
+			out[outputKey] = element.toISOString()
+		}
+		// If it's an array, process each element
+		else if (Array.isArray(element)) {
+			out[outputKey] = element.map((elm) => (hasToJSON ? elm.toJSON() : flatten(elm)))
+		}
+		// If it has a toJSON method, use its result
+		else if (hasToJSON) {
+			out[outputKey] = element.toJSON()
+		}
+		// If it's an object with a primitive valueOf, use that
+		else if (elemIsObj && valueOf && typeof valueOf !== 'object') {
+			out[outputKey] = valueOf
+		}
+		// If it's a non-primitive object, flatten it
+		else if (elemIsObj) {
+			out[outputKey] = flatten(element)
+		}
+		// Otherwise, use the element as-is (primitives)
+		else {
+			out[outputKey] = element
 		}
 	}
 
-	return result as Flattened<T>
+	return out as Flattened<T>
 }
 
 export { flatten, type Flattened }