Skip to content

Commit 404d4a1

Browse files
committed
Introduce "*Tc" constrained versions of the ArrayOr* type variables
1 parent f307a04 commit 404d4a1

File tree

1 file changed

+20
-1
lines changed

1 file changed

+20
-1
lines changed

arraycontext/context.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -226,20 +226,39 @@ def __rtruediv__(self, other: Self | ScalarLike) -> Self: ...
226226
Scalar = ScalarLike
227227
ScalarLikeT = TypeVar("ScalarLikeT", bound=ScalarLike)
228228

229+
# NOTE: I'm kind of not sure about the *Tc versions of these type variables.
230+
# mypy seems better at understanding arithmetic performed on the *Tc versions
231+
# than the *T, versions, whereas pyright doesn't seem to care.
232+
#
233+
# This issue seems to be part of it:
234+
# https://github.com/python/mypy/issues/18203
235+
# but there is likely other stuff lurking.
236+
#
237+
# For now, they're purposefully not in the main arraycontext.* name space.
229238
ArrayT = TypeVar("ArrayT", bound=Array)
230239
ArrayOrScalar: TypeAlias = "Array | ScalarLike"
231240
ArrayOrContainer: TypeAlias = "Array | ArrayContainer"
232241
ArrayOrArithContainer: TypeAlias = "Array | ArithArrayContainer"
233242
ArrayOrContainerT = TypeVar("ArrayOrContainerT", bound=ArrayOrContainer)
243+
ArrayOrContainerTc = TypeVar("ArrayOrContainerTc",
244+
Array, "ArrayContainer", "ArithArrayContainer")
234245
ArrayOrArithContainerT = TypeVar("ArrayOrArithContainerT", bound=ArrayOrArithContainer)
246+
ArrayOrArithContainerTc = TypeVar("ArrayOrArithContainerTc",
247+
Array, "ArithArrayContainer")
235248
ArrayOrContainerOrScalar: TypeAlias = "Array | ArrayContainer | ScalarLike"
236249
ArrayOrArithContainerOrScalar: TypeAlias = "Array | ArithArrayContainer | ScalarLike"
237250
ArrayOrContainerOrScalarT = TypeVar(
238251
"ArrayOrContainerOrScalarT",
239252
bound=ArrayOrContainerOrScalar)
240253
ArrayOrArithContainerOrScalarT = TypeVar(
241254
"ArrayOrArithContainerOrScalarT",
242-
bound=ArrayOrContainerOrScalar)
255+
bound=ArrayOrArithContainerOrScalar)
256+
ArrayOrContainerOrScalarTc = TypeVar(
257+
"ArrayOrContainerOrScalarTc",
258+
ScalarLike, Array, "ArrayContainer", "ArithArrayContainer")
259+
ArrayOrArithContainerOrScalarTc = TypeVar(
260+
"ArrayOrArithContainerOrScalarTc",
261+
ScalarLike, Array, "ArithArrayContainer")
243262

244263

245264
ContainerOrScalarT = TypeVar("ContainerOrScalarT", bound="ArrayContainer | ScalarLike")

0 commit comments

Comments
 (0)