diff --git a/format/format.go b/format/format.go index 18875a7679..7095f79feb 100644 --- a/format/format.go +++ b/format/format.go @@ -60,7 +60,7 @@ func Ast(x interface{}) ([]byte, error) { // or internal.member_3, it will sugarize them into usage of the `in` // operator. It has to ensure that the proper future keyword import is // present. - extraFutureKeywordImports := map[string]bool{} + extraFutureKeywordImports := map[string]struct{}{} // Preprocess the AST. Set any required defaults and calculate // values required for printing the formatted output. @@ -76,9 +76,9 @@ func Ast(x interface{}) ([]byte, error) { case *ast.Expr: switch { case n.IsCall() && ast.Member.Ref().Equal(n.Operator()) || ast.MemberWithKey.Ref().Equal(n.Operator()): - extraFutureKeywordImports["in"] = true + extraFutureKeywordImports["in"] = struct{}{} case n.IsEvery(): - extraFutureKeywordImports["every"] = true + extraFutureKeywordImports["every"] = struct{}{} } } if x.Loc() == nil { @@ -1239,11 +1239,22 @@ func (w *writer) down() { func ensureFutureKeywordImport(imps []*ast.Import, kw string) []*ast.Import { allKeywords := ast.MustParseTerm("future.keywords") - kwPath := ast.MustParseTerm("future.keywords." + kw) + kwPath := keyword(kw) + every := keyword("every") for _, imp := range imps { - if allKeywords.Equal(imp.Path) || imp.Path.Equal(kwPath) { + if allKeywords.Equal(imp.Path) || + imp.Path.Equal(kwPath) || + (imp.Path.Equal(every) && kw == "in") { // "every" implies "in", so we don't need to add both return imps } } - return append(imps, &ast.Import{Path: kwPath}) + imp := &ast.Import{ + Path: kwPath, + } + imp.Location = defaultLocation(imp) + return append(imps, imp) +} + +func keyword(kw string) *ast.Term { + return ast.MustParseTerm("future.keywords." + kw) } diff --git a/format/testfiles/test_in_operator_without_import.rego.formatted b/format/testfiles/test_in_operator_without_import.rego.formatted index efd6f755e7..040c28dfdb 100644 --- a/format/testfiles/test_in_operator_without_import.rego.formatted +++ b/format/testfiles/test_in_operator_without_import.rego.formatted @@ -1,6 +1,7 @@ package p import future.keywords.in + import input.foo r { diff --git a/format/testfiles/test_issue_4606.rego b/format/testfiles/test_issue_4606.rego new file mode 100644 index 0000000000..a71142ed97 --- /dev/null +++ b/format/testfiles/test_issue_4606.rego @@ -0,0 +1,6 @@ +package p + +# +import future.keywords.every + +r { 1 in [] } \ No newline at end of file diff --git a/format/testfiles/test_issue_4606.rego.formatted b/format/testfiles/test_issue_4606.rego.formatted new file mode 100644 index 0000000000..ce00ead8c9 --- /dev/null +++ b/format/testfiles/test_issue_4606.rego.formatted @@ -0,0 +1,8 @@ +package p + +# +import future.keywords.every + +r { + 1 in [] +}