Define precise overloads for Array.reshape

Author: inducer

pytato/array.py

@@ -199,6 +199,7 @@
     TYPE_CHECKING,
     Any,
     ClassVar,
+    Literal,
     Protocol,
     TypeAlias,
     TypeVar,
@@ -246,6 +247,8 @@
 ShapeType: TypeAlias = tuple[ShapeComponent, ...]
 ConvertibleToShape: TypeAlias = "ShapeComponent | Sequence[ShapeComponent]"
 
+OrderCF: TypeAlias = Literal["C"] | Literal["F"]
+
 # }}}
 
 
@@ -971,7 +974,13 @@ def astype(self, dtype: DTypeLike) -> Array:
             dtype=dtype,
         )
 
-    def reshape(self, *shape: int | Sequence[int], order: str = "C") -> Array:
+    @overload
+    def reshape(self, *shape: int, order: OrderCF = "C") -> Array: ...
+
+    @overload
+    def reshape(self, shape: tuple[int, ...], /, *, order: OrderCF = "C") -> Array: ...
+
+    def reshape(self, *shape: int | Sequence[int], order: OrderCF = "C") -> Array:
         import pytato as pt
         if len(shape) == 0:
             raise TypeError("reshape takes at least one argument (0 given)")