Make a type-checking-only ObjectArray type

[inducer]

@alexfikl @majosm What do you think of this?

Object arrays are super-duper in the way of the type checker understanding hierarchies of array containers. I can't for the life of me figure out how to teach numpy to take any old object as its dtype:

x: np.ndarray[tuple[int], np.dtype[ArithmeticExpression]]

results in

  /home/andreas/src/pytools/pytools/obj_array.py:90:36 - error: Type "ArithmeticExpression" cannot be assigned to type variable "_ScalarT_co@dtype"
    Type "ArithmeticExpression" is not assignable to upper bound "generic[Any]" for type variable "_ScalarT_co@dtype"
      Type "ArithmeticExpression" is not assignable to type "generic[Any]"
        "ExpressionNode" is not assignable to "generic[Any]" (reportInvalidTypeArguments)

In addition, shape-precise typing for numpy remains a "next year, for sure" type of thing. Where as in this reality, we can have all of that, today! 🪄

I somehow don't see much downside. Especially since it's type-checking only, it can be ripped out without damaging anybody's ability to run code. But maybe I'm too high on type checker juice here. Reality check please? 🙂

https://youtu.be/_oNgyUAEv0Q?si=4YId-5UMWjmQ4MbE&t=52

[inducer]

This would also help with some ugly bits in inducer/arraycontext#322, such as the overlap between array containers and numpy arrays that's causing overlap in from_numpy's signature.

[inducer]

Alright, since nobody jumped at the chance to provide a voice of reason, I'm plotting to proceed with this.

[alexfikl]

Left a few nitpick comments, but overall this looks reasonable to me!

[inducer]

I'm hitting a bit of a conundrum with this.

On the one hand, I would like ObjectArray[ShapeT, T] to be covariant on T (so that if Cat is a subtype of Animal, ObjectArray[ShapeT, Cat] is a subtype of ObjectArray[ShapeT, Animal]). This is relevant in arraycontext where, without this, ArithArrayContainer is no longer a subtype of ArrayContainer (i.e. you can't pass Arith where general is allowed).

On the other hand, covariance prohibits arithmetic of the form

def __mul__(self, other: T, /) -> Self: ...

because other is a sink, and ObjectArray[ShapeT, Cat] could not honestly be a subtype of ObjectArray[ShapeT, Animal] because it would need to keep accepting Animals.

Options:

• Shut up pyright about the non-covariance of the arithmetic. Unless I'm missing something, this means pyright will accept less downstream code than prescribed by strict subtyping. I think that's OK.
• Declare ArrayContainer to consist of ObjectArray[ShapeT, ...ArithContainer...]. This just looks decidedly weird to me.
• Others I have not thought of?

As you can tell, I'm leaning towards the first. I'd be happy to hear dissenting comments though.

[alexfikl]

Shut up pyright about the non-covariance of the arithmetic. Unless I'm missing something, this means pyright will accept less downstream code than prescribed by strict subtyping. I think that's OK.

I'm not sure I understand all the ramifications of this, since I haven't been steeped into the current typing run, but this also sounds like the best option to me. It seems more useful for callers to be able to pass in cats when animals are expected in various functions than have the arithmetic methods be technically correct.

When would pyright complain downstream for the arithmetic operations? I can see it complains here because the covariant T isn't allowed in the method argument like that.

Others I have not thought of?

Very stupid (untested) thought, but would something like

U = TypeVar("U", bound=T)

def __add__(self, other: U) -> Self: ...

do anything useful?

EDIT: A more believable hack:

class SupportsArithmetic(Protocol):
	def __add__(self, other: Any) -> Any: ...
	
class ObjectArray(Generic[ShapeT, T]):
	def __add__(self, other: SupportsArithmetic) -> Self: ...

? Nicely duplicates everything and not sure the Any there does anyone any favors..

[inducer]

(untested) thought, but would something like

Interesting thought, tried it out on the pyright playground. In short: no dice.

EDIT: A more believable hack:

Nah, not loving the Anys.

---

In the meantime, I've gathered a bit of experience now using this in pymbolic and arraycontext, and it seems to do its job as hoped. I think I'm reasonably comfortable with this now.