Add option to marshal location text (#6234)

ast/annotations.go

@@ -11,6 +11,7 @@ import (
 	"sort"
 	"strings"
 
+	astJSON "github.com/open-policy-agent/opa/ast/json"
 	"github.com/open-policy-agent/opa/internal/deepcopy"
 	"github.com/open-policy-agent/opa/util"
 )
@@ -39,7 +40,7 @@ type (
 
 		comments    []*Comment
 		node        Node
-		jsonOptions JSONOptions
+		jsonOptions astJSON.Options
 	}
 
 	// SchemaAnnotation contains a schema declaration for the document identified by the path.
@@ -76,7 +77,7 @@ type (
 		Annotations *Annotations `json:"annotations,omitempty"`
 		Location    *Location    `json:"location,omitempty"` // The location of the node the annotations are applied to
 
-		jsonOptions JSONOptions
+		jsonOptions astJSON.Options
 
 		node Node // The node the annotations are applied to
 	}
@@ -180,8 +181,11 @@ func (a *Annotations) GetTargetPath() Ref {
 	}
 }
 
-func (a *Annotations) setJSONOptions(opts JSONOptions) {
+func (a *Annotations) setJSONOptions(opts astJSON.Options) {
 	a.jsonOptions = opts
+	if a.Location != nil {
+		a.Location.JSONOptions = opts
+	}
 }
 
 func (a *Annotations) MarshalJSON() ([]byte, error) {

ast/json/json.go

@@ -0,0 +1,33 @@
+package json
+
+// Options defines the options for JSON operations,
+// currently only marshaling can be configured
+type Options struct {
+	MarshalOptions MarshalOptions
+}
+
+// MarshalOptions defines the options for JSON marshaling,
+// currently only toggling the marshaling of location information is supported
+type MarshalOptions struct {
+	// IncludeLocation toggles the marshaling of location information
+	IncludeLocation NodeToggle
+	// IncludeLocationText additionally/optionally includes the text of the location
+	IncludeLocationText bool
+}
+
+// NodeToggle is a generic struct to allow the toggling of
+// settings for different ast node types
+type NodeToggle struct {
+	Term           bool
+	Package        bool
+	Comment        bool
+	Import         bool
+	Rule           bool
+	Head           bool
+	Expr           bool
+	SomeDecl       bool
+	Every          bool
+	With           bool
+	Annotations    bool
+	AnnotationsRef bool
+}

ast/location/location.go

@@ -3,8 +3,11 @@ package location
 
 import (
 	"bytes"
+	"encoding/json"
 	"errors"
 	"fmt"
+
+	astJSON "github.com/open-policy-agent/opa/ast/json"
 )
 
 // Location records a position in source code
@@ -14,6 +17,9 @@ type Location struct {
 	Row    int    `json:"row"`  // The line in the source.
 	Col    int    `json:"col"`  // The column in the row.
 	Offset int    `json:"-"`    // The byte offset for the location in the source.
+
+	// JSONOptions specifies options for marshaling and unmarshaling of locations
+	JSONOptions astJSON.Options
 }
 
 // NewLocation returns a new Location object.
@@ -87,3 +93,23 @@ func (loc *Location) Compare(other *Location) int {
 	}
 	return 0
 }
+
+func (loc *Location) MarshalJSON() ([]byte, error) {
+	// structs are used here to preserve the field ordering of the original Location struct
+	data := struct {
+		File string `json:"file"`
+		Row  int    `json:"row"`
+		Col  int    `json:"col"`
+		Text []byte `json:"text,omitempty"`
+	}{
+		File: loc.File,
+		Row:  loc.Row,
+		Col:  loc.Col,
+	}
+
+	if loc.JSONOptions.MarshalOptions.IncludeLocationText {
+		data.Text = loc.Text
+	}
+
+	return json.Marshal(data)
+}

ast/location/location_test.go

@@ -1,8 +1,10 @@
 package location
 
 import (
+	"encoding/json"
 	"testing"
 
+	astJSON "github.com/open-policy-agent/opa/ast/json"
 	"github.com/open-policy-agent/opa/util"
 )
 
@@ -85,3 +87,46 @@ func TestLocationCompare(t *testing.T) {
 		}
 	}
 }
+
+func TestLocationMarshal(t *testing.T) {
+	testCases := map[string]struct {
+		loc *Location
+		exp string
+	}{
+		"default json options": {
+			loc: &Location{
+				Text: []byte("text"),
+				File: "file",
+				Row:  1,
+				Col:  1,
+			},
+			exp: `{"file":"file","row":1,"col":1}`,
+		},
+		"including text": {
+			loc: &Location{
+				Text: []byte("text"),
+				File: "file",
+				Row:  1,
+				Col:  1,
+				JSONOptions: astJSON.Options{
+					MarshalOptions: astJSON.MarshalOptions{
+						IncludeLocationText: true,
+					},
+				},
+			},
+			exp: `{"file":"file","row":1,"col":1,"text":"dGV4dA=="}`,
+		},
+	}
+
+	for id, tc := range testCases {
+		t.Run(id, func(t *testing.T) {
+			bs, err := json.Marshal(tc.loc)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if string(bs) != tc.exp {
+				t.Fatalf("Expected %v but got %v", tc.exp, string(bs))
+			}
+		})
+	}
+}

ast/marshal.go

@@ -1,7 +1,11 @@
 package ast
 
+import (
+	astJSON "github.com/open-policy-agent/opa/ast/json"
+)
+
 // customJSON is an interface that can be implemented by AST nodes that
 // allows the parser to set options for JSON operations on that node.
 type customJSON interface {
-	setJSONOptions(JSONOptions)
+	setJSONOptions(astJSON.Options)
 }