[STLIB] Add the unified "add" in Sofa.Core.Node (#503)

* Add Sofa.Core.Node.add with preliminary test * rename the "add" function and make it private * Apply suggestions from code review Co-authored-by: Hugo <hugo.talbot@sofa-framework.org> --------- Co-authored-by: Hugo <hugo.talbot@sofa-framework.org>
damienmarchal · Apr 14, 2025 · Apr 14, 2025 · Apr 14, 2025 · Apr 14, 2025 · Apr 14, 2025
commit aad2a5e4c4d3f54ff412d06d55f8b3f0a49a0be5
diff --git a/stlib/__init__.py b/stlib/__init__.py
@@ -1 +1,56 @@
 __all__ = ["core","entities","prefabs","shapes"]
+
+import Sofa.Core
+def __genericAdd(self : Sofa.Core.Node, typeName, **kwargs):
+    def findName(cname, names):
+        """Compute a working unique name in the node"""
+        rname = cname
+        for i in range(0, len(names)):
+            if rname not in names:
+                return rname
+            rname = cname + str(i+1)
+        return rname
+
+    # Check if a name is provided, if not, use the one of the class
+    params = kwargs.copy()
+    isNode = False
+    if "name" not in params:
+        if isinstance(typeName, str):
+            params["name"] = typeName
+            isNode=True
+        elif isinstance(typeName, type) and issubclass(typeName, Sofa.Core.Node):
+            params["name"] = "Node"
+            isNode=True
+        elif isinstance(typeName, Sofa.Core.Node):
+            params["name"] = "Node"
+            isNode=True
+        elif isinstance(typeName, type) and issubclass(typeName, Sofa.Core.Object):
+            params["name"] = typeName.name.value
+        elif isinstance(typeName, type) and issubclass(typeName, Sofa.Core.ObjectDeclaration):
+            params["name"] = typeName.__name__
+        else:
+            raise RuntimeError("Invalid argument ", typeName)
+
+    # Check if the name already exists, if this happens, create a new one.
+    if params["name"] in self.children or params["name"] in self.objects:
+        names = {node.name.value for node in self.children}
+        names = names.union({object.name.value for object in self.objects})
+        params["name"] = findName(params["name"], names)
+
+    # Dispatch the creation to either addObject or addChild
+    if isinstance(typeName, type) and issubclass(typeName, Sofa.Core.Node):
+        pref = self.addChild(typeName(params["name"]))
+    elif isinstance(typeName, Sofa.Core.Node):
+        pref = self.addChild(typeName)
+    elif isinstance(typeName, type) and issubclass(typeName, Sofa.Core.Object):
+        pref = self.addObject(typeName(**params))
+    elif isinstance(typeName, type) and issubclass(typeName, Sofa.Core.ObjectDeclaration):
+        pref = self.addObject(typeName.__name__, **params)
+    elif isinstance(typeName, str):
+        pref = self.addObject(typeName, **params)
+    else:
+        raise RuntimeError("Invalid argument", typeName)
+    return pref
+
+# Inject the method so it become available as if it was part of Sofa.Core.Node
+Sofa.Core.Node.add = __genericAdd
diff --git a/stlib/core/basePrefab.py b/stlib/core/basePrefab.py
@@ -3,33 +3,6 @@
 import Sofa.Core 
 from stlib.core.basePrefabParameters import BasePrefabParameters
 
-
-def addFromTypeName(self : Sofa.Core.Node, typeName, params = BasePrefabParameters, **kwargs):  
-    def findName(cname, node):
-        """Compute a working unique name in the node"""
-        rname = cname 
-        for i in range(0, len(node.children)):
-            if rname not in node.children:
-                return rname
-            rname = cname + str(i+1)
-        return rname
-
-    for k,v in kwargs.items():
-        if hasattr(params, k):
-            setattr(params, k, v)
-
-    params = copy.copy(params)
-    if params.name in self.children:
-        params.name = findName(params.name, self)
-
-    pref = self.addChild(typeName(params)) 
-    pref.init() 
-
-    return pref
-
-Sofa.Core.Node.add = addFromTypeName
-
-
 class BasePrefab(Sofa.Core.Node):
     """
     A Prefab is a Sofa.Node that assembles a set of components and nodes

diff --git a/tests/test_new_add.py b/tests/test_new_add.py
@@ -0,0 +1,64 @@
+import unittest 
+import Sofa
+import SofaRuntime
+import Sofa.Core
+import stlib
+
+class ObjectDeclaration(object):
+    ...
+
+Sofa.Core.ObjectDeclaration = ObjectDeclaration
+
+class MechanicalObject(ObjectDeclaration):
+    pass
+
+class TestNewAdd(unittest.TestCase):
+    def test_add_node_with_node_type(self):
+        root = Sofa.Core.Node("root")
+        root.add(Sofa.Core.Node, name="aNodeA")
+        self.assertEqual(len(root.children), 1)
+        self.assertEqual(root.children[0].name.value, "aNodeA")
+
+    def test_add_node_with_node_instance(self):
+        root = Sofa.Core.Node("root")
+        root.add(Sofa.Core.Node("aNodeB"))
+        self.assertEqual(len(root.children), 1)
+        self.assertEqual(root.children[0].name.value, "aNodeB")
+
+    def test_add_object_with_string_type(self):
+        root = Sofa.Core.Node("root")
+        root.add("MechanicalObject", name="anObject1", position=[[1,2,3]])
+        self.assertEqual(len(root.objects), 1)
+        self.assertEqual(root.objects[0].name.value, "anObject1")
+        self.assertEqual(root.objects[0].position.value.shape, (1,3))
+
+    def test_add_object_with_object_type(self):
+        root = Sofa.Core.Node("root")
+        root.add(MechanicalObject, name="anObject2", position=[[1,2,3]])
+        self.assertEqual(len(root.objects), 1)
+        self.assertEqual(root.objects[0].name.value, "anObject2")
+        self.assertEqual(root.objects[0].position.value.shape, (1,3))
+
+    def test_automatic_name_generation(self):
+        root = Sofa.Core.Node("root")
+        root.add(MechanicalObject, position=[[1,2,3]])
+        root.add(MechanicalObject, position=[[1,2,3]])
+        root.add(MechanicalObject, position=[[1,2,3]])
+        self.assertEqual(root.objects[0].name.value, "MechanicalObject")
+        self.assertEqual(root.objects[1].name.value, "MechanicalObject1")
+        self.assertEqual(root.objects[2].name.value, "MechanicalObject2")
+
+        root.add(Sofa.Core.Node, name="TestNode")
+        root.add(Sofa.Core.Node, name="TestNode")
+        self.assertEqual(root.children[0].name.value, "TestNode")
+        self.assertEqual(root.children[1].name.value, "TestNode1")
+
+        root.add(Sofa.Core.Node)
+        root.add(Sofa.Core.Node)
+        self.assertEqual(root.children[2].name.value, "Node")
+        self.assertEqual(root.children[3].name.value, "Node1")
+
+if __name__ == '__main__':
+
+    SofaRuntime.importPlugin("Sofa.Component.StateContainer")
+    unittest.main()