Maya MPxNode multiple outputs - python

I'm trying to create an MPxNode with multiple outputs, but I can only get one to work properly. The other output doesn't set properly after connecting the node and during undos.
Is it possible to set both outputs at the same time in compute like how I'm trying to? It does work if I change the first line in compute to if plug != self.output1 and plug != self.output2, but that means it would calculate twice which is a waste of memory. And you could imagine how bad this would be if there were even more outputs.
I managed to minimize the code to this simple example. I'm scripting it in Python on Maya 2018:
import maya.OpenMayaMPx as OpenMayaMPx
import maya.OpenMaya as OpenMaya
class MyAwesomeNode(OpenMayaMPx.MPxNode):
# Define node properties.
kname = "myAwesomeNode"
kplugin_id = OpenMaya.MTypeId(0x90000005)
# Define node attributes.
in_val = OpenMaya.MObject()
output1 = OpenMaya.MObject()
output2 = OpenMaya.MObject()
def __init__(self):
OpenMayaMPx.MPxNode.__init__(self)
def compute(self, plug, data):
# Only operate on output1 attribute.
if plug != self.output1:
return OpenMaya.kUnknownParameter
# Get input value.
val = data.inputValue(MyAwesomeNode.in_val).asFloat()
# Set output 2.
# This fails when setting up the node and during undos.
out_plug_2 = data.outputValue(self.output2)
if val > 0:
out_plug_2.setFloat(1)
else:
out_plug_2.setFloat(0)
out_plug_2.setClean()
# Set output 1.
# This works as expected.
out_plug_1 = data.outputValue(self.output1)
out_plug_1.setFloat(val)
out_plug_1.setClean()
data.setClean(plug)
return True
def creator():
return OpenMayaMPx.asMPxPtr(MyAwesomeNode())
def initialize():
nattr = OpenMaya.MFnNumericAttribute()
MyAwesomeNode.output2 = nattr.create("output2", "output2", OpenMaya.MFnNumericData.kFloat)
nattr.setWritable(False)
nattr.setStorable(False)
MyAwesomeNode.addAttribute(MyAwesomeNode.output2)
MyAwesomeNode.output1 = nattr.create("output1", "output1", OpenMaya.MFnNumericData.kFloat)
nattr.setWritable(False)
nattr.setStorable(False)
MyAwesomeNode.addAttribute(MyAwesomeNode.output1)
MyAwesomeNode.in_val = nattr.create("input", "input", OpenMaya.MFnNumericData.kFloat, 1)
nattr.setKeyable(True)
MyAwesomeNode.addAttribute(MyAwesomeNode.in_val)
MyAwesomeNode.attributeAffects(MyAwesomeNode.in_val, MyAwesomeNode.output2)
MyAwesomeNode.attributeAffects(MyAwesomeNode.in_val, MyAwesomeNode.output1)
def initializePlugin(obj):
plugin = OpenMayaMPx.MFnPlugin(obj, "Me", "1.0", "Any")
try:
plugin.registerNode(MyAwesomeNode.kname, MyAwesomeNode.kplugin_id, creator, initialize)
except:
raise RuntimeError, "Failed to register node: '{}'".format(MyAwesomeNode.kname)
def uninitializePlugin(obj):
plugin = OpenMayaMPx.MFnPlugin(obj)
try:
plugin.deregisterNode(MyAwesomeNode.kplugin_id)
except:
raise RuntimeError, "Failed to register node: '{}'".format(MyAwesomeNode.kname)
# Example usage of node
if __name__ == "__main__":
import maya.cmds as cmds
cmds.createNode("transform", name="result")
cmds.setAttr("result.displayLocalAxis", True)
cmds.createNode("myAwesomeNode", name="myAwesomeNode")
cmds.connectAttr("myAwesomeNode.output1", "result.translateX")
# This output fails.
cmds.polyCube(name="cube")
cmds.setAttr("cube.translate", 0, 3, 0)
cmds.connectAttr("myAwesomeNode.output2", "cube.scaleX")
cmds.connectAttr("myAwesomeNode.output2", "cube.scaleY")
cmds.connectAttr("myAwesomeNode.output2", "cube.scaleZ")

I have a solution which is working as expected. All outputs still must go through compute() but only one output will do the actual heavy calculations.
When going through compute, it checks all output plugs if any are clean. If all are dirty, then we need to re-calculate, otherwise if we do find one clean plug we can just use the cached values we saved earlier.
Here's an example:
import maya.OpenMayaMPx as OpenMayaMPx
import maya.OpenMaya as OpenMaya
class MyAwesomeNode(OpenMayaMPx.MPxNode):
# Define node properties.
kname = "myAwesomeNode"
kplugin_id = OpenMaya.MTypeId(0x90000005)
# Define node attributes.
in_val = OpenMaya.MObject()
output1 = OpenMaya.MObject()
output2 = OpenMaya.MObject()
def __init__(self):
OpenMayaMPx.MPxNode.__init__(self)
# Store value here.
self.cached_value = 0
def compute(self, plug, data):
# Include all outputs here.
if plug != self.output1 and plug != self.output2:
return OpenMaya.kUnknownParameter
# Get plugs.
val = data.inputValue(MyAwesomeNode.in_val).asFloat()
out_plug_1 = data.outputValue(self.output1)
out_plug_2 = data.outputValue(self.output2)
dep_node = OpenMaya.MFnDependencyNode(self.thisMObject())
# Determine if this output needs to recalculate or simply use cached values.
use_cache_values = False
for name in ["output1", "output2"]:
mplug = dep_node.findPlug(name)
if data.isClean(mplug):
# If we find a clean plug then just use cached values.
use_cache_values = True
break
if use_cache_values:
# Use cached value.
value = self.cached_value
else:
# Calculate value.
# We potentially can make big computations here.
self.cached_value = val
value = val
# Set output 1.
if plug == self.output1:
out_plug_1.setFloat(value)
out_plug_1.setClean()
# Set output 2.
if plug == self.output2:
if value > 0:
out_plug_2.setFloat(1)
else:
out_plug_2.setFloat(0)
out_plug_2.setClean()
data.setClean(plug)
return True
def creator():
return OpenMayaMPx.asMPxPtr(MyAwesomeNode())
def initialize():
nattr = OpenMaya.MFnNumericAttribute()
MyAwesomeNode.output2 = nattr.create("output2", "output2", OpenMaya.MFnNumericData.kFloat)
nattr.setWritable(False)
nattr.setStorable(False)
MyAwesomeNode.addAttribute(MyAwesomeNode.output2)
MyAwesomeNode.output1 = nattr.create("output1", "output1", OpenMaya.MFnNumericData.kFloat)
nattr.setWritable(False)
nattr.setStorable(False)
MyAwesomeNode.addAttribute(MyAwesomeNode.output1)
MyAwesomeNode.in_val = nattr.create("input", "input", OpenMaya.MFnNumericData.kFloat, -1)
nattr.setKeyable(True)
MyAwesomeNode.addAttribute(MyAwesomeNode.in_val)
# Include both outputs.
MyAwesomeNode.attributeAffects(MyAwesomeNode.in_val, MyAwesomeNode.output1)
MyAwesomeNode.attributeAffects(MyAwesomeNode.in_val, MyAwesomeNode.output2)
def initializePlugin(obj):
plugin = OpenMayaMPx.MFnPlugin(obj, "Me", "1.0", "Any")
try:
plugin.registerNode(MyAwesomeNode.kname, MyAwesomeNode.kplugin_id, creator, initialize)
except:
raise RuntimeError, "Failed to register node: '{}'".format(MyAwesomeNode.kname)
def uninitializePlugin(obj):
plugin = OpenMayaMPx.MFnPlugin(obj)
try:
plugin.deregisterNode(MyAwesomeNode.kplugin_id)
except:
raise RuntimeError, "Failed to register node: '{}'".format(MyAwesomeNode.kname)
# Example usage of node
if __name__ == "__main__":
import maya.cmds as cmds
cmds.createNode("transform", name="result")
cmds.setAttr("result.displayLocalAxis", True)
cmds.createNode("myAwesomeNode", name="myAwesomeNode")
cmds.connectAttr("myAwesomeNode.output1", "result.translateX")
# This output fails.
cmds.polyCube(name="cube")
cmds.setAttr("cube.translate", 0, 3, 0)
cmds.connectAttr("myAwesomeNode.output2", "cube.scaleX")
cmds.connectAttr("myAwesomeNode.output2", "cube.scaleY")
cmds.connectAttr("myAwesomeNode.output2", "cube.scaleZ")
The outputs seem to be reacting ok when re-opening the file, when importing it to a new scene, and referencing it. I just need to transfer the same idea to c++ then it'll be golden.

Related

Python Runtime Profiler?

Most python profilers are made for python programs or scripts, in my case I'm working with a python plugin for a third-party app (blender 3d), therefore the profiling needs to be sampled in real-time while the user is interacting with the plugin.
I'm currently trying an injection strategy, which consists of procedurally searching through all plugin modules, & injecting a profiler wrapper to every single function.
see below, this is what my current profiler looks like
I'm wondering if there are other profilers out there that can be used for run-time scenarios such as plugins
class ModuleProfiler:
#profiler is running?
allow = False #must be True in order to start the profiler
activated = False #read only indication if profiler has been activated
#please define your plugin main module here
plugin_main_module = "MyBlenderPlugin"
#function calls registry
registry = {}
#ignore parameters, typically ui functions/modules
ignore_fcts = [
"draw",
"foo",
]
ignore_module = [
"interface_drawing",
]
event_prints = True #print all event?
#classmethod
def print_registry(cls):
"""print all registered benchmark"""
#generate total
for k,v in cls.registry.copy().items():
cls.registry[k]["averagetime"] = v["runtime"]/v["calls"]
print("")
print("PROFILER: PRINTING OUTCOME")
sorted_registry = dict(sorted(cls.registry.items(), key=lambda item:item[1]["runtime"], reverse=False))
for k,v in sorted_registry.items():
print("\n",k,":")
for a,val in v.items():
print(" "*6,a,":",val)
return None
#classmethod
def update_registry(cls, fct, exec_time=0):
"""update internal benchmark with new data"""
key = f"{fct.__module__}.{fct.__name__}"
r = cls.registry.get(key)
if (r is None):
cls.registry[key] = {}
cls.registry[key]["calls"] = 0
cls.registry[key]["runtime"] = 0
r = cls.registry[key]
r["calls"] +=1
r["runtime"] += exec_time
return None
#classmethod
def profile_wrap(cls, fct):
"""wrap any functions with our benchmark & call-counter"""
#ignore some function?
if (fct.__name__ in cls.ignore_fcts):
return fct
import functools
import time
#functools.wraps(fct)
def inner(*args,**kwargs):
t = time.time()
r = fct(*args,**kwargs)
cls.update_registry(fct, exec_time=time.time()-t)
if cls.event_prints:
print(f"PROFILER : {fct.__module__}.{fct.__name__} : {time.time()-t}")
return r
return inner
#classmethod
def start(cls):
"""inject the wrapper for every functions of every sub-modules of our plugin
used for benchmark or debugging purpose only"""
if (not cls.allow):
return None
cls.activated = True
import types
import sys
def is_function(obj):
"""check if given object is a function"""
return isinstance(obj, types.FunctionType)
print("")
#for all modules in sys.modules
for mod_k,mod in sys.modules.copy().items():
#separate module componments names
mod_list = mod_k.split('.')
#fileter what isn't ours
if (mod_list[0]!=cls.plugin_main_module):
continue
#ignore some modules?
if any([m in cls.ignore_module for m in mod_list]):
continue
print("PROFILER_SEARCH : ",mod_k)
#for each objects found in module
for ele_k,ele in mod.__dict__.items():
#if it does not have a name, skip
if (not hasattr(ele,"__name__")):
continue
#we have a global function
elif is_function(ele):
print(f" INJECT LOCAL_FUNCTION: {mod_k}.{ele_k}")
mod.__dict__[ele_k] = cls.profile_wrap(ele)
#then we have a homebrewed class? search for class.fcts
#class.fcts implementation is not flawless, need to investigate issue(s)
elif repr(ele).startswith(f"<class '{cls.plugin_main_module}."):
for class_k,class_e in ele.__dict__.items():
if is_function(class_e):
print(f" INJECT CLASS_FUNCTION: {mod_k}.{ele_k}.{class_k}")
setattr( mod.__dict__[ele_k], class_k, cls.profile_wrap(class_e),) #class.__dict__ are mapping proxies, need to assign this way,
continue
print("")
return None
ModuleProfiler.allow = True
ModuleProfiler.plugin_main_module = "MyModule"
ModuleProfiler.start()

Is there a way to give a function access to the (external-scope) name of the variable being passed in? [duplicate]

This question already has answers here:
Getting the name of a variable as a string
(32 answers)
Closed 4 months ago.
Is it possible to get the original variable name of a variable passed to a function? E.g.
foobar = "foo"
def func(var):
print var.origname
So that:
func(foobar)
Returns:
>>foobar
EDIT:
All I was trying to do was make a function like:
def log(soup):
f = open(varname+'.html', 'w')
print >>f, soup.prettify()
f.close()
.. and have the function generate the filename from the name of the variable passed to it.
I suppose if it's not possible I'll just have to pass the variable and the variable's name as a string each time.
EDIT: To make it clear, I don't recommend using this AT ALL, it will break, it's a mess, it won't help you in any way, but it's doable for entertainment/education purposes.
You can hack around with the inspect module, I don't recommend that, but you can do it...
import inspect
def foo(a, f, b):
frame = inspect.currentframe()
frame = inspect.getouterframes(frame)[1]
string = inspect.getframeinfo(frame[0]).code_context[0].strip()
args = string[string.find('(') + 1:-1].split(',')
names = []
for i in args:
if i.find('=') != -1:
names.append(i.split('=')[1].strip())
else:
names.append(i)
print names
def main():
e = 1
c = 2
foo(e, 1000, b = c)
main()
Output:
['e', '1000', 'c']
To add to Michael Mrozek's answer, you can extract the exact parameters versus the full code by:
import re
import traceback
def func(var):
stack = traceback.extract_stack()
filename, lineno, function_name, code = stack[-2]
vars_name = re.compile(r'\((.*?)\).*$').search(code).groups()[0]
print vars_name
return
foobar = "foo"
func(foobar)
# PRINTS: foobar
Looks like Ivo beat me to inspect, but here's another implementation:
import inspect
def varName(var):
lcls = inspect.stack()[2][0].f_locals
for name in lcls:
if id(var) == id(lcls[name]):
return name
return None
def foo(x=None):
lcl='not me'
return varName(x)
def bar():
lcl = 'hi'
return foo(lcl)
bar()
# 'lcl'
Of course, it can be fooled:
def baz():
lcl = 'hi'
x='hi'
return foo(lcl)
baz()
# 'x'
Moral: don't do it.
Another way you can try if you know what the calling code will look like is to use traceback:
def func(var):
stack = traceback.extract_stack()
filename, lineno, function_name, code = stack[-2]
code will contain the line of code that was used to call func (in your example, it would be the string func(foobar)). You can parse that to pull out the argument
You can't. It's evaluated before being passed to the function. All you can do is pass it as a string.
#Ivo Wetzel's answer works in the case of function call are made in one line, like
e = 1 + 7
c = 3
foo(e, 100, b=c)
In case that function call is not in one line, like:
e = 1 + 7
c = 3
foo(e,
1000,
b = c)
below code works:
import inspect, ast
def foo(a, f, b):
frame = inspect.currentframe()
frame = inspect.getouterframes(frame)[1]
string = inspect.findsource(frame[0])[0]
nodes = ast.parse(''.join(string))
i_expr = -1
for (i, node) in enumerate(nodes.body):
if hasattr(node, 'value') and isinstance(node.value, ast.Call)
and hasattr(node.value.func, 'id') and node.value.func.id == 'foo' # Here goes name of the function:
i_expr = i
break
i_expr_next = min(i_expr + 1, len(nodes.body)-1)
lineno_start = nodes.body[i_expr].lineno
lineno_end = nodes.body[i_expr_next].lineno if i_expr_next != i_expr else len(string)
str_func_call = ''.join([i.strip() for i in string[lineno_start - 1: lineno_end]])
params = str_func_call[str_func_call.find('(') + 1:-1].split(',')
print(params)
You will get:
[u'e', u'1000', u'b = c']
But still, this might break.
You can use python-varname package
from varname import nameof
s = 'Hey!'
print (nameof(s))
Output:
s
Package below:
https://github.com/pwwang/python-varname
For posterity, here's some code I wrote for this task, in general I think there is a missing module in Python to give everyone nice and robust inspection of the caller environment. Similar to what rlang eval framework provides for R.
import re, inspect, ast
#Convoluted frame stack walk and source scrape to get what the calling statement to a function looked like.
#Specifically return the name of the variable passed as parameter found at position pos in the parameter list.
def _caller_param_name(pos):
#The parameter name to return
param = None
#Get the frame object for this function call
thisframe = inspect.currentframe()
try:
#Get the parent calling frames details
frames = inspect.getouterframes(thisframe)
#Function this function was just called from that we wish to find the calling parameter name for
function = frames[1][3]
#Get all the details of where the calling statement was
frame,filename,line_number,function_name,source,source_index = frames[2]
#Read in the source file in the parent calling frame upto where the call was made
with open(filename) as source_file:
head=[source_file.next() for x in xrange(line_number)]
source_file.close()
#Build all lines of the calling statement, this deals with when a function is called with parameters listed on each line
lines = []
#Compile a regex for matching the start of the function being called
regex = re.compile(r'\.?\s*%s\s*\(' % (function))
#Work backwards from the parent calling frame line number until we see the start of the calling statement (usually the same line!!!)
for line in reversed(head):
lines.append(line.strip())
if re.search(regex, line):
break
#Put the lines we have groked back into sourcefile order rather than reverse order
lines.reverse()
#Join all the lines that were part of the calling statement
call = "".join(lines)
#Grab the parameter list from the calling statement for the function we were called from
match = re.search('\.?\s*%s\s*\((.*)\)' % (function), call)
paramlist = match.group(1)
#If the function was called with no parameters raise an exception
if paramlist == "":
raise LookupError("Function called with no parameters.")
#Use the Python abstract syntax tree parser to create a parsed form of the function parameter list 'Name' nodes are variable names
parameter = ast.parse(paramlist).body[0].value
#If there were multiple parameters get the positional requested
if type(parameter).__name__ == 'Tuple':
#If we asked for a parameter outside of what was passed complain
if pos >= len(parameter.elts):
raise LookupError("The function call did not have a parameter at postion %s" % pos)
parameter = parameter.elts[pos]
#If there was only a single parameter and another was requested raise an exception
elif pos != 0:
raise LookupError("There was only a single calling parameter found. Parameter indices start at 0.")
#If the parameter was the name of a variable we can use it otherwise pass back None
if type(parameter).__name__ == 'Name':
param = parameter.id
finally:
#Remove the frame reference to prevent cyclic references screwing the garbage collector
del thisframe
#Return the parameter name we found
return param
If you want a Key Value Pair relationship, maybe using a Dictionary would be better?
...or if you're trying to create some auto-documentation from your code, perhaps something like Doxygen (http://www.doxygen.nl/) could do the job for you?
I wondered how IceCream solves this problem. So I looked into the source code and came up with the following (slightly simplified) solution. It might not be 100% bullet-proof (e.g. I dropped get_text_with_indentation and I assume exactly one function argument), but it works well for different test cases. It does not need to parse source code itself, so it should be more robust and simpler than previous solutions.
#!/usr/bin/env python3
import inspect
from executing import Source
def func(var):
callFrame = inspect.currentframe().f_back
callNode = Source.executing(callFrame).node
source = Source.for_frame(callFrame)
expression = source.asttokens().get_text(callNode.args[0])
print(expression, '=', var)
i = 1
f = 2.0
dct = {'key': 'value'}
obj = type('', (), {'value': 42})
func(i)
func(f)
func(s)
func(dct['key'])
func(obj.value)
Output:
i = 1
f = 2.0
s = string
dct['key'] = value
obj.value = 42
Update: If you want to move the "magic" into a separate function, you simply have to go one frame further back with an additional f_back.
def get_name_of_argument():
callFrame = inspect.currentframe().f_back.f_back
callNode = Source.executing(callFrame).node
source = Source.for_frame(callFrame)
return source.asttokens().get_text(callNode.args[0])
def func(var):
print(get_name_of_argument(), '=', var)
If you want to get the caller params as in #Matt Oates answer answer without using the source file (ie from Jupyter Notebook), this code (combined from #Aeon answer) will do the trick (at least in some simple cases):
def get_caller_params():
# get the frame object for this function call
thisframe = inspect.currentframe()
# get the parent calling frames details
frames = inspect.getouterframes(thisframe)
# frame 0 is the frame of this function
# frame 1 is the frame of the caller function (the one we want to inspect)
# frame 2 is the frame of the code that calls the caller
caller_function_name = frames[1][3]
code_that_calls_caller = inspect.findsource(frames[2][0])[0]
# parse code to get nodes of abstract syntact tree of the call
nodes = ast.parse(''.join(code_that_calls_caller))
# find the node that calls the function
i_expr = -1
for (i, node) in enumerate(nodes.body):
if _node_is_our_function_call(node, caller_function_name):
i_expr = i
break
# line with the call start
idx_start = nodes.body[i_expr].lineno - 1
# line with the end of the call
if i_expr < len(nodes.body) - 1:
# next expression marks the end of the call
idx_end = nodes.body[i_expr + 1].lineno - 1
else:
# end of the source marks the end of the call
idx_end = len(code_that_calls_caller)
call_lines = code_that_calls_caller[idx_start:idx_end]
str_func_call = ''.join([line.strip() for line in call_lines])
str_call_params = str_func_call[str_func_call.find('(') + 1:-1]
params = [p.strip() for p in str_call_params.split(',')]
return params
def _node_is_our_function_call(node, our_function_name):
node_is_call = hasattr(node, 'value') and isinstance(node.value, ast.Call)
if not node_is_call:
return False
function_name_correct = hasattr(node.value.func, 'id') and node.value.func.id == our_function_name
return function_name_correct
You can then run it as this:
def test(*par_values):
par_names = get_caller_params()
for name, val in zip(par_names, par_values):
print(name, val)
a = 1
b = 2
string = 'text'
test(a, b,
string
)
to get the desired output:
a 1
b 2
string text
Since you can have multiple variables with the same content, instead of passing the variable (content), it might be safer (and will be simpler) to pass it's name in a string and get the variable content from the locals dictionary in the callers stack frame. :
def displayvar(name):
import sys
return name+" = "+repr(sys._getframe(1).f_locals[name])
If it just so happens that the variable is a callable (function), it will have a __name__ property.
E.g. a wrapper to log the execution time of a function:
def time_it(func, *args, **kwargs):
start = perf_counter()
result = func(*args, **kwargs)
duration = perf_counter() - start
print(f'{func.__name__} ran in {duration * 1000}ms')
return result

How to write test cases for a init function by Unit test in python?

I'm new in Python. I have a __init__ function with information about data sink config and I have to write some test cases fot it. But init is a function which doesn't return any result so confuses me. Can you show me some ways to implement test case for it ? Thank you.
My function :
class DataSinkConfig():
DEFAULT_ENTRIES_TABLE = 'entries'
DEFAULT_OUTLIERS_TABLE = 'outliers'
DEFAULT_OUTLIERCAUSES_TABLE = 'outlier_causes'
# Note
def __init__(self, dbconf):
insopt = dbconf.get('insopt')
if insopt is None:
self.entries_table_name = self.DEFAULT_ENTRIES_TABLE
self.outliers_table_name = self.DEFAULT_OUTLIERS_TABLE
self.outlier_causes_table_name = self.DEFAULT_OUTLIERCAUSES_TABLE
else:
try:
dict(insopt)
except Exception as e:
raise Exception(
"Invalid value {} for parameter 'insopt'.".format(insopt))
self.entries_table_name = self.__complement_item(insopt, 'entry_tbl', self.DEFAULT_ENTRIES_TABLE)
self.outliers_table_name = self.__complement_item(insopt, 'outlier_tbl', self.DEFAULT_OUTLIERS_TABLE)
self.outlier_causes_table_name = self.__complement_item(insopt, 'cause_tbl', self.DEFAULT_OUTLIERCAUSES_TABLE)
My test is stuck :
import unittest
import DataSinkConfig, DataSourceConfig
class TestDataSink(unittest.TestCase):
def setUp(self):
#pass
if __name__ == "__main__":
unittest.main()
You can do stuff like check the state on the object.
Obviously you can move DataSink instantiation to the setUp method.
import unittest
import DataSinkConfig, DataSourceConfig
class TestDataSink(unittest.TestCase):
def setUp(self):
#pass
def test_data_sink(self):
dsc = DataSinkConfig(db_conf)
self.assertEqual(dsc.entries_table_name, "WHATREVER COMPLIMENT ITEM WAS SUPPOSE TO SET")
if __name__ == "__main__":
unittest.main()

How do you obtain the body of a function?

I would need to separate the signature from the body of a function in Python, i.e. :
def mult(a, b):
"Multiplication"
if a or b is None:
# border case
return None
else:
return a * b
(the function is there for demonstration only).
I know that inspect.getsourcelines(mult) will get me the full code, but I would like to have only the body, i.e. less the signature and the docstring:
if a or b is None:
# border case
return None
else:
return a * b
Is there an elegant way to obtain that, ideally using Python3 built-in parsing tools, rather than string manipulation?
As I don't know any function that does this, here is a homemade function that should work.
I don't have time for more in-depth use-cases. So i'll leave my incomplete answer here.
def get_body(func):
""" Using the magic method __doc__, we KNOW the size of the docstring.
We then, just substract this from the total length of the function
"""
try:
lines_to_skip = len(func.__doc__.split('\n'))
except AttributeError:
lines_to_skip = 0
lines = getsourcelines(func)[0]
return ''.join( lines[lines_to_skip+1:] )
def ex_0(a, b):
""" Please !
Don't
Show
This
"""
if a or b is None:
# border case
return None
else:
return a * b
def ex_1(a, b):
''' Please !Don'tShowThis'''
if a or b is None:
# border case
return None
else:
return a * b
def ex_2(a, b):
if a or b is None:
# border case
return None
else:
return a * b
def ex_3(bullshit, hello):
pass
get_body(ex_0)
# if a or b is None:
# # border case
# return None
# else:
# return a * b
get_body(ex_1)
# if a or b is None:
# # border case
# return None
# else:
# return a * b
get_body(ex_2)
# if a or b is None:
# # border case
# return None
# else:
# return a * b
get_body(ex_3)
# pass
you can use this also
i = inspect.getsource(fun_name).index('\n')
j = inspect.getsource(fun_name).rindex(':',0,i)
print(inspect.getsource(fun_name)[j+1:])
If you know your code will always have a non-empty docstring like a above, You can do this with the built in inspect library.
The below snippet should get you acquainted with how to see the source code body of a function.
def hello_world():
"""Python port of original Fortran code"""
print("hello_world")
import inspect
source = inspect.getsource(hello_world)
doc = hello_world.__doc__
code = source.split(doc)
body = code[1]
print(body)
EDIT:
I thought about this some more, you can first remove the defintion line and then strip out the docstring if it exists:
def hello_world():
"""Python port of original Fortran code"""
print("hello_world")
import inspect
source = inspect.getsource(hello_world)
doc = hello_world.__doc__
code = source.split(':',1)
body= code[1].replace(doc, "")
body = body.replace('""""""',"")
print(body)
Thanks a lot for the inspiring answers! I tried to take this further and I found a solution that uses Python's ast facilities.
import inspect, ast
from astunparse import unparse
def mult(a,
b): # comment
"""
This is a 'doctstring', "hello"
"""
if not a is None or b is None:
# border case
return None
else:
return a * b
def get_body(f, doc_string=False):
"""
Get the body text of a function, i.e. without the signature.
NOTE: Comments are stripped.
"""
COMMENT_START = ("'", '"')
code = inspect.getsource(mult)
# print("Function's code:\n", code)
module_tree = ast.parse(code) # this creates a module
# print("Dump:\n", ast.dump(module_tree))
# the first element of the module is evidently the function:
function_body = module_tree.body[0].body
# strip the code lines to remove the additional lines:
lines = [unparse(code_line).strip() for code_line in function_body]
# for i, line in enumerate(lines): print("%s: %s" % (i, line.strip()))
# in principle the notion of body contains the docstring:
if not doc_string:
lines = (line for line in lines if not line.startswith(COMMENT_START))
return '\n'.join(lines).strip()
s =get_body(mult)
print("---------")
print(s)
Here is the result:
$ python3 function_body.py
---------
if ((not (a is None)) or (b is None)):
return None
else:
return (a * b)
You can choose if you want the docstring or not. The downside of this approach (which should not be a downside in my use case) is that comments are stripped.
I also left some commented print statements, in case someone would like to explore the various steps.

Getting Python's nosetests results in a tearDown() method

I want to be able to get the result of a particular test method and output it inside the teardown method, while using the nose test runner.
There is a very good example here.
But unfortunately, running nosetests example.py does not work, since nose doesn't seem to like the fact that the run method in the superclass is being overridden:
AttributeError: 'ResultProxy' object has no attribute 'wasSuccessful'
Caveat: the following doesn't actually access the test during the tearDown, but it does access each result.
You might want to write a nose plugin (see the API documentation here). The method that you are probably interested in is afterTest(), which is run... after the test. :) Though, depending on your exact application, handleError()/handleFailure() or finalize() might actually be more useful.
Here is an example plugin that accesses the result of a test immediately after it is executed.
from nose.plugins import Plugin
import logging
log = logging.getLogger('nose.plugins.testnamer')
class ReportResults(Plugin):
def __init__(self, *args, **kwargs):
super(ReportResults, self).__init__(*args, **kwargs)
self.passes = 0
self.failures = 0
def afterTest(self, test):
if test.passed:
self.passes += 1
else:
self.failures += 1
def finalize(self, result):
print "%d successes, %d failures" % (self.passes, self.failures)
This trivial example merely reports the number of passes and failures (like the link you included, but I'm sure you can extend it to do something more interesting (here's another fun idea). To use this, make sure that it is installed in Nose (or load it into a custom runner), and then activate it with --with-reportresults.
If you are OK with adding some boilerplate code to the tests, something like the following might work.
In MyTest1, tearDown is called at the end of each test, and the value of self.result has been set to a tuple containing the method name and a dictionary (but you could set that to whatever you like). The inspect module is used to get the method name, so tearDown knows which test just ran.
In MyTest2, all the results are saved in a dictionary (results), which you can do with what you like in the tearDownClass method.
import inspect
import unittest
class MyTest1(unittest.TestCase):
result = None
def tearDown(self):
print "tearDown:", self.result
def test_aaa(self):
frame = inspect.currentframe()
name = inspect.getframeinfo(frame).function
del frame
self.result = (name, None)
x = 1 + 1
self.assertEqual(x, 2)
self.result = (name, dict(x=x))
def test_bbb(self):
frame = inspect.currentframe()
name = inspect.getframeinfo(frame).function
del frame
self.result = (name, None)
# Intentional fail.
x = -1
self.assertEqual(x, 0)
self.result = (name, dict(x=x))
class MyTest2(unittest.TestCase):
results = {}
#classmethod
def tearDownClass(cls):
print "tearDownClass:", cls.results
def test_aaa(self):
frame = inspect.currentframe()
name = inspect.getframeinfo(frame).function
del frame
self.results[name] = None
x = 1 + 1
self.assertEqual(x, 2)
self.results[name] = dict(x=x)
def test_bbb(self):
frame = inspect.currentframe()
name = inspect.getframeinfo(frame).function
del frame
self.results[name] = None
x = -1
self.assertEqual(x, 0)
self.results[name] = dict(x=x)
if __name__ == '__main__':
unittest.main()

Categories