sqlc/internal/python/printer/printer_test.go

package printer

import (
	"strings"
	"testing"

	"github.com/google/go-cmp/cmp"

	"github.com/kyleconroy/sqlc/internal/python/ast"
)

type testcase struct {
	Node     *ast.Node
	Expected string
}

func TestPrinter(t *testing.T) {
	for name, tc := range map[string]testcase{
		"assign": {
			Node: &ast.Node{
				Node: &ast.Node_Assign{
					Assign: &ast.Assign{
						Targets: []*ast.Node{
							{
								Node: &ast.Node_Name{
									Name: &ast.Name{Id: "FICTION"},
								},
							},
						},
						Value: &ast.Node{
							Node: &ast.Node_Constant{
								Constant: &ast.Constant{
									Value: &ast.Constant_Str{
										Str: "FICTION",
									},
								},
							},
						},
					},
				},
			},
			Expected: `FICTION = "FICTION"`,
		},
		"class-base": {
			Node: &ast.Node{
				Node: &ast.Node_ClassDef{
					ClassDef: &ast.ClassDef{
						Name: "Foo",
						Bases: []*ast.Node{
							{
								Node: &ast.Node_Name{
									Name: &ast.Name{Id: "str"},
								},
							},
							{
								Node: &ast.Node_Attribute{
									Attribute: &ast.Attribute{
										Value: &ast.Node{
											Node: &ast.Node_Name{
												Name: &ast.Name{Id: "enum"},
											},
										},
										Attr: "Enum",
									},
								},
							},
						},
					},
				},
			},
			Expected: `class Foo(str, enum.Enum):`,
		},
		"dataclass": {
			Node: &ast.Node{
				Node: &ast.Node_ClassDef{
					ClassDef: &ast.ClassDef{
						Name: "Foo",
						DecoratorList: []*ast.Node{
							{
								Node: &ast.Node_Name{
									Name: &ast.Name{
										Id: "dataclass",
									},
								},
							},
						},
						Body: []*ast.Node{
							{
								Node: &ast.Node_AnnAssign{
									AnnAssign: &ast.AnnAssign{
										Target: &ast.Name{Id: "bar"},
										Annotation: &ast.Node{
											Node: &ast.Node_Name{
												Name: &ast.Name{Id: "int"},
											},
										},
									},
								},
							},
							{
								Node: &ast.Node_AnnAssign{
									AnnAssign: &ast.AnnAssign{
										Target: &ast.Name{Id: "bat"},
										Annotation: &ast.Node{
											Node: &ast.Node_Subscript{
												Subscript: &ast.Subscript{
													Value: &ast.Name{Id: "Optional"},
													Slice: &ast.Node{
														Node: &ast.Node_Name{
															Name: &ast.Name{Id: "int"},
														},
													},
												},
											},
										},
									},
								},
							},
						},
					},
				},
			},
			Expected: `
@dataclass
class Foo:
    bar: int
    bat: Optional[int]
`,
		},
		"call": {
			Node: &ast.Node{
				Node: &ast.Node_Call{
					Call: &ast.Call{
						Func: &ast.Node{
							Node: &ast.Node_Alias{
								Alias: &ast.Alias{
									Name: "foo",
								},
							},
						},
					},
				},
			},
			Expected: `foo()`,
		},

		"import": {
			Node: &ast.Node{
				Node: &ast.Node_Import{
					Import: &ast.Import{
						Names: []*ast.Node{
							{
								Node: &ast.Node_Alias{
									Alias: &ast.Alias{
										Name: "foo",
									},
								},
							},
						},
					},
				},
			},
			Expected: `import foo`,
		},
		"import-from": {
			Node: &ast.Node{
				Node: &ast.Node_ImportFrom{
					ImportFrom: &ast.ImportFrom{
						Module: "pkg",
						Names: []*ast.Node{
							{
								Node: &ast.Node_Alias{
									Alias: &ast.Alias{
										Name: "foo",
									},
								},
							},
							{
								Node: &ast.Node_Alias{
									Alias: &ast.Alias{
										Name: "bar",
									},
								},
							},
						},
					},
				},
			},
			Expected: `from pkg import foo, bar`,
		},
	} {
		tc := tc
		t.Run(name, func(t *testing.T) {
			result := Print(tc.Node, Options{})
			if diff := cmp.Diff(strings.TrimSpace(tc.Expected), strings.TrimSpace(string(result.Python))); diff != "" {
				t.Errorf("print mismatch (-want +got):\n%s", diff)
			}
		})
	}
}