123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719 |
- package build
- import (
- "bytes"
- "fmt"
- "strings"
- )
- const nestedIndentation = 2
- const listIndentation = 4
- func Format(f *File) []byte {
- pr := &printer{}
- pr.file(f)
- return pr.Bytes()
- }
- func FormatString(x Expr) string {
- pr := &printer{}
- switch x := x.(type) {
- case *File:
- pr.file(x)
- default:
- pr.expr(x, precLow)
- }
- return pr.String()
- }
- type printer struct {
- bytes.Buffer
- comment []Comment
- margin int
- depth int
- }
- func (p *printer) printf(format string, args ...interface{}) {
- fmt.Fprintf(p, format, args...)
- }
- func (p *printer) indent() int {
- b := p.Bytes()
- n := 0
- for n < len(b) && b[len(b)-1-n] != '\n' {
- n++
- }
- return n
- }
- func (p *printer) newline() {
- if len(p.comment) > 0 {
- p.printf(" ")
- for i, com := range p.comment {
- if i > 0 {
- p.trim()
- p.printf("\n%*s", p.margin, "")
- }
- p.printf("%s", strings.TrimSpace(com.Token))
- }
- p.comment = p.comment[:0]
- }
- p.trim()
- p.printf("\n%*s", p.margin, "")
- }
- func (p *printer) breakline() {
- if p.depth == 0 {
-
- p.printf(" \\\n%*s", p.margin, "")
- return
- }
-
- p.newline()
- }
- func (p *printer) trim() {
-
- b := p.Bytes()
- n := len(b)
- for n > 0 && b[n-1] == ' ' {
- n--
- }
- p.Truncate(n)
- }
- func (p *printer) file(f *File) {
- for _, com := range f.Before {
- p.printf("%s", strings.TrimSpace(com.Token))
- p.newline()
- }
- p.statements(f.Stmt)
- for _, com := range f.After {
- p.printf("%s", strings.TrimSpace(com.Token))
- p.newline()
- }
-
- p.trim()
- }
- func (p *printer) statements(stmts []Expr) {
- for i, stmt := range stmts {
- switch stmt := stmt.(type) {
- case *CommentBlock:
-
- case *PythonBlock:
- for _, com := range stmt.Before {
- p.printf("%s", strings.TrimSpace(com.Token))
- p.newline()
- }
- p.printf("%s", stmt.Token)
- p.newline()
- default:
- p.expr(stmt, precLow)
-
-
- if !isCodeBlock(stmt) {
- p.newline()
- }
- }
- for _, com := range stmt.Comment().After {
- p.printf("%s", strings.TrimSpace(com.Token))
- p.newline()
- }
- if i+1 < len(stmts) && !compactStmt(stmt, stmts[i+1], p.margin == 0) {
- p.newline()
- }
- }
- }
- func compactStmt(s1, s2 Expr, isTopLevel bool) bool {
- if len(s2.Comment().Before) > 0 {
- return false
- }
- if isTopLevel {
- return isCall(s1, "load") && isCall(s2, "load")
- } else {
- return !(isCodeBlock(s1) || isCodeBlock(s2))
- }
- }
- func isCall(x Expr, name string) bool {
- c, ok := x.(*CallExpr)
- if !ok {
- return false
- }
- nam, ok := c.X.(*LiteralExpr)
- if !ok {
- return false
- }
- return nam.Token == name
- }
- func isCodeBlock(x Expr) bool {
- switch x.(type) {
- case *FuncDef:
- return true
- case *ForLoop:
- return true
- case *IfElse:
- return true
- default:
- return false
- }
- }
- const (
- precLow = iota
- precAssign
- precComma
- precColon
- precIn
- precOr
- precAnd
- precCmp
- precAdd
- precMultiply
- precSuffix
- precUnary
- precConcat
- )
- var opPrec = map[string]int{
- "=": precAssign,
- "+=": precAssign,
- "-=": precAssign,
- "*=": precAssign,
- "/=": precAssign,
- "//=": precAssign,
- "%=": precAssign,
- "or": precOr,
- "and": precAnd,
- "<": precCmp,
- ">": precCmp,
- "==": precCmp,
- "!=": precCmp,
- "<=": precCmp,
- ">=": precCmp,
- "+": precAdd,
- "-": precAdd,
- "*": precMultiply,
- "/": precMultiply,
- "//": precMultiply,
- "%": precMultiply,
- }
- func (p *printer) expr(v Expr, outerPrec int) {
-
-
-
-
-
-
-
-
-
-
-
- if before := v.Comment().Before; len(before) > 0 {
-
-
- p.trim()
- if p.indent() > 0 {
-
- p.printf("\n")
- }
-
- p.printf("%*s", p.margin, "")
- for _, com := range before {
- p.printf("%s", strings.TrimSpace(com.Token))
- p.newline()
- }
- }
-
-
-
-
-
-
-
- parenthesized := false
- addParen := func(prec int) {
- if prec < outerPrec {
- p.printf("(")
- p.depth++
- parenthesized = true
- }
- }
- switch v := v.(type) {
- default:
- panic(fmt.Errorf("printer: unexpected type %T", v))
- case *LiteralExpr:
- p.printf("%s", v.Token)
- case *StringExpr:
-
-
-
-
- if strings.HasPrefix(v.Token, `"`) {
- s, triple, err := unquote(v.Token)
- if s == v.Value && triple == v.TripleQuote && err == nil {
- p.printf("%s", v.Token)
- break
- }
- }
- p.printf("%s", quote(v.Value, v.TripleQuote))
- case *DotExpr:
- addParen(precSuffix)
- p.expr(v.X, precSuffix)
- p.printf(".%s", v.Name)
- case *IndexExpr:
- addParen(precSuffix)
- p.expr(v.X, precSuffix)
- p.printf("[")
- p.expr(v.Y, precLow)
- p.printf("]")
- case *KeyValueExpr:
- p.expr(v.Key, precLow)
- p.printf(": ")
- p.expr(v.Value, precLow)
- case *SliceExpr:
- addParen(precSuffix)
- p.expr(v.X, precSuffix)
- p.printf("[")
- if v.From != nil {
- p.expr(v.From, precLow)
- }
- p.printf(":")
- if v.To != nil {
- p.expr(v.To, precLow)
- }
- if v.SecondColon.Byte != 0 {
- p.printf(":")
- if v.Step != nil {
- p.expr(v.Step, precLow)
- }
- }
- p.printf("]")
- case *UnaryExpr:
- addParen(precUnary)
- if v.Op == "not" {
- p.printf("not ")
- } else {
- p.printf("%s", v.Op)
- }
- p.expr(v.X, precUnary)
- case *LambdaExpr:
- addParen(precColon)
- p.printf("lambda ")
- for i, name := range v.Var {
- if i > 0 {
- p.printf(", ")
- }
- p.expr(name, precLow)
- }
- p.printf(": ")
- p.expr(v.Expr, precColon)
- case *BinaryExpr:
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
- prec := opPrec[v.Op]
- addParen(prec)
- m := p.margin
- if v.LineBreak {
- p.margin = p.indent()
- if v.Op == "=" {
- p.margin += listIndentation
- }
- }
- p.expr(v.X, prec)
- p.printf(" %s", v.Op)
- if v.LineBreak {
- p.breakline()
- } else {
- p.printf(" ")
- }
- p.expr(v.Y, prec+1)
- p.margin = m
- case *ParenExpr:
- p.seq("()", []Expr{v.X}, &v.End, modeParen, false, v.ForceMultiLine)
- case *CallExpr:
- addParen(precSuffix)
- p.expr(v.X, precSuffix)
- p.seq("()", v.List, &v.End, modeCall, v.ForceCompact, v.ForceMultiLine)
- case *ListExpr:
- p.seq("[]", v.List, &v.End, modeList, false, v.ForceMultiLine)
- case *SetExpr:
- p.seq("{}", v.List, &v.End, modeList, false, v.ForceMultiLine)
- case *TupleExpr:
- p.seq("()", v.List, &v.End, modeTuple, v.ForceCompact, v.ForceMultiLine)
- case *DictExpr:
- var list []Expr
- for _, x := range v.List {
- list = append(list, x)
- }
- p.seq("{}", list, &v.End, modeDict, false, v.ForceMultiLine)
- case *ListForExpr:
- p.listFor(v)
- case *ConditionalExpr:
- addParen(precSuffix)
- p.expr(v.Then, precSuffix)
- p.printf(" if ")
- p.expr(v.Test, precSuffix)
- p.printf(" else ")
- p.expr(v.Else, precSuffix)
- case *ReturnExpr:
- p.printf("return")
- if v.X != nil {
- p.printf(" ")
- p.expr(v.X, precSuffix)
- }
- case *FuncDef:
- p.printf("def ")
- p.printf(v.Name)
- p.seq("()", v.Args, &v.End, modeCall, v.ForceCompact, v.ForceMultiLine)
- p.printf(":")
- p.margin += nestedIndentation
- p.newline()
- p.statements(v.Body.Statements)
- p.margin -= nestedIndentation
- case *ForLoop:
- p.printf("for ")
- for i, loopVar := range v.LoopVars {
- if i > 0 {
- p.printf(", ")
- }
- p.expr(loopVar, precLow)
- }
- p.printf(" in ")
- p.expr(v.Iterable, precLow)
- p.printf(":")
- p.margin += nestedIndentation
- p.newline()
- p.statements(v.Body.Statements)
- p.margin -= nestedIndentation
- case *IfElse:
- for i, block := range v.Conditions {
- if i == 0 {
- p.printf("if ")
- } else if block.If == nil {
- p.newline()
- p.printf("else")
- } else {
- p.newline()
- p.printf("elif ")
- }
- if block.If != nil {
- p.expr(block.If, precLow)
- }
- p.printf(":")
- p.margin += nestedIndentation
- p.newline()
- p.statements(block.Then.Statements)
- p.margin -= nestedIndentation
- }
- }
-
- if parenthesized {
- p.depth--
- p.printf(")")
- }
-
-
- p.comment = append(p.comment, v.Comment().Suffix...)
- }
- type seqMode int
- const (
- _ seqMode = iota
- modeCall
- modeList
- modeTuple
- modeParen
- modeDict
- modeSeq
- )
- func (p *printer) seq(brack string, list []Expr, end *End, mode seqMode, forceCompact, forceMultiLine bool) {
- p.printf("%s", brack[:1])
- p.depth++
-
-
- for _, x := range list {
- if len(x.Comment().Before) > 0 {
- forceMultiLine = true
- }
- }
- if len(end.Before) > 0 {
- forceMultiLine = true
- }
-
-
- if forceMultiLine {
- forceCompact = false
- }
- switch {
- case len(list) == 0 && !forceMultiLine:
-
- case len(list) == 1 && !forceMultiLine:
-
- p.expr(list[0], precLow)
-
- if mode == modeTuple {
- p.printf(",")
- }
- case forceCompact:
-
- for i, x := range list {
- if i > 0 {
- p.printf(", ")
- }
- p.expr(x, precLow)
- }
- default:
-
- p.margin += listIndentation
- for i, x := range list {
-
-
-
-
-
-
-
- if i == 0 && len(p.comment) > 0 {
- p.printf("\n%*s", p.margin-2, "")
- }
- p.newline()
- p.expr(x, precLow)
- if mode != modeParen || i+1 < len(list) {
- p.printf(",")
- }
- }
-
- for _, com := range end.Before {
- p.newline()
- p.printf("%s", strings.TrimSpace(com.Token))
- }
- p.margin -= listIndentation
- p.newline()
- }
- p.depth--
- p.printf("%s", brack[1:])
- }
- func (p *printer) listFor(v *ListForExpr) {
- multiLine := v.ForceMultiLine || len(v.End.Before) > 0
-
-
- space := func() {
- if multiLine {
- p.breakline()
- } else {
- p.printf(" ")
- }
- }
- if v.Brack != "" {
- p.depth++
- p.printf("%s", v.Brack[:1])
- }
- if multiLine {
- if v.Brack != "" {
- p.margin += listIndentation
- }
- p.newline()
- }
- p.expr(v.X, precLow)
- for _, c := range v.For {
- space()
- p.printf("for ")
- for i, name := range c.For.Var {
- if i > 0 {
- p.printf(", ")
- }
- p.expr(name, precLow)
- }
- p.printf(" in ")
- p.expr(c.For.Expr, precLow)
- p.comment = append(p.comment, c.For.Comment().Suffix...)
- for _, i := range c.Ifs {
- space()
- p.printf("if ")
- p.expr(i.Cond, precLow)
- p.comment = append(p.comment, i.Comment().Suffix...)
- }
- p.comment = append(p.comment, c.Comment().Suffix...)
- }
- if multiLine {
- for _, com := range v.End.Before {
- p.newline()
- p.printf("%s", strings.TrimSpace(com.Token))
- }
- if v.Brack != "" {
- p.margin -= listIndentation
- }
- p.newline()
- }
- if v.Brack != "" {
- p.printf("%s", v.Brack[1:])
- p.depth--
- }
- }
- func (p *printer) isTopLevel() bool {
- return p.margin == 0
- }
|