How to infer type from a schema generated by a generic function in zod?

I have a generic function which generated a very simple schema.

const LeafSchema = <T extends z.ZodTypeAny>(valueSchema: T) { 
    return z.object({ value: valueSchema })
}

Now I want to infer the type of this generated Schema, with using ReturnType, but also I have to use a generic, if not value will be any.

type Leaf<T extends z.ZodTypeAny> = z.infer<ReturnType<typeof LeafSchema<T>>>
const getLeafValue = <T extends z.ZodTypeAny>(leaf: Leaf<T>) => leaf.value

However, I still get any when calling getLeafValue with { 'value': 123 }. I understand that something wrong with my generics, but i'm not sure how to fix it. Sadly I haven't found much about generics with zod...

Heres the Playground

Solution

I got it working by using <T> with ZodType<T> instead of <T extends ZodAnyType> with <T>. Almost works as intended, as undefined will be return beside number in the getLeafValue function.

type Leaf<T> = z.infer<ReturnType<typeof LeafSchema<z.ZodType<T>>>>
type NumberLeaf = Leaf<number>
//    ^? { value: number }

const getLeafValue = <T>(leaf: Leaf<T>) => leaf.value
const test = { 'value': 123 }
const val = getLeafValue(test)
//     ^? { value: number | undefined }

Wrapping the z.infer with a Required<> would fix it in this example, but would make optional value (when present) also required. So, that only half the solution...

Playground