add onp.NDim0, which, unlike NDim, excludes 0

jorenham · Feb 5, 2025 · 003b3d8 · 003b3d8
1 parent 7bf9e72
commit 003b3d8
Showing 1 changed file with 10 additions and 20 deletions.
diff --git a/optype/numpy/_shape.py b/optype/numpy/_shape.py
@@ -6,12 +6,11 @@
 else:
     from typing_extensions import TypeAliasType, TypeVar, Unpack
 
-from ._compat import NP20
 
 __all__ = [
     "AtLeast0D", "AtLeast1D", "AtLeast2D", "AtLeast3D",
     "AtMost0D", "AtMost1D", "AtMost2D", "AtMost3D",
-    "NDim",
+    "NDim", "NDim0",
 ]  # fmt: skip
 
 
@@ -80,21 +79,12 @@ def __dir__() -> list[str]:
 
 # ND
 
-if NP20:
-    NDim = TypeAliasType(
-        "NDim",
-        Literal[
-            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
-            17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
-            33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48,
-            49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64,
-        ],
-    )  # fmt: skip
-else:
-    NDim = TypeAliasType(
-        "NDim",
-        Literal[
-            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
-            17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
-        ],
-    )  # fmt: skip
+# NOTE: on `numpy<2` this was at most 32
+_NDim0: TypeAlias = Literal[
+    1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
+    17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32,
+    33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48,
+    49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64,
+]  # fmt: skip
+NDim0 = TypeAliasType("NDim0", _NDim0)
+NDim = TypeAliasType("NDim", Literal[0, _NDim0])