updated ipython notebook format to work with nbformat == 4

This commit is contained in:
Ekaterina Tuzova
2014-11-14 18:42:37 +03:00
parent f47b874ff3
commit 8501fbbecd
@@ -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.<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 +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<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 +129,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 +148,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 +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;
}
}