starlark/typed/list.go - infra/luci/luci-go - Git at Google

 // Copyright 2019 The LUCI Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 package typed

 import (
 	"fmt"

 	"go.starlark.net/starlark"
 	"go.starlark.net/syntax"
 )

 // List is a Starlark value that looks like a regular list, but implements type
 // checks and implicit conversions when assigning elements.
 //
 // Differences from regular lists:
 //   * Not comparable by value to regular lists, only to other typed lists.
 //     go.starlark.net has no way to express that typed lists should be
 //     comparable to regular lists.
 //   * Only == and != comparison operators are supported.
 //   * `l += ...` is same as `l = l + ...`, not `l.extend(...)`
 type List struct {
 	itemT Converter      // defines the type of elements
 	list  *starlark.List // the underlying untyped list
 }

 // Note that these interfaces intersect, so we can't mush them into a single
 // interface to check *List conforms it.
 var (
 	_ starlark.Value       = (*List)(nil)
 	_ starlark.Sequence    = (*List)(nil)
 	_ starlark.Indexable   = (*List)(nil)
 	_ starlark.HasBinary   = (*List)(nil)
 	_ starlark.HasAttrs    = (*List)(nil)
 	_ starlark.HasSetIndex = (*List)(nil)
 	_ starlark.Comparable  = (*List)(nil)
 )

 // NewList returns a list with given elements, converted to necessary type
 // through 'item' converter.
 //
 // All subsequently added elements will be converter through this converter too.
 func NewList(item Converter, elems []starlark.Value) (*List, error) {
 	l := &List{itemT: item}
 	vals := make([]starlark.Value, len(elems))
 	for i, e := range elems {
 		var err error
 		if vals[i], err = l.conv(e, i); err != nil {
 			return nil, err
 		}
 	}
 	l.list = starlark.NewList(vals)
 	return l, nil
 }

 // AsTypedList allocates a new list<t>, and copies it from 'x'.
 //
 // Returns an error if 'x' is not an iterable or some of its elements can't be
 // converted to 't'.
 func AsTypedList(t Converter, x starlark.Value) (*List, error) {
 	it := starlark.Iterate(x)
 	if it == nil {
 		return nil, fmt.Errorf("got %s, want an iterable", x.Type())
 	}
 	defer it.Done()

 	var vals []starlark.Value
 	if l := starlark.Len(x); l > 0 {
 		vals = make([]starlark.Value, 0, l)
 	}
 	var itm starlark.Value
 	for it.Next(&itm) {
 		vals = append(vals, itm)
 	}

 	return NewList(t, vals)
 }

 // conv calls Convert, annotating the error with item's index if idx >= 0.
 func (l *List) conv(v starlark.Value, idx int) (starlark.Value, error) {
 	switch v, err := l.itemT.Convert(v); {
 	case err == nil:
 		return v, nil
 	case idx >= 0:
 		return nil, fmt.Errorf("item #%d: %s", idx, err)
 	default:
 		return nil, err
 	}
 }

 // Converter returns the type converter used by this list.
 func (l *List) Converter() Converter { return l.itemT }

 // Binary implements 'l <op> y' (or 'y <op> l', depending on 'side').
 //
 // Note that *starlark.List has its binary operators implement natively by the
 // Starlark evaluator, not via HasBinary interface.
 func (l *List) Binary(op syntax.Token, y starlark.Value, side starlark.Side) (starlark.Value, error) {
 	switch op {
 	case syntax.IN, syntax.NOT_IN:
 		return starlark.Binary(op, y, l.list)

 	case syntax.STAR: // 2*[1,2] => [1,2,1,2], also typed
 		if _, ok := y.(starlark.Int); !ok {
 			goto unsupported
 		}
 		return l.asTypedList(starlark.Binary(op, l.list, y))

 	case syntax.PLUS:
 		// Convert regular lists to typed lists first.
 		if yl, ok := y.(*starlark.List); ok {
 			vals := make([]starlark.Value, yl.Len())
 			for i := 0; i < yl.Len(); i++ {
 				vals[i] = yl.Index(i)
 			}
 			asTyped, err := NewList(l.itemT, vals)
 			if err != nil {
 				return nil, err
 			}
 			return l.Binary(op, asTyped, side)
 		}

 		// When adding typed lists their types must match.
 		if yl, ok := y.(*List); ok {
 			if yl.itemT != l.itemT {
 				goto unsupported
 			}
 			lhs, rhs := maybeSwap(l.list, yl.list, side)
 			return l.asTypedList(starlark.Binary(op, lhs, rhs))
 		}
 	}

 unsupported:
 	return nil, fmt.Errorf("unknown binary op: %s %s %s", l.Type(), op, y.Type())
 }

 func (l *List) asTypedList(v starlark.Value, err error) (starlark.Value, error) {
 	if err != nil {
 		return nil, err
 	}
 	if asL, ok := v.(*starlark.List); ok {
 		return &List{itemT: l.itemT, list: asL}, nil
 	}
 	return nil, fmt.Errorf("expected a list, but got %q", v.Type())
 }

 func maybeSwap(l, r starlark.Value, side starlark.Side) (starlark.Value, starlark.Value) {
 	if side == starlark.Left {
 		return l, r
 	}
 	return r, l
 }

 // starlark.List methods we reimplement or proxy.

 func (l *List) Append(v starlark.Value) error {
 	v, err := l.conv(v, -1)
 	if err != nil {
 		return err
 	}
 	return l.list.Append(v)
 }

 func (l *List) AttrNames() []string {
 	// Note: we hardcode the list of methods (rather than call l.list.AttrNames)
 	// to be able to notice (via the failing test) when go.starlark.net adds some
 	// new methods that we might need to "wrap" in Attr(...).
 	return []string{
 		"append",
 		"clear",
 		"extend",
 		"index",
 		"insert",
 		"pop",
 		"remove",
 	}
 }

 func (l *List) Attr(name string) (starlark.Value, error) {
 	if name == "append" {
 		return l.implAppend(), nil // fully reimplemented, no need for orig
 	}
 	orig, err := l.list.Attr(name)
 	if err != nil {
 		return nil, err
 	}
 	switch name {
 	case "extend":
 		return l.implExtend(orig.(*starlark.Builtin)), nil
 	case "insert":
 		return l.implInsert(orig.(*starlark.Builtin)), nil
 	default:
 		return orig, nil // no need to wrap
 	}
 }

 func (l *List) SetIndex(i int, v starlark.Value) error {
 	// Note: 'i' is never negative here, the Starlark evaluator deals with
 	// negative indexes before calling SetIndex.
 	v, err := l.conv(v, i)
 	if err != nil {
 		return err
 	}
 	return l.list.SetIndex(i, v)
 }

 func (l *List) Slice(start, end, step int) starlark.Value {
 	return &List{
 		itemT: l.itemT,
 		list:  l.list.Slice(start, end, step).(*starlark.List),
 	}
 }

 func (l *List) Type() string {
 	return fmt.Sprintf("list<%s>", l.itemT.Type())
 }

 func (l *List) String() string {
 	return fmt.Sprintf("list<%s>(%s)", l.itemT.Type(), l.list.String())
 }

 func (l *List) CompareSameType(op syntax.Token, y starlark.Value, depth int) (bool, error) {
 	switch op {
 	case syntax.EQL:
 		return listsEqual(l, y.(*List), depth)
 	case syntax.NEQ:
 		eq, err := listsEqual(l, y.(*List), depth)
 		return !eq, err
 	default:
 		return false, fmt.Errorf("%q is not implemented for %s", op, l.Type())
 	}
 }

 func listsEqual(l, r *List, depth int) (bool, error) {
 	if l.itemT != r.itemT {
 		return false, nil
 	}
 	return l.list.CompareSameType(syntax.EQL, r.list, depth)
 }

 // starlark.List methods delegated without any modifications.

 func (l *List) Clear() error               { return l.list.Clear() }
 func (l *List) Freeze()                    { l.list.Freeze() }
 func (l *List) Hash() (uint32, error)      { return l.list.Hash() }
 func (l *List) Index(i int) starlark.Value { return l.list.Index(i) }
 func (l *List) Iterate() starlark.Iterator { return l.list.Iterate() }
 func (l *List) Len() int                   { return l.list.Len() }
 func (l *List) Truth() starlark.Bool       { return l.list.Truth() }

 // Reimplemented builtins.

 // implAppend reimplements 'append' via Append(...).
 func (l *List) implAppend() *starlark.Builtin {
 	return starlark.NewBuiltin("append", func(th *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 		var object starlark.Value
 		if err := starlark.UnpackPositionalArgs("append", args, kwargs, 1, &object); err != nil {
 			return nil, err
 		}
 		if err := l.Append(object); err != nil {
 			return nil, fmt.Errorf("append: %s", err)
 		}
 		return starlark.None, nil
 	})
 }

 // implExtend reimplements 'extend' via Append(...).
 //
 // Uses original extend() for the fast path to optimize memory allocations.
 func (l *List) implExtend(orig *starlark.Builtin) *starlark.Builtin {
 	return starlark.NewBuiltin("extend", func(th *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 		var iterable starlark.Iterable
 		if err := starlark.UnpackPositionalArgs("extend", args, kwargs, 1, &iterable); err != nil {
 			return nil, err
 		}

 		// Fast path: the iterable has the exact same type as us, we can skip the
 		// conversion completely, since per Convert idempotency property this will
 		// be a noop anyway. By calling original 'extend' implementation we are
 		// basically doing one big memcpy(...) instead of adding items one by one.
 		if it, ok := iterable.(*List); ok && it.itemT == l.itemT {
 			return orig.CallInternal(th, starlark.Tuple{it.list}, nil)
 		}

 		// Slow path: append items one by one.
 		iter := iterable.Iterate()
 		defer iter.Done()
 		var x starlark.Value
 		for i := 0; iter.Next(&x); i++ {
 			if err := l.Append(x); err != nil {
 				return nil, fmt.Errorf("extend: item #%d: %s", i, err)
 			}
 		}
 		return starlark.None, nil
 	})
 }

 // implInsert implements 'insert' by wrapping the original 'insert'.
 //
 // It's not easily reimplementable, since starlark.List doesn't expose its
 // internal list in a simple way.
 func (l *List) implInsert(orig *starlark.Builtin) *starlark.Builtin {
 	return starlark.NewBuiltin("insert", func(th *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 		var index starlark.Value
 		var object starlark.Value
 		if err := starlark.UnpackPositionalArgs("insert", args, kwargs, 2, &index, &object); err != nil {
 			return nil, err
 		}
 		x, err := l.conv(object, -1)
 		if err != nil {
 			return nil, fmt.Errorf("insert: %s", err)
 		}
 		return orig.CallInternal(th, starlark.Tuple{index, x}, nil)
 	})
 }
	// Copyright 2019 The LUCI Authors.
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	package typed

	import (
	"fmt"

	"go.starlark.net/starlark"
	"go.starlark.net/syntax"
	)

	// List is a Starlark value that looks like a regular list, but implements type
	// checks and implicit conversions when assigning elements.
	//
	// Differences from regular lists:
	// * Not comparable by value to regular lists, only to other typed lists.
	// go.starlark.net has no way to express that typed lists should be
	// comparable to regular lists.
	// * Only == and != comparison operators are supported.
	// * `l += ...` is same as `l = l + ...`, not `l.extend(...)`
	type List struct {
	itemT Converter // defines the type of elements
	list *starlark.List // the underlying untyped list
	}

	// Note that these interfaces intersect, so we can't mush them into a single
	// interface to check *List conforms it.
	var (
	_ starlark.Value = (*List)(nil)
	_ starlark.Sequence = (*List)(nil)
	_ starlark.Indexable = (*List)(nil)
	_ starlark.HasBinary = (*List)(nil)
	_ starlark.HasAttrs = (*List)(nil)
	_ starlark.HasSetIndex = (*List)(nil)
	_ starlark.Comparable = (*List)(nil)
	)

	// NewList returns a list with given elements, converted to necessary type
	// through 'item' converter.
	//
	// All subsequently added elements will be converter through this converter too.
	func NewList(item Converter, elems []starlark.Value) (*List, error) {
	l := &List{itemT: item}
	vals := make([]starlark.Value, len(elems))
	for i, e := range elems {
	var err error
	if vals[i], err = l.conv(e, i); err != nil {
	return nil, err
	}
	}
	l.list = starlark.NewList(vals)
	return l, nil
	}

	// AsTypedList allocates a new list<t>, and copies it from 'x'.
	//
	// Returns an error if 'x' is not an iterable or some of its elements can't be
	// converted to 't'.
	func AsTypedList(t Converter, x starlark.Value) (*List, error) {
	it := starlark.Iterate(x)
	if it == nil {
	return nil, fmt.Errorf("got %s, want an iterable", x.Type())
	}
	defer it.Done()

	var vals []starlark.Value
	if l := starlark.Len(x); l > 0 {
	vals = make([]starlark.Value, 0, l)
	}
	var itm starlark.Value
	for it.Next(&itm) {
	vals = append(vals, itm)
	}

	return NewList(t, vals)
	}

	// conv calls Convert, annotating the error with item's index if idx >= 0.
	func (l *List) conv(v starlark.Value, idx int) (starlark.Value, error) {
	switch v, err := l.itemT.Convert(v); {
	case err == nil:
	return v, nil
	case idx >= 0:
	return nil, fmt.Errorf("item #%d: %s", idx, err)
	default:
	return nil, err
	}
	}

	// Converter returns the type converter used by this list.
	func (l *List) Converter() Converter { return l.itemT }

	// Binary implements 'l <op> y' (or 'y <op> l', depending on 'side').
	//
	// Note that *starlark.List has its binary operators implement natively by the
	// Starlark evaluator, not via HasBinary interface.
	func (l *List) Binary(op syntax.Token, y starlark.Value, side starlark.Side) (starlark.Value, error) {
	switch op {
	case syntax.IN, syntax.NOT_IN:
	return starlark.Binary(op, y, l.list)

	case syntax.STAR: // 2*[1,2] => [1,2,1,2], also typed
	if _, ok := y.(starlark.Int); !ok {
	goto unsupported
	}
	return l.asTypedList(starlark.Binary(op, l.list, y))

	case syntax.PLUS:
	// Convert regular lists to typed lists first.
	if yl, ok := y.(*starlark.List); ok {
	vals := make([]starlark.Value, yl.Len())
	for i := 0; i < yl.Len(); i++ {
	vals[i] = yl.Index(i)
	}
	asTyped, err := NewList(l.itemT, vals)
	if err != nil {
	return nil, err
	}
	return l.Binary(op, asTyped, side)
	}

	// When adding typed lists their types must match.
	if yl, ok := y.(*List); ok {
	if yl.itemT != l.itemT {
	goto unsupported
	}
	lhs, rhs := maybeSwap(l.list, yl.list, side)
	return l.asTypedList(starlark.Binary(op, lhs, rhs))
	}
	}

	unsupported:
	return nil, fmt.Errorf("unknown binary op: %s %s %s", l.Type(), op, y.Type())
	}

	func (l *List) asTypedList(v starlark.Value, err error) (starlark.Value, error) {
	if err != nil {
	return nil, err
	}
	if asL, ok := v.(*starlark.List); ok {
	return &List{itemT: l.itemT, list: asL}, nil
	}
	return nil, fmt.Errorf("expected a list, but got %q", v.Type())
	}

	func maybeSwap(l, r starlark.Value, side starlark.Side) (starlark.Value, starlark.Value) {
	if side == starlark.Left {
	return l, r
	}
	return r, l
	}

	// starlark.List methods we reimplement or proxy.

	func (l *List) Append(v starlark.Value) error {
	v, err := l.conv(v, -1)
	if err != nil {
	return err
	}
	return l.list.Append(v)
	}

	func (l *List) AttrNames() []string {
	// Note: we hardcode the list of methods (rather than call l.list.AttrNames)
	// to be able to notice (via the failing test) when go.starlark.net adds some
	// new methods that we might need to "wrap" in Attr(...).
	return []string{
	"append",
	"clear",
	"extend",
	"index",
	"insert",
	"pop",
	"remove",
	}
	}

	func (l *List) Attr(name string) (starlark.Value, error) {
	if name == "append" {
	return l.implAppend(), nil // fully reimplemented, no need for orig
	}
	orig, err := l.list.Attr(name)
	if err != nil {
	return nil, err
	}
	switch name {
	case "extend":
	return l.implExtend(orig.(*starlark.Builtin)), nil
	case "insert":
	return l.implInsert(orig.(*starlark.Builtin)), nil
	default:
	return orig, nil // no need to wrap
	}
	}

	func (l *List) SetIndex(i int, v starlark.Value) error {
	// Note: 'i' is never negative here, the Starlark evaluator deals with
	// negative indexes before calling SetIndex.
	v, err := l.conv(v, i)
	if err != nil {
	return err
	}
	return l.list.SetIndex(i, v)
	}

	func (l *List) Slice(start, end, step int) starlark.Value {
	return &List{
	itemT: l.itemT,
	list: l.list.Slice(start, end, step).(*starlark.List),
	}
	}

	func (l *List) Type() string {
	return fmt.Sprintf("list<%s>", l.itemT.Type())
	}

	func (l *List) String() string {
	return fmt.Sprintf("list<%s>(%s)", l.itemT.Type(), l.list.String())
	}

	func (l *List) CompareSameType(op syntax.Token, y starlark.Value, depth int) (bool, error) {
	switch op {
	case syntax.EQL:
	return listsEqual(l, y.(*List), depth)
	case syntax.NEQ:
	eq, err := listsEqual(l, y.(*List), depth)
	return !eq, err
	default:
	return false, fmt.Errorf("%q is not implemented for %s", op, l.Type())
	}
	}

	func listsEqual(l, r *List, depth int) (bool, error) {
	if l.itemT != r.itemT {
	return false, nil
	}
	return l.list.CompareSameType(syntax.EQL, r.list, depth)
	}

	// starlark.List methods delegated without any modifications.

	func (l *List) Clear() error { return l.list.Clear() }
	func (l *List) Freeze() { l.list.Freeze() }
	func (l *List) Hash() (uint32, error) { return l.list.Hash() }
	func (l *List) Index(i int) starlark.Value { return l.list.Index(i) }
	func (l *List) Iterate() starlark.Iterator { return l.list.Iterate() }
	func (l *List) Len() int { return l.list.Len() }
	func (l *List) Truth() starlark.Bool { return l.list.Truth() }

	// Reimplemented builtins.

	// implAppend reimplements 'append' via Append(...).
	func (l List) implAppend() starlark.Builtin {
	return starlark.NewBuiltin("append", func(th starlark.Thread, _ starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var object starlark.Value
	if err := starlark.UnpackPositionalArgs("append", args, kwargs, 1, &object); err != nil {
	return nil, err
	}
	if err := l.Append(object); err != nil {
	return nil, fmt.Errorf("append: %s", err)
	}
	return starlark.None, nil
	})
	}

	// implExtend reimplements 'extend' via Append(...).
	//
	// Uses original extend() for the fast path to optimize memory allocations.
	func (l List) implExtend(orig starlark.Builtin) *starlark.Builtin {
	return starlark.NewBuiltin("extend", func(th starlark.Thread, _ starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var iterable starlark.Iterable
	if err := starlark.UnpackPositionalArgs("extend", args, kwargs, 1, &iterable); err != nil {
	return nil, err
	}

	// Fast path: the iterable has the exact same type as us, we can skip the
	// conversion completely, since per Convert idempotency property this will
	// be a noop anyway. By calling original 'extend' implementation we are
	// basically doing one big memcpy(...) instead of adding items one by one.
	if it, ok := iterable.(*List); ok && it.itemT == l.itemT {
	return orig.CallInternal(th, starlark.Tuple{it.list}, nil)
	}

	// Slow path: append items one by one.
	iter := iterable.Iterate()
	defer iter.Done()
	var x starlark.Value
	for i := 0; iter.Next(&x); i++ {
	if err := l.Append(x); err != nil {
	return nil, fmt.Errorf("extend: item #%d: %s", i, err)
	}
	}
	return starlark.None, nil
	})
	}

	// implInsert implements 'insert' by wrapping the original 'insert'.
	//
	// It's not easily reimplementable, since starlark.List doesn't expose its
	// internal list in a simple way.
	func (l List) implInsert(orig starlark.Builtin) *starlark.Builtin {
	return starlark.NewBuiltin("insert", func(th starlark.Thread, _ starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var index starlark.Value
	var object starlark.Value
	if err := starlark.UnpackPositionalArgs("insert", args, kwargs, 2, &index, &object); err != nil {
	return nil, err
	}
	x, err := l.conv(object, -1)
	if err != nil {
	return nil, fmt.Errorf("insert: %s", err)
	}
	return orig.CallInternal(th, starlark.Tuple{index, x}, nil)
	})
	}