diff --git a/stlib/__init__.py b/stlib/__init__.py index 70108a41..01c21a40 100644 --- 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 index c2e10148..1dd5665a 100644 --- 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 new file mode 100644 index 00000000..fcadda6f --- /dev/null +++ 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()