Skip to content

Commit d8b1cdb

Browse files
committed
TYP: improve type annotations of arrays to also express constraints on dimensionality
1 parent f3cc0a5 commit d8b1cdb

File tree

7 files changed

+171
-165
lines changed

7 files changed

+171
-165
lines changed

pyproject.toml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -167,6 +167,9 @@ convention = "numpy"
167167
[tool.ruff.lint.per-file-ignores]
168168
"tests/**" = ["D"]
169169

170+
[tool.ty.rules]
171+
unused-type-ignore-comment = "ignore"
172+
170173
[tool.mypy]
171174
python_version = "3.11"
172175
show_error_codes = true

src/gpgi/_boundaries.py

Lines changed: 50 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,23 @@
11
from collections.abc import Callable
22
from threading import Lock
3-
from typing import Any, Literal, cast
3+
from typing import Any, Literal
44

5-
from numpy.typing import NDArray
6-
7-
from gpgi._typing import FloatT
5+
from gpgi._typing import DH, F, FArray
86

97
BoundaryRecipeT = Callable[
108
[
11-
NDArray[FloatT],
12-
NDArray[FloatT],
13-
NDArray[FloatT],
14-
NDArray[FloatT],
15-
NDArray[FloatT],
16-
NDArray[FloatT],
17-
NDArray[FloatT],
18-
NDArray[FloatT],
9+
FArray[DH, F],
10+
FArray[DH, F],
11+
FArray[DH, F],
12+
FArray[DH, F],
13+
FArray[DH, F],
14+
FArray[DH, F],
15+
FArray[DH, F],
16+
FArray[DH, F],
1917
Literal["left", "right"],
2018
dict[str, Any],
2119
],
22-
NDArray[FloatT],
20+
FArray[DH, F],
2321
]
2422

2523

@@ -125,64 +123,64 @@ def __contains__(self, key: str) -> bool:
125123

126124
# basic recipes
127125
def open_boundary(
128-
same_side_active_layer: NDArray[FloatT],
129-
same_side_ghost_layer: NDArray[FloatT],
130-
opposite_side_active_layer: NDArray[FloatT],
131-
opposite_side_ghost_layer: NDArray[FloatT],
132-
weight_same_side_active_layer: NDArray[FloatT],
133-
weight_same_side_ghost_layer: NDArray[FloatT],
134-
weight_opposite_side_active_layer: NDArray[FloatT],
135-
weight_opposite_side_ghost_layer: NDArray[FloatT],
126+
same_side_active_layer: FArray[DH, F],
127+
same_side_ghost_layer: FArray[DH, F],
128+
opposite_side_active_layer: FArray[DH, F],
129+
opposite_side_ghost_layer: FArray[DH, F],
130+
weight_same_side_active_layer: FArray[DH, F],
131+
weight_same_side_ghost_layer: FArray[DH, F],
132+
weight_opposite_side_active_layer: FArray[DH, F],
133+
weight_opposite_side_ghost_layer: FArray[DH, F],
136134
side: Literal["left", "right"],
137135
metadata: dict[str, Any],
138-
) -> NDArray[FloatT]:
136+
) -> FArray[DH, F]:
139137
# return the active layer unchanged
140138
return same_side_active_layer
141139

142140

143141
def wall_boundary(
144-
same_side_active_layer: NDArray[FloatT],
145-
same_side_ghost_layer: NDArray[FloatT],
146-
opposite_side_active_layer: NDArray[FloatT],
147-
opposite_side_ghost_layer: NDArray[FloatT],
148-
weight_same_side_active_layer: NDArray[FloatT],
149-
weight_same_side_ghost_layer: NDArray[FloatT],
150-
weight_opposite_side_active_layer: NDArray[FloatT],
151-
weight_opposite_side_ghost_layer: NDArray[FloatT],
142+
same_side_active_layer: FArray[DH, F],
143+
same_side_ghost_layer: FArray[DH, F],
144+
opposite_side_active_layer: FArray[DH, F],
145+
opposite_side_ghost_layer: FArray[DH, F],
146+
weight_same_side_active_layer: FArray[DH, F],
147+
weight_same_side_ghost_layer: FArray[DH, F],
148+
weight_opposite_side_active_layer: FArray[DH, F],
149+
weight_opposite_side_ghost_layer: FArray[DH, F],
152150
side: Literal["left", "right"],
153151
metadata: dict[str, Any],
154-
) -> NDArray[FloatT]:
155-
return cast("NDArray[FloatT]", same_side_active_layer + same_side_ghost_layer)
152+
) -> FArray[DH, F]:
153+
return same_side_active_layer + same_side_ghost_layer
156154

157155

158156
def antisymmetric_boundary(
159-
same_side_active_layer: NDArray[FloatT],
160-
same_side_ghost_layer: NDArray[FloatT],
161-
opposite_side_active_layer: NDArray[FloatT],
162-
opposite_side_ghost_layer: NDArray[FloatT],
163-
weight_same_side_active_layer: NDArray[FloatT],
164-
weight_same_side_ghost_layer: NDArray[FloatT],
165-
weight_opposite_side_active_layer: NDArray[FloatT],
166-
weight_opposite_side_ghost_layer: NDArray[FloatT],
157+
same_side_active_layer: FArray[DH, F],
158+
same_side_ghost_layer: FArray[DH, F],
159+
opposite_side_active_layer: FArray[DH, F],
160+
opposite_side_ghost_layer: FArray[DH, F],
161+
weight_same_side_active_layer: FArray[DH, F],
162+
weight_same_side_ghost_layer: FArray[DH, F],
163+
weight_opposite_side_active_layer: FArray[DH, F],
164+
weight_opposite_side_ghost_layer: FArray[DH, F],
167165
side: Literal["left", "right"],
168166
metadata: dict[str, Any],
169-
) -> NDArray[FloatT]:
170-
return cast("NDArray[FloatT]", same_side_active_layer - same_side_ghost_layer)
167+
) -> FArray[DH, F]:
168+
return same_side_active_layer - same_side_ghost_layer
171169

172170

173171
def periodic_boundary(
174-
same_side_active_layer: NDArray[FloatT],
175-
same_side_ghost_layer: NDArray[FloatT],
176-
opposite_side_active_layer: NDArray[FloatT],
177-
opposite_side_ghost_layer: NDArray[FloatT],
178-
weight_same_side_active_layer: NDArray[FloatT],
179-
weight_same_side_ghost_layer: NDArray[FloatT],
180-
weight_opposite_side_active_layer: NDArray[FloatT],
181-
weight_opposite_side_ghost_layer: NDArray[FloatT],
172+
same_side_active_layer: FArray[DH, F],
173+
same_side_ghost_layer: FArray[DH, F],
174+
opposite_side_active_layer: FArray[DH, F],
175+
opposite_side_ghost_layer: FArray[DH, F],
176+
weight_same_side_active_layer: FArray[DH, F],
177+
weight_same_side_ghost_layer: FArray[DH, F],
178+
weight_opposite_side_active_layer: FArray[DH, F],
179+
weight_opposite_side_ghost_layer: FArray[DH, F],
182180
side: Literal["left", "right"],
183181
metadata: dict[str, Any],
184-
) -> NDArray[FloatT]:
185-
return cast("NDArray[FloatT]", same_side_active_layer + opposite_side_ghost_layer)
182+
) -> FArray[DH, F]:
183+
return same_side_active_layer + opposite_side_ghost_layer
186184

187185

188186
_base_registry: dict[str, BoundaryRecipeT] = {

0 commit comments

Comments
 (0)