erp5 jupyter kernel: print capturing using AST processor to modify print calls

Before, we we're redirecting `sys.stdout` and this doesn't play nice with the
distribute architecture of ERP5 and our Jupyter kernel needs to be adjusted for this.

So, we're now using an AST processor to fix print calls. It will modify the print
and make it write to a different file-like object. All the writes are collected
after code execution and sent to Jupyter.

It's still necesasry though to fix print inside other libraries. But for this
deeper investigation is necessary because we cannot replace print as a statement
inside `exec` contetx, it needs to be used as a function. Code can be compiled to
run with `print` as a function, but then external libraries calls will be
broken.

bt5/erp5_data_notebook/ExtensionTemplateItem/portal_components/extension.erp5.JupyterCompile.py

@@ -176,6 +176,8 @@ def Base_runJupyterCode(self, jupyter_code, old_notebook_context):
   # Something like -- user_context = globals(); user_context['context'] = self;
   user_context = {}
   
+  output = ''
+
   # Saving the initial globals dict so as to compare it after code execution
   globals_dict = globals()
   result_string = ''
@@ -193,8 +195,11 @@ def Base_runJupyterCode(self, jupyter_code, old_notebook_context):
     
     # Fixing "normal" imports and detecting environment object usage
     import_fixer = ImportFixer()
+    print_fixer = PrintFixer()
     environment_collector = EnvironmentParser()
     ast_node = import_fixer.visit(ast_node)
+    ast_node = print_fixer.visit(ast_node)
+    ast.fix_missing_locations(ast_node)
     
     # The collector also raises errors when environment.define and undefine
     # calls are made incorrectly, so we need to capture them to propagate
@@ -205,7 +210,6 @@ def Base_runJupyterCode(self, jupyter_code, old_notebook_context):
       transaction.abort()
       return getErrorMessageForException(self, e, notebook_context)
     
-    
     # Get the node list from the parsed tree
     nodelist = ast_node.body
 
@@ -225,13 +229,6 @@ def Base_runJupyterCode(self, jupyter_code, old_notebook_context):
       elif interactivity == 'last':
         to_run_exec, to_run_interactive = nodelist[:-1], nodelist[-1:]
       
-      # TODO: fix this global handling by replacing the print statement with
-      # a custom print function. Tip: create an ast.NodeTransformer, like the
-      # one used to fix imports.
-      old_stdout = sys.stdout
-      result = StringIO()
-      sys.stdout = result
-      
       # Variables used at the display hook to get the proper form to display
       # the last returning variable of any code cell.
       display_data = {'result': '', 'mime_type': None}
@@ -255,7 +252,8 @@ def Base_runJupyterCode(self, jupyter_code, old_notebook_context):
         'environment': Environment(),
         '_display_data': display_data,
         '_processor_list': processor_list,
-        '_volatile_variable_list': []
+        '_volatile_variable_list': [],
+        '_print': CustomPrint()
       }
       user_context.update(inject_variable_dict)
       user_context.update(notebook_context['variables'])
@@ -337,6 +335,7 @@ def Base_runJupyterCode(self, jupyter_code, old_notebook_context):
           "func_name": func_name,
           "code": setup_string
         }
+        inject_variable_dict['_print'].write(setup_string)
 
       # Iterating over envinronment.define calls captured by the environment collector
       # that are simple variables and saving them in the setup.
@@ -349,7 +348,7 @@ def Base_runJupyterCode(self, jupyter_code, old_notebook_context):
         user_context['_volatile_variable_list'] += variable
         
       if environment_collector.showEnvironmentSetup():
-        result_string += "%s\n" % str(notebook_context['setup'])
+        inject_variable_dict.write("%s\n" % str(notebook_context['setup']))
 
       # Execute the nodes with 'exec' mode
       for node in to_run_exec:
@@ -381,9 +380,9 @@ def Base_runJupyterCode(self, jupyter_code, old_notebook_context):
           # Clear the portal cache from previous transaction
           return getErrorMessageForException(self, e, notebook_context)
 
-      sys.stdout = old_stdout
       mime_type = display_data['mime_type'] or mime_type
-      result_string += "\n".join(removed_setup_message_list) + result.getvalue() + display_data['result']
+      inject_variable_dict['_print'].write("\n".join(removed_setup_message_list) + display_data['result'])
+    
     
     # Saves a list of all the variables we injected into the user context and
     # shall be deleted before saving the context.
@@ -396,12 +395,13 @@ def Base_runJupyterCode(self, jupyter_code, old_notebook_context):
           notebook_context['variables'][key] = val
         else:
           del user_context[key]
-          result_string += (
+          message = (
             "Cannot serialize the variable named %s whose value is %s, "
             "thus it will not be stored in the context. "
             "You should move it's definition to a function and " 
             "use the environment object to load it.\n"
           ) % (key, val)
+          inject_variable_dict['_print'].write(message)
     
     # Deleting from the variable storage the keys that are not in the user 
     # context anymore (i.e., variables that are deleted by the user).
@@ -409,8 +409,10 @@ def Base_runJupyterCode(self, jupyter_code, old_notebook_context):
       if not key in user_context:
         del notebook_context['variables'][key]
     
+    output = inject_variable_dict['_print'].getCapturedOutputString()
+  
   result = {
-    'result_string': result_string,
+    'result_string': output,
     'notebook_context': notebook_context,
     'status': status,
     'mime_type': mime_type,
@@ -486,6 +488,26 @@ def canSerialize(obj):
       return True
   
   
+class CustomPrint(object):
+  
+  def __init__(self):
+    self.captured_output_list = []
+    
+  def write(self, *args):
+    self.captured_output_list += args
+    
+  def getCapturedOutputString(self):
+    return ''.join(self.captured_output_list)
+    
+
+class PrintFixer(ast.NodeTransformer):
+    
+  def visit_Print(self, node):
+    _print_name_node = ast.Name(id="_print", ctx=ast.Load())
+    node.dest = _print_name_node
+    return node
+  
+
 class EnvironmentParser(ast.NodeTransformer):
   """
     EnvironmentParser class is an AST transformer that walks in the abstract
@@ -741,7 +763,7 @@ def renderAsHtml(self, renderable_object):
   compile_jupyter_locals = compile_jupyter_frame.f_locals
   processor = compile_jupyter_locals['processor_list'].getProcessorFor(renderable_object)
   result, mime_type = processor(renderable_object).process()
-  compile_jupyter_locals['result'].write(result)
+  compile_jupyter_locals['inject_variable_dict']['_print'].write(result)
   compile_jupyter_locals['display_data']['mime_type'] = 'text/html'
 
 def getErrorMessageForException(self, exception, notebook_context):
@@ -963,3 +985,4 @@ def erp5PivotTableUI(self, df):
   iframe_host = self.REQUEST['HTTP_X_FORWARDED_HOST'].split(',')[0]
   url = "https://%s/erp5/Base_displayPivotTableFrame?key=%s" % (iframe_host, key)
   return IFrame(src=url, width='100%', height='500')
+

bt5/erp5_data_notebook/TestTemplateItem/portal_components/test.erp5.testExecuteJupyter.py

@@ -575,7 +575,7 @@ environment.define(x='couscous')
     self.tic()
     self.assertEquals(json.loads(result)['status'], 'ok')
     
-    jupyter_code = "'x' in locals()"
+    jupyter_code = "print 'x' in locals()"
     result = self.portal.Base_executeJupyter(
       reference=reference,
       python_expression=jupyter_code
@@ -704,5 +704,6 @@ context.Base_renderAsHtml(iframe)
     # The big hash in this string was previous calculated using the expect hash
     # of the pivot table page's html.
     pivottable_frame_display_path = 'Base_displayPivotTableFrame?key=853524757258b19805d13beb8c6bd284a7af4a974a96a3e5a4847885df069a74d3c8c1843f2bcc4d4bb3c7089194b57c90c14fe8dd0c776d84ce0868e19ac411'
+    
     self.assertTrue(pivottable_frame_display_path in json_result['code_result'])