add support for Enum types in Python, based on const values of the same type in Go. Also generates the raw const values at global package scope, as accessed in Go.

Author: Randall C. O'Reilly

bind/gen.go

@@ -250,6 +250,7 @@ os.chdir(cwd)
 # %[2]s
 
 import _%[1]s, collections
+from enum import Enum
 
 # to use this code in your end-user python file, import it as follows:
 # from %[1]s import %[3]s
@@ -757,6 +758,11 @@ func (g *pyGen) genAll() {
 		g.genType(sym, false, false) // not exttypes
 	}
 
+	g.pywrap.Printf("\n\n#---- Enums from Go (collections of consts with same type) ---\n")
+	for _, e := range g.pkg.enums {
+		g.genEnum(e)
+	}
+
 	g.pywrap.Printf("\n\n#---- Constants from Go: Python can only ask that you please don't change these! ---\n")
 	for _, c := range g.pkg.consts {
 		g.genConst(c)

bind/gen_varconst.go

@@ -6,6 +6,7 @@ package bind
 
 import (
 	"fmt"
+	"strings"
 )
 
 func (g *pyGen) genConst(c *Const) {
@@ -110,7 +111,7 @@ func (g *pyGen) genVarSetter(v *Var) {
 
 func (g *pyGen) genConstValue(c *Const) {
 	// constants go directly into wrapper as-is
-	val := c.obj.Val().ExactString()
+	val := c.val
 	switch val {
 	case "true":
 		val = "True"
@@ -119,3 +120,34 @@ func (g *pyGen) genConstValue(c *Const) {
 	}
 	g.pywrap.Printf("%s = %s\n", c.GoName(), val)
 }
+
+func (g *pyGen) genEnum(e *Enum) {
+	g.pywrap.Printf("class %s(Enum):\n", e.typ.Obj().Name())
+	g.pywrap.Indent()
+	doc := e.Doc()
+	if doc != "" {
+		lns := strings.Split(doc, "\n")
+		g.pywrap.Printf(`"""`)
+		g.pywrap.Printf("\n")
+		for _, l := range lns {
+			g.pywrap.Printf("%s\n", l)
+		}
+		g.pywrap.Printf(`"""`)
+		g.pywrap.Printf("\n")
+	}
+	e.SortConsts()
+	for _, c := range e.items {
+		g.genConstValue(c)
+	}
+	g.pywrap.Outdent()
+
+	// Go has each const value globally available within a given package
+	// so to keep the code consistent, we redundantly generate the consts
+	// again here.  The Enum organization however is critical for organizing
+	// the values under the type (making them accessible programmatically)
+	g.pywrap.Printf("\n")
+	for _, c := range e.items {
+		g.genConstValue(c)
+	}
+	g.pywrap.Printf("\n")
+}

bind/package.go

@@ -10,6 +10,7 @@ import (
 	"go/types"
 	"path/filepath"
 	"reflect"
+	"strconv"
 	"strings"
 )
 
@@ -24,6 +25,7 @@ type Package struct {
 	syms      *symtab // note: this is now *always* = symbols.current
 	objs      map[string]Object
 	consts    []*Const
+	enums     []*Enum
 	vars      []*Var
 	structs   []*Struct
 	ifaces    []*Interface
@@ -484,7 +486,34 @@ func (p *Package) process() error {
 	return err
 }
 
+func (p *Package) findEnum(ntyp *types.Named) *Enum {
+	for _, enm := range p.enums {
+		if enm.typ == ntyp {
+			return enm
+		}
+	}
+	return nil
+}
+
 func (p *Package) addConst(obj *types.Const) {
+	if ntyp, ok := obj.Type().(*types.Named); ok {
+		enm := p.findEnum(ntyp)
+		if enm != nil {
+			enm.AddConst(p, obj)
+			return
+		} else {
+			val := obj.Val().String()
+			_, err := strconv.Atoi(val)
+			if err == nil {
+				enm, err := newEnum(p, obj)
+				if err == nil {
+					p.enums = append(p.enums, enm)
+					return
+				}
+			}
+		}
+	}
+
 	nc, err := newConst(p, obj)
 	if err == nil {
 		p.consts = append(p.consts, nc)

bind/types.go

@@ -7,6 +7,8 @@ package bind
 import (
 	"fmt"
 	"go/types"
+	"sort"
+	"strconv"
 )
 
 type Object interface {
@@ -461,20 +463,23 @@ type Const struct {
 	obj *types.Const
 	id  string
 	doc string
+	val string
 }
 
 func newConst(p *Package, o *types.Const) (*Const, error) {
 	pkg := o.Pkg()
 	sym := p.syms.symtype(o.Type())
 	id := pkg.Name() + "_" + o.Name()
 	doc := p.getDoc("", o)
+	val := o.Val().String()
 
 	return &Const{
 		pkg: p,
 		sym: sym,
 		obj: o,
 		id:  id,
 		doc: doc,
+		val: val,
 	}, nil
 }
 
@@ -483,6 +488,59 @@ func (c *Const) Doc() string        { return c.doc }
 func (c *Const) GoName() string     { return c.obj.Name() }
 func (c *Const) GoType() types.Type { return c.obj.Type() }
 
+///////////////////////////////////////////////////////////////////////////////////
+//  Enum
+
+type Enum struct {
+	pkg   *Package
+	sym   *symbol
+	obj   *types.Const // first one -- random..
+	typ   *types.Named
+	id    string
+	doc   string
+	items []*Const
+}
+
+func newEnum(p *Package, o *types.Const) (*Enum, error) {
+	pkg := o.Pkg()
+	sym := p.syms.symtype(o.Type())
+	id := pkg.Name() + "_" + o.Name()
+	typ := o.Type().(*types.Named)
+	doc := p.getDoc("", typ.Obj())
+
+	e := &Enum{
+		pkg: p,
+		sym: sym,
+		obj: o,
+		typ: typ,
+		id:  id,
+		doc: doc,
+	}
+	e.AddConst(p, o)
+	return e, nil
+}
+
+func (e *Enum) ID() string         { return e.id }
+func (e *Enum) Doc() string        { return e.doc }
+func (e *Enum) GoName() string     { return e.obj.Name() }
+func (e *Enum) GoType() types.Type { return e.obj.Type() }
+
+func (e *Enum) AddConst(p *Package, o *types.Const) (*Const, error) {
+	c, err := newConst(p, o)
+	if err == nil {
+		e.items = append(e.items, c)
+	}
+	return c, err
+}
+
+func (e *Enum) SortConsts() {
+	sort.Slice(e.items, func(i, j int) bool {
+		iv, _ := strconv.Atoi(e.items[i].val)
+		jv, _ := strconv.Atoi(e.items[j].val)
+		return iv < jv
+	})
+}
+
 ///////////////////////////////////////////////////////////////////////////////////
 //  Var