steve/databricks-cli
steve
/
databricks-cli
mirror of https://github.com/databricks/cli.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add more typing for ipynb

This commit is contained in:
kartikgupta-db 2023-09-13 13:04:21 +02:00
parent f45935f8ae
commit b7a2b23ba9
No known key found for this signature in database
GPG Key ID: 6AD5FA11FACDEA39
1 changed files with 17 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
bundle/python/notebook_task_wrappers.go
View File

@ -99,23 +99,32 @@ func getDbnbTemplate(s string) (string, error) {
`, notebookTrampolineData, s), nil `, notebookTrampolineData, s), nil
} }
type IpynbData struct {
Cells []IpynbCell `json:"cells"`
}
type IpynbCell struct {
CellType string `json:"cell_type"`
Source []string `json:"source"`
}
func getIpynbTemplate(s string) (string, error) { func getIpynbTemplate(s string) (string, error) {
var data map[string]any var data IpynbData
err := json.Unmarshal([]byte(s), &data) err := json.Unmarshal([]byte(s), &data)
if err != nil { if err != nil {
return "", err return "", err
} }
if data["cells"] == nil { if data.Cells == nil {
data["cells"] = []any{} data.Cells = []IpynbCell{}
} }
data["cells"] = append([]any{ data.Cells = append([]IpynbCell{
map[string]any{ {
"cell_type": "code", CellType: "code",
"source": []string{notebookTrampolineData}, Source: []string{"# Databricks notebook source"},
}, },
}, data["cells"].([]any)...) }, data.Cells...)
bytes, err := json.Marshal(data) bytes, err := json.Marshal(data)
if err != nil { if err != nil {