Fix RUF052 and numpy 2.2 typecheck failures

inducer · Dec 17, 2024 · b29b045 · b29b045
1 parent 96d5e06
commit b29b045
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 18 deletions.
diff --git a/grudge/array_context.py b/grudge/array_context.py
@@ -246,15 +246,15 @@ def _dag_to_compiled_func(self, dict_of_named_arrays,
 
         # type-ignore-reason: 'PytatoPyOpenCLArrayContext' has no 'mpi_communicator'
         # pylint: disable=no-member
-        distributed_partition, _new_mpi_base_tag = number_distributed_tags(
+        distributed_partition, new_mpi_base_tag = number_distributed_tags(
                 self.actx.mpi_communicator,
                 distributed_partition,
                 base_tag=prev_mpi_base_tag)
 
         assert prev_mpi_base_tag == self.actx.mpi_base_tag
         # FIXME: Updating stuff inside the array context from here is *cough*
         # not super pretty.
-        self.actx.mpi_base_tag = _new_mpi_base_tag
+        self.actx.mpi_base_tag = new_mpi_base_tag
 
         self.actx._compile_trace_callback(self.f, "post_number_distributed_tags",
                 distributed_partition)

diff --git a/grudge/dof_desc.py b/grudge/dof_desc.py
@@ -389,8 +389,11 @@ def _normalize_domain_and_discr_tag(
         *, _contextual_volume_tag: VolumeTag | None = None
         ) -> tuple[DomainTag, DiscretizationTag]:
 
-    if _contextual_volume_tag is None:
-        _contextual_volume_tag = VTAG_ALL
+    contextual_volume_tag = _contextual_volume_tag
+    del _contextual_volume_tag
+
+    if contextual_volume_tag is None:
+        contextual_volume_tag = VTAG_ALL
 
     if domain == "scalar":
         domain = DTAG_SCALAR
@@ -399,13 +402,13 @@ def _normalize_domain_and_discr_tag(
     elif domain in [VTAG_ALL, "vol"]:
         domain = DTAG_VOLUME_ALL
     elif domain in [FACE_RESTR_ALL, "all_faces"]:
-        domain = BoundaryDomainTag(FACE_RESTR_ALL, _contextual_volume_tag)
+        domain = BoundaryDomainTag(FACE_RESTR_ALL, contextual_volume_tag)
     elif domain in [FACE_RESTR_INTERIOR, "int_faces"]:
-        domain = BoundaryDomainTag(FACE_RESTR_INTERIOR, _contextual_volume_tag)
+        domain = BoundaryDomainTag(FACE_RESTR_INTERIOR, contextual_volume_tag)
     elif isinstance(domain, BTAG_PARTITION):
-        domain = BoundaryDomainTag(domain, _contextual_volume_tag)
+        domain = BoundaryDomainTag(domain, contextual_volume_tag)
     elif domain in [BTAG_ALL, BTAG_REALLY_ALL, BTAG_NONE]:
-        domain = BoundaryDomainTag(domain, _contextual_volume_tag)
+        domain = BoundaryDomainTag(domain, contextual_volume_tag)
     else:
         raise ValueError(f"domain tag not understood: {domain}")
 

diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py
@@ -330,7 +330,8 @@ def inverse_metric_derivative_mat(
     result = np.zeros((ambient_dim, dim), dtype=object)
     for i in range(dim):
         for j in range(ambient_dim):
-            result[i, j] = inverse_metric_derivative(
+            # type-ignore because numpy typing doesn't like object arrays
+            result[i, j] = inverse_metric_derivative(  # type: ignore[call-overload]
                 actx, dcoll, i, j, dd=dd,
                 _use_geoderiv_connection=_use_geoderiv_connection
             )
@@ -564,17 +565,17 @@ def _signed_face_ones(
         actx, dtype=dcoll.real_dtype
     ) + 1
 
-    _signed_face_ones_numpy = actx.to_numpy(signed_ones)
+    signed_face_ones_numpy = actx.to_numpy(signed_ones)
 
     for igrp, grp in enumerate(all_faces_conn.groups):
         for batch in grp.batches:
             assert batch.to_element_face is not None
             i = actx.to_numpy(actx.thaw(batch.to_element_indices))
-            grp_field = _signed_face_ones_numpy[igrp].reshape(-1)
+            grp_field = signed_face_ones_numpy[igrp].reshape(-1)
             grp_field[i] = \
                 (2.0 * (batch.to_element_face % 2) - 1.0) * grp_field[i]
 
-    return actx.from_numpy(_signed_face_ones_numpy)
+    return actx.from_numpy(signed_face_ones_numpy)
 
 
 def parametrization_derivative(
@@ -724,10 +725,13 @@ def mv_normal(
     """
     dd = dof_desc.as_dofdesc(dd)
 
-    if _use_geoderiv_connection is None:
-        _use_geoderiv_connection = actx.supports_nonscalar_broadcasting
+    use_geoderiv_connection = _use_geoderiv_connection
+    del _use_geoderiv_connection
 
-    @memoize_in(dcoll, (mv_normal, dd, _use_geoderiv_connection))
+    if use_geoderiv_connection is None:
+        use_geoderiv_connection = actx.supports_nonscalar_broadcasting
+
+    @memoize_in(dcoll, (mv_normal, dd, use_geoderiv_connection))
     def _normal():
         dim = dcoll.discr_from_dd(dd).dim
         ambient_dim = dcoll.ambient_dim
@@ -740,7 +744,7 @@ def _normal():
         if dim == ambient_dim - 1:
             result = rel_mv_normal(
                 actx, dcoll, dd=dd,
-                _use_geoderiv_connection=_use_geoderiv_connection)
+                use_geoderiv_connection=use_geoderiv_connection)
         else:
             # NOTE: In the case of (d - 2)-dimensional curves, we don't really have
             # enough information on the face to decide what an "exterior face normal"
@@ -759,12 +763,12 @@ def _normal():
                         rel_mv_normal(
                             actx, dcoll,
                             dd=dd.untrace(),
-                            _use_geoderiv_connection=_use_geoderiv_connection
+                            use_geoderiv_connection=use_geoderiv_connection
                         ).as_vector(dtype=object))
             )
             pder = pseudoscalar(
                 actx, dcoll, dd=dd,
-                _use_geoderiv_connection=_use_geoderiv_connection)
+                use_geoderiv_connection=use_geoderiv_connection)
 
             mv = -(volm_normal ^ pder) << volm_normal.I.inv()