github.com/qri-io/qri@v0.10.1-0.20220104210721-c771715036cb/transform/startf/qri/entry_writer.go (about)

     1  package qri
     2  
     3  import (
     4  	"github.com/qri-io/dataset"
     5  	"github.com/qri-io/dataset/dsio"
     6  	"github.com/qri-io/starlib/util"
     7  	starlark "go.starlark.net/starlark"
     8  )
     9  
    10  // StarlarkEntryWriter creates a starlark.Value as an EntryWriter
    11  type StarlarkEntryWriter struct {
    12  	IsDict bool
    13  	Struct *dataset.Structure
    14  	Object starlark.Value
    15  }
    16  
    17  // WriteEntry adds an entry to the underlying starlark.Value
    18  func (w *StarlarkEntryWriter) WriteEntry(ent dsio.Entry) error {
    19  	if w.IsDict {
    20  		dict := w.Object.(*starlark.Dict)
    21  		key, err := util.Marshal(ent.Key)
    22  		if err != nil {
    23  			return err
    24  		}
    25  		val, err := util.Marshal(ent.Value)
    26  		if err != nil {
    27  			return err
    28  		}
    29  		dict.SetKey(key, val)
    30  	} else {
    31  		list := w.Object.(*starlark.List)
    32  		val, err := util.Marshal(ent.Value)
    33  		if err != nil {
    34  			return err
    35  		}
    36  		list.Append(val)
    37  	}
    38  	return nil
    39  }
    40  
    41  // Structure returns the EntryWriter's dataset structure
    42  func (w *StarlarkEntryWriter) Structure() *dataset.Structure {
    43  	return w.Struct
    44  }
    45  
    46  // Close is a no-op
    47  func (w *StarlarkEntryWriter) Close() error {
    48  	return nil
    49  }
    50  
    51  // Value returns the underlying starlark.Value
    52  func (w *StarlarkEntryWriter) Value() starlark.Value {
    53  	return w.Object
    54  }
    55  
    56  // NewStarlarkEntryWriter returns a new StarlarkEntryWriter
    57  func NewStarlarkEntryWriter(st *dataset.Structure) (*StarlarkEntryWriter, error) {
    58  	mode, err := schemaScanMode(st)
    59  	if err != nil {
    60  		return nil, err
    61  	}
    62  	if mode == smObject {
    63  		return &StarlarkEntryWriter{IsDict: true, Struct: st, Object: &starlark.Dict{}}, nil
    64  	}
    65  	return &StarlarkEntryWriter{IsDict: false, Struct: st, Object: &starlark.List{}}, nil
    66  }
    67  
    68  // TODO: Refactor everything below this so that jsonschema returns this in a simple way
    69  type scanMode int
    70  
    71  const (
    72  	smArray scanMode = iota
    73  	smObject
    74  )
    75  
    76  // schemaScanMode determines weather the top level is an array or object
    77  func schemaScanMode(st *dataset.Structure) (scanMode, error) {
    78  	tlt, err := dsio.GetTopLevelType(st)
    79  	if err != nil {
    80  		return smArray, err
    81  	}
    82  	if tlt == "array" {
    83  		return smArray, nil
    84  	}
    85  	return smObject, nil
    86  }