diff --git a/internal/codegen/codegen.go b/internal/codegen/codegen.go
index c4161fa..2c9aecc 100644
--- a/internal/codegen/codegen.go
+++ b/internal/codegen/codegen.go
@@ -11,14 +11,6 @@ import (
 	"strings"
 )
 
-const headerFmt = `// GENERATED BY gitlab.codemonkeysoftware.net/b/hatmill/internal/codegen
-// DO NOT EDIT!
-
-package %s
-
-import "gitlab.codemonkeysoftware.net/b/hatmill"
-`
-
 func identifier(s string) string {
 	words := strings.Split(s, "-")
 	for i, word := range words {
@@ -27,84 +19,67 @@ func identifier(s string) string {
 	return strings.Join(words, "")
 }
 
-func fileHeader(packageName string) string {
-	return fmt.Sprintf(headerFmt, packageName)
+var attribTypes = map[string]string {
+	"string":  `// %[1]s creates a "%[2]s" attribute
+	func %[1]s(value string) hatmill.Attrib {
+		return hatmill.Attrib{
+			Key: "%[2]s",
+			Value: value,
+		}
+	}
+	`,
+
+	"bool": `// %[1]s creates a "%[2]s" attribute
+	func %[1]s() hatmill.Attrib {
+		return hatmill.Attrib{
+			Key: "%[2]s",
+		}
+	}
+	`,
 }
 
-type AttribType int
-
-const (
-	String AttribType = iota
-	Bool
-)
-
-func (t *AttribType) UnmarshalJSON(data []byte) error {
-	if string(data) == "null" {
-		return nil
-	}
-
-	var typeName string
-	err := json.Unmarshal(data, &typeName)
-	if err != nil {
-		return fmt.Errorf("type property must be a string")
-	}
-	switch typeName {
-	case "bool":
-		*t = Bool
-	case "string":
-		*t = String
-	default:
-		return fmt.Errorf("unrecognized attribute type %s", typeName)
-	}
-	return nil
-}
-
-type AttribDef struct {
+type AttribSpec struct {
 	Name string     `json:"name"`
-	Type AttribType `json:"type"`
+	Type string `json:"type"`
 }
 
-func (def AttribDef) Generate() string {
-	const (
-		boolType   = "bool"
-		stringType = "string"
+type AttribSpecs []AttribSpec
 
-		stringTemplate = `// %[1]s creates a "%[2]s" attribute
-        func %[1]s(value string) hatmill.Attrib {
-            return hatmill.Attrib{
-                Key: "%[2]s",
-                Value: value,
-            }
-        }
-        `
-		boolTemplate = `// %[1]s creates a "%[2]s" attribute
-        func %[1]s() hatmill.Attrib {
-            return hatmill.Attrib{
-                Key: "%[2]s",
-            }
-        }
-        `
-	)
+func (a AttribSpecs) Code() Code {
+	var c Code
+	const hatmillImport = "gitlab.codemonkeysoftware.net/b/hatmill"
 
-	var template string
-	switch def.Type {
-	case Bool:
-		template = boolTemplate
-	case String:
-		template = stringTemplate
-	default:
-		panic(fmt.Errorf("unknown attribute type: %v", def.Type))
+	c.Imports = append(c.Imports, hatmillImport)
+	for _, spec := range a {
+		c.Defs = append(c.Defs, spec.Generate())
 	}
+	return c
+}
 
+type ElemSpecs []ElemSpec
+
+func (e ElemSpecs) Code() Code {
+	var c Code
+	const hatmillImport = "gitlab.codemonkeysoftware.net/b/hatmill"
+
+	c.Imports = append(c.Imports, hatmillImport)
+	for _, spec := range e {
+		c.Defs = append(c.Defs, spec.Generate())
+	}
+	return c
+}
+
+func (def AttribSpec) Generate() string {
+	template := attribTypes[def.Type]
 	return fmt.Sprintf(template, identifier(def.Name), def.Name)
 }
 
-type ElemDef struct {
+type ElemSpec struct {
 	Name string `json:"name"`
 	Void bool   `json:"void"`
 }
 
-func (def ElemDef) Generate() string {
+func (def ElemSpec) Generate() string {
 	const (
 		parentTemplate = `// %[1]s creates a <%[2]s> element.
         func %[1]s(attribs ...hatmill.Attrib) func(children ...hatmill.Term) hatmill.ParentElement {
@@ -136,20 +111,46 @@ func (def ElemDef) Generate() string {
 	return fmt.Sprintf(template, identifier(def.Name), def.Name)
 }
 
-type Defs struct {
-	Attributes []AttribDef `json:"attributes"`
-	Elements   []ElemDef   `json:"elements"`
+type Specs struct {
+	Attributes AttribSpecs `json:"attributes"`
+	Elements   ElemSpecs   `json:"elements"`
 }
 
-func writeFormatted(fileName, packageName string, writeDefs func(*bytes.Buffer)) error {
-	src := new(bytes.Buffer)
-	src.WriteString(fileHeader(packageName))
-	writeDefs(src)
-	formatted, err := format.Source(src.Bytes())
+// Code is a very slight abstraction of a source code file.
+type Code struct {
+	Package string
+	Imports []string
+	Defs []string
+}
+
+func (c Code) Bytes() ([]byte, error) {
+	buf := new(bytes.Buffer)
+	buf.WriteString(`// GENERATED BY gitlab.codemonkeysoftware.net/b/hatmill/internal/codegen
+	// DO NOT EDIT!
+
+	`)
+	fmt.Fprintln(buf, "package ", c.Package)
+	for _, imp := range c.Imports {
+		fmt.Fprintf(buf, "import \"%s\"\n", imp)
+	}
+	for _, def := range c.Defs {
+		buf.WriteString(def)
+	}
+	return format.Source(buf.Bytes())
+}
+
+type Coder interface {
+	Code() Code
+}
+
+func WriteCodeFile(c Coder, path, pkg string) error {
+	code := c.Code()
+	code.Package = pkg
+	b, err := code.Bytes()
 	if err != nil {
 		return err
 	}
-	return ioutil.WriteFile(fileName, formatted, 0644)
+	return ioutil.WriteFile(path, b, 0644)
 }
 
 func main() {
@@ -165,26 +166,18 @@ func main() {
 		log.Fatal(err)
 	}
 
-	var defs Defs
-	err = json.Unmarshal(input, &defs)
+	var specs Specs
+	err = json.Unmarshal(input, &specs)
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	err = writeFormatted(*attribPath, *attribPkg, func(buf *bytes.Buffer) {
-		for _, attribDef := range defs.Attributes {
-			buf.WriteString(attribDef.Generate())
-		}
-	})
+	err = WriteCodeFile(specs.Attributes, *attribPath, *attribPkg)
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	err = writeFormatted(*elemPath, *elemPkg, func(buf *bytes.Buffer) {
-		for _, elemDef := range defs.Elements {
-			buf.WriteString(elemDef.Generate())
-		}
-	})
+	err = WriteCodeFile(specs.Elements, *elemPath, *elemPkg)
 	if err != nil {
 		log.Fatal(err)
 	}