Handle comments in multi-line cell outputs

panel/io/handlers.py

@@ -165,6 +165,8 @@ def capture_code_cell(cell):
     parses = False
     while not parses:
         try:
+            if not cell_out.strip():
+                raise SyntaxError
             ast.parse(cell_out)
             parses = True
         except SyntaxError:
@@ -185,7 +187,7 @@ def capture_code_cell(cell):
         return code
 
     # Remove code comments
-    if '#' in cell_out:
+    if '#' in cell_out and not cell_out.count('\n'):
         try:
             # To not remove "#000000"
             cell_tmp = cell_out[:cell_out.index('#')].rstrip()

panel/tests/io/test_handlers.py

@@ -72,6 +72,17 @@
     'test.csv'
 )"""
 
+code_function = """
+def foo():
+    # Comment
+    return 1+1
+"""
+
+code_expr_multi_line_with_comment = """
+(
+  1 + 1 # Comment
+)
+"""
 
 def test_extract_panel_block():
     f = StringIO(md1)
@@ -110,6 +121,26 @@ def test_capture_code_cell_expression():
     _pn__state._cell_outputs['foo'].append(_fig__out)
 """]
 
+def test_capture_code_cell_function():
+    assert capture_code_cell({'id': 'foo', 'source': code_function}) == [
+        '', 'def foo():\n    # Comment\n    return 1+1\n'
+   ]
+
+def test_capture_code_expression_multi_line_with_comment():
+    assert capture_code_cell({'id': 'foo', 'source': code_expr_multi_line_with_comment}) == [
+        '', """\
+_pn__state._cell_outputs['foo'].append(((
+  1 + 1 # Comment
+)
+))
+for _cell__out in _CELL__DISPLAY:
+    _pn__state._cell_outputs['foo'].append(_cell__out)
+_CELL__DISPLAY.clear()
+_fig__out = _get__figure()
+if _fig__out:
+    _pn__state._cell_outputs['foo'].append(_fig__out)
+"""]
+
 
 def test_capture_code_cell_multi_line_expression():
     assert capture_code_cell({'id': 'foo', 'source': code_multi_line}) == [