Merge pull request #761 from merkys/fix_tcodexporter

Fix tcodexporter
aiidateam · Oct 13, 2017 · b5ae353 · b5ae353
2 parents 28eb294 + 4881332
commit b5ae353
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git 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
@@ -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