Refactor Notebook cell parsing

zasper-io · Dec 16, 2024 · f4fe5d7 · f4fe5d7
1 parent 7c59e7f
commit f4fe5d7
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 10 deletions.
diff --git a/content/content_manager.go b/content/content_manager.go
@@ -73,14 +73,14 @@ func getNotebookModel(path string) models.ContentModel {
 	return output
 }
 
-func nbformatReads(data string, version int, capture_validation_error bool) Notebook {
+func nbformatReads(data string, version int, capture_validation_error bool) OutNotebook {
 	// output := make(map[string]interface{})
 	var nb Notebook
 	_ = json.Unmarshal([]byte(data), &nb)
-	rejoinLines(&nb)
+	output := rejoinLines(nb)
 	stripTransient(&nb)
 
-	return nb
+	return output
 }
 
 func getDirectoryModel(relativePath string) models.ContentModel {

diff --git a/content/notebook.go b/content/notebook.go
@@ -32,13 +32,15 @@ func _rejoinMimeBundle(data map[string]interface{}) map[string]interface{} {
 }
 
 // rejoinLines rejoins multi-line text into strings.
-func rejoinLines(nb *Notebook) {
+func rejoinLines(nb Notebook) OutNotebook {
+	outNb := OutNotebook{
+		Metadata: nb.Metadata,
+	}
 	for _, cell := range nb.Cells {
+		data := ""
 		if cell.Source != nil {
-			sourceList := cell.Source
-			// if sourceList, ok := cell.Source.([]string); ok {
-			cell.Source[0] = strings.Join(sourceList, "")
-			// }
+			sourceList := toStringSlice(cell.Source)
+			data = strings.Join(sourceList, "")
 		}
 
 		for _, attachment := range cell.Attachments {
@@ -57,7 +59,15 @@ func rejoinLines(nb *Notebook) {
 				}
 			}
 		}
+		outNb.Cells = append(outNb.Cells,
+			OutCell{Source: data,
+				CellType:       cell.CellType,
+				ExecutionCount: cell.ExecutionCount,
+				Metadata:       cell.Metadata,
+				Attachments:    cell.Attachments,
+				Outputs:        cell.Outputs})
 	}
+	return outNb
 }
 
 // _splitMimeBundle splits multi-line string fields in a mimebundle.
@@ -130,9 +140,24 @@ type Notebook struct {
 	Metadata map[string]interface{} `json:"metadata"`
 }
 
+// Notebook struct for handling notebook cells
+type OutNotebook struct {
+	Cells    []OutCell              `json:"cells"`
+	Metadata map[string]interface{} `json:"metadata"`
+}
+
 // Cell struct for handling individual cells in a notebook
 type Cell struct {
-	Source         []string                          `json:"source"`
+	Source         []interface{}                     `json:"source"`
+	ExecutionCount int                               `json:"execution_count"`
+	CellType       string                            `json:"cell_type"`
+	Attachments    map[string]map[string]interface{} `json:"attachments"`
+	Outputs        []Output                          `json:"outputs"`
+	Metadata       map[string]interface{}            `json:"metadata"`
+}
+
+type OutCell struct {
+	Source         string                            `json:"source"`
 	ExecutionCount int                               `json:"execution_count"`
 	CellType       string                            `json:"cell_type"`
 	Attachments    map[string]map[string]interface{} `json:"attachments"`

diff --git a/ui/src/ide/editor/notebook/Cell.tsx b/ui/src/ide/editor/notebook/Cell.tsx
@@ -50,7 +50,7 @@ export interface CodeMirrorRef {
 const Cell = React.forwardRef((props: ICellProps, ref) => {
   const cell = props.cell
   const [theme] = useAtom(themeAtom)
-  const [cellContents, setCellContents] = useState(cell.source[0])
+  const [cellContents, setCellContents] = useState(cell.source)
   const [cursorPosition, setCursorPosition] = useState(0)
   const [totalLines, setTotalLines] = useState(0)