[ast/opa parse] Support marshalling of all ast location data (#5576)

Signed-off-by: Charlie Egan <[email protected]>

Author: Charlie Egan

ast/annotations.go

@@ -26,7 +26,6 @@ const (
 type (
 	// Annotations represents metadata attached to other AST nodes such as rules.
 	Annotations struct {
-		Location         *Location                    `json:"-"`
 		Scope            string                       `json:"scope"`
 		Title            string                       `json:"title,omitempty"`
 		Entrypoint       bool                         `json:"entrypoint,omitempty"`
@@ -36,8 +35,11 @@ type (
 		Authors          []*AuthorAnnotation          `json:"authors,omitempty"`
 		Schemas          []*SchemaAnnotation          `json:"schemas,omitempty"`
 		Custom           map[string]interface{}       `json:"custom,omitempty"`
-		node             Node
-		comments         []*Comment
+		Location         *Location                    `json:"location,omitempty"`
+
+		comments    []*Comment
+		node        Node
+		jsonOptions JSONOptions
 	}
 
 	// SchemaAnnotation contains a schema declaration for the document identified by the path.
@@ -70,10 +72,13 @@ type (
 	}
 
 	AnnotationsRef struct {
-		Location    *Location    `json:"location"` // The location of the node the annotations are applied to
-		Path        Ref          `json:"path"`     // The path of the node the annotations are applied to
+		Path        Ref          `json:"path"` // The path of the node the annotations are applied to
 		Annotations *Annotations `json:"annotations,omitempty"`
-		node        Node         // The node the annotations are applied to
+		Location    *Location    `json:"location,omitempty"` // The location of the node the annotations are applied to
+
+		jsonOptions JSONOptions
+
+		node Node // The node the annotations are applied to
 	}
 
 	AnnotationsRefSet []*AnnotationsRef
@@ -82,7 +87,7 @@ type (
 )
 
 func (a *Annotations) String() string {
-	bs, _ := json.Marshal(a)
+	bs, _ := a.MarshalJSON()
 	return string(bs)
 }
 
@@ -175,6 +180,60 @@ func (a *Annotations) GetTargetPath() Ref {
 	}
 }
 
+func (a *Annotations) setJSONOptions(opts JSONOptions) {
+	a.jsonOptions = opts
+}
+
+func (a *Annotations) MarshalJSON() ([]byte, error) {
+	if a == nil {
+		return []byte(`{"scope":""}`), nil
+	}
+
+	data := map[string]interface{}{
+		"scope": a.Scope,
+	}
+
+	if a.Title != "" {
+		data["title"] = a.Title
+	}
+
+	if a.Description != "" {
+		data["description"] = a.Description
+	}
+
+	if a.Entrypoint {
+		data["entrypoint"] = a.Entrypoint
+	}
+
+	if len(a.Organizations) > 0 {
+		data["organizations"] = a.Organizations
+	}
+
+	if len(a.RelatedResources) > 0 {
+		data["related_resources"] = a.RelatedResources
+	}
+
+	if len(a.Authors) > 0 {
+		data["authors"] = a.Authors
+	}
+
+	if len(a.Schemas) > 0 {
+		data["schemas"] = a.Schemas
+	}
+
+	if len(a.Custom) > 0 {
+		data["custom"] = a.Custom
+	}
+
+	if a.jsonOptions.MarshalOptions.IncludeLocation.Annotations {
+		if a.Location != nil {
+			data["location"] = a.Location
+		}
+	}
+
+	return json.Marshal(data)
+}
+
 func NewAnnotationsRef(a *Annotations) *AnnotationsRef {
 	var loc *Location
 	if a.node != nil {
@@ -209,6 +268,24 @@ func (ar *AnnotationsRef) GetRule() *Rule {
 	}
 }
 
+func (ar *AnnotationsRef) MarshalJSON() ([]byte, error) {
+	data := map[string]interface{}{
+		"path": ar.Path,
+	}
+
+	if ar.Annotations != nil {
+		data["annotations"] = ar.Annotations
+	}
+
+	if ar.jsonOptions.MarshalOptions.IncludeLocation.AnnotationsRef {
+		if ar.Location != nil {
+			data["location"] = ar.Location
+		}
+	}
+
+	return json.Marshal(data)
+}
+
 func scopeCompare(s1, s2 string) int {
 
 	o1 := scopeOrder(s1)

ast/annotations_test.go

@@ -52,8 +52,8 @@ p := 7`},
 	}
 
 	// Output:
-	// data.foo at foo.rego:5 has annotations {"scope":"subpackages","organizations":["Acme Corp."]}
-	// data.foo.bar at mod:3 has annotations {"scope":"package","description":"A couple of useful rules"}
+	// data.foo at foo.rego:5 has annotations {"organizations":["Acme Corp."],"scope":"subpackages"}
+	// data.foo.bar at mod:3 has annotations {"description":"A couple of useful rules","scope":"package"}
 	// data.foo.bar.p at mod:7 has annotations {"scope":"rule","title":"My Rule P"}
 }
 
@@ -102,8 +102,8 @@ p := 7`},
 
 	// Output:
 	// data.foo.bar.p at mod:7 has annotations {"scope":"rule","title":"My Rule P"}
-	// data.foo.bar at mod:3 has annotations {"scope":"package","description":"A couple of useful rules"}
-	// data.foo at foo.rego:5 has annotations {"scope":"subpackages","organizations":["Acme Corp."]}
+	// data.foo.bar at mod:3 has annotations {"description":"A couple of useful rules","scope":"package"}
+	// data.foo at foo.rego:5 has annotations {"organizations":["Acme Corp."],"scope":"subpackages"}
 }
 
 func TestAnnotationSet_Flatten(t *testing.T) {

ast/marshal.go

@@ -0,0 +1,7 @@
+package ast
+
+// 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)
+}