Fixup for TensorProductElement

python/dune/codegen/pdelab/driver/gridfunctionspace.py

@@ -120,15 +120,33 @@ def name_leafview():
     return name
 
 
+def get_short_name(element):
+    if isinstance(element, TensorProductElement):
+        assert len(set(subel._short_name for subel in element.sub_elements())) == 1
+        return get_short_name(element.sub_elements()[0])
+
+    return element._short_name
+
+
 @preamble(section="fem")
 def typedef_fem(element, name):
     gv = type_leafview()
     df = type_domainfield()
     r = type_range()
     dim = get_dimension()
-    short = element._short_name
     cell = element.cell()
     degree = element.degree()
+    short = get_short_name(element)
+
+    # We currently only support TensorProductElement from UFL if it aliases another finite element
+    # available from UFL. Here, we check this condition and recover the aliases element
+    if isinstance(element, TensorProductElement):
+        subels = set(subel._short_name for subel in element.sub_elements())
+        if len(subels) != 1 or len(set(element.degree())) != 1:
+            raise CodegenUnsupportedFiniteElementError(element)
+
+        degree = element.degree()[0]
+        cell = TensorProductCell(*tuple(subel.cell() for subel in element.sub_elements()))
 
     # The blockstructured code branch has its own handling of finite element selection
     if get_form_option("blockstructured"):
@@ -141,17 +159,6 @@ def typedef_fem(element, name):
     if short == "DG" and isSimplical(cell):
         short = "OPB"
 
-    # We currently only support TensorProductElement from UFL if it aliases another finite element
-    # available from UFL. Here, we check this condition and recover the aliases element
-    if isinstance(element, TensorProductElement):
-        subels = set(subel._short_name for subel in element.sub_elements())
-        if len(subels) != 1 or len(set(element.degrees())) != 1:
-            raise CodegenUnsupportedFiniteElementError(element)
-
-        short = element.sub_elements()[0]._short_name
-        degree = element.degrees()[0]
-        cell = TensorProductCell(*tuple(subel.cell() for subel in element.sub_elements()))
-
     # Choose the correct finite element implementation
     if short == "CG":
         if isSimplical(cell):
@@ -218,7 +225,7 @@ def define_fem(element, name):
     femtype = type_fem(element)
 
     # Determine whether the FEM is grid-dependent - currently on the Lagrange elements are
-    if element._short_name == "CG":
+    if get_short_name(element) == "CG":
         gv = name_leafview()
         return "{} {}({});".format(femtype, name, gv)
     else: