github.com/attic-labs/noms@v0.0.0-20210827224422-e5fa29d95e8b/samples/go/csv/write.go (about)

     1  // Copyright 2016 Attic Labs, Inc. All rights reserved.
     2  // Licensed under the Apache License, version 2.0:
     3  // http://www.apache.org/licenses/LICENSE-2.0
     4  
     5  package csv
     6  
     7  import (
     8  	"encoding/csv"
     9  	"fmt"
    10  	"io"
    11  
    12  	"github.com/attic-labs/noms/go/d"
    13  	"github.com/attic-labs/noms/go/types"
    14  )
    15  
    16  func getElemDesc(s types.Collection, index int) types.StructDesc {
    17  	t := types.TypeOf(s).Desc.(types.CompoundDesc).ElemTypes[index]
    18  	if types.StructKind != t.TargetKind() {
    19  		d.Panic("Expected StructKind, found %s", t.Kind())
    20  	}
    21  	return t.Desc.(types.StructDesc)
    22  }
    23  
    24  // GetListElemDesc ensures that l is a types.List of structs, pulls the types.StructDesc that describes the elements of l out of vr, and returns the StructDesc.
    25  func GetListElemDesc(l types.List, vr types.ValueReader) types.StructDesc {
    26  	return getElemDesc(l, 0)
    27  }
    28  
    29  // GetMapElemDesc ensures that m is a types.Map of structs, pulls the types.StructDesc that describes the elements of m out of vr, and returns the StructDesc.
    30  // If m is a nested types.Map of types.Map, then GetMapElemDesc will descend the levels of the enclosed types.Maps to get to a types.Struct
    31  func GetMapElemDesc(m types.Map, vr types.ValueReader) types.StructDesc {
    32  	t := types.TypeOf(m).Desc.(types.CompoundDesc).ElemTypes[1]
    33  	if types.StructKind == t.TargetKind() {
    34  		return t.Desc.(types.StructDesc)
    35  	} else if types.MapKind == t.TargetKind() {
    36  		_, v := m.First()
    37  		return GetMapElemDesc(v.(types.Map), vr)
    38  	}
    39  	panic(fmt.Sprintf("Expected StructKind or MapKind, found %s", t.Kind().String()))
    40  }
    41  
    42  func writeValuesFromChan(structChan chan types.Struct, sd types.StructDesc, comma rune, output io.Writer) {
    43  	fieldNames := getFieldNamesFromStruct(sd)
    44  	csvWriter := csv.NewWriter(output)
    45  	csvWriter.Comma = comma
    46  	if csvWriter.Write(fieldNames) != nil {
    47  		d.Panic("Failed to write header %v", fieldNames)
    48  	}
    49  	record := make([]string, len(fieldNames))
    50  	for s := range structChan {
    51  		i := 0
    52  		s.WalkValues(func(v types.Value) {
    53  			record[i] = fmt.Sprintf("%v", v)
    54  			i++
    55  		})
    56  		if csvWriter.Write(record) != nil {
    57  			d.Panic("Failed to write record %v", record)
    58  		}
    59  	}
    60  
    61  	csvWriter.Flush()
    62  	if csvWriter.Error() != nil {
    63  		d.Panic("error flushing csv")
    64  	}
    65  }
    66  
    67  // WriteList takes a types.List l of structs (described by sd) and writes it to output as comma-delineated values.
    68  func WriteList(l types.List, sd types.StructDesc, comma rune, output io.Writer) {
    69  	structChan := make(chan types.Struct, 1024)
    70  	go func() {
    71  		l.IterAll(func(v types.Value, index uint64) {
    72  			structChan <- v.(types.Struct)
    73  		})
    74  		close(structChan)
    75  	}()
    76  	writeValuesFromChan(structChan, sd, comma, output)
    77  }
    78  
    79  func sendMapValuesToChan(m types.Map, structChan chan<- types.Struct) {
    80  	m.IterAll(func(k, v types.Value) {
    81  		if subMap, ok := v.(types.Map); ok {
    82  			sendMapValuesToChan(subMap, structChan)
    83  		} else {
    84  			structChan <- v.(types.Struct)
    85  		}
    86  	})
    87  }
    88  
    89  // WriteMap takes a types.Map m of structs (described by sd) and writes it to output as comma-delineated values.
    90  func WriteMap(m types.Map, sd types.StructDesc, comma rune, output io.Writer) {
    91  	structChan := make(chan types.Struct, 1024)
    92  	go func() {
    93  		sendMapValuesToChan(m, structChan)
    94  		close(structChan)
    95  	}()
    96  	writeValuesFromChan(structChan, sd, comma, output)
    97  }
    98  
    99  func getFieldNamesFromStruct(structDesc types.StructDesc) (fieldNames []string) {
   100  	structDesc.IterFields(func(name string, t *types.Type, optional bool) {
   101  		if !types.IsPrimitiveKind(t.TargetKind()) {
   102  			d.Panic("Expected primitive kind, found %s", t.TargetKind().String())
   103  		}
   104  		fieldNames = append(fieldNames, name)
   105  	})
   106  	return
   107  }