Merge pull request #14 from dave/space-type

Space type

README.md

@@ -13,51 +13,63 @@ comments don't remain attached to the correct nodes:
 ```go
 code := `package a
 
-var a int    // foo
-var b string // bar
-`
+func main() { 
+	var a int    // foo
+	var b string // bar
+}`
 fset := token.NewFileSet()
 f, err := parser.ParseFile(fset, "a.go", code, parser.ParseComments)
 if err != nil {
 	panic(err)
 }
-f.Decls = []ast.Decl{f.Decls[1], f.Decls[0]}
+
+body := f.Decls[0].(*ast.FuncDecl).Body
+body.List = []ast.Stmt{body.List[1], body.List[0]}
+
 if err := format.Node(os.Stdout, fset, f); err != nil {
 	panic(err)
 }
 
 //Output:
 //package a
 //
-//// foo
-//var b string
-//var a int
-//
-//// bar
+//func main() {
+//	// foo
+//	var b string
+//	var a int
+//	// bar
+//}
 ```
 
 Here's the same example using `dst`:
 
 ```go
 code := `package a
 
-var a int    // foo
-var b string // bar
+func main() { 
+	var a int    // foo
+	var b string // bar
+}
 `
 f, err := decorator.Parse(code)
 if err != nil {
 	panic(err)
 }
-f.Decls = []dst.Decl{f.Decls[1], f.Decls[0]}
+
+body := f.Decls[0].(*dst.FuncDecl).Body
+body.List = []dst.Stmt{body.List[1], body.List[0]}
+
 if err := decorator.Print(f); err != nil {
 	panic(err)
 }
 
 //Output:
 //package a
 //
-//var b string // bar
-//var a int    // foo
+//func main() {
+//	var b string // bar
+//	var a int    // foo
+//}
 ```
 
 ### Example

decorations.go

@@ -21,3 +21,11 @@ type DecorationStmtDecorations struct {
 type DecorationDeclDecorations struct {
 	Start Decorations
 }
+
+type SpaceType int
+
+const (
+	None      SpaceType = 0
+	NewLine   SpaceType = 1
+	EmptyLine SpaceType = 2
+)

decorations_test.go

@@ -57,47 +57,59 @@ func ExampleDecorations() {
 func ExampleAstBroken() {
 	code := `package a
 
-	var a int    // foo
-	var b string // bar
-	`
+	func main() { 
+		var a int    // foo
+		var b string // bar
+	}`
 	fset := token.NewFileSet()
 	f, err := parser.ParseFile(fset, "a.go", code, parser.ParseComments)
 	if err != nil {
 		panic(err)
 	}
-	f.Decls = []ast.Decl{f.Decls[1], f.Decls[0]}
+
+	body := f.Decls[0].(*ast.FuncDecl).Body
+	body.List = []ast.Stmt{body.List[1], body.List[0]}
+
 	if err := format.Node(os.Stdout, fset, f); err != nil {
 		panic(err)
 	}
 
 	//Output:
 	//package a
 	//
-	//// foo
-	//var b string
-	//var a int
-	//
-	//// bar
+	//func main() {
+	//	// foo
+	//	var b string
+	//	var a int
+	//	// bar
+	//}
 }
 
 func ExampleDstFixed() {
 	code := `package a
 
-	var a int    // foo
-	var b string // bar
+	func main() { 
+		var a int    // foo
+		var b string // bar
+	}
 	`
 	f, err := decorator.Parse(code)
 	if err != nil {
 		panic(err)
 	}
-	f.Decls = []dst.Decl{f.Decls[1], f.Decls[0]}
+
+	body := f.Decls[0].(*dst.FuncDecl).Body
+	body.List = []dst.Stmt{body.List[1], body.List[0]}
+
 	if err := decorator.Print(f); err != nil {
 		panic(err)
 	}
 
 	//Output:
 	//package a
 	//
-	//var b string // bar
-	//var a int    // foo
+	//func main() {
+	//	var b string // bar
+	//	var a int    // foo
+	//}
 }