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

     1  package ds
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  
     7  	"github.com/qri-io/dataset"
     8  	"github.com/qri-io/dataset/dsio"
     9  	"github.com/qri-io/starlib/util"
    10  	"go.starlark.net/starlark"
    11  )
    12  
    13  // EntryReader implements the dsio.EntryReader interface for starlark.Iterable's
    14  type EntryReader struct {
    15  	i    int
    16  	st   *dataset.Structure
    17  	iter starlark.Iterator
    18  	data starlark.Value
    19  }
    20  
    21  var _ dsio.EntryReader = (*EntryReader)(nil)
    22  
    23  // NewEntryReader creates a new Entry Reader
    24  func NewEntryReader(st *dataset.Structure, iter starlark.Iterable) *EntryReader {
    25  	return &EntryReader{
    26  		st:   st,
    27  		data: iter.(starlark.Value),
    28  		iter: iter.Iterate(),
    29  	}
    30  }
    31  
    32  // Structure gives this reader's structure
    33  func (r *EntryReader) Structure() *dataset.Structure {
    34  	return r.st
    35  }
    36  
    37  // ReadEntry reads one entry from the reader
    38  func (r *EntryReader) ReadEntry() (e dsio.Entry, err error) {
    39  	// Read next element (key for object, value for array).
    40  	var next starlark.Value
    41  	if !r.iter.Next(&next) {
    42  		r.iter.Done()
    43  		return e, io.EOF
    44  	}
    45  
    46  	// Handle array entry.
    47  	tlt, err := dsio.GetTopLevelType(r.st)
    48  	if err != nil {
    49  		return
    50  	}
    51  
    52  	if tlt == "array" {
    53  		e.Index = r.i
    54  		r.i++
    55  		e.Value, err = util.Unmarshal(next)
    56  		if err != nil {
    57  			fmt.Printf("reading error: %s\n", err.Error())
    58  		}
    59  		return
    60  	}
    61  
    62  	// Handle object entry. Assume key is a string.
    63  	var ok bool
    64  	e.Key, ok = starlark.AsString(next)
    65  	if !ok {
    66  		fmt.Printf("key error: %s\n", next)
    67  	}
    68  	// Lookup the corresponding value for the key.
    69  	dict := r.data.(*starlark.Dict)
    70  	value, ok, err := dict.Get(next)
    71  	if err != nil {
    72  		fmt.Printf("reading error: %s\n", err.Error())
    73  	}
    74  	e.Value, err = util.Unmarshal(value)
    75  	if err != nil {
    76  		fmt.Printf("reading error: %s\n", err.Error())
    77  	}
    78  	return
    79  }
    80  
    81  // Close finalizes the reader
    82  func (r *EntryReader) Close() error {
    83  	return nil
    84  }