Change to int | None

Author: schnellerhase

test/test_form.py

@@ -116,6 +116,12 @@ def test_form_arity(functional, mass, stiffness, convection, load) -> None:
     assert convection.arity == 2
     assert load.arity == 1
 
+    assert (functional + load).arity is None
+    assert (functional + mass).arity is None
+    assert (functional + stiffness).arity is None
+    assert (load + mass).arity is None
+    assert (load + stiffness).arity is None
+
 
 def test_form_coefficients(element, domain):
     space = FunctionSpace(domain, element)

ufl/form.py

@@ -341,9 +341,28 @@ def empty(self):
         return len(self.integrals()) == 0
 
     @property
-    def arity(self):
-        """Arity of the form."""
-        return len(self.arguments())
+    def arity(self) -> int | None:
+        """Arity of the form.
+
+        Returns:
+            Number of arguments if all integrals share the argument count, otherwise None.
+        """
+        if not self._integrals:
+            return 0
+
+        from ufl.algorithms.analysis import extract_terminals_with_domain
+
+        arity = None
+        for integral in self._integrals:
+            args, _, _ = extract_terminals_with_domain(integral.integrand())
+            _arity = len(set(args))
+
+            if arity is None:
+                arity = _arity
+            elif arity != _arity:
+                return None
+
+        return arity
 
     def ufl_domains(self):
         """Return the geometric integration domains occuring in the form.