Is there a smart way to pass a class variable using a function?

Question

I have a class to organize my datasets that contains several variables. I read in multiple datasets that's I'd like to compare, so I would like to creat a function that takes one of the class variables as an argument, if that's possible. Currently I have the setup below which I think should work, but I was wondering if there's a more elegant/pythonic way to go about it?

import numpy as np
import matplotlib.pyplot as plt

class MyClass:
    def __init__(self, input=None):
        _tmp = np.genfromtxt(input)
        self.var1 = _tmp['f0']
        self.var2 = _tmp['f1']
        self.var3 = var1 + var2
        self.var4 = var1 * var2

    def var(self, v='var1'):
        if v=='var1':
            return self.var1
        if v=='var2':
            return self.var2
        if v=='var3':
            return self.var3
        if v=='var4':
            return self.var4

def plot_MyClass(data1, data2, var='var1'):
    plt.plot(data1.var(var), data2.var(var), 'k-')
    return

###########

data1 = MyClass(file1.txt)
data2 = MyClass(file2.txt)
plot_MyClass(data1, data2, var='var3')
plt.show()

What's the point of `def var(self, v='var1'):`?? Why not just `data.var1` or `data.var2`? — juanpa.arrivillaga, Jul 11 '22 at 17:41
You can use `getattr(self, v)`, but then you should consider whether you want 4 separate instance attributes in the first place, or a single `dict`-valued attribute. — chepner, Jul 11 '22 at 17:42
@juanpa.arrivillaga In my actual code, I have other details in MyClass that the plotting function requires, so I need to pass the whole object and I wanted an argument to say which of the variables to graph. But you're correct, in the provided code, plt.plot(data1.var1, data2.var2) would work just fine — VioChemist, Jul 11 '22 at 18:37
@chepner So, I could do away with the "var" function in the class and in my plotting fuction I could just do something like d1 = getattr(data1, var); d2 = getattr(data2, var) ?[sorry, I don't know how to format code in comments] — VioChemist, Jul 11 '22 at 18:45
@chepner This worked very nicely and is exactly what I was looking for. I just forgot entirely about 'getattr'. If you want to post it as an answer, I can select it. — VioChemist, Jul 11 '22 at 18:52

Basil · Answer 1 · 2022-07-11T18:06:07.470

Simplify it using a dictionary. There is no special magic required:

import numpy as np
import matplotlib.pyplot as plt

class MyClass:
    def __init__(self, input=None):
        _tmp = np.genfromtxt(input)
        # OPTION 1 ===========================
        self.var = {}
        self.var['var1'] = var1 = _tmp['f0']
        self.var['var2'] = var2 = _tmp['f1']
        self.var['var3'] = var3 = var1 + var2
        self.var['var4'] = var4 = var1 * var2
        # And so on...use varX variables from above if needed...
        # OPTION 2 ====================================================
        # The following is identical to option 1, but may be easier to read:
        var1 = _tmp['f0']
        var2 = _tmp['f1']
        self.var = {
            'var1': var1,
            'var2': var2,
            'var3': var1 + var2,
            'var4': var1 * var2
        }


def plot_MyClass(data1, data2, var='var1'):
    plt.plot(data1.var[var], data2.var[var], 'k-')
    return

###########

data1 = MyClass('file1.txt')
data2 = MyClass('file2.txt')
plot_MyClass(data1, data2, var='var3')
plt.show()

Note that var1, var2, etc., in the constructor are just local variables that make it easier to construct your expressions for the different plot variables of interest.

This is ok, and I will consider it, but would require a fair amount of refactoring to my actual code (already 500 lines). My hack is doing ok, it just doesn't seem like the "right" way to do it — VioChemist, Jul 11 '22 at 18:42

Is there a smart way to pass a class variable using a function?

1 Answers1