Rearranged generated code

internal/codegen/codegen.go

@@ -133,10 +133,23 @@ type Defs struct {
     Elements []ElemDef `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())
+	if err != nil {
+		return err
+	}
+	return ioutil.WriteFile(fileName, formatted, 0644)
+}
+
 func main() {
     inputPath := flag.String("input", "", "JSON input file")
-    outputPath := flag.String("output", "", ".go output file")
-    packageName := flag.String("package", "", "output package name")
+	elemPath := flag.String("elemfile", "", "generated element .go file")
+	elemPkg := flag.String("elempkg", "", "generated element package name")
+	attribPath := flag.String("attribfile", "", "generated attribute .go file")
+	attribPkg := flag.String("attribpkg", "", "generated attribute package name")
     flag.Parse()
 
     input, err := ioutil.ReadFile(*inputPath)
@@ -150,22 +163,21 @@ func main() {
         log.Fatal(err)
     }
 
-    output := new(bytes.Buffer)
-    output.WriteString(fileHeader(*packageName))
-    for _, elemDef := range defs.Elements {
-        output.WriteString(elemDef.Generate())
-    }
-    for _, attribDef := range defs.Attributes {
-        output.WriteString(attribDef.Generate())
-    }
+	err = writeFormatted(*attribPath, *attribPkg, func(buf *bytes.Buffer) {
+		for _, attribDef := range defs.Attributes {
+	        buf.WriteString(attribDef.Generate())
+	    }
+	})
+	if err != nil {
+		log.Fatal(err)
+	}
 
-    formatted, err := format.Source(output.Bytes())
-    if err != nil {
-        log.Fatal(err)
-    }
-
-    err = ioutil.WriteFile(*outputPath, formatted, 0644)
-    if err != nil {
-        log.Fatal(err)
-    }
+	err = writeFormatted(*elemPath, *elemPkg, func(buf *bytes.Buffer) {
+		for _, elemDef := range defs.Elements {
+	        buf.WriteString(elemDef.Generate())
+	    }
+	})
+	if err != nil {
+		log.Fatal(err)
+	}
 }