From 8501fbbecd003c0c84b4cbd5b53dc72480a17832 Mon Sep 17 00:00:00 2001 From: Ekaterina Tuzova Date: Fri, 14 Nov 2014 18:42:37 +0300 Subject: [PATCH] updated ipython notebook format to work with nbformat == 4 --- .../plugins/ipnb/format/IpnbParser.java | 103 ++++++++++++++---- 1 file changed, 80 insertions(+), 23 deletions(-) diff --git a/python/ipnb/src/org/jetbrains/plugins/ipnb/format/IpnbParser.java b/python/ipnb/src/org/jetbrains/plugins/ipnb/format/IpnbParser.java index 0f1f001dc3b0..9e553c7c4906 100644 --- a/python/ipnb/src/org/jetbrains/plugins/ipnb/format/IpnbParser.java +++ b/python/ipnb/src/org/jetbrains/plugins/ipnb/format/IpnbParser.java @@ -3,6 +3,8 @@ 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 org.jetbrains.annotations.NotNull; @@ -14,6 +16,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 +38,19 @@ public class IpnbParser { if (rawFile == null) return new IpnbFile(new IpnbFileRaw(), Lists.newArrayList(), path); List cells = new ArrayList(); final IpnbWorksheet[] worksheets = rawFile.worksheets; - for (IpnbWorksheet worksheet : worksheets) { - final List 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 rawCells = worksheet.cells; + for (IpnbCellRaw rawCell : rawCells) { + cells.add(rawCell.createCell()); + } + } + } return new IpnbFile(rawFile, cells, path); } @@ -65,12 +75,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 +116,11 @@ public class IpnbParser { } public static class IpnbFileRaw { + IpnbWorksheet[] worksheets; + List cells = new ArrayList(); Map metadata = new HashMap(); int nbformat = 3; int nbformat_minor; - IpnbWorksheet[] worksheets; } private static class IpnbWorksheet { @@ -105,14 +129,16 @@ public class IpnbParser { private static class IpnbCellRaw { String cell_type; + Integer execution_count; + Map metadata = new HashMap(); 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 +148,19 @@ public class IpnbParser { raw.cell_type = "code"; final ArrayList outputRaws = new ArrayList(); 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 +183,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 +201,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; @@ -181,7 +217,7 @@ public class IpnbParser { String[] traceback; - 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 IpnbPngOutputCell) { @@ -203,7 +239,12 @@ public class IpnbParser { 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(); } @@ -218,9 +259,18 @@ public class IpnbParser { 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; } @@ -239,8 +289,8 @@ public class IpnbParser { else if (latex != null) { outputCell = new IpnbLatexOutputCell(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); @@ -251,10 +301,17 @@ public class IpnbParser { else if ("pyout".equals(output_type)) { outputCell = new IpnbOutOutputCell(text, prompt_number); } + else if ("execute_result".equals(output_type)) { + 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; + } }