Enable mypy type-checking across the full icon4py codebase - #1382
Draft
msimberg wants to merge 179 commits into
Draft
Enable mypy type-checking across the full icon4py codebase#1382msimberg wants to merge 179 commits into
msimberg wants to merge 179 commits into
Conversation
The icon4py mypy plugin registers a type-analyze hook for icon4py.model.common.type_alias.vpfloat, resolving it to builtins.float in type positions. This mirrors the existing GT4Py plugin's blur_float_precision hook and eliminates the need for the disable_error_code = [valid-type] override that was applied to advection.stencils, diffusion.stencils, dycore.dycore_states, and dycore.stencils. The vpfloat variable in type_alias.py is a runtime-mutable global (reassigned by set_precision() to switch between single and double precision). While this works at runtime, mypy does not recognize it as a valid type. The plugin resolves this without changing any runtime behavior or touching any of the 695+ call sites that use vpfloat in type positions. Additionally fixes pre-existing valid-type errors in two dycore stencil files that used gtx.Field[[dims.X, dims.Y], T] (bare list) instead of gtx.Field[gtx.Dims[dims.X, dims.Y], T]. These were previously masked by the valid-type override. Removes two now-unused type: ignore[return-value] comments in the solve_tridiagonal scan operators — the return type errors are resolved by the vpfloat plugin.
Fix 4 pre-existing errors in type_alias.py: - Change set_precision parameter from Literal['double','mixed','single'] to str (the function already validates internally via match/case + ValueError; the only caller passes str from os.environ.get()) - Add type: ignore[assignment] on 3 lines where gtx.float32 is assigned to vpfloat/wpfloat. The GT4Py type plugin blurs float32/float64 to builtins.float in type annotations only, not in value positions. This is a consequence of the runtime precision-switching design (using type objects as values), which is preserved per design decision D1.
Adds 8 level-0 files (no unchecked internal dependencies) to the mypy files list: dimension.py, exceptions.py, model_backends.py, grid/geometry_config.py, math/projection.py, math/stencils/generic_math_operations_array_ns.py, metrics/compute_coeff_gradekin.py, metrics/compute_diffusion_metrics.py. Fixes 9 type errors in model_backends.py: - get_allocator: change return type from Backend to FieldBufferAllocatorProtocol (the function returns an allocator, not a Backend). When input is already a Backend, return backend.allocator instead of backend itself. - is_cpu_backend/is_gpu_backend: use __gt_device_type__ instead of device_type (GT4Py stubs only expose __gt_device_type__ on FieldBufferAllocatorProtocol, not device_type). - make_custom_gtfn_backend: change return type from Backend to GTFNBackendFactory (the function returns a factory, not a Backend). - make_custom_gtfn_backend/make_custom_dace_backend: type **kwargs instead of untyped **_. Follow-on fix in standalone_driver.py: update _allocator property type from gtx.typing.Backend to FieldBufferAllocatorProtocol to match the new return type of get_allocator.
Both files are level-1 (depend only on type_alias.py and dimension.py, both already checked) and require no source changes.
Fix 2 errors: - attrs: MutableMapping[str, ...] is invalid (Ellipsis not a valid type argument); change to MutableMapping[str, Any] - Add missing return type annotation -> xa.DataArray
Add missing type annotation for e2c parameter (data_alloc.NDArray).
Fix 5 errors: - Add type: ignore[override] on Zone.value property (intentional override of Enum.value to return combined name_level string) - Add type annotation for get_last_nudging parameter and return type - Remove 3 stale type: ignore comments (name-defined, attr-defined) that are no longer needed
Fix 6 errors: - Remove stale type: ignore[return-value] on line 99 (no longer needed after model_backends.py fixes) - Add explicit Any type annotation for backend_factory to prevent DeviceType being inferred as callable - Add type: ignore[arg-type] on **variants passed to compile() and **constant_args/**horizontal_sizes/**vertical_sizes passed to functools.partial(). GT4Py Program.compile and __call__ use dynamic **kwargs that mypy cannot verify against strict method signatures. Also removes unused import gt_coredefs from states/model.py (ruff).
Contributor
Author
|
Addressed the outstanding review comments in commits b74c39b through d7aa7ca. Fixed:
Already in place / not applicable:
Discussion answers:
🤖 Posted by an agent on behalf of @msimberg |
- Remove now-unused NDArrayObject __setitem__ ignores in initial_condition modules (GT4Py 1.2.0 no longer emits index errors there). - Replace per-line ghex/viztracer import ignores with module-level ignore_missing_imports overrides in pyproject.toml, matching the existing dace pattern. This makes type-checking consistent whether the optional distributed/profiling dependencies are installed.
…ns [R30,R31,R32-R36,R39,N10]
…s [R40,R41,R45,R46]
Contributor
Author
|
cscs-ci run default |
Contributor
Author
|
cscs-ci run default |
Contributor
Author
|
cscs-ci run default |
msimberg
commented
Aug 10, 2026
| from types import TracebackType | ||
|
|
||
| import viztracer # type: ignore[import-not-found] | ||
| import viztracer # optional profiling dependency, not in typing env |
Contributor
Author
There was a problem hiding this comment.
If it doesn't need to be ignored anymore, I'd remove the comment.
Suggested change
| import viztracer # optional profiling dependency, not in typing env | |
| import viztracer |
msimberg
commented
Aug 10, 2026
Comment on lines
-34
to
+56
| grid_functionality = collections.defaultdict(dict) | ||
| grid_functionality: collections.defaultdict[str, dict[str, Any]] = collections.defaultdict(dict) | ||
|
|
||
|
|
||
| def get_grid_for_experiment(experiment: test_defs.Experiment, backend: gtx_typing.Backend): | ||
| def get_grid_for_experiment(experiment: test_defs.Experiment, backend: gtx_typing.Backend) -> Any: | ||
| return _get_or_initialize(experiment, backend, "grid") | ||
|
|
||
|
|
||
| def get_edge_geometry_for_experiment(experiment: test_defs.Experiment, backend: gtx_typing.Backend): | ||
| def get_edge_geometry_for_experiment( | ||
| experiment: test_defs.Experiment, backend: gtx_typing.Backend | ||
| ) -> Any: | ||
| return _get_or_initialize(experiment, backend, "edge_geometry") | ||
|
|
||
|
|
||
| def get_cell_geometry_for_experiment(experiment: test_defs.Experiment, backend: gtx_typing.Backend): | ||
| def get_cell_geometry_for_experiment( | ||
| experiment: test_defs.Experiment, backend: gtx_typing.Backend | ||
| ) -> Any: | ||
| return _get_or_initialize(experiment, backend, "cell_geometry") | ||
|
|
||
|
|
||
| def _get_or_initialize(experiment: test_defs.Experiment, backend: gtx_typing.Backend, name: str): | ||
| def _get_or_initialize( | ||
| experiment: test_defs.Experiment, backend: gtx_typing.Backend, name: str | ||
| ) -> Any: |
Contributor
Author
There was a problem hiding this comment.
Can these return types not be more tight than Any? If yes, change them.
Comment on lines
-188
to
+211
| theta_v_at_cells_on_half_levels=None, | ||
| perturbed_exner_at_cells_on_model_levels=None, | ||
| rho_at_cells_on_half_levels=None, | ||
| exner_tendency_due_to_slow_physics=None, | ||
| grf_tend_rho=None, | ||
| grf_tend_thv=None, | ||
| grf_tend_w=None, | ||
| mass_flux_at_edges_on_model_levels=None, | ||
| normal_wind_tendency_due_to_slow_physics_process=None, | ||
| grf_tend_vn=None, | ||
| theta_v_at_cells_on_half_levels=data_alloc.zero_field( | ||
| icon_grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, allocator=backend | ||
| ), | ||
| perturbed_exner_at_cells_on_model_levels=data_alloc.zero_field( | ||
| icon_grid, dims.CellDim, dims.KDim, allocator=backend | ||
| ), | ||
| rho_at_cells_on_half_levels=data_alloc.zero_field( | ||
| icon_grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, allocator=backend | ||
| ), | ||
| exner_tendency_due_to_slow_physics=data_alloc.zero_field( | ||
| icon_grid, dims.CellDim, dims.KDim, allocator=backend | ||
| ), | ||
| grf_tend_rho=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend), | ||
| grf_tend_thv=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend), | ||
| grf_tend_w=data_alloc.zero_field( | ||
| icon_grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, allocator=backend | ||
| ), | ||
| mass_flux_at_edges_on_model_levels=data_alloc.zero_field( | ||
| icon_grid, dims.EdgeDim, dims.KDim, allocator=backend | ||
| ), | ||
| normal_wind_tendency_due_to_slow_physics_process=data_alloc.zero_field( | ||
| icon_grid, dims.EdgeDim, dims.KDim, allocator=backend | ||
| ), | ||
| grf_tend_vn=data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim, allocator=backend), |
Contributor
Author
There was a problem hiding this comment.
Hmm, why are these allocated now instead of being None? It might be the right thing to do, but can you explain the change and why it's not unnecessary allocations?
Contributor
Author
There was a problem hiding this comment.
DiagnosticStateNonHydro and PrognosticState constructors require typed gtx.Field values. None would violate their annotations and need # type: ignore[arg-type] suppressions. Zero fields satisfy the types and act as neutral test values.\n\n🤖 Posted by an agent on behalf of @msimberg
Comment on lines
-358
to
+405
| theta_v=None, | ||
| rho=None, | ||
| exner=None, | ||
| theta_v=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend), | ||
| rho=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend), | ||
| exner=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend), |
Contributor
Author
There was a problem hiding this comment.
Same here, and just above: why are these now allocated to zero fields?
Contributor
Author
There was a problem hiding this comment.
Same reason: the state constructors require non-optional Field values per their annotations. data_alloc.zero_field gives a typed neutral value instead of None with a suppression.\n\n🤖 Posted by an agent on behalf of @msimberg
Comment on lines
-86
to
+124
| virtual_temperature=None, | ||
| virtual_temperature=virtual_temperature, | ||
| pressure=entry_savepoint.pressure(), | ||
| pressure_ifc=None, | ||
| u=None, | ||
| v=None, | ||
| pressure_ifc=pressure_ifc, | ||
| u=u, | ||
| v=v, |
Contributor
Author
There was a problem hiding this comment.
Can you explain why these fields aren't None anymore?
Contributor
Author
There was a problem hiding this comment.
PrognosticState requires non-optional vn, w, exner, and theta_v fields under the current annotations. The previous None values were only accepted via # type: ignore[arg-type]. Using data_alloc.zero_field satisfies the types without suppression.\n\n🤖 Posted by an agent on behalf of @msimberg
Contributor
Author
|
cscs-ci run default |
Contributor
Author
|
cscs-ci run default 🤖 Posted by an agent on behalf of @msimberg |
|
When developing, you can test your changes on CSCS CI before merge with the You can pass options to override pipeline variables, for example:
Available options are:
For each option, See The Merging Once your PR is approved and ready for merging, add it to the merge queue. The Optional Tests To run benchmarks you can use:
For more detailed information please look at CI in the EXCLAIM universe. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Opened by an agent.
Enable mypy type-checking across the full icon4py codebase. Until now, most Python files were excluded from mypy through an opt-in whitelist, so type errors were only caught at runtime. This branch replaces that whitelist with an opt-out configuration, fixes the resulting type errors, and annotates the previously untyped test code.
The main pieces of work are: adding a small icon4py mypy plugin that treats
vpfloatasfloatin type positions, mirroring the GT4Py plugin; enabling type-checking for the remaining source and test files inmodel/common,model/atmosphere,model/driver,model/standalone_driver,model/testing,tools, andbindings; fully annotating test functions; and replacing broaddisable_error_codeoverrides with per-line# type: ignore[code]comments that include brief justifications. The only remaining module-level suppression isignore_errors = trueforicon4py.model.testing.serialbox, becauseserialbox4pyis a C extension without type stubs.A few design changes were required to expose and break circular imports that type-checking surfaced:
initial_conditionand related state helpers moved frommodel/commonintomodel/standalone_driver, andread_experiment_config_from_fortranwas renamed toread_configwith a simplified signature. These are intentional and limited to what was needed to make the dependency graph type-safe.All
# type: ignorecomments carry an error code and a short justification; there are no bare suppressions.