Is this the Pythonic way to add a partial method to an existing class?

Question

I am using the daft package to draw graphical models. I want to create new methods for the daft.PGM class to simplify the amount of arguments I have to write when adding nodes for constructing the output. Specifically, I seek to reduce the arguments to the add_node method of the daft.PGM class.

Here is what I have done to create new methods for creating different types of nodes (ie obsNode , decNode , detNode , and latNode ).

import matplotlib.pyplot as plt
import daft   ### %pip install -U git+https://github.com/daft-dev/daft.git
from functools import partial, partialmethod

class dag(daft.PGM):
    def __init__(self, *args, **kwargs):
        daft.PGM.__init__(self, *args, **kwargs)
    
    obsNode = partialmethod(dag.add_node, aspect = 2.2, fontsize = 10, plot_params = {'facecolor': 'cadetblue'})
    decNode = partialmethod(dag.add_node, aspect = 2.2, fontsize = 10, shape = "rectangle", plot_params = {'facecolor': 'thistle'})
    detNode = partialmethod(dag.add_node, aspect = 2.2, fontsize = 10, alternate = True, plot_params = {'facecolor': 'aliceblue'})
    latNode = partialmethod(dag.add_node, aspect = 2.2, fontsize = 10, plot_params = {'facecolor': 'aliceblue'})



pgm = dag(node_fc="aliceblue", dpi = 150, alternate_style="outer")
pgm.obsNode("sb","Start\nBalance", 1, 4)
pgm.decNode("ba","Bet\nAmount", 1, 3)
pgm.detNode("w","Winnings", 2.7, 3)
pgm.latNode("cf","Coin\nFlip",2.7,2)
pgm.detNode("nb","New\nBalance",2.7,4)

pgm.add_edge("sb", "ba")
pgm.add_edge("ba","w")
pgm.add_edge("cf","w")
pgm.add_edge("w", "nb")
pgm.add_edge("sb", "nb")
pgm.render()

And the output is as follows:

I suspect I am doing many things wrong here, but am happy to have output that works. Please help me improve my code with any suggestions you think might help. Thanks!

Answer 1

Maybe this helps a little to reduce code redundancy:

edges = [["sb", "ba"], ["ba","w"], ["cf","w"], ["w", "nb"], ["sb", "nb"]]

for i in edges:
    pgm.add_edge(i[0], i[1])