ofxgo/generate_constants.py

#!/usr/bin/env python

enums = {
    # Bank/general
    "AcctType": ["Checking", "Savings", "MoneyMrkt", "CreditLine", "CD"],
    "TrnType": ["Credit", "Debit", "Int", "Div", "Fee", "SrvChg", "Dep", "ATM", "POS", "Xfer", "Check", "Payment", "Cash", "DirectDep", "DirectDebit", "RepeatPmt", "Hold", "Other"],
    "ImageType": ["Statement", "Transaction", "Tax"],
    "ImageRefType": ["Opaque", "URL", "FormURL"],
    "CheckSup": ["FrontOnly", "BackOnly", "FrontAndBack"],
    "CorrectAction": ["Delete", "Replace"],
    "BalType": ["Dollar", "Percent", "Number"],

    # InvStmt
    "Inv401kSource": ["PreTax", "AfterTax", "Match", "ProfitSharing", "Rollover", "OtherVest", "OtherNonVest"],
    "SubAcctType": ["Cash", "Margin", "Short", "Other"], # used in fields named: SubAcctSec, HeldInAcct, SubAcctFund
    "BuyType": ["Buy", "BuyToCover"],
    "OptAction": ["Exercise", "Assign", "Expire"],
    "TferAction": ["In", "Out"],
    "PosType": ["Long", "Short"],
    "Secured": ["Naked", "Covered"],
    "Duration": ["Day", "GoodTilCancel", "Immediate"],
    "Restriction": ["AllOrNone", "MinUnits", "None"],
    "UnitType": ["Shares", "Currency"],
    "OptBuyType": ["BuyToOpen", "BuyToClose"],
    "SellType": ["Sell", "SellShort"],
    "LoanPmtFreq": ["Weekly", "Biweekly", "TwiceMonthly", "Monthly", "FourWeeks", "BiMonthly", "Quarterly", "Semiannually", "Annually", "Other"],
    "IncomeType": ["CGLong", "CGShort", "Div", "Interest", "Misc"],
    "SellReason": ["Call", "Sell", "Maturity"],
    "OptSellType": ["SellToClose", "SellToOpen"],
    "RelType": ["Spread", "Straddle", "None", "Other"],

    # Prof
    "CharType": ["AlphaOnly", "NumericOnly", "AlphaOrNumeric", "AlphaAndNumeric"],
    "SyncMode": ["Full", "Lite"],

    # SecList
    "DebtType": ["Coupon", "Zero"],
    "DebtClass": ["Treasury", "Municipal", "Corporate", "Other"],
    "CouponFreq": ["Monthly", "Quarterly", "Semiannual", "Annual", "Other"],
    "CallType": ["Call", "Put", "Prefund", "Maturity"],
    "AssetClass": ["DomesticBond", "IntlBond", "LargeStock", "SmallStock", "IntlStock", "MoneyMrkt", "Other"],
    "MfType": ["OpenEnd", "CloseEnd", "Other"],
    "OptType": ["Put", "Call"],
    "StockType": ["Common", "Preferred", "Convertible", "Other"],
    "OfxSec": ["None", "Type 1"],

    # Signup
    "HolderType": ["Individual", "Joint", "Custodial", "Trust", "Other"],
    "AcctClassification": ["Personal", "Business", "Corporate", "Other"],
    "SvcStatus": ["Avail", "Pend", "Active"],
    "UsProductType": ["401K", "403B", "IRA", "KEOGH", "Other", "SARSEP", "Simple", "Normal", "TDA", "Trust", "UGMA"],
}

header = """package ofxgo

/*
 * Do not edit this file by hand. It is auto-generated by calling `go generate`.
 * To make changes, edit generate_constants.py, re-run `go generate`, and check
 * in the result.
 */

import (
	"errors"
	"fmt"
	"github.com/aclindsa/go/src/encoding/xml"
	"strings"
)
"""

template = """
type {enumLower} uint

const (
{constNames})

var {enumLower}s = [...]string{{"{upperValueString}"}}

func (e {enumLower}) Valid() bool {{
	return e >= {firstValue} && e <= {lastValue}
}}

func (e {enumLower}) String() string {{
	if e.Valid() {{
		return {enumLower}s[e-1]
	}}
	return fmt.Sprintf("invalid {enumLower} (%d)", e)
}}

func (e *{enumLower}) FromString(in string) error {{
	value := strings.TrimSpace(in)

	for i, s := range {enumLower}s {{
		if s == value {{
			*e = {enumLower}(i + 1)
			return nil
		}}
	}}
	*e = 0
	return errors.New("Invalid {enum}: \\\"" + in + "\\\"")
}}

func (e *{enumLower}) UnmarshalXML(d *xml.Decoder, start xml.StartElement) error {{
	var value string
	err := d.DecodeElement(&value, &start)
	if err != nil {{
		return err
	}}

	return e.FromString(value)
}}

func (e *{enumLower}) MarshalXML(enc *xml.Encoder, start xml.StartElement) error {{
	if *e == 0 {{
		return nil
	}} else if !e.Valid() {{
		return errors.New("Invalid {enum}")
	}}
	enc.EncodeElement({enumLower}s[*e-1], start)
	return nil
}}

func New{enum}(s string) ({enumLower}, error) {{
	var e {enumLower}
	err := e.FromString(s)
	if err != nil {{
		return 0, err
	}}
	return e, nil
}}
"""

with open("constants.go", 'w') as f:
    f.write(header)

    for enum in enums:
        enumLower = enum[:1].lower() + enum[1:].replace(" ", "")
        firstValue = enum+enums[enum][0].replace(" ", "")
        lastValue = enum+enums[enum][-1].replace(" ", "")

        constNames = "\t{firstValue} {enumLower} = 1 + iota\n".format(
            enum=enum,
            firstValue=firstValue,
            enumLower=enumLower)
        for value in enums[enum][1:]:
            constNames += "\t{enum}{value}\n".format(
                enum=enum,
                value=value.replace(" ", ""))

        upperValueString = "\", \"".join([s.upper() for s in enums[enum]])

        f.write(template.format(enum=enum,
                              enumLower=enumLower,
                              firstValue=firstValue,
                              lastValue=lastValue,
                              constNames=constNames,
                              upperValueString=upperValueString))

test_header = """package ofxgo_test

/*
 * Do not edit this file by hand. It is auto-generated by calling `go generate`.
 * To make changes, edit generate_constants.py, re-run `go generate`, and check
 * in the result.
 */

import (
	"github.com/aclindsa/go/src/encoding/xml"
	"github.com/aclindsa/ofxgo"
	"testing"
)
"""

test_template = """
func Test{enum}(t *testing.T) {{
	e, err := ofxgo.New{enum}("{firstValueUpper}")
	if err != nil {{
		t.Fatalf("Unexpected error creating new {enum} from string \\\"{firstValueUpper}\\\"\\n")
	}}
	if !e.Valid() {{
		t.Fatalf("{enum} unexpectedly invalid\\n")
	}}
	err = e.FromString("{lastValueUpper}")
	if err != nil {{
		t.Fatalf("Unexpected error on {enum}.FromString(\\\"{lastValueUpper}\\\")\\n")
	}}
	if e.String() != "{lastValueUpper}" {{
		t.Fatalf("{enum}.String() expected to be \\\"{lastValueUpper}\\\"\\n")
	}}

	marshalHelper(t, "{lastValueUpper}", &e)

	overwritten, err := ofxgo.New{enum}("THISWILLNEVERBEAVALIDENUMSTRING")
	if err == nil {{
		t.Fatalf("Expected error creating new {enum} from string \\\"THISWILLNEVERBEAVALIDENUMSTRING\\\"\\n")
	}}
	if overwritten.Valid() {{
		t.Fatalf("{enum} created with string \\\"THISWILLNEVERBEAVALIDENUMSTRING\\\" should not be valid\\n")
	}}

	b, err := xml.Marshal(&overwritten)
	if err != nil {{
		t.Fatalf("Unexpected error on xml.Marshal({enum}): %s\\n", err)
	}}
	if string(b) != "" {{
		t.Fatalf("Expected empty string, got '%s'\\n", string(b))
	}}

	unmarshalHelper(t, "{lastValueUpper}", &e, &overwritten)
}}
"""

with open("constants_test.go", 'w') as f:
    f.write(test_header)

    for enum in enums:
        firstValueUpper = enums[enum][0].upper()
        lastValueUpper = enums[enum][-1].upper()
        f.write(test_template.format(enum=enum,
                              firstValueUpper=firstValueUpper,
                              lastValueUpper=lastValueUpper))