Merge remote-tracking branch 'origin/master'
authorEkaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Mon, 17 Nov 2014 11:18:00 +0000 (14:18 +0300)
committerEkaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Mon, 17 Nov 2014 11:18:00 +0000 (14:18 +0300)
python/ipnb/src/org/jetbrains/plugins/ipnb/format/IpnbParser.java

index 0f1f001dc3b05c9b44f190c797d73077fa50a6d4..15a12fd5dff0a3017d4aba1ddeda20cd057790af 100644 (file)
@@ -3,8 +3,11 @@ package org.jetbrains.plugins.ipnb.format;
 import com.google.common.collect.Lists;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
+import com.google.gson.annotations.SerializedName;
+import com.google.gson.stream.JsonWriter;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.editor.Document;
+import com.intellij.openapi.util.text.StringUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.plugins.ipnb.editor.panels.IpnbEditablePanel;
 import org.jetbrains.plugins.ipnb.editor.panels.IpnbFilePanel;
@@ -14,6 +17,7 @@ import org.jetbrains.plugins.ipnb.format.cells.output.*;
 import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
+import java.io.StringWriter;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -35,12 +39,19 @@ public class IpnbParser {
     if (rawFile == null) return new IpnbFile(new IpnbFileRaw(), Lists.<IpnbCell>newArrayList(), path);
     List<IpnbCell> cells = new ArrayList<IpnbCell>();
     final IpnbWorksheet[] worksheets = rawFile.worksheets;
-    for (IpnbWorksheet worksheet : worksheets) {
-      final List<IpnbCellRaw> rawCells = worksheet.cells;
-      for (IpnbCellRaw rawCell : rawCells) {
+    if (worksheets == null) {
+      for (IpnbCellRaw rawCell : rawFile.cells) {
         cells.add(rawCell.createCell());
       }
     }
+    else {
+      for (IpnbWorksheet worksheet : worksheets) {
+        final List<IpnbCellRaw> rawCells = worksheet.cells;
+        for (IpnbCellRaw rawCell : rawCells) {
+          cells.add(rawCell.createCell());
+        }
+      }
+    }
     return new IpnbFile(rawFile, cells, path);
   }
 
@@ -65,12 +76,25 @@ public class IpnbParser {
     }
 
     final IpnbFileRaw fileRaw = ipnbFile.getRawFile();
-    final IpnbWorksheet worksheet = new IpnbWorksheet();
-    for (IpnbCell cell: ipnbFile.getCells()) {
-      worksheet.cells.add(IpnbCellRaw.fromCell(cell));
+    if (fileRaw.nbformat == 4) {
+      fileRaw.cells.clear();
+      for (IpnbCell cell: ipnbFile.getCells()) {
+        fileRaw.cells.add(IpnbCellRaw.fromCell(cell, fileRaw.nbformat));
+      }
     }
-    fileRaw.worksheets = new IpnbWorksheet[]{worksheet};
-    return gson.toJson(fileRaw);
+    else {
+      final IpnbWorksheet worksheet = new IpnbWorksheet();
+      worksheet.cells.clear();
+      for (IpnbCell cell : ipnbFile.getCells()) {
+        worksheet.cells.add(IpnbCellRaw.fromCell(cell, fileRaw.nbformat));
+      }
+      fileRaw.worksheets = new IpnbWorksheet[]{worksheet};
+    }
+    final StringWriter stringWriter = new StringWriter();
+    final JsonWriter writer = new JsonWriter(stringWriter);
+    writer.setIndent(" ");
+    gson.toJson(fileRaw, fileRaw.getClass(), writer);
+    return stringWriter.toString();
   }
 
   private static void writeToFile(@NotNull final String path, @NotNull final String json) {
@@ -93,10 +117,11 @@ public class IpnbParser {
   }
 
   public static class IpnbFileRaw {
+    IpnbWorksheet[] worksheets;
+    List<IpnbCellRaw> cells = new ArrayList<IpnbCellRaw>();
     Map<String, Object> metadata = new HashMap<String, Object>();
     int nbformat = 3;
     int nbformat_minor;
-    IpnbWorksheet[] worksheets;
   }
 
   private static class IpnbWorksheet {
@@ -105,14 +130,16 @@ public class IpnbParser {
 
   private static class IpnbCellRaw {
     String cell_type;
+    Integer execution_count;
+    Map<String, Object> metadata = new HashMap<String, Object>();
     Integer level;
+    CellOutputRaw[] outputs;
     String[] source;
     String[] input;
     String language;
-    CellOutputRaw[] outputs;
     Integer prompt_number;
 
-    public static IpnbCellRaw fromCell(@NotNull final IpnbCell cell) {
+    public static IpnbCellRaw fromCell(@NotNull final IpnbCell cell, int nbformat) {
       final IpnbCellRaw raw = new IpnbCellRaw();
       if (cell instanceof IpnbMarkdownCell) {
         raw.cell_type = "markdown";
@@ -122,13 +149,19 @@ public class IpnbParser {
         raw.cell_type = "code";
         final ArrayList<CellOutputRaw> outputRaws = new ArrayList<CellOutputRaw>();
         for (IpnbOutputCell outputCell : ((IpnbCodeCell)cell).getCellOutputs()) {
-          outputRaws.add(CellOutputRaw.fromOutput(outputCell));
+          outputRaws.add(CellOutputRaw.fromOutput(outputCell, nbformat));
         }
         raw.outputs = outputRaws.toArray(new CellOutputRaw[outputRaws.size()]);
-        raw.language = ((IpnbCodeCell)cell).getLanguage();
-        raw.input = ((IpnbCodeCell)cell).getSource();
         final Integer promptNumber = ((IpnbCodeCell)cell).getPromptNumber();
-        raw.prompt_number = promptNumber != null && promptNumber >= 0 ? promptNumber : null;
+        if (nbformat == 4) {
+          raw.execution_count = promptNumber != null && promptNumber >= 0 ? promptNumber : null;
+          raw.source = ((IpnbCodeCell)cell).getSource();
+        }
+        else {
+          raw.prompt_number = promptNumber != null && promptNumber >= 0 ? promptNumber : null;
+          raw.language = ((IpnbCodeCell)cell).getLanguage();
+          raw.input = ((IpnbCodeCell)cell).getSource();
+        }
       }
       else if (cell instanceof IpnbRawCell) {
         raw.cell_type = "raw";
@@ -151,7 +184,8 @@ public class IpnbParser {
         for (CellOutputRaw outputRaw : outputs) {
           outputCells.add(outputRaw.createOutput());
         }
-        cell = new IpnbCodeCell(language, input, prompt_number, outputCells);
+        cell = new IpnbCodeCell(language == null ? "python" : language, input == null ? source : input,
+                                prompt_number == null ? execution_count : prompt_number, outputCells);
       }
       else if (cell_type.equals("raw")) {
         cell = new IpnbRawCell();
@@ -168,7 +202,10 @@ public class IpnbParser {
 
   private static class CellOutputRaw {
     String ename;
+    String name;
     String evalue;
+    OutputDataRaw data;
+    Integer execution_count;
     String output_type;
     String png;
     String stream;
@@ -179,82 +216,146 @@ public class IpnbParser {
     Integer prompt_number;
     String[] text;
     String[] traceback;
+    Map<String, Object> metadata;
 
-
-    public static CellOutputRaw fromOutput(@NotNull final IpnbOutputCell outputCell) {
+    public static CellOutputRaw fromOutput(@NotNull final IpnbOutputCell outputCell, int nbformat) {
       final CellOutputRaw raw = new CellOutputRaw();
+      if (!(outputCell instanceof IpnbStreamOutputCell) && !(outputCell instanceof IpnbErrorOutputCell)) {
+        raw.metadata = new HashMap<String, Object>();
+      }
 
       if (outputCell instanceof IpnbPngOutputCell) {
-        raw.png = ((IpnbPngOutputCell)outputCell).getBase64String();
+        if (nbformat == 4) {
+          final OutputDataRaw dataRaw = new OutputDataRaw();
+          dataRaw.png = new String[]{((IpnbPngOutputCell)outputCell).getBase64String()};
+          raw.data = dataRaw;
+        }
+        else {
+          raw.png = ((IpnbPngOutputCell)outputCell).getBase64String();
+        }
         raw.text = outputCell.getText();
         raw.output_type = "display_data";
       }
       else if (outputCell instanceof IpnbSvgOutputCell) {
-        raw.svg = ((IpnbSvgOutputCell)outputCell).getSvg();
+        if (nbformat == 4) {
+          final OutputDataRaw dataRaw = new OutputDataRaw();
+          dataRaw.svg = ((IpnbSvgOutputCell)outputCell).getSvg();
+          raw.data = dataRaw;
+        }
+        else {
+          raw.svg = ((IpnbSvgOutputCell)outputCell).getSvg();
+        }
         raw.text = outputCell.getText();
       }
       else if (outputCell instanceof IpnbJpegOutputCell) {
-        raw.jpeg = ((IpnbJpegOutputCell)outputCell).getBase64String();
+        if (nbformat == 4) {
+          final OutputDataRaw dataRaw = new OutputDataRaw();
+          dataRaw.jpeg = new String[]{((IpnbJpegOutputCell)outputCell).getBase64String()};
+          raw.data = dataRaw;
+        }
+        else {
+          raw.jpeg = ((IpnbJpegOutputCell)outputCell).getBase64String();
+        }
         raw.text = outputCell.getText();
       }
       else if (outputCell instanceof IpnbLatexOutputCell) {
-        raw.latex = ((IpnbLatexOutputCell)outputCell).getLatex();
+        if (nbformat == 4) {
+          final OutputDataRaw dataRaw = new OutputDataRaw();
+          dataRaw.latex = ((IpnbLatexOutputCell)outputCell).getLatex();
+          raw.data = dataRaw;
+        }
+        else {
+          raw.latex = ((IpnbLatexOutputCell)outputCell).getLatex();
+        }
         raw.prompt_number = outputCell.getPromptNumber();
         raw.text = outputCell.getText();
       }
       else if (outputCell instanceof IpnbStreamOutputCell) {
-        raw.stream = ((IpnbStreamOutputCell)outputCell).getStream();
+        if (nbformat == 4) {
+          raw.name = ((IpnbStreamOutputCell)outputCell).getStream();
+        }
+        else {
+          raw.stream = ((IpnbStreamOutputCell)outputCell).getStream();
+        }
         raw.output_type = "stream";
         raw.text = outputCell.getText();
       }
       else if (outputCell instanceof IpnbHtmlOutputCell) {
-        raw.html = ((IpnbHtmlOutputCell)outputCell).getHtmls();
+        if (nbformat == 4) {
+          final OutputDataRaw dataRaw = new OutputDataRaw();
+          dataRaw.html = ((IpnbHtmlOutputCell)outputCell).getHtmls();
+          raw.data = dataRaw;
+        }
+        else {
+          raw.html = ((IpnbHtmlOutputCell)outputCell).getHtmls();
+        }
         raw.text = outputCell.getText();
       }
       else if (outputCell instanceof IpnbErrorOutputCell) {
-        raw.output_type = "pyerr";
+        raw.output_type = nbformat == 4 ? "error" : "pyerr";
         raw.evalue = ((IpnbErrorOutputCell)outputCell).getEvalue();
         raw.ename = ((IpnbErrorOutputCell)outputCell).getEname();
         raw.traceback = outputCell.getText();
       }
       else if (outputCell instanceof IpnbOutOutputCell) {
-        raw.output_type = "pyout";
-        raw.text = outputCell.getText();
-        raw.prompt_number = outputCell.getPromptNumber();
+        if (nbformat == 4) {
+          raw.execution_count = outputCell.getPromptNumber();
+          raw.output_type = "execute_result";
+          final OutputDataRaw dataRaw = new OutputDataRaw();
+          dataRaw.text = outputCell.getText();
+          raw.data = dataRaw;
+        }
+        else {
+          raw.output_type = "pyout";
+          raw.prompt_number = outputCell.getPromptNumber();
+          raw.text = outputCell.getText();
+        }
       }
       return raw;
     }
 
     public IpnbOutputCell createOutput() {
       final IpnbOutputCell outputCell;
-      if (png != null) {
-        outputCell = new IpnbPngOutputCell(png, text, prompt_number);
+      if (png != null || (data != null && data.png != null)) {
+        outputCell = new IpnbPngOutputCell(png == null ? StringUtil.join(data.png) : png, text, prompt_number);
       }
-      else if (jpeg != null) {
-        outputCell = new IpnbJpegOutputCell(jpeg, text, prompt_number);
+      else if (jpeg != null || (data != null && data.jpeg != null)) {
+        outputCell = new IpnbJpegOutputCell(jpeg == null ? StringUtil.join(data.jpeg) : jpeg, text, prompt_number);
       }
-      else if (svg != null) {
-        outputCell = new IpnbSvgOutputCell(svg, text, prompt_number);
+      else if (svg != null || (data != null && data.svg != null)) {
+        outputCell = new IpnbSvgOutputCell(svg == null ? data.svg : svg, text, prompt_number);
       }
-      else if (latex != null) {
-        outputCell = new IpnbLatexOutputCell(latex, prompt_number, text);
+      else if (latex != null || (data != null && data.latex != null)) {
+        outputCell = new IpnbLatexOutputCell(latex == null ? data.latex : latex, prompt_number, text);
       }
-      else if (stream != null) {
-        outputCell = new IpnbStreamOutputCell(stream, text, prompt_number);
+      else if (stream != null || name != null) {
+        outputCell = new IpnbStreamOutputCell(stream == null ? name : stream, text, prompt_number);
       }
-      else if (html != null) {
-        outputCell = new IpnbHtmlOutputCell(html, text, prompt_number);
+      else if (html != null || (data != null && data.html != null)) {
+        outputCell = new IpnbHtmlOutputCell(html == null ? data.html : html, text, prompt_number);
       }
-      else if ("pyerr".equals(output_type)) {
+      else if ("pyerr".equals(output_type) || "error".equals(output_type)) {
         outputCell = new IpnbErrorOutputCell(evalue, ename, traceback, prompt_number);
       }
       else if ("pyout".equals(output_type)) {
         outputCell = new IpnbOutOutputCell(text, prompt_number);
       }
+      else if ("execute_result".equals(output_type) && data != null) {
+        outputCell = new IpnbOutOutputCell(data.text, execution_count);
+      }
       else {
         outputCell = new IpnbOutputCell(text, prompt_number);
       }
       return outputCell;
     }
   }
+
+  private static class OutputDataRaw {
+    @SerializedName("text/plain") String[] text;
+    @SerializedName("text/html") String[] html;
+    @SerializedName("image/svg+xml") String[] svg;
+    @SerializedName("image/png") String[] png;
+    @SerializedName("image/jpeg") String[] jpeg;
+    @SerializedName("text/latex") String[] latex;
+  }
 }