diff --git a/aiida/tools/dbexporters/tcod.py b/aiida/tools/dbexporters/tcod.py index 8dfbaaf400..ee108d8044 100644 --- a/aiida/tools/dbexporters/tcod.py +++ b/aiida/tools/dbexporters/tcod.py @@ -385,16 +385,18 @@ def _inline_to_standalone_script(calc): args = ["{}=load_node('{}')".format(x, input_dict[x].uuid) for x in input_dict.keys()] args_string = ",\n ".join(sorted(args)) + code_string = calc.get_attr('source_file').encode('utf-8') + if calc.get_attr('namespace', '__main__').startswith('aiida.'): + code_string = "from {} import {}".format(calc.get_attr('namespace', '__main__'), + calc.get_attr('function_name','f')) return """#!/usr/bin/env runaiida {} for key, value in {}( {} - )[1].iteritems(): + ).iteritems(): value.store() -""".format(calc.get_attr('source_file').encode('utf-8'), - calc.get_attr('function_name','f'), - args_string) +""".format(code_string, calc.get_attr('function_name','f'), args_string) def _collect_calculation_data(calc): diff --git a/aiida/utils/calculation.py b/aiida/utils/calculation.py index 7d913b462b..e0fe8acda8 100644 --- a/aiida/utils/calculation.py +++ b/aiida/utils/calculation.py @@ -39,3 +39,7 @@ def add_source_info(node, func): node._set_attr("source_file", source) except (IOError, OSError): pass + try: + node._set_attr("namespace", func.func_globals["__name__"]) + except Exception: + pass