mapiter.go

Documentation: github.com/lestrrat-go/iter/mapiter

     1  package mapiter
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"reflect"
     7  	"sync"
     8  )
     9  
    10  // Iterate creates an iterator from arbitrary map types. This is not
    11  // the most efficient tool, but it's the quickest way to create an
    12  // iterator for maps.
    13  // Also, note that you cannot make any assumptions on the order of
    14  // pairs being returned.
    15  func Iterate(ctx context.Context, m interface{}) (Iterator, error) {
    16  	mrv := reflect.ValueOf(m)
    17  
    18  	if mrv.Kind() != reflect.Map {
    19  		return nil, fmt.Errorf(`argument must be a map (%s)`, mrv.Type())
    20  	}
    21  
    22  	ch := make(chan *Pair)
    23  	go func(ctx context.Context, ch chan *Pair, mrv reflect.Value) {
    24  		defer close(ch)
    25  		for _, key := range mrv.MapKeys() {
    26  			value := mrv.MapIndex(key)
    27  			pair := &Pair{
    28  				Key:   key.Interface(),
    29  				Value: value.Interface(),
    30  			}
    31  			select {
    32  			case <-ctx.Done():
    33  				return
    34  			case ch <- pair:
    35  			}
    36  		}
    37  	}(ctx, ch, mrv)
    38  
    39  	return New(ch), nil
    40  }
    41  
    42  // Source represents a map that knows how to create an iterator
    43  type Source interface {
    44  	Iterate(context.Context) Iterator
    45  }
    46  
    47  // Pair represents a single pair of key and value from a map
    48  type Pair struct {
    49  	Key   interface{}
    50  	Value interface{}
    51  }
    52  
    53  // Iterator iterates through keys and values of a map
    54  type Iterator interface {
    55  	Next(context.Context) bool
    56  	Pair() *Pair
    57  }
    58  
    59  type iter struct {
    60  	ch   chan *Pair
    61  	mu   sync.RWMutex
    62  	next *Pair
    63  }
    64  
    65  // Visitor represents an object that handles each pair in a map
    66  type Visitor interface {
    67  	Visit(interface{}, interface{}) error
    68  }
    69  
    70  // VisitorFunc is a type of Visitor based on a function
    71  type VisitorFunc func(interface{}, interface{}) error
    72  
    73  func (fn VisitorFunc) Visit(s interface{}, v interface{}) error {
    74  	return fn(s, v)
    75  }
    76  
    77  func New(ch chan *Pair) Iterator {
    78  	return &iter{
    79  		ch: ch,
    80  	}
    81  }
    82  
    83  // Next returns true if there are more items to read from the iterator
    84  func (i *iter) Next(ctx context.Context) bool {
    85  	i.mu.RLock()
    86  	if i.ch == nil {
    87  		i.mu.RUnlock()
    88  		return false
    89  	}
    90  	i.mu.RUnlock()
    91  
    92  	i.mu.Lock()
    93  	defer i.mu.Unlock()
    94  	select {
    95  	case <-ctx.Done():
    96  		i.ch = nil
    97  		return false
    98  	case v, ok := <-i.ch:
    99  		if !ok {
   100  			i.ch = nil
   101  			return false
   102  		}
   103  		i.next = v
   104  		return true
   105  	}
   106  
   107  	//nolint:govet
   108  	return false // never reached
   109  }
   110  
   111  // Pair returns the currently buffered Pair. Calling Next() will reset its value
   112  func (i *iter) Pair() *Pair {
   113  	i.mu.RLock()
   114  	defer i.mu.RUnlock()
   115  	return i.next
   116  }
   117  
   118  // Walk walks through each element in the map
   119  func Walk(ctx context.Context, s Source, v Visitor) error {
   120  	for i := s.Iterate(ctx); i.Next(ctx); {
   121  		pair := i.Pair()
   122  		if err := v.Visit(pair.Key, pair.Value); err != nil {
   123  			return fmt.Errorf(`failed to visit key %s: %w`, pair.Key, err)
   124  		}
   125  	}
   126  	return nil
   127  }
   128  
   129  // AsMap returns the values obtained from the source as a map
   130  func AsMap(ctx context.Context, s interface{}, v interface{}) error {
   131  	var iter Iterator
   132  	switch reflect.ValueOf(s).Kind() {
   133  	case reflect.Map:
   134  		x, err := Iterate(ctx, s)
   135  		if err != nil {
   136  			return fmt.Errorf(`failed to iterate over map type: %w`, err)
   137  		}
   138  		iter = x
   139  	default:
   140  		ssrc, ok := s.(Source)
   141  		if !ok {
   142  			return fmt.Errorf(`cannot iterate over %T: not a mapiter.Source type`, s)
   143  		}
   144  		iter = ssrc.Iterate(ctx)
   145  	}
   146  
   147  	dst := reflect.ValueOf(v)
   148  
   149  	// dst MUST be a pointer to a map type
   150  	if kind := dst.Kind(); kind != reflect.Ptr {
   151  		return fmt.Errorf(`dst must be a pointer to a map (%s)`, dst.Type())
   152  	}
   153  
   154  	dst = dst.Elem()
   155  	if dst.Kind() != reflect.Map {
   156  		return fmt.Errorf(`dst must be a pointer to a map (%s)`, dst.Type())
   157  	}
   158  
   159  	if dst.IsNil() {
   160  		dst.Set(reflect.MakeMap(dst.Type()))
   161  	}
   162  
   163  	// dst must be assignable
   164  	if !dst.CanSet() {
   165  		return fmt.Errorf(`dst is not writeable`)
   166  	}
   167  
   168  	keytyp := dst.Type().Key()
   169  	valtyp := dst.Type().Elem()
   170  
   171  	for iter.Next(ctx) {
   172  		pair := iter.Pair()
   173  
   174  		rvkey := reflect.ValueOf(pair.Key)
   175  		rvvalue := reflect.ValueOf(pair.Value)
   176  
   177  		if !rvkey.Type().AssignableTo(keytyp) {
   178  			return fmt.Errorf(`cannot assign key of type %s to map key of type %s`, rvkey.Type(), keytyp)
   179  		}
   180  
   181  		switch rvvalue.Kind() {
   182  		// we can only check if we can assign to rvvalue to valtyp if it's non-nil
   183  		case reflect.Invalid:
   184  			rvvalue = reflect.New(valtyp).Elem()
   185  		default:
   186  			if !rvvalue.Type().AssignableTo(valtyp) {
   187  				return fmt.Errorf(`cannot assign value of type %s to map value of type %s`, rvvalue.Type(), valtyp)
   188  			}
   189  		}
   190  
   191  		dst.SetMapIndex(rvkey, rvvalue)
   192  	}
   193  
   194  	return nil
   195  }
   196
View as plain text