Use AST parser to find python multilines with function handling

docs/sphinxext/ipython_directive.py

@@ -455,14 +455,16 b' class EmbeddedSphinxShell(object):'
         output = []
         savefig = False # keep up with this to clear figure
         multiline = False # to handle line continuation
+        multiline_start = None
         fmtin = self.promptin
 
+        ct = 0
+
         for lineno, line in enumerate(content):
 
             line_stripped = line.strip()
-
             if not len(line):
-                output.append(line) # preserve empty lines in output
+                output.append(line)
                 continue
 
             # handle decorators
@@ -477,50 +479,43 b' class EmbeddedSphinxShell(object):'
                 output.extend([line])
                 continue
 
-            # deal with multilines
-            if not multiline: # not currently on a multiline
-
-                if line_stripped.endswith('\\'): # now we are
+            # deal with lines checking for multiline
+            continuation  = u'   %s:'% ''.join(['.']*(len(str(ct))+2))
+            if not multiline:
+                modified = u"%s %s" % (fmtin % ct, line_stripped)
+                output.append(modified)
+                ct += 1
+                try:
+                    ast.parse(line_stripped)
+                    output.append(u'')
+                except Exception: # on a multiline
                     multiline = True
-                    cont_len = len(str(lineno)) + 2
-                    line_to_process = line.strip('\\')
-                    output.extend([u"%s %s" % (fmtin%lineno,line)])
-                    continue
-                else: # no we're still not
-                    line_to_process = line.strip('\\')
-            else: # we are currently on a multiline
-                line_to_process += line.strip('\\')
-                if line_stripped.endswith('\\'): # and we still are
-                    continuation = '.' * cont_len
-                    output.extend([(u'   %s: '+line_stripped) % continuation])
-                    continue
-                # else go ahead and run this multiline then carry on
-
-            # get output of line
-            self.process_input_line(unicode(line_to_process.strip()),
-                                    store_history=False)
-            out_line = self.cout.getvalue()
-            self.clear_cout()
-
-            # clear current figure if plotted
-            if savefig:
+                    multiline_start = lineno
+                    if line_stripped.startswith('def '):
+                        is_function = True
+            else: # still on a multiline
+                modified = u'%s %s' % (continuation, line)
+                output.append(modified)
+                try:
+                    mod = ast.parse(
+                            '\n'.join(content[multiline_start:lineno+1]))
+                    if isinstance(mod.body[0], ast.FunctionDef):
+                        # check to see if we have the whole function
+                        for element in mod.body[0].body:
+                            if isinstance(element, ast.Return):
+                                multiline = False
+                    else:
+                        output.append(u'')
+                        multiline = False
+                except Exception:
+                    pass
+
+            if savefig: # clear figure if plotted
                 self.ensure_pyplot()
                 self.process_input_line('plt.clf()', store_history=False)
                 self.clear_cout()
                 savefig = False
 
-            # line numbers don't actually matter, they're replaced later
-            if not multiline:
-                in_line = u"%s %s" % (fmtin%lineno,line)
-
-                output.extend([in_line])
-            else:
-                output.extend([(u'   %s: '+line_stripped) % continuation])
-                multiline = False
-            if len(out_line):
-                output.extend([out_line])
-            output.extend([u''])
-
         return output
 
 class IpythonDirective(Directive):