add topology.WithIdentifierTopology

Author: joostvanzwieten

nutils/topology.py

@@ -1986,6 +1986,44 @@ def refined(self):
   def basis(self, name, *args, **kwargs):
     return function.asarray([1.])
 
+class WithIdentifierTopology(Topology):
+  '''A topology that appends an :class:`nutils.transform.Identifier` to the ``transforms`` and ``opposites`` of another topology.
+
+  Parameters
+  ----------
+  parent : :class:`Topology`
+      The parent topology.
+  token : :class:`object`
+      An immutable token that will be used to create the
+      :class:`nutils.transform.Identifier`.
+  '''
+
+  @types.apply_annotations
+  def __init__(self, parent:stricttopology, token):
+    self._parent = parent
+    self._token = token
+    super().__init__(parent.references,
+                     transformseq.WithIdentifierTransforms(parent.transforms, token),
+                     transformseq.WithIdentifierTransforms(parent.opposites, token))
+
+  def basis(self, *args, **kwargs):
+    return function.WithTransformsBasis(self._parent.basis(*args, **kwargs), self.transforms)
+
+  @property
+  def refined(self):
+    return WithIdentifierTopology(self._parent.refined, self._token)
+
+  @property
+  def boundary(self):
+    return WithIdentifierTopology(self._parent.boundary, self._token)
+
+  @property
+  def interfaces(self):
+    return WithIdentifierTopology(self._parent.interfaces, self._token)
+
+  def getitem(self, item):
+    return WithIdentifierTopology(self._parent.getitem(item), self._token)
+
 class PatchBoundary(types.Singleton):
 
   __slots__ = 'id', 'dim', 'side', 'reverse', 'transpose'

nutils/transformseq.py

@@ -807,6 +807,72 @@ def index_with_tail(self, trans):
     iderived = self._derived_transforms.index(tail[0])
     return iparent*len(self._derived_transforms) + iderived, tail[1:]
 
+class WithIdentifierTransforms(Transforms):
+  '''A sequence that appends a :class:`nutils.transform.Identifier` to all transforms of another sequence.
+
+  Parameters
+  ----------
+  parent : :class:`Transforms`
+      The parent transforms.
+  token : :class:`object`
+      An immutable token that will be used to create the
+      :class:`nutils.transform.Identifier`.
+  '''
+
+  @types.apply_annotations
+  def __init__(self, parent:stricttransforms, token):
+    self._parent = parent
+    self._token = token
+    super().__init__(parent.fromdims)
+
+  def __len__(self):
+    return len(self._parent)
+
+  def __getitem__(self, index):
+    parent_trans = self._parent[index]
+    if isinstance(parent_trans, Transforms):
+      return WithIdentifierTransforms(parent_trans, self._token)
+    else:
+      return parent_trans + (transform.Identifier(parent_trans[-1].fromdims, self._token),)
+
+  def __iter__(self):
+    for parent_trans in self._parent:
+      yield parent_trans + (transform.Identifier(parent_trans[-1].fromdims, self._token),)
+
+  def _remove_identifier(self, trans):
+    if not trans:
+      raise ValueError
+    for i, item in enumerate(trans):
+      if isinstance(item, transform.Identifier) and item.token == self._token:
+        return trans[:i] + trans[i+1:]
+    raise ValueError
+
+  def index_with_tail(self, trans):
+    return self._parent.index_with_tail(self._remove_identifier(trans))
+
+  def index(self, trans):
+    return self._parent.index(self._remove_identifier(trans))
+
+  def contains_with_tail(self, trans):
+    try:
+      head = self._remove_identifier(trans)
+    except ValueError:
+      return False
+    return self._parent.contains_with_tail(head)
+
+  def contains(self, trans):
+    try:
+      head = self._remove_identifier(trans)
+    except ValueError:
+      return False
+    return self._parent.contains(head)
+
+  def refined(self, references):
+    return WithIdentifierTransforms(self._parent.refined(references), self._token)
+
+  def edges(self, references):
+    return WithIdentifierTransforms(self._parent.edges(references), self._token)
+
 class ProductTransforms(Transforms):
   '''The product of two :class:`Transforms` objects.
 

tests/test_transformseq.py

@@ -216,6 +216,16 @@ def setUp(self):
     self.checkrefs = nutils.elementseq.asreferences([line]*4, 1)
     self.checkfromdims = 1
 
+class WithIdentifierTransforms(TestCase, Common, Edges):
+  def setUp(self):
+    parent = nutils.transformseq.PlainTransforms([(x1,s0),(x1,s1),(x1,s2),(x1,s3)], fromdims=1)
+    self.seq = nutils.transformseq.WithIdentifierTransforms(parent, 'token')
+    token = nutils.transform.Identifier(1, 'token')
+    self.check = (x1,s0,token),(x1,s1,token),(x1,s2,token),(x1,s3,token)
+    self.checkmissing = (x1,s0,nutils.transform.Identifier(1, 'other')),
+    self.checkrefs = nutils.elementseq.asreferences([line]*4, 1)
+    self.checkfromdims = 1
+
 class ChainedTransforms(TestCase, Common, Edges):
   def setUp(self):
     self.seq = nutils.transformseq.ChainedTransforms([nutils.transformseq.PlainTransforms([(x1,s0),(x1,s1)], fromdims=1), nutils.transformseq.PlainTransforms([(x1,s2),(x1,s3)], fromdims=1)])