PEP 696 input type default

[nstarman]

Thinking about this more, we could even target PEP 696 and define:

ArrayT = TypeVar('ArrayT', bound=Array)
ArrayLikeT = TypeVar('ArrayLikeT', default=float | ArrayT)

class Cosmology(Generic[ArrayT, ArrayLikeT], Protocol): ...

Cosmologies using numpy arrays could then accept numpy.typing.ArrayLike inputs, and e.g. cosmo.comoving_distance([z1, z2]) would remain legal, which is a big plus in my mind.

Originally posted by @ntessore in https://github.com/cosmology-api/cosmology-api/issues/27#issuecomment-1421654691

[nstarman]

I think the way users will expect is something like this:

ArrayT = OutputT = TypeVar('ArrayT', bound=Array)
InputT = TypeVar("InputT", default=float | ArrayT)

class Cosmology(Protocol[InputT, OutputT): ...

Unfortunately, I think this will always require both type inputs since InputT gets bound before OutputT. We can fix this with https://mypy.readthedocs.io/en/stable/extending_mypy.html by having Cosmology[OutputT] map to Cosmology[OutputT, OutputT]. But we might need a plugin for pyright as well... does pyright even support plugins?

[ntessore]

I'm not sure I follow. What we need to prescribe in the first instance is the array type that the cosmology instance works with: Cosmology[ArrayT].

Then we have a situation where the input type can be array-like instead of being the exact array type. So we add a second type to the generic for that: Cosmology[ArrayT, ArrayLikeT].

Isn't that a good interface up to this point?

The point of the issue is then merely that we could go one step further and leverage the PEP 696 semantics to give a sensible default value to ArrayLikeT.

We could also not do that at all, because defining cosmology protocol implementations is really only something that vendor libraries do, not users in their day to day work.

[nstarman]

My point is that the expectation is probably

Protocol[InputT, OutputT] not Protocol[OutputT, InputT].

But InputT with a default can't precede OutputT, without. At least not normally. We can make Protocol[InputT, OutputT] work as expected for mypy with https://mypy.readthedocs.io/en/stable/extending_mypy.html, but I'm not sure if that will work for pyright.

[ntessore]

Yeah, I can see that if we phrased it in terms of input and output types, but I think in terms of array and array-like types it's fine.

In any case I think we can separate this into three issues and discussions

• [x] support float input
• [x] support array-like input
• [ ] support PEP 696 for default array-like type

What I meant to say above is that I think it would even be fine to support array-like input without PEP 696 (or other) default behaviour, because it's something that needs to be specified once per cosmology implementation.