Go back to LBYL programming style (instead of EAFP) and several mypy fixes in fem.petsc and la.petsc

python/dolfinx/fem/assemble.py

@@ -7,11 +7,10 @@
 
 from __future__ import annotations
 
-import collections
 import functools
 import typing
 import warnings
-from collections.abc import Iterable
+from collections.abc import Sequence
 
 import numpy as np
 import numpy.typing as npt
@@ -28,7 +27,7 @@
 
 def pack_constants(
     form: typing.Union[Form, typing.Sequence[Form]],
-) -> typing.Union[np.ndarray, typing.Sequence[np.ndarray]]:
+) -> typing.Union[npt.NDArray, typing.Sequence[npt.NDArray]]:
     """Pack form constants for use in assembly.
 
     Pack the 'constants' that appear in forms. The packed constants can
@@ -50,7 +49,7 @@ def pack_constants(
     def _pack(form):
         if form is None:
             return None
-        elif isinstance(form, collections.abc.Iterable):
+        elif isinstance(form, Sequence):
             return list(map(lambda sub_form: _pack(sub_form), form))
         else:
             return _pack_constants(form._cpp_object)
@@ -84,7 +83,7 @@ def pack_coefficients(
     def _pack(form):
         if form is None:
             return {}
-        elif isinstance(form, collections.abc.Iterable):
+        elif isinstance(form, Sequence):
             return list(map(lambda sub_form: _pack(sub_form), form))
         else:
             return _pack_coefficients(form._cpp_object)
@@ -106,7 +105,7 @@ def create_vector(L: Form) -> la.Vector:
     """
     # Can just take the first dofmap here, since all dof maps have the same
     # index map in mixed-topology meshes
-    dofmap = L.function_spaces[0].dofmaps(0)  # type: ignore
+    dofmap = L.function_spaces[0].dofmaps(0)  # type: ignore[attr-defined]
     return la.vector(dofmap.index_map, dofmap.index_map_bs, dtype=L.dtype)
 
 
@@ -133,7 +132,11 @@ def create_matrix(a: Form, block_mode: typing.Optional[la.BlockMode] = None) ->
 # -- Scalar assembly ------------------------------------------------------
 
 
-def assemble_scalar(M: Form, constants: typing.Optional[np.ndarray] = None, coeffs=None):
+def assemble_scalar(
+    M: Form,
+    constants: typing.Optional[npt.NDArray] = None,
+    coeffs: typing.Optional[dict[tuple[IntegralType, int], npt.NDArray]] = None,
+):
     """Assemble functional. The returned value is local and not
     accumulated across processes.
 
@@ -155,22 +158,28 @@ def assemble_scalar(M: Form, constants: typing.Optional[np.ndarray] = None, coef
         To compute the functional value on the whole domain, the output
         of this function is typically summed across all MPI ranks.
     """
-    constants = pack_constants(M) if constants is None else constants  # type: ignore
-    coeffs = coeffs or pack_coefficients(M)
+    constants = pack_constants(M) if constants is None else constants  # type: ignore[assignment]
+    coeffs = pack_coefficients(M) if coeffs is None else constants  # type: ignore[assignment]
     return _cpp.fem.assemble_scalar(M._cpp_object, constants, coeffs)
 
 
 # -- Vector assembly ------------------------------------------------------
 
 
 @functools.singledispatch
-def assemble_vector(L: typing.Any, constants: typing.Optional[np.ndarray] = None, coeffs=None):
+def assemble_vector(
+    L: typing.Any,
+    constants: typing.Optional[npt.NDArray] = None,
+    coeffs: typing.Optional[dict[tuple[IntegralType, int], npt.NDArray]] = None,
+):
     return _assemble_vector_form(L, constants, coeffs)
 
 
-@assemble_vector.register(Form)
+@assemble_vector.register
 def _assemble_vector_form(
-    L: Form, constants: typing.Optional[np.ndarray] = None, coeffs=None
+    L: Form,
+    constants: typing.Optional[npt.NDArray] = None,
+    coeffs: typing.Optional[dict[tuple[IntegralType, int], npt.NDArray]] = None,
 ) -> la.Vector:
     """Assemble linear form into a new Vector.
 
@@ -197,15 +206,18 @@ def _assemble_vector_form(
     """
     b = create_vector(L)
     b.array[:] = 0
-    constants = constants or pack_constants(L)  # type: ignore
-    coeffs = coeffs or pack_coefficients(L)
+    constants = pack_constants(L) if constants is None else constants  # type: ignore[assignment]
+    coeffs = pack_coefficients(L) if coeffs is None else coeffs  # type: ignore[assignment]
     _assemble_vector_array(b.array, L, constants, coeffs)
     return b
 
 
 @assemble_vector.register(np.ndarray)
 def _assemble_vector_array(
-    b: np.ndarray, L: Form, constants: typing.Optional[np.ndarray] = None, coeffs=None
+    b: npt.NDArray,
+    L: Form,
+    constants: typing.Optional[npt.NDArray] = None,
+    coeffs: typing.Optional[dict[tuple[IntegralType, int], npt.NDArray]] = None,
 ):
     """Assemble linear form into an existing array.
 
@@ -229,8 +241,8 @@ def _assemble_vector_array(
         :func:`dolfinx.la.Vector.scatter_reverse` on the return vector
         can accumulate ghost contributions.
     """
-    constants = pack_constants(L) if constants is None else constants  # type: ignore
-    coeffs = pack_coefficients(L) if coeffs is None else coeffs
+    constants = pack_constants(L) if constants is None else constants  # type: ignore[assignment]
+    coeffs = pack_coefficients(L) if coeffs is None else coeffs  # type: ignore[assignment]
     _cpp.fem.assemble_vector(b, L._cpp_object, constants, coeffs)
     return b
 
@@ -241,10 +253,10 @@ def _assemble_vector_array(
 @functools.singledispatch
 def assemble_matrix(
     a: typing.Any,
-    bcs: typing.Optional[list[DirichletBC]] = None,
-    diagonal: float = 1.0,
-    constants: typing.Optional[np.ndarray] = None,
-    coeffs=None,
+    bcs: typing.Optional[Sequence[DirichletBC]] = None,
+    diag: float = 1.0,
+    constants: typing.Optional[npt.NDArray] = None,
+    coeffs: typing.Optional[dict[tuple[IntegralType, int], npt.NDArray]] = None,
     block_mode: typing.Optional[la.BlockMode] = None,
 ):
     """Assemble bilinear form into a matrix.
@@ -274,18 +286,18 @@ def assemble_matrix(
     """
     bcs = [] if bcs is None else bcs
     A: la.MatrixCSR = create_matrix(a, block_mode)
-    _assemble_matrix_csr(A, a, bcs, diagonal, constants, coeffs)
+    _assemble_matrix_csr(A, a, bcs, diag, constants, coeffs)
     return A
 
 
-@assemble_matrix.register(la.MatrixCSR)
+@assemble_matrix.register
 def _assemble_matrix_csr(
     A: la.MatrixCSR,
     a: Form,
-    bcs: typing.Optional[list[DirichletBC]] = None,
-    diagonal: float = 1.0,
-    constants: typing.Optional[np.ndarray] = None,
-    coeffs=None,
+    bcs: typing.Optional[Sequence[DirichletBC]] = None,
+    diag: float = 1.0,
+    constants: typing.Optional[npt.NDArray] = None,
+    coeffs: typing.Optional[dict[tuple[IntegralType, int], npt.NDArray]] = None,
 ) -> la.MatrixCSR:
     """Assemble bilinear form into a matrix.
 
@@ -311,28 +323,28 @@ def _assemble_matrix_csr(
         accumulated.
     """
     bcs = [] if bcs is None else [bc._cpp_object for bc in bcs]
-    constants = pack_constants(a) if constants is None else constants  # type: ignore
-    coeffs = pack_coefficients(a) if coeffs is None else coeffs
+    constants = pack_constants(a) if constants is None else constants  # type: ignore[assignment]
+    coeffs = pack_coefficients(a) if coeffs is None else coeffs  # type: ignore[assignment]
     _cpp.fem.assemble_matrix(A._cpp_object, a._cpp_object, constants, coeffs, bcs)
 
     # If matrix is a 'diagonal'block, set diagonal entry for constrained
     # dofs
     if a.function_spaces[0] is a.function_spaces[1]:
-        _cpp.fem.insert_diagonal(A._cpp_object, a.function_spaces[0], bcs, diagonal)
+        _cpp.fem.insert_diagonal(A._cpp_object, a.function_spaces[0], bcs, diag)
     return A
 
 
 # -- Modifiers for Dirichlet conditions -----------------------------------
 
 
 def apply_lifting(
-    b: np.ndarray,
-    a: Iterable[Form],
-    bcs: Iterable[Iterable[DirichletBC]],
-    x0: typing.Optional[Iterable[np.ndarray]] = None,
+    b: npt.NDArray,
+    a: Sequence[Form],
+    bcs: Sequence[Sequence[DirichletBC]],
+    x0: typing.Optional[Sequence[npt.NDArray]] = None,
     alpha: float = 1,
-    constants: typing.Optional[Iterable[np.ndarray]] = None,
-    coeffs=None,
+    constants: typing.Optional[npt.NDArray] = None,
+    coeffs: typing.Optional[dict[tuple[IntegralType, int], npt.NDArray]] = None,
 ) -> None:
     """Modify right-hand side vector ``b`` for lifting of Dirichlet
     boundary conditions.
@@ -430,12 +442,12 @@ def apply_lifting(
     """
     x0 = [] if x0 is None else x0
     constants = (
-        [pack_constants(form) if form is not None else np.array([], dtype=b.dtype) for form in a]  # type: ignore
+        [pack_constants(form) if form is not None else np.array([], dtype=b.dtype) for form in a]  # type: ignore[assignment]
         if constants is None
         else constants
     )
     coeffs = (
-        [{} if form is None else pack_coefficients(form) for form in a]
+        [{} if form is None else pack_coefficients(form) for form in a]  # type: ignore[assignment]
         if coeffs is None
         else coeffs
     )
@@ -445,10 +457,10 @@ def apply_lifting(
 
 
 def set_bc(
-    b: np.ndarray,
-    bcs: list[DirichletBC],
-    x0: typing.Optional[np.ndarray] = None,
-    scale: float = 1,
+    b: npt.NDArray,
+    bcs: Sequence[DirichletBC],
+    x0: typing.Optional[npt.NDArray] = None,
+    alpha: float = 1,
 ) -> None:
     """Insert boundary condition values into vector.
 
@@ -464,4 +476,4 @@ def set_bc(
         DeprecationWarning,
     )
     for bc in bcs:
-        bc.set(b, x0, scale)
+        bc.set(b, x0, alpha)

python/dolfinx/fem/bcs.py

@@ -11,16 +11,14 @@
 
 import numbers
 import typing
-
-import numpy.typing as npt
-
-if typing.TYPE_CHECKING:
-    from dolfinx.fem.function import Constant, Function
+from collections.abc import Iterable
 
 import numpy as np
+import numpy.typing as npt
 
 import dolfinx
 from dolfinx import cpp as _cpp
+from dolfinx.fem.function import Constant, Function
 
 
 def locate_dofs_geometrical(
@@ -48,11 +46,11 @@ def locate_dofs_geometrical(
         Returned degree-of-freedom indices are unique and ordered by the
         first column.
     """
-    try:
-        return _cpp.fem.locate_dofs_geometrical(V._cpp_object, marker)  # type: ignore
-    except AttributeError:
-        _V = [space._cpp_object for space in V]  # type: ignore
-        return _cpp.fem.locate_dofs_geometrical(_V, marker)
+    if not isinstance(V, Iterable):
+        return _cpp.fem.locate_dofs_geometrical(V._cpp_object, marker)
+
+    _V = [space._cpp_object for space in V]
+    return _cpp.fem.locate_dofs_geometrical(_V, marker)
 
 
 def locate_dofs_topological(
@@ -84,11 +82,11 @@ def locate_dofs_topological(
         first column.
     """
     _entities = np.asarray(entities, dtype=np.int32)
-    try:
-        return _cpp.fem.locate_dofs_topological(V._cpp_object, entity_dim, _entities, remote)  # type: ignore
-    except AttributeError:
-        _V = [space._cpp_object for space in V]  # type: ignore
-        return _cpp.fem.locate_dofs_topological(_V, entity_dim, _entities, remote)
+    if not isinstance(V, Iterable):
+        return _cpp.fem.locate_dofs_topological(V._cpp_object, entity_dim, _entities, remote)
+
+    _V = [space._cpp_object for space in V]
+    return _cpp.fem.locate_dofs_topological(_V, entity_dim, _entities, remote)
 
 
 class DirichletBC:
@@ -223,9 +221,9 @@ def dirichletbc(
         _value = value
     else:
         try:
-            _value = value._cpp_object  # type: ignore
+            _value = value._cpp_object
         except AttributeError:
-            _value = value  # type: ignore
+            _value = value  # type: ignore[assignment]
 
     if V is not None:
         try: