handle union set in prettify function (#198)

arr-ai · Jun 16, 2021 · c98c79e · c98c79e
1 parent 33856cd
commit c98c79e
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 7 deletions.
diff --git a/rel/value_repr.go b/rel/value_repr.go
@@ -103,13 +103,13 @@ func reprDict(d Dict, w io.Writer) {
 	fmt.Fprint(w, "}")
 }
 
-// orderableSet is a type used to repr GenericSet and UnionSet to avoid duplications.
-type orderableSet interface {
+// OrderableSet is a type used to repr GenericSet and UnionSet to avoid duplications.
+type OrderableSet interface {
 	Set
 	OrderedValues() []Value
 }
 
-func reprOrderableSet(s orderableSet, w io.Writer) {
+func reprOrderableSet(s OrderableSet, w io.Writer) {
 	if s.Equal(True) {
 		fmt.Fprint(w, sTrue)
 		return
@@ -178,7 +178,7 @@ func reprValue(v Value, w io.Writer) {
 		reprDict(v, w)
 	case Relation:
 		fmt.Fprint(w, v.String())
-	case orderableSet:
+	case OrderableSet:
 		reprOrderableSet(v, w)
 	case Closure:
 		reprClosure(v, w)

diff --git a/syntax/prettify.go b/syntax/prettify.go
@@ -60,8 +60,8 @@ func PrettifyString(val interface{}, indentsNum int) (string, error) {
 		return prettifyArray(t, indentsNum+1)
 	case rel.Dict: // {'a': 1}
 		return prettifyDict(t, indentsNum+1)
-	case rel.GenericSet: // {1, 2}
-		return prettifySet(t, indentsNum+1)
+	case rel.OrderableSet: // {1, 2}
+		return prettifyOrderableSet(t, indentsNum+1)
 	case rel.Relation:
 		return prettifyRelation(t, indentsNum+1)
 	case rel.String:
@@ -107,7 +107,7 @@ func prettifyRelation(r rel.Relation, indentsNum int) (string, error) {
 	return sb.String(), nil
 }
 
-func prettifySet(arr rel.GenericSet, indentsNum int) (string, error) {
+func prettifyOrderableSet(arr rel.OrderableSet, indentsNum int) (string, error) {
 	content, err := prettifyItems(arr.OrderedValues(), indentsNum)
 	if err != nil {
 		return "", err

diff --git a/syntax/std_fmt_test.go b/syntax/std_fmt_test.go
@@ -86,6 +86,67 @@ func TestFmtPrettySet(t *testing.T) {
 	AssertCodesEvalToSameValue(t, `"{1, 2, {3, 4, 5}}"`, `//fmt.pretty({1,2,{3,4,5}})`)
 }
 
+func TestFmtPrettyUnionSet(t *testing.T) {
+	t.Parallel()
+	AssertCodesEvalToSameValue(t,
+		`"{
+  (
+    a: 'abc',
+  ),
+  (
+    b: 2,
+  ),
+}"`,
+		`//fmt.pretty({(a: 'abc'), (b: 2)})`,
+	)
+	AssertCodesEvalToSameValue(t,
+		`"{
+  1,
+  2,
+  3,
+  'a',
+  [1, 2, 3],
+  (
+    a: 1,
+  ),
+}"`,
+		`//fmt.pretty({(a: 1), 1, 2, 3, "a", [1, 2, 3]})`,
+	)
+	AssertCodesEvalToSameValue(t,
+		`"{
+  (
+    a: (
+      b: (
+        c: [
+          (
+            d: 1,
+          ),
+        ],
+      ),
+    ),
+  ),
+  (
+    e: {
+      (
+        a: {
+          (
+            a: 1,
+          ),
+          (
+            b: 1,
+          ),
+        },
+      ),
+      (
+        b: 2,
+      ),
+    },
+  ),
+}"`,
+		`//fmt.pretty({(a: (b: (c: [(d: 1)]))), (e: {(a: {(a: 1), (b: 1)}), (b: 2)})})`,
+	)
+}
+
 func TestFmtPrettyString(t *testing.T) {
 	t.Parallel()
 	AssertCodesEvalToSameValue(t, `"{}"`, `//fmt.pretty('')`)