import com.google.common.base.Predicate; import com.google.common.collect.Iterables; import com.intellij.openapi.application.PathManager; import com.intellij.openapi.util.text.StringUtil; import com.intellij.testFramework.LightVirtualFile; import junit.framework.TestCase; import org.jetbrains.annotations.NotNull; import org.jetbrains.plugins.ipnb.format.IpnbFile; import org.jetbrains.plugins.ipnb.format.IpnbParser; import org.jetbrains.plugins.ipnb.format.cells.IpnbCell; import org.jetbrains.plugins.ipnb.format.cells.IpnbCodeCell; import org.jetbrains.plugins.ipnb.format.cells.IpnbMarkdownCell; import org.jetbrains.plugins.ipnb.format.cells.output.IpnbOutputCell; import java.io.BufferedReader; import java.io.FileReader; import java.io.IOException; import java.util.List; public class JsonParserTest extends TestCase { public void testFile() throws IOException { final String fileName = "testData/SymPy.ipynb"; final String fileText = getFileText(fileName); final IpnbFile ipnbFile = IpnbParser.parseIpnbFile(fileText, new LightVirtualFile()); assertNotNull(ipnbFile); assertEquals(31, ipnbFile.getCells().size()); } public void testMarkdownCells() throws IOException { final String fileName = "testData/SymPy.ipynb"; final String fileText = getFileText(fileName); final IpnbFile ipnbFile = IpnbParser.parseIpnbFile(fileText, new LightVirtualFile()); assertNotNull(ipnbFile); final List cells = ipnbFile.getCells(); Iterables.removeIf(cells, new Predicate() { @Override public boolean apply(IpnbCell cell) { return !(cell instanceof IpnbMarkdownCell); } }); assertEquals(7, cells.size()); } public void testMarkdownCell() throws IOException { final String fileName = "testData/markdown.ipynb"; final String fileText = getFileText(fileName); final IpnbFile ipnbFile = IpnbParser.parseIpnbFile(fileText, new LightVirtualFile()); assertNotNull(ipnbFile); final List cells = ipnbFile.getCells(); assertEquals(1, cells.size()); final IpnbCell cell = cells.get(0); assertTrue(cell instanceof IpnbMarkdownCell); final List source = ((IpnbMarkdownCell)cell).getSource(); final String joined = StringUtil.join(source, ""); assertEquals("", joined); } public void testCodeCell() throws IOException { final String fileName = "testData/code.ipynb"; final String fileText = getFileText(fileName); final IpnbFile ipnbFile = IpnbParser.parseIpnbFile(fileText, new LightVirtualFile()); assertNotNull(ipnbFile); final List cells = ipnbFile.getCells(); assertEquals(1, cells.size()); final IpnbCell cell = cells.get(0); assertTrue(cell instanceof IpnbCodeCell); final List outputs = ((IpnbCodeCell)cell).getCellOutputs(); assertEquals(0, outputs.size()); final List source = ((IpnbCodeCell)cell).getSource(); final String joined = StringUtil.join(source, ""); assertEquals("e = x + 2*y", joined); final String language = ((IpnbCodeCell)cell).getLanguage(); assertEquals("python", language); final Integer number = ((IpnbCodeCell)cell).getPromptNumber(); assertEquals(new Integer(4), number); } public void testOutputs() throws IOException { final String fileName = "testData/outputs.ipynb"; final String fileText = getFileText(fileName); final IpnbFile ipnbFile = IpnbParser.parseIpnbFile(fileText, new LightVirtualFile()); assertNotNull(ipnbFile); final List cells = ipnbFile.getCells(); assertEquals(1, cells.size()); final IpnbCell cell = cells.get(0); assertTrue(cell instanceof IpnbCodeCell); final List outputs = ((IpnbCodeCell)cell).getCellOutputs(); assertEquals(1, outputs.size()); final IpnbOutputCell output = outputs.get(0); final List text = output.getText(); assertNotNull(text); final String joined = StringUtil.join(text, ""); assertEquals("\"Add(Symbol('x'), Mul(Integer(2), Symbol('y')))\"", joined); } private static String getFileText(@NotNull final String fileName) throws IOException { String fullPath = PathManager.getHomePath() + "/community/python/ipnb/" + fileName; final BufferedReader br = new BufferedReader(new FileReader(fullPath)); try { final StringBuilder sb = new StringBuilder(); String line = br.readLine(); while (line != null) { sb.append(line); sb.append("\n"); line = br.readLine(); } return sb.toString(); } finally { br.close(); } } }