Make systems/inputs/outputs/connections editable

pyxdsm/XDSM.py

@@ -3,7 +3,7 @@
 import numpy as np
 import json
 import subprocess
-from collections import namedtuple
+from recordclass import recordclass
 
 from pyxdsm import __version__ as pyxdsm_version
 
@@ -110,10 +110,10 @@ def _label_to_spec(label, spec):
             spec.add(var)
 
 
-System = namedtuple("System", "node_name style label stack faded label_width spec_name")
-Input = namedtuple("Input", "node_name label label_width style stack faded")
-Output = namedtuple("Output", "node_name label label_width style stack faded side")
-Connection = namedtuple("Connection", "src target label label_width style stack faded")
+System = recordclass("System", "node_name style label stack faded label_width spec_name")
+Input = recordclass("Input", "node_name label label_width style stack faded")
+Output = recordclass("Output", "node_name label label_width style stack faded side")
+Connection = recordclass("Connection", "src target label label_width style stack faded")
 
 
 class XDSM(object):
@@ -196,6 +196,7 @@ def add_system(
 
         sys = System(node_name, style, label, stack, faded, label_width, spec_name)
         self.systems.append(sys)
+        return sys
 
     def add_input(self, name, label, label_width=None, style="DataIO", stack=False, faded=False):
         """
@@ -229,7 +230,9 @@ def add_input(self, name, label, label_width=None, style="DataIO", stack=False,
         faded : bool
             If true, the component will be faded, in order to highlight some other system.
         """
-        self.ins[name] = Input("output_" + name, label, label_width, style, stack, faded)
+        inp = Input("output_" + name, label, label_width, style, stack, faded)
+        self.ins[name] = inp
+        return inp
 
     def add_output(self, name, label, label_width=None, style="DataIO", stack=False, faded=False, side="left"):
         """
@@ -268,11 +271,14 @@ def add_output(self, name, label, label_width=None, style="DataIO", stack=False,
             is placed on the left-most column or the right-most column of the diagram.
         """
         if side == "left":
-            self.left_outs[name] = Output("left_output_" + name, label, label_width, style, stack, faded, side)
+            outp = Output("left_output_" + name, label, label_width, style, stack, faded, side)
+            self.left_outs[name] = outp
         elif side == "right":
-            self.right_outs[name] = Output("right_output_" + name, label, label_width, style, stack, faded, side)
+            outp = Output("right_output_" + name, label, label_width, style, stack, faded, side)
+            self.right_outs[name] = outp
         else:
             raise ValueError("The option 'side' must be given as either 'left' or 'right'!")
+        return outp
 
     def connect(
         self,
@@ -325,7 +331,9 @@ def connect(
         if (not isinstance(label_width, int)) and (label_width is not None):
             raise ValueError("label_width argument must be an integer")
 
-        self.connections.append(Connection(src, target, label, label_width, style, stack, faded))
+        connection = Connection(src, target, label, label_width, style, stack, faded)
+        self.connections.append(connection)
+        return connection
 
     def add_process(self, systems, arrow=True):
         """

setup.py

@@ -26,7 +26,7 @@
         "pyxdsm",
     ],
     package_data={"pyxdsm": ["*.tex"]},
-    install_requires=["numpy>=1.16"],
+    install_requires=["numpy>=1.16", "recordclass>=0.16"],
     python_requires=">=3",
     classifiers=[
         "Operating System :: OS Independent",

tests/test_xdsm.py

@@ -3,7 +3,7 @@
 import shutil
 import tempfile
 import subprocess
-from pyxdsm.XDSM import XDSM, OPT, FUNC, SOLVER, LEFT, RIGHT
+from pyxdsm.XDSM import XDSM, OPT, FUNC, SOLVER, LEFT, RIGHT, IFUNC
 from numpy.distutils.exec_command import find_executable
 
 basedir = os.path.dirname(os.path.abspath(__file__))
@@ -323,6 +323,15 @@ def test_write_outdir(self):
         # no files outside the subdirs
         self.assertFalse(any(os.path.isfile(fp) for fp in os.listdir(self.tempdir)))
 
+    def test_mutability_of_systems(self):
+        x = XDSM()
+        f = x.add_system("func", FUNC, "func")
+        g = x.add_system("cons", IFUNC, "cons")
+        f.stack = True
+        g.label = "cons_new"
+        self.assertEqual(x.systems[0].stack, True)
+        self.assertEqual(x.systems[1].label, "cons_new")
+
 
 if __name__ == "__main__":
     unittest.main()