Generate a Go Enumeration

generate_enum

#!/usr/bin/env ruby

require 'fileutils'
require 'optparse'

def snake_case(str)
  return str.downcase if str =~ /\A[A-Z]+\z/
  str.gsub(/([A-Z]+)([A-Z][a-z])/, '\1_\2')
     .gsub(/([a-z])([A-Z])/, '\1_\2')
     .downcase
end

def make_identifier(str)
  str.split(/[_-](?!\d)/).collect(&:capitalize).join
end

def generate_enumeration(file_name, package_name, states, iota_offset)
  File.open(file_name, 'w') do |f|
    f << file_header(package_name)
    f << generate_states(package_name, states, iota_offset)
    f << generate_valid_strings(states)
    f << stringer
    f << generate_parser(package_name, states)
    f << generate_marshal(package_name, states)
    f << unmarshal_text
  end
end

def file_header(package_name)
  <<-ENDFUNC
// Code generated by generate_enum; DO NOT EDIT.

package #{package_name}

import (
    "fmt"
    "encoding/json"
    "reflect"
)

  ENDFUNC
end

def generate_states(package_name, states, iota_offset)
  first = make_identifier(states[0])
  result = <<-FIRST
// Enum is an enumerated type
type Enum uint8

const (
  // #{first} is an enumeration for #{package_name}.Enum
	#{first} Enum = iota + #{iota_offset}
  FIRST
  states[1..-1].reduce(result) do |memo, state|
    identifier = make_identifier(state)
    memo << <<-NEXT
	// #{identifier} is an enumeration for #{package_name}.Enum
	#{identifier}
    NEXT
  end << <<-FINAL
)

FINAL
end

def generate_valid_strings(states)
  result = <<-FIRST
// Set of strings that are valid inputs for ParseEnum
var ValidStrings = []string{
  FIRST
  states.reduce(result) do |result, state|
    result << <<-NEXT
    #{make_identifier(state)}.String(),
    NEXT
  end << <<-FINAL
	}

FINAL
end

def stringer
  <<-ENDFUNC
// String makes Enum satisfy the Stringer interface
func (i Enum) String() string {
	tmp, err := i.MarshalText()
	if err == nil {
		return string(tmp)
	}
	return ""
}

  ENDFUNC
end

def generate_parser(package_name, states)
  result = <<-FIRST
// ParseEnum attempts to convert a string into a Enum
func ParseEnum(name string) (Enum, error) {
	switch name {
  FIRST
  states.reduce(result) do |result, state|
    result << <<-NEXT
	case "#{state.downcase}":
		return #{make_identifier(state)}, nil
    NEXT
  end << <<-FINAL
	}
  var zero Enum
  return zero, fmt.Errorf("%s is not a valid #{package_name}.Enum", name)
}

FINAL
end

def generate_marshal(package_name, states)
  result = <<-FIRST
// MarshalText implements the text marshaller method
func (i Enum) MarshalText() ([]byte, error) {
	switch i {
  FIRST
  states.reduce(result) do |result, state|
    result << <<-NEXT
	case #{make_identifier(state)}:
		return []byte("#{state.downcase}"), nil
    NEXT
  end << <<-FINAL
	}
  return nil, fmt.Errorf("%d is not a valid #{package_name}.Enum", i)
}

FINAL
end

def unmarshal_text
  <<-ENDFUNC
// UnmarshalText implements the text unmarshaller method
func (i *Enum) UnmarshalText(text []byte) error {
	name := string(text)
	tmp, err := ParseEnum(name)
	if err != nil {
		return &json.UnmarshalTypeError{
			Value: name,
			Type:  reflect.TypeOf(*i),
		}
	}
	*i = tmp
	return nil
}

  ENDFUNC
end

# Check if the zero value is to be used in this enumeration
iota = 1
OptionParser.new do |opts|
  opts.on('-z', 'Make first enumeration the default') do |default|
    iota = default ? 0 : 1
  end
end.parse!

# First argument is the package name
package_name = ARGV.shift
if package_name.to_s.empty? || ARGV.empty?
  STDERR.puts('Need a package name and at least one enumeration')
  exit 1
end
enum_dir = File.join('enums', package_name)
file_name = File.join(enum_dir, "#{package_name}.go")
FileUtils.mkdir_p(enum_dir)

# Generate the enumeration from the rest of the arguments
generate_enumeration(file_name, package_name, ARGV, iota)
exec('gofmt', '-w', file_name)