github.com/emcfarlane/larking@v0.0.0-20220605172417-1704b45ee6c3/starlib/encoding/starlarkproto/proto.go (about) 1 // Copyright 2021 Edward McFarlane. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 // Package starlarkproto provides support for protocol buffers. 6 package starlarkproto 7 8 import ( 9 "fmt" 10 "sort" 11 "strings" 12 13 "go.starlark.net/starlark" 14 "go.starlark.net/syntax" 15 "google.golang.org/protobuf/encoding/protojson" 16 "google.golang.org/protobuf/encoding/prototext" 17 "google.golang.org/protobuf/proto" 18 "google.golang.org/protobuf/reflect/protodesc" 19 "google.golang.org/protobuf/reflect/protoreflect" 20 "google.golang.org/protobuf/reflect/protoregistry" 21 "google.golang.org/protobuf/types/dynamicpb" 22 23 "github.com/emcfarlane/larking/starlib/starext" 24 "github.com/emcfarlane/larking/starlib/starlarkstruct" 25 ) 26 27 const ( 28 resolverKey = "protodescResolver" 29 ) 30 31 func SetProtodescResolver(thread *starlark.Thread, resolver protodesc.Resolver) { 32 thread.SetLocal(resolverKey, resolver) 33 } 34 35 func GetProtodescResolver(thread *starlark.Thread) protodesc.Resolver { 36 if resolver, ok := thread.Local(resolverKey).(protodesc.Resolver); ok { 37 return resolver 38 } 39 return protoregistry.GlobalFiles 40 } 41 42 func NewModule() *starlarkstruct.Module { 43 p := NewProto() 44 return &starlarkstruct.Module{ 45 Name: "proto", 46 Members: starlark.StringDict{ 47 "file": starext.MakeBuiltin("proto.file", p.File), 48 //"load": starext.MakeBuiltin("proto.load", p.Load), 49 "new": starext.MakeBuiltin("proto.new", p.New), 50 "marshal": starext.MakeBuiltin("proto.marshal", p.Marshal), 51 "unmarshal": starext.MakeBuiltin("proto.unmarshal", p.Unmarshal), 52 "marshal_json": starext.MakeBuiltin("proto.marshal_json", p.MarshalJSON), 53 "unmarshal_json": starext.MakeBuiltin("proto.unmarshal_json", p.UnmarshalJSON), 54 "marshal_text": starext.MakeBuiltin("proto.marshal_text", p.MarshalText), 55 "unmarshal_text": starext.MakeBuiltin("proto.unmarshal_text", p.UnmarshalText), 56 }, 57 } 58 } 59 60 type Proto struct { 61 //resolver protodesc.Resolver 62 types protoregistry.Types // TODO: wrap resolver to register extensions. 63 } 64 65 func NewProto() *Proto { 66 return &Proto{} 67 } 68 69 func (p *Proto) File(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 70 var name string 71 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 1, &name); err != nil { 72 return nil, err 73 } 74 75 fileDesc, err := GetProtodescResolver(thread).FindFileByPath(name) 76 if err != nil { 77 return nil, err 78 } 79 return &Descriptor{desc: fileDesc}, nil 80 } 81 82 //func (p *Proto) Load(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 83 // var data string 84 // if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 1, &data); err != nil { 85 // return nil, err 86 // } 87 // 88 // resolver := GetProtodescResolver(thread) 89 // 90 // var file descriptorpb.FieldDescriptorProto 91 // if err := proto.Unmarshal([]byte(data), &file); err != nil { 92 // return nil, err 93 // } 94 // 95 // fileDesc, err := 96 // //fileDesc, err := .FindFileByPath(name) 97 // //if err != nil { 98 // // return nil, err 99 // //} 100 // //return &Descriptor{desc: fileDesc}, nil 101 //} 102 103 func (p *Proto) New(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 104 var name string 105 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 1, &name); err != nil { 106 return nil, err 107 } 108 fullname := protoreflect.FullName(name) 109 110 desc, err := GetProtodescResolver(thread).FindDescriptorByName(fullname) 111 if err != nil { 112 return nil, err 113 } 114 return &Descriptor{desc: desc}, nil 115 } 116 117 func (p *Proto) Marshal(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 118 var msg *Message 119 var options proto.MarshalOptions 120 if err := starlark.UnpackPositionalArgs( 121 fnname, args, kwargs, 1, &msg, 122 "allow_partial?", &options.AllowPartial, 123 "deterministic?", &options.Deterministic, 124 "use_cache_size?", &options.UseCachedSize, 125 ); err != nil { 126 return nil, err 127 } 128 data, err := options.Marshal(msg) 129 if err != nil { 130 return nil, err 131 } 132 return starlark.String(string(data)), nil 133 } 134 135 func (p *Proto) Unmarshal(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 136 var str string 137 var msg *Message 138 options := proto.UnmarshalOptions{ 139 Resolver: &p.types, // TODO: types... 140 } 141 if err := starlark.UnpackPositionalArgs( 142 fnname, args, kwargs, 2, &str, &msg, 143 "merge?", &options.Merge, 144 "allow_partial?", &options.AllowPartial, 145 "discard_unknown?", &options.DiscardUnknown, 146 ); err != nil { 147 return nil, err 148 } 149 if err := msg.checkMutable(fnname); err != nil { 150 return nil, err 151 } 152 if err := proto.Unmarshal([]byte(str), msg); err != nil { 153 return nil, err 154 } 155 return starlark.None, nil 156 } 157 158 func (p *Proto) MarshalJSON(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 159 var msg *Message 160 var options protojson.MarshalOptions 161 if err := starlark.UnpackPositionalArgs( 162 fnname, args, kwargs, 1, &msg, 163 "multiline?", &options.Multiline, 164 "indent?", &options.Indent, 165 "allow_partial?", &options.AllowPartial, 166 "use_proto_names?", &options.UseProtoNames, 167 "use_enum_numbers?", &options.UseEnumNumbers, 168 "emit_unpopulated?", &options.EmitUnpopulated, 169 ); err != nil { 170 return nil, err 171 } 172 data, err := options.Marshal(msg) 173 if err != nil { 174 return nil, err 175 } 176 return starlark.String(string(data)), nil 177 } 178 179 func (p *Proto) UnmarshalJSON(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 180 var str string 181 var msg *Message 182 options := protojson.UnmarshalOptions{ 183 Resolver: &p.types, // TODO: types... 184 } 185 if err := starlark.UnpackPositionalArgs( 186 fnname, args, kwargs, 2, &str, &msg, 187 "allow_partial?", &options.AllowPartial, 188 "discard_unknown?", &options.DiscardUnknown, 189 ); err != nil { 190 return nil, err 191 } 192 if err := msg.checkMutable(fnname); err != nil { 193 return nil, err 194 } 195 if err := proto.Unmarshal([]byte(str), msg); err != nil { 196 return nil, err 197 } 198 return starlark.None, nil 199 } 200 201 func (p *Proto) MarshalText(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 202 var msg *Message 203 var options prototext.MarshalOptions 204 if err := starlark.UnpackPositionalArgs( 205 fnname, args, kwargs, 1, &msg, 206 "multiline?", &options.Multiline, 207 "indent?", &options.Indent, 208 "allow_partial?", &options.AllowPartial, 209 "emit_unknown?", &options.EmitUnknown, 210 ); err != nil { 211 return nil, err 212 } 213 data, err := options.Marshal(msg) 214 if err != nil { 215 return nil, err 216 } 217 return starlark.String(string(data)), nil 218 } 219 220 func (p *Proto) UnmarshalText(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 221 var str string 222 var msg *Message 223 options := prototext.UnmarshalOptions{ 224 Resolver: &p.types, // TODO: types... 225 } 226 if err := starlark.UnpackPositionalArgs( 227 fnname, args, kwargs, 2, &str, &msg, 228 "allow_partial?", &options.AllowPartial, 229 "discard_unknown?", &options.DiscardUnknown, 230 ); err != nil { 231 return nil, err 232 } 233 if err := msg.checkMutable(fnname); err != nil { 234 return nil, err 235 } 236 if err := proto.Unmarshal([]byte(str), msg); err != nil { 237 return nil, err 238 } 239 return starlark.None, nil 240 } 241 242 func equalFullName(a, b protoreflect.FullName) error { 243 if a != b { 244 return fmt.Errorf("type mismatch %s != %s", a, b) 245 } 246 return nil 247 } 248 249 type Descriptor struct { 250 desc protoreflect.Descriptor 251 252 frozen bool 253 attrs map[string]protoreflect.Descriptor 254 } 255 256 func NewDescriptor(desc protoreflect.Descriptor) *Descriptor { return &Descriptor{desc: desc} } 257 258 // Descriptor exports proto.Descriptor 259 func (d *Descriptor) Descriptor() protoreflect.Descriptor { return d.desc } 260 261 func (d *Descriptor) String() string { return string(d.desc.Name()) } 262 func (d *Descriptor) Type() string { return "proto.desc" } 263 func (d *Descriptor) Freeze() { d.frozen = true } 264 func (d *Descriptor) Truth() starlark.Bool { return d.desc != nil } 265 func (d *Descriptor) Hash() (uint32, error) { return 0, fmt.Errorf("unhashable type: proto.desc") } 266 func (d *Descriptor) Name() string { return string(d.desc.Name()) } // TODO 267 func (d *Descriptor) CallInternal(thread *starlark.Thread, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 268 switch v := d.desc.(type) { 269 case protoreflect.FileDescriptor: 270 return nil, fmt.Errorf("proto: file descriptor not callable") 271 272 case protoreflect.EnumDescriptor: 273 if len(kwargs) > 0 { 274 return nil, fmt.Errorf("unexpected kwargs") 275 } 276 if len(args) != 1 { 277 return nil, fmt.Errorf("unexpected number of args") 278 } 279 vals := v.Values() 280 return NewEnum(vals, args[0]) 281 282 case protoreflect.MessageDescriptor: 283 // Create the msg, try to use to Go type if avaliable. 284 var msg protoreflect.Message 285 if mt, err := protoregistry.GlobalTypes.FindMessageByName( 286 v.FullName(), 287 ); err == nil { 288 msg = mt.New() 289 fmt.Println("go type", msg) 290 } else { 291 // Fallback to dynamic meessages. 292 msg = dynamicpb.NewMessage(v) 293 fmt.Println("dynamic type", msg) 294 } 295 return NewMessage(msg, args, kwargs) 296 297 default: 298 return nil, fmt.Errorf("proto: desc missing call type %T", v) 299 } 300 } 301 302 func (d *Descriptor) getAttrs() map[string]protoreflect.Descriptor { 303 if d.attrs != nil { 304 return d.attrs 305 } 306 m := make(map[string]protoreflect.Descriptor) 307 308 switch v := d.desc.(type) { 309 case protoreflect.FileDescriptor: 310 for i, eds := 0, v.Enums(); i < eds.Len(); i++ { 311 ed := eds.Get(i) 312 m[string(ed.Name())] = ed 313 } 314 for i, mds := 0, v.Messages(); i < mds.Len(); i++ { 315 md := mds.Get(i) 316 m[string(md.Name())] = md 317 } 318 for i, eds := 0, v.Extensions(); i < eds.Len(); i++ { 319 ed := eds.Get(i) 320 m[string(ed.Name())] = ed 321 } 322 for i, sds := 0, v.Services(); i < sds.Len(); i++ { 323 sd := sds.Get(i) 324 m[string(sd.Name())] = sd 325 } 326 327 case protoreflect.EnumDescriptor: 328 for i, eds := 0, v.Values(); i < eds.Len(); i++ { 329 evd := eds.Get(i) 330 m[string(evd.Name())] = evd 331 } 332 333 case protoreflect.MessageDescriptor: 334 for i, eds := 0, v.Enums(); i < eds.Len(); i++ { 335 ed := eds.Get(i) 336 m[string(ed.Name())] = ed 337 } 338 for i, mds := 0, v.Messages(); i < mds.Len(); i++ { 339 md := mds.Get(i) 340 m[string(md.Name())] = md 341 } 342 for i, ods := 0, v.Oneofs(); i < ods.Len(); i++ { 343 od := ods.Get(i) 344 m[string(od.Name())] = od 345 } 346 347 case protoreflect.ServiceDescriptor: 348 for i, mds := 0, v.Methods(); i < mds.Len(); i++ { 349 md := mds.Get(i) 350 m[string(md.Name())] = md 351 } 352 353 default: 354 panic(fmt.Sprintf("proto: desc missing attr type %T", v)) 355 } 356 357 if !d.frozen { 358 d.attrs = m 359 } 360 return m 361 } 362 363 func (d *Descriptor) Attr(name string) (starlark.Value, error) { 364 // TODO: can this just use the resolver? 365 attrs := d.getAttrs() 366 desc, ok := attrs[name] 367 if !ok { 368 return nil, nil 369 } 370 // Special descriptor type handling 371 switch v := desc.(type) { 372 case protoreflect.EnumValueDescriptor: 373 return Enum{edesc: v}, nil 374 default: 375 return &Descriptor{desc: desc}, nil 376 } 377 } 378 379 func (d *Descriptor) AttrNames() []string { 380 var names []string 381 for name := range d.getAttrs() { 382 names = append(names, name) 383 } 384 sort.Strings(names) 385 return names 386 } 387 388 // Message represents a proto.Message as a starlark.Value. 389 type Message struct { 390 msg protoreflect.Message 391 frozen *bool 392 } 393 394 // ProtoReflect implements proto.Message 395 func (m *Message) ProtoReflect() protoreflect.Message { return m.msg } 396 397 // Type conversions rules: 398 // 399 // ═══════════════╤════════════════════════════════════ 400 // Starlark type │ Protobuf Type 401 // ═══════════════╪════════════════════════════════════ 402 // NoneType │ MessageKind, GroupKind 403 // Bool │ BoolKind 404 // Int │ Int32Kind, Sint32Kind, Sfixed32Kind, 405 // │ Int64Kind, Sint64Kind, Sfixed64Kind, 406 // │ Uint32Kind, Fixed32Kind, 407 // │ Uint64Kind, Fixed64Kind 408 // Float │ FloatKind, DoubleKind 409 // String │ StringKind, BytesKind 410 // *List │ List<Kind> 411 // Tuple │ n/a 412 // *Dict │ Map<Kind><Kind> 413 // *Set │ n/a 414 // 415 func toStarlark(v protoreflect.Value, fd protoreflect.FieldDescriptor, frozen *bool) starlark.Value { 416 switch v := v.Interface().(type) { 417 case nil: 418 return starlark.None 419 case bool: 420 return starlark.Bool(v) 421 case int32: 422 return starlark.MakeInt(int(v)) 423 case int64: 424 return starlark.MakeInt(int(v)) 425 case uint32: 426 return starlark.MakeInt(int(v)) 427 case uint64: 428 return starlark.MakeInt(int(v)) 429 case float32: 430 return starlark.Float(float64(v)) 431 case float64: 432 return starlark.Float(v) 433 case string: 434 return starlark.String(v) 435 case []byte: 436 return starlark.String(v) 437 case protoreflect.EnumNumber: 438 evdesc := fd.Enum().Values().ByNumber(v) 439 if evdesc == nil { 440 evdesc = fd.DefaultEnumValue() // TODO: error? 441 } 442 return Enum{edesc: evdesc} 443 case protoreflect.List: 444 return &List{list: v, fd: fd, frozen: frozen} 445 case protoreflect.Message: 446 return &Message{msg: v, frozen: frozen} 447 case protoreflect.Map: 448 return &Map{m: v, keyfd: fd.MapKey(), valfd: fd.MapValue(), frozen: frozen} 449 default: 450 panic(fmt.Sprintf("unhandled proto type %s %T", v, v)) 451 } 452 } 453 454 func allocField(parent protoreflect.Value, fd protoreflect.FieldDescriptor) protoreflect.Value { 455 switch v := parent.Interface().(type) { 456 case protoreflect.List: 457 return v.NewElement() 458 case protoreflect.Map: 459 return v.NewValue() 460 case protoreflect.Message: 461 return v.NewField(fd) 462 default: 463 panic(fmt.Sprintf("unhandled parent value type: %T", v)) 464 } 465 } 466 467 func toProtobuf(v starlark.Value, fd protoreflect.FieldDescriptor, parent protoreflect.Value) (protoreflect.Value, error) { 468 switch kind := fd.Kind(); kind { 469 case protoreflect.BoolKind: 470 if b, ok := v.(starlark.Bool); ok { 471 return protoreflect.ValueOfBool(bool(b)), nil 472 } 473 case protoreflect.Int32Kind, protoreflect.Sint32Kind, protoreflect.Sfixed32Kind: 474 if x, err := starlark.NumberToInt(v); err == nil { 475 v, err := starlark.AsInt32(x) 476 if err != nil { 477 return protoreflect.Value{}, err 478 } 479 return protoreflect.ValueOfInt32(int32(v)), nil 480 } 481 case protoreflect.Int64Kind, protoreflect.Sint64Kind, protoreflect.Sfixed64Kind: 482 if x, err := starlark.NumberToInt(v); err == nil { 483 v, _ := x.Int64() 484 return protoreflect.ValueOfInt64(v), nil 485 } 486 case protoreflect.Uint32Kind, protoreflect.Fixed32Kind: 487 if x, err := starlark.NumberToInt(v); err == nil { 488 v, _ := x.Uint64() 489 return protoreflect.ValueOfUint32(uint32(v)), nil 490 } 491 case protoreflect.Uint64Kind, protoreflect.Fixed64Kind: 492 if x, err := starlark.NumberToInt(v); err == nil { 493 v, _ := x.Uint64() 494 return protoreflect.ValueOfUint64(v), nil 495 } 496 case protoreflect.FloatKind: 497 if x, ok := starlark.AsFloat(v); ok { 498 return protoreflect.ValueOfFloat32(float32(x)), nil 499 } 500 case protoreflect.DoubleKind: 501 if x, ok := starlark.AsFloat(v); ok { 502 return protoreflect.ValueOfFloat64(float64(x)), nil 503 } 504 case protoreflect.StringKind: 505 if x, ok := v.(starlark.String); ok { 506 return protoreflect.ValueOfString(string(x)), nil 507 } 508 case protoreflect.BytesKind: 509 if x, ok := v.(starlark.String); ok { 510 return protoreflect.ValueOfBytes([]byte(x)), nil 511 } 512 case protoreflect.EnumKind: 513 switch v := v.(type) { 514 case starlark.String: 515 enumVal := fd.Enum().Values().ByName(protoreflect.Name(string(v))) 516 if enumVal == nil { 517 return protoreflect.Value{}, fmt.Errorf("proto: enum has no %s value", v) 518 } 519 return protoreflect.ValueOfEnum(enumVal.Number()), nil 520 case starlark.Int, starlark.Float: 521 i, err := starlark.NumberToInt(v) 522 if err != nil { 523 return protoreflect.Value{}, err 524 } 525 x, ok := i.Int64() 526 if !ok { 527 return protoreflect.Value{}, fmt.Errorf("proto: enum has no %s value", v) 528 } 529 return protoreflect.ValueOfEnum(protoreflect.EnumNumber(int32(x))), nil 530 case Enum: 531 return protoreflect.ValueOfEnum(v.edesc.Number()), nil 532 } 533 case protoreflect.MessageKind: 534 if fd.IsMap() { 535 switch v := v.(type) { 536 case *Map: 537 return protoreflect.ValueOfMap(v.m), nil 538 case starlark.IterableMapping: 539 val := allocField(parent, fd) 540 541 mm := val.Map() 542 kfd := fd.MapKey() 543 vfd := fd.MapValue() 544 545 items := v.Items() 546 for _, item := range items { 547 // can only be scalar. 548 kval, err := toProtobuf(item[0], kfd, val) 549 if err != nil { 550 return protoreflect.Value{}, err 551 } 552 mkey := kval.MapKey() 553 554 mval, err := toProtobuf(item[1], vfd, val) 555 if err != nil { 556 return protoreflect.Value{}, err 557 } 558 559 mm.Set(mkey, mval) 560 } 561 return val, nil 562 //return protoreflect.ValueOfMap(mm), nil 563 } 564 } else { 565 switch v := v.(type) { 566 case *Message: 567 return protoreflect.ValueOfMessage(v.msg), nil 568 case starlark.NoneType: 569 msg := parent.Message() 570 msg.Clear(fd) 571 return msg.Get(fd), nil // RO 572 case starlark.IterableMapping: 573 val := allocField(parent, fd) 574 m := Message{msg: val.Message(), frozen: new(bool)} // wrap for set 575 576 for _, kv := range v.Items() { 577 key, ok := kv[0].(starlark.String) 578 if !ok { 579 return protoreflect.Value{}, fmt.Errorf("proto: invalid key type %s", kv[0].Type()) 580 } 581 if err := m.SetField(string(key), kv[1]); err != nil { 582 return protoreflect.Value{}, err 583 } 584 } 585 return val, nil 586 case starlark.HasAttrs: 587 val := allocField(parent, fd) 588 m := Message{msg: val.Message(), frozen: new(bool)} // wrap for set 589 590 for _, name := range v.AttrNames() { 591 val, err := v.Attr(name) 592 if err != nil { 593 return protoreflect.Value{}, err 594 } 595 if err := m.SetField(name, val); err != nil { 596 return protoreflect.Value{}, err 597 } 598 } 599 return val, nil 600 } 601 } 602 default: 603 panic(fmt.Sprintf("unknown kind %q", kind)) 604 } 605 return protoreflect.Value{}, fmt.Errorf( 606 "proto: unknown type conversion %s<%T> to %s", v, v, fd.Kind(), 607 ) 608 } 609 610 func (m *Message) encodeField(v starlark.Value, fd protoreflect.FieldDescriptor) (protoreflect.Value, error) { 611 // It is equivalent to checking whether Cardinality is Repeated and 612 // that IsMap reports false. 613 if !fd.IsList() { 614 return toProtobuf(v, fd, protoreflect.ValueOfMessage(m.msg)) 615 } 616 617 switch v := v.(type) { 618 case *List: 619 // Starlark type is wrapped in ref by caller. 620 return protoreflect.ValueOfList(v.list), nil 621 622 case starlark.Indexable: 623 val := m.msg.NewField(fd) 624 l := val.List() 625 626 for i := 0; i < v.Len(); i++ { 627 val, err := toProtobuf(v.Index(i), fd, val) 628 if err != nil { 629 return protoreflect.Value{}, err 630 } 631 l.Append(val) 632 } 633 return val, nil 634 635 case starlark.Iterable: 636 val := m.msg.NewField(fd) 637 l := val.List() 638 639 iter := v.Iterate() 640 defer iter.Done() 641 642 var p starlark.Value 643 for iter.Next(&p) { 644 val, err := toProtobuf(p, fd, val) 645 if err != nil { 646 return protoreflect.Value{}, err 647 } 648 l.Append(val) 649 } 650 return val, nil 651 } 652 return protoreflect.Value{}, fmt.Errorf("proto: unknown repeated type conversion %s", v.Type()) 653 } 654 655 func (m *Message) checkMutable(verb string) error { 656 if *m.frozen { 657 return fmt.Errorf("cannot %s frozen message", verb) 658 } 659 if !m.msg.IsValid() { 660 return fmt.Errorf("cannot %s non mutable message", verb) 661 } 662 return nil 663 } 664 665 func NewMessage(msg protoreflect.Message, args starlark.Tuple, kwargs []starlark.Tuple) (*Message, error) { 666 hasArgs := len(args) > 0 667 hasKwargs := len(kwargs) > 0 668 669 if hasArgs && len(args) > 1 { 670 return nil, fmt.Errorf("unexpected number of args") 671 } 672 673 if hasArgs && hasKwargs { 674 return nil, fmt.Errorf("unxpected args and kwargs") 675 } 676 677 if hasArgs { 678 switch v := args[0].(type) { 679 case *Message: 680 return v, nil 681 case starlark.NoneType: 682 return &Message{msg: msg.Type().Zero(), frozen: new(bool)}, nil // RO 683 case starlark.IterableMapping: 684 m := &Message{msg: msg, frozen: new(bool)} 685 for _, kv := range v.Items() { 686 key, ok := kv[0].(starlark.String) 687 if !ok { 688 return nil, fmt.Errorf("proto: invalid key type %s", kv[0].Type()) 689 } 690 if err := m.SetField(string(key), kv[1]); err != nil { 691 return nil, err 692 } 693 } 694 return m, nil 695 case starlark.HasAttrs: 696 m := &Message{msg: msg, frozen: new(bool)} 697 for _, name := range v.AttrNames() { 698 val, err := v.Attr(name) 699 if err != nil { 700 return nil, err 701 } 702 if err := m.SetField(name, val); err != nil { 703 return nil, err 704 } 705 } 706 return m, nil 707 default: 708 return nil, fmt.Errorf("proto: unknown type conversion %s<%T> to proto.message", v, v) 709 } 710 } 711 712 m := &Message{msg: msg, frozen: new(bool)} 713 for _, kwarg := range kwargs { 714 k := string(kwarg[0].(starlark.String)) 715 v := kwarg[1] 716 717 if err := m.SetField(k, v); err != nil { 718 return nil, err 719 } 720 } 721 return m, nil 722 } 723 724 func (m *Message) String() string { 725 desc := m.msg.Descriptor() 726 buf := new(strings.Builder) 727 buf.WriteString(string(desc.Name())) 728 729 buf.WriteByte('(') 730 if m.msg.IsValid() { 731 fds := desc.Fields() 732 for i := 0; i < fds.Len(); i++ { 733 if i > 0 { 734 buf.WriteString(", ") 735 } 736 fd := fds.Get(i) 737 buf.WriteString(string(fd.Name())) 738 buf.WriteString(" = ") 739 v := m.msg.Get(fd) 740 buf.WriteString(v.String()) 741 } 742 } else { 743 buf.WriteString("None") 744 } 745 buf.WriteByte(')') 746 return buf.String() 747 } 748 749 func (m *Message) Type() string { return "proto.message" } 750 func (m *Message) Truth() starlark.Bool { return starlark.Bool(m.msg.IsValid()) } 751 func (m *Message) Hash() (uint32, error) { 752 return 0, fmt.Errorf("unhashable type: proto.message") 753 } 754 func (m *Message) Freeze() { *m.frozen = true } 755 756 // Attr returns the value of the specified field. 757 func (m *Message) Attr(name string) (starlark.Value, error) { 758 fd, err := m.fieldDesc(name) 759 if err != nil { 760 return nil, err 761 } 762 // Get mutable references if we can. 763 if fd.IsMap() || fd.IsList() || (fd.Kind() == protoreflect.MessageKind && m.msg.Has(fd)) { 764 return toStarlark(m.msg.Mutable(fd), fd, m.frozen), nil 765 } 766 return toStarlark(m.msg.Get(fd), fd, m.frozen), nil 767 } 768 769 func (x *Message) Binary(op syntax.Token, y starlark.Value, side starlark.Side) (starlark.Value, error) { 770 return nil, nil // unhandled 771 } 772 773 // AttrNames returns a new sorted list of the message fields. 774 func (m *Message) AttrNames() []string { 775 desc := m.msg.Descriptor() 776 fds := desc.Fields() 777 ods := desc.Oneofs() 778 names := make([]string, fds.Len()+ods.Len()) 779 for i := 0; i < fds.Len(); i++ { 780 fd := fds.Get(i) 781 names[i] = string(fd.Name()) 782 } 783 offset := fds.Len() 784 for i := 0; i < ods.Len(); i++ { 785 od := ods.Get(i) 786 names[offset+i] = string(od.Name()) 787 } 788 sort.Strings(names) // TODO: sort by protobuf number 789 return names 790 } 791 792 func (m *Message) fieldDesc(name string) (protoreflect.FieldDescriptor, error) { 793 desc := m.msg.Descriptor() 794 if fd := desc.Fields().ByName(protoreflect.Name(name)); fd != nil { 795 return fd, nil 796 } 797 798 if od := desc.Oneofs().ByName(protoreflect.Name(name)); od != nil { 799 return m.msg.WhichOneof(od), nil 800 } 801 return nil, starlark.NoSuchAttrError( 802 fmt.Sprintf("%s has no .%s attribute", desc.Name(), name), 803 ) 804 } 805 806 func (m *Message) SetField(name string, val starlark.Value) error { 807 if err := m.checkMutable("set field"); err != nil { 808 return err 809 } 810 fd, err := m.fieldDesc(name) 811 if err != nil { 812 return err 813 } 814 815 if val == starlark.None { 816 m.msg.Clear(fd) 817 return nil 818 } 819 820 v, err := m.encodeField(val, fd) 821 if err != nil { 822 return err 823 } 824 825 m.msg.Set(fd, v) 826 return nil 827 } 828 829 func (x *Message) CompareSameType(op syntax.Token, y_ starlark.Value, depth int) (bool, error) { 830 y := y_.(*Message) 831 switch op { 832 case syntax.EQL: 833 return proto.Equal(x, y), nil 834 case syntax.NEQ: 835 return !proto.Equal(x, y), nil 836 case syntax.LE, syntax.LT, syntax.GE, syntax.GT: 837 return false, fmt.Errorf("%v not implemented", op) 838 default: 839 panic(op) 840 } 841 } 842 843 // List represents a repeated field as a starlark.List. 844 type List struct { 845 list protoreflect.List 846 fd protoreflect.FieldDescriptor 847 848 frozen *bool 849 itercount uint32 850 } 851 852 type listAttr func(l *List) starlark.Value 853 854 // methods from starlark/library.go 855 var listAttrs = map[string]listAttr{ 856 "append": func(l *List) starlark.Value { return starext.MakeMethod(l, "append", l.append) }, 857 "clear": func(l *List) starlark.Value { return starext.MakeMethod(l, "clear", l.clear) }, 858 "extend": func(l *List) starlark.Value { return starext.MakeMethod(l, "extend", l.extend) }, 859 "index": func(l *List) starlark.Value { return starext.MakeMethod(l, "index", l.index) }, 860 "insert": func(l *List) starlark.Value { return starext.MakeMethod(l, "insert", l.insert) }, 861 "pop": func(l *List) starlark.Value { return starext.MakeMethod(l, "pop", l.pop) }, 862 "remove": func(l *List) starlark.Value { return starext.MakeMethod(l, "remove", l.remove) }, 863 } 864 865 func (l *List) Attr(name string) (starlark.Value, error) { 866 if a := listAttrs[name]; a != nil { 867 return a(l), nil 868 } 869 return nil, nil 870 871 } 872 func (l *List) AttrNames() []string { 873 names := make([]string, 0, len(listAttrs)) 874 for name := range listAttrs { 875 names = append(names, name) 876 } 877 sort.Strings(names) 878 return names 879 } 880 881 func (l *List) String() string { 882 buf := new(strings.Builder) 883 buf.WriteByte('[') 884 if l.list.IsValid() { 885 for i := 0; i < l.Len(); i++ { 886 if i > 0 { 887 buf.WriteString(", ") 888 } 889 buf.WriteString(l.Index(i).String()) 890 } 891 } 892 buf.WriteByte(']') 893 return buf.String() 894 } 895 896 func (l *List) Freeze() { *l.frozen = true } 897 898 func (l *List) Hash() (uint32, error) { 899 return 0, fmt.Errorf("unhashable type: proto.list") 900 } 901 902 func (l *List) checkMutable(verb string) error { 903 if *l.frozen { 904 return fmt.Errorf("cannot %s frozen list", verb) 905 } 906 if l.itercount > 0 { 907 return fmt.Errorf("cannot %s list during iteration", verb) 908 } 909 if !l.list.IsValid() { 910 return fmt.Errorf("cannot %s non mutable list", verb) 911 } 912 return nil 913 } 914 915 func (l *List) Index(i int) starlark.Value { 916 return toStarlark(l.list.Get(i), l.fd, l.frozen) 917 } 918 919 type listIterator struct { 920 l *List 921 i int 922 } 923 924 func (it *listIterator) Next(p *starlark.Value) bool { 925 if it.i < it.l.Len() { 926 v := it.l.list.Get(it.i) 927 *p = toStarlark(v, it.l.fd, it.l.frozen) 928 return true 929 } 930 return false 931 } 932 933 func (it *listIterator) Done() { 934 if !*it.l.frozen { 935 it.l.itercount-- 936 } 937 } 938 939 func (l *List) Iterate() starlark.Iterator { 940 if !*l.frozen { 941 l.itercount++ 942 } 943 return &listIterator{l: l} 944 } 945 946 // From Hacker's Delight, section 2.8. 947 func signum(x int64) int { return int(uint64(x>>63) | uint64(-x)>>63) } 948 949 // Slice copies values to a starlark.List 950 func (l *List) Slice(start, end, step int) starlark.Value { 951 sign := signum(int64(step)) 952 953 var elems []starlark.Value 954 for i := start; signum(int64(end-i)) == sign; i += step { 955 elems = append(elems, l.Index(i)) 956 } 957 return starlark.NewList(elems) 958 } 959 960 func (l *List) Clear() error { 961 if err := l.checkMutable("clear"); err != nil { 962 return err 963 } 964 if l.list.Len() > 0 { 965 l.list.Truncate(0) 966 } 967 return nil 968 } 969 970 func (l *List) Type() string { return l.fd.Kind().String() } 971 func (l *List) Len() int { return l.list.Len() } 972 func (l *List) Truth() starlark.Bool { return l.Len() > 0 } 973 974 func (l *List) SetIndex(i int, v starlark.Value) error { 975 if err := l.checkMutable("assign to element of"); err != nil { 976 return err 977 } 978 979 val, err := toProtobuf(v, l.fd, protoreflect.ValueOfList(l.list)) 980 if err != nil { 981 return err 982 } 983 984 l.list.Set(i, val) 985 return nil 986 } 987 988 func (l *List) Append(v starlark.Value) error { 989 if err := l.checkMutable("append to"); err != nil { 990 return err 991 } 992 993 val, err := toProtobuf(v, l.fd, protoreflect.ValueOfList(l.list)) 994 if err != nil { 995 return err 996 } 997 998 l.list.Append(val) 999 return nil 1000 } 1001 1002 func (l *List) Pop(i int) (starlark.Value, error) { 1003 v := l.Index(i) 1004 n := l.Len() 1005 1006 // shift list after index 1007 for j := i; j < n-1; j++ { 1008 val := l.list.Get(j + 1) 1009 l.list.Set(j, val) 1010 } 1011 l.list.Truncate(n - 1) 1012 1013 return v, nil 1014 1015 } 1016 1017 func (v *List) append(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1018 var object starlark.Value 1019 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 1, &object); err != nil { 1020 return nil, err 1021 } 1022 if err := v.Append(object); err != nil { 1023 return nil, err 1024 } 1025 return starlark.None, nil 1026 } 1027 1028 func (v *List) clear(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1029 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 0); err != nil { 1030 return nil, err 1031 } 1032 if err := v.Clear(); err != nil { 1033 return nil, err 1034 } 1035 return starlark.None, nil 1036 } 1037 1038 func (v *List) extend(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1039 var iterable starlark.Iterable 1040 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 1, &iterable); err != nil { 1041 return nil, err 1042 } 1043 iter := iterable.Iterate() 1044 var p starlark.Value 1045 for iter.Next(&p) { 1046 if err := v.Append(p); err != nil { 1047 return nil, err 1048 } 1049 } 1050 return starlark.None, nil 1051 } 1052 1053 func outOfRange(i, n int, x starlark.Value) error { 1054 if n == 0 { 1055 return fmt.Errorf("index %d out of range: empty %s", i, x.Type()) 1056 } else { 1057 return fmt.Errorf("%s index %d out of range [%d:%d]", x.Type(), i, -n, n-1) 1058 } 1059 } 1060 1061 func absIndex(i, len int) int { 1062 if i < 0 { 1063 i += len // negative offset 1064 } 1065 // clamp [0:len] 1066 if i < 0 { 1067 i = 0 1068 } else if i > len { 1069 i = len 1070 } 1071 return i 1072 } 1073 1074 func asIndex(v starlark.Value, len int, result *int) (err error) { 1075 if v != nil && v != starlark.None { 1076 *result, err = starlark.AsInt32(v) 1077 if err != nil { 1078 return err 1079 } 1080 *result = absIndex(*result, len) 1081 } 1082 return nil 1083 } 1084 1085 func (v *List) index(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1086 var value, start_, end_ starlark.Value 1087 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 1, &value, &start_, &end_); err != nil { 1088 return nil, err 1089 } 1090 1091 len := v.Len() 1092 start := 0 1093 if err := asIndex(start_, len, &start); err != nil { 1094 return nil, err 1095 } 1096 1097 end := len 1098 if err := asIndex(end_, len, &end); err != nil { 1099 return nil, err 1100 } 1101 1102 // find 1103 for i := start; i < end; i++ { 1104 if ok, err := starlark.Equal(v.Index(i), value); ok { 1105 return starlark.MakeInt(i), nil 1106 } else if err != nil { 1107 return nil, fmt.Errorf("%s: %w", fnname, err) 1108 } 1109 } 1110 return nil, fmt.Errorf("%s: value not in list", fnname) 1111 } 1112 1113 func (v *List) insert(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1114 var ( 1115 index int 1116 object starlark.Value 1117 ) 1118 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 2, &index, &object); err != nil { 1119 return nil, err 1120 } 1121 if err := v.checkMutable("insert into"); err != nil { 1122 return nil, fmt.Errorf("%s: %w", v.Type(), err) 1123 } 1124 1125 len := v.Len() 1126 index = absIndex(index, len) 1127 if index >= len { 1128 if err := v.Append(object); err != nil { 1129 return nil, err 1130 } 1131 return starlark.None, nil 1132 } 1133 1134 val, err := toProtobuf(object, v.fd, protoreflect.ValueOfList(v.list)) 1135 if err != nil { 1136 return nil, err 1137 } 1138 1139 for i := index; i < len; i++ { 1140 swap := v.list.Get(i) 1141 v.list.Set(i, val) 1142 val = swap 1143 } 1144 1145 v.list.Append(val) 1146 return starlark.None, nil 1147 } 1148 1149 func (v *List) pop(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1150 n := v.Len() 1151 i := n - 1 1152 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 0, &i); err != nil { 1153 return nil, err 1154 } 1155 if err := v.checkMutable("pop from"); err != nil { 1156 return nil, fmt.Errorf("%s: %w", fnname, err) 1157 } 1158 origI := i 1159 if i < 0 { 1160 i += n 1161 } 1162 if i < 0 || i >= n { 1163 return nil, fmt.Errorf("%s: %w", fnname, outOfRange(origI, n, v)) 1164 } 1165 return v.Pop(i) 1166 } 1167 1168 func (v *List) remove(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1169 var value starlark.Value 1170 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 1, &value); err != nil { 1171 return nil, err 1172 } 1173 if err := v.checkMutable("remove from"); err != nil { 1174 return nil, fmt.Errorf("%s: %w", v.Type(), err) 1175 } 1176 1177 // find 1178 for i := 0; i < v.Len(); i++ { 1179 if ok, err := starlark.Equal(v.Index(i), value); ok { 1180 // pop 1181 if _, err := v.Pop(i); err != nil { 1182 return nil, err 1183 } 1184 return starlark.None, nil 1185 1186 } else if err != nil { 1187 return nil, fmt.Errorf("%s: %w", fnname, err) 1188 } 1189 } 1190 return nil, fmt.Errorf("%s: element not found", fnname) 1191 } 1192 1193 // Enum is the type of a protobuf enum. 1194 type Enum struct { 1195 edesc protoreflect.EnumValueDescriptor 1196 } 1197 1198 func NewEnum(enum protoreflect.EnumValueDescriptors, arg starlark.Value) (Enum, error) { 1199 switch v := arg.(type) { 1200 case starlark.String: 1201 edesc := enum.ByName(protoreflect.Name(v)) 1202 if edesc == nil { 1203 return Enum{}, fmt.Errorf("proto: enum not found") 1204 } 1205 return Enum{edesc: edesc}, nil 1206 1207 case starlark.Int: 1208 n, _ := v.Int64() // TODO: checks? 1209 edesc := enum.ByNumber(protoreflect.EnumNumber(n)) 1210 return Enum{edesc: edesc}, nil 1211 1212 case Enum: 1213 return Enum{edesc: v.edesc}, nil 1214 1215 default: 1216 return Enum{}, fmt.Errorf("unsupported type %s", arg.Type()) 1217 } 1218 } 1219 1220 func (e Enum) String() string { return string(e.edesc.Name()) } 1221 func (e Enum) Type() string { return "proto.enum" } 1222 func (e Enum) Freeze() {} // immutable 1223 func (e Enum) Truth() starlark.Bool { return e.edesc.Number() > 0 } 1224 func (e Enum) Hash() (uint32, error) { return uint32(e.edesc.Number()), nil } 1225 func (x Enum) CompareSameType(op syntax.Token, y_ starlark.Value, depth int) (bool, error) { 1226 y := y_.(Enum) 1227 if err := equalFullName(x.edesc.Parent().FullName(), y.edesc.Parent().FullName()); err != nil { 1228 return false, err 1229 } 1230 i, j := x.edesc.Number(), y.edesc.Number() 1231 switch op { 1232 case syntax.EQL: 1233 return i == j, nil 1234 case syntax.NEQ: 1235 return i != j, nil 1236 case syntax.LE: 1237 return i <= j, nil 1238 case syntax.LT: 1239 return i < j, nil 1240 case syntax.GE: 1241 return i >= j, nil 1242 case syntax.GT: 1243 return i > j, nil 1244 default: 1245 panic(op) 1246 } 1247 } 1248 1249 type Map struct { 1250 m protoreflect.Map 1251 keyfd protoreflect.FieldDescriptor 1252 valfd protoreflect.FieldDescriptor 1253 1254 frozen *bool 1255 itercount uint32 1256 } 1257 1258 func (m *Map) Clear() error { 1259 m.m.Range(func(key protoreflect.MapKey, _ protoreflect.Value) bool { 1260 m.m.Clear(key) 1261 return true 1262 }) 1263 return nil 1264 } 1265 func (m *Map) parseKey(k starlark.Value) (protoreflect.MapKey, error) { 1266 keyval, err := toProtobuf(k, m.keyfd, protoreflect.ValueOfMap(m.m)) 1267 if err != nil { 1268 return protoreflect.MapKey{}, err 1269 } 1270 return keyval.MapKey(), nil 1271 } 1272 func (m *Map) toValue(key protoreflect.MapKey) (starlark.Value, bool) { 1273 val := m.m.Get(key) 1274 if !val.IsValid() { 1275 return starlark.None, false 1276 } 1277 return toStarlark(val, m.valfd, m.frozen), true 1278 } 1279 func (m *Map) Delete(k starlark.Value) (v starlark.Value, found bool, err error) { 1280 key, err := m.parseKey(k) 1281 if err != nil { 1282 return nil, false, err 1283 } 1284 1285 v, found = m.toValue(key) 1286 if found { 1287 m.m.Clear(key) 1288 } 1289 return v, found, nil 1290 } 1291 func (m *Map) Get(k starlark.Value) (v starlark.Value, found bool, err error) { 1292 key, err := m.parseKey(k) 1293 if err != nil { 1294 return nil, false, err 1295 } 1296 1297 v, found = m.toValue(key) 1298 return v, found, nil 1299 } 1300 1301 type byTuple []starlark.Tuple 1302 1303 func (a byTuple) Len() int { return len(a) } 1304 func (a byTuple) Swap(i, j int) { a[i], a[j] = a[j], a[i] } 1305 func (a byTuple) Less(i, j int) bool { 1306 c := a[i][0].(starlark.Comparable) 1307 ok, err := c.CompareSameType(syntax.LT, a[j][0], 1) 1308 if err != nil { 1309 panic(err) 1310 } 1311 return ok 1312 } 1313 1314 func (m *Map) Items() []starlark.Tuple { 1315 v := make([]starlark.Tuple, 0, m.Len()) 1316 m.m.Range(func(key protoreflect.MapKey, val protoreflect.Value) bool { 1317 v = append(v, starlark.Tuple{ 1318 toStarlark(key.Value(), m.keyfd, m.frozen), 1319 toStarlark(val, m.valfd, m.frozen), 1320 }) 1321 return true 1322 }) 1323 sort.Sort(byTuple(v)) 1324 return v 1325 } 1326 1327 type byValue []starlark.Value 1328 1329 func (a byValue) Len() int { return len(a) } 1330 func (a byValue) Swap(i, j int) { a[i], a[j] = a[j], a[i] } 1331 func (a byValue) Less(i, j int) bool { 1332 c := a[i].(starlark.Comparable) 1333 ok, err := c.CompareSameType(syntax.LT, a[j], 1) 1334 if err != nil { 1335 panic(err) 1336 } 1337 return ok 1338 } 1339 1340 func (m *Map) Keys() []starlark.Value { 1341 v := make([]starlark.Value, 0, m.Len()) 1342 m.m.Range(func(key protoreflect.MapKey, _ protoreflect.Value) bool { 1343 v = append(v, toStarlark(key.Value(), m.keyfd, m.frozen)) 1344 return true 1345 }) 1346 sort.Sort(byValue(v)) 1347 return v 1348 } 1349 func (m *Map) Len() int { 1350 return m.m.Len() 1351 } 1352 1353 type keyIterator struct { 1354 m *Map 1355 keys []starlark.Value // copy 1356 i int 1357 } 1358 1359 func (ki *keyIterator) Next(k *starlark.Value) bool { 1360 if ki.i < len(ki.keys) { 1361 *k = ki.keys[ki.i] 1362 ki.i++ 1363 return true 1364 } 1365 return false 1366 } 1367 1368 func (ki *keyIterator) Done() { 1369 if !*ki.m.frozen { 1370 ki.m.itercount-- 1371 } 1372 } 1373 1374 func (m *Map) Iterate() starlark.Iterator { 1375 if !*m.frozen { 1376 m.itercount-- 1377 } 1378 return &keyIterator{m: m, keys: m.Keys()} 1379 } 1380 func (m *Map) SetKey(k, v starlark.Value) error { 1381 if err := m.checkMutable("set"); err != nil { 1382 return err 1383 } 1384 1385 keyval, err := toProtobuf(k, m.keyfd, protoreflect.Value{}) 1386 if err != nil { 1387 return err 1388 } 1389 key := keyval.MapKey() 1390 1391 val, err := toProtobuf(k, m.valfd, protoreflect.ValueOfMap(m.m)) 1392 if err != nil { 1393 return err 1394 } 1395 m.m.Set(key, val) 1396 return nil 1397 } 1398 func (m *Map) String() string { 1399 buf := new(strings.Builder) 1400 buf.WriteByte('{') 1401 if m.m.IsValid() { 1402 for i, item := range m.Items() { 1403 if i > 0 { 1404 buf.WriteString(", ") 1405 } 1406 k, v := item[0], item[1] 1407 1408 buf.WriteString(k.String()) 1409 buf.WriteString(": ") 1410 buf.WriteString(v.String()) 1411 } 1412 } 1413 buf.WriteByte('}') 1414 return buf.String() 1415 } 1416 func (m *Map) Type() string { return "proto.map" } // TODO 1417 func (m *Map) Freeze() { *m.frozen = true } 1418 func (m *Map) Truth() starlark.Bool { return m.Len() > 0 } 1419 func (m *Map) Hash() (uint32, error) { return 0, fmt.Errorf("unhashable type: map") } 1420 func (m *Map) checkMutable(verb string) error { 1421 if *m.frozen { 1422 return fmt.Errorf("cannot %s frozen map", verb) 1423 } 1424 if m.itercount > 0 { 1425 return fmt.Errorf("cannot %s map during iteration", verb) 1426 } 1427 return nil 1428 } 1429 1430 type mapAttr func(m *Map) starlark.Value 1431 1432 // methods from starlark/library.go 1433 var mapAttrs = map[string]mapAttr{ 1434 "clear": func(m *Map) starlark.Value { return starext.MakeMethod(m, "clear", m.clear) }, 1435 "get": func(m *Map) starlark.Value { return starext.MakeMethod(m, "get", m.get) }, 1436 "items": func(m *Map) starlark.Value { return starext.MakeMethod(m, "items", m.items) }, 1437 "keys": func(m *Map) starlark.Value { return starext.MakeMethod(m, "keys", m.keys) }, 1438 "pop": func(m *Map) starlark.Value { return starext.MakeMethod(m, "pop", m.pop) }, 1439 //"popitem": starlark.NewBuiltin("popitem", dict_popitem), // TODO: list based? 1440 "setdefault": func(m *Map) starlark.Value { return starext.MakeMethod(m, "setdefault", m.setdefault) }, 1441 //"update": starlark.NewBuiltin("update", dict_update), // TODO: update list. 1442 "values": func(m *Map) starlark.Value { return starext.MakeMethod(m, "values", m.values) }, 1443 } 1444 1445 func (m *Map) Attr(name string) (starlark.Value, error) { 1446 if a := mapAttrs[name]; a != nil { 1447 return a(m), nil 1448 } 1449 return nil, nil 1450 } 1451 func (m *Map) AttrNames() []string { 1452 names := make([]string, 0, len(mapAttrs)) 1453 for name := range mapAttrs { 1454 names = append(names, name) 1455 } 1456 sort.Strings(names) 1457 return names 1458 } 1459 1460 func (m *Map) clear(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1461 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 0); err != nil { 1462 return nil, err 1463 } 1464 if err := m.Clear(); err != nil { 1465 return nil, err 1466 } 1467 return starlark.None, nil 1468 } 1469 1470 func (m *Map) get(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1471 var key, dflt starlark.Value 1472 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 1, &key, &dflt); err != nil { 1473 return nil, err 1474 } 1475 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 0); err != nil { 1476 return nil, err 1477 } 1478 if v, ok, err := m.Get(key); err != nil { 1479 return nil, err 1480 } else if ok { 1481 return v, nil 1482 } else if dflt != nil { 1483 return dflt, nil 1484 } 1485 return starlark.None, nil 1486 } 1487 1488 func (m *Map) items(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1489 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 0); err != nil { 1490 return nil, err 1491 } 1492 items := m.Items() 1493 res := make([]starlark.Value, len(items)) 1494 for i, item := range items { 1495 res[i] = item // convert [2]Value to Value 1496 } 1497 return starlark.NewList(res), nil 1498 } 1499 1500 func (m *Map) keys(_ *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1501 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 0); err != nil { 1502 return nil, err 1503 } 1504 return starlark.NewList(m.Keys()), nil 1505 } 1506 1507 func (m *Map) pop(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1508 var k, d starlark.Value 1509 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 1, &k, &d); err != nil { 1510 return nil, err 1511 } 1512 if v, found, err := m.Delete(k); err != nil { 1513 return nil, err 1514 } else if found { 1515 return v, nil 1516 } else if d != nil { 1517 return d, nil 1518 } 1519 return nil, fmt.Errorf("%s: missing key", fnname) 1520 } 1521 1522 func (m *Map) setdefault(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1523 var key, dflt starlark.Value = nil, starlark.None 1524 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 1, &key, &dflt); err != nil { 1525 return nil, err 1526 } 1527 if v, ok, err := m.Get(key); err != nil { 1528 return nil, err 1529 } else if ok { 1530 return v, nil 1531 } else if err := m.SetKey(key, dflt); err != nil { 1532 return nil, err 1533 } 1534 return dflt, nil 1535 } 1536 1537 /*func (m *Map) update(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1538 if len(args) > 1 { 1539 return nil, fmt.Errorf("update: got %d arguments, want at most 1", len(args)) 1540 } 1541 // TODO: update 1542 return starlark.None, nil 1543 }*/ 1544 1545 func (m *Map) values(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { 1546 if err := starlark.UnpackPositionalArgs(fnname, args, kwargs, 0); err != nil { 1547 return nil, err 1548 } 1549 items := m.Items() 1550 res := make([]starlark.Value, len(items)) 1551 for i, item := range items { 1552 res[i] = item[1] 1553 } 1554 return starlark.NewList(res), nil 1555 }