Python: can I change function behaviour whether output is assigned or not?

Question

In Matlab, nargout is a variable that tells you if the output is assigned, so

x = f(2);

and

f(2);

can behave differently.

Is it possible to do similar in Python?

I have a function that plots to screen and returns a matplotlib figure object. I want that if output is assigned to a variable then do not plot to screen.

Answer 1

Here's a way you can do it (not that i'd advise it), but it has many cases where it won't work - to make it work you'd essentially need to parse the python code in the line and see what it is doing, which would be possible down to some level, but there are likely always going to be ways to get around it.

import inspect, re


def func(x, noCheck=False):
  if not noCheck:
      #Get the line the function was called on.
      _, _, _, _, lines, _ = inspect.getouterframes(inspect.currentframe())[1]

      #Now we need to search through `line` to see how the function is called.

      line = lines[0].split("#")[0] #Get rid of any comments at the end of the line.

      match = re.search(r"[a-zA-Z0-9]+ *= *func\(.*\)", line) #Search for instances of `func` being called after an equals sign

      try:
        variable, functioncall = match.group(0).split("=")
        print variable, "=", functioncall, "=", eval(functioncall.strip()[:-1] + ", noCheck=True)")

      except:
        pass #print "not assigned to a variable"

  #Actually make the function do something
  return 3*x**2 + 2*x + 1

func(1) # x = func(1)
x = func(1)

Another way to do it would be to examine all of the set local variables when you call the code, and check if any of them have been set to the result of your function, then use that information to help parse the python.

Or you could look at object IDs, and try and do things that way, but that's not goign to be straightforward, as not all objects work the same way (i.e. do a=10 and c=10 and then have a look at each object's IDs, they're the same ven though a and c are seperate. The same happens with short strings too)

If you can think up a way to do this that would work universally, i'd be interested to know how you do it, I'd pressume it will need to be done by digging around in inspect though, rather than through parsing the actual code.

Answer 2

Others have mentioned that this is complex, but can be done with inspect. You may want a simple approach by having a separate function to plot it, or pass an extra variable that says to plot.

def create_plot(x):
    return plot

def display(plot):
    # show the plot

x = create_plot(2)
display(x)

Plot variable

def plot(x, show=False)
    # create the plot
    if show:
        # show the plot

plot(2, True)
x = plot(2)

It is probably not worth the time and easier to just create the two functions.

Personally, I think this is ugly, nasty, and I do not believe that functionality should be based on something catching the return value. However, I was curious, and I found a way. You could probably turn this into a decorator if you want to use it in the future, but I still suggest that you use two separate methods instead of checking for an output.

import inspect

def f(val):
    has_output = False
    frame = inspect.currentframe()
    name = frame.f_code.co_name
    outer = inspect.getouterframes(frame)[1] # may want to loop through available frames.
    for i in range(len(outer)):
        item = str(outer[i]).replace(" ", "")
        check = "="+name+"("
        if check in item and "="+check not in item: # also check assignment vs equality
            # Your method has an output
            has_output = True
            break

    if has_output:
        print("Something catches the output")
    return val*val
# end f

In many cases this will not work either. You will have to make really good regex for the check if you always want it to work.

import my_lib

x = my_lib.f(2)