diff --git a/pysb/core.py b/pysb/core.py index 65f726de1..75297cc49 100644 --- a/pysb/core.py +++ b/pysb/core.py @@ -191,11 +191,14 @@ def __init__(self, name, _export=True): mod_name = frame.f_globals.get('__name__', '__unnamed__') if mod_name in ['IPython.core.interactiveshell', '__main__']: break - if mod_name not in ['pysb.core', 'pysb.macros'] and not \ + if mod_name != 'pysb.core' and not \ mod_name.startswith('importlib.'): self._modules.append(mod_name) if self._function is None: - self._function = frame.f_code.co_name + if mod_name == 'pysb.macros': + self._function = frame.f_back.f_code.co_name + else: + self._function = frame.f_code.co_name frame = frame.f_back def __getstate__(self): diff --git a/pysb/tests/test_core.py b/pysb/tests/test_core.py index 8ca4af54d..ab3f8be3f 100644 --- a/pysb/tests/test_core.py +++ b/pysb/tests/test_core.py @@ -32,6 +32,32 @@ def test_component_names_invalid(): for name in 'a!', '!B', 'A!bC~`\\', '_!', '_!7', '__a01b 999x_x___!': assert_raises(InvalidComponentNameError, Component, name, _export=False) +@with_model +def test_function_introspection(): + # Case 1: Component defined inside function + Monomer('A') + assert A._function == 'test_function_introspection' + + # Case 2: Component defined inside nested function + def define_monomer_b(): + Monomer('B') + define_monomer_b() + assert B._function == 'define_monomer_b' + + # Case 3: Component defined by macro + from pysb.macros import equilibrate + equilibrate(A(), B(), [1, 1]) + + assert model.rules['equilibrate_A_to_B']._function == 'equilibrate' + + # Case 4: Component defined by macro inside function + def define_macro_inside_function(): + Monomer('C') + equilibrate(A(), C(), [2, 2]) + define_macro_inside_function() + assert model.rules['equilibrate_A_to_C']._function == 'equilibrate' + + def test_monomer(): sites = ['x', 'y', 'z'] states = {'y': ['foo', 'bar', 'baz'], 'x': ['e']}