Port Import provider to Rego (#853)

Signed-off-by: Anders Eknert <[email protected]> Signed-off-by: Charlie Egan <[email protected]> Co-authored-by: Anders Eknert <[email protected]>
StyraInc · Jun 19, 2024 · 0536573 · 0536573
1 parent e9a2ab6
commit 0536573
Show file tree

Hide file tree

Showing 9 changed files with 131 additions and 170 deletions.
diff --git a/bundle/regal/lsp/completion/providers/import/import.rego b/bundle/regal/lsp/completion/providers/import/import.rego
@@ -0,0 +1,28 @@
+package regal.lsp.completion.providers["import"]
+
+import rego.v1
+
+import data.regal.lsp.completion.kind
+import data.regal.lsp.completion.location
+
+items contains item if {
+	position := location.to_position(input.regal.context.location)
+	line := input.regal.file.lines[position.line]
+	word := location.word_at(line, input.regal.context.location.col)
+
+	invoke_suggestion(line)
+
+	item := {
+		"label": "import",
+		"kind": kind.keyword,
+		"detail": "import <path>",
+		"textEdit": {
+			"range": location.word_range(word, position),
+			"newText": "import ",
+		},
+	}
+}
+
+invoke_suggestion("")
+
+invoke_suggestion(line) if startswith("import", line)
diff --git a/bundle/regal/lsp/completion/providers/import/import_test.rego b/bundle/regal/lsp/completion/providers/import/import_test.rego
@@ -0,0 +1,65 @@
+package regal.lsp.completion.providers.import_test
+
+import rego.v1
+
+import data.regal.lsp.completion.providers["import"] as provider
+
+test_import_completion_empty_line if {
+	policy := `package policy
+
+import rego.v1
+
+`
+
+	regal_module := {"regal": {
+		"file": {
+			"name": "p.rego",
+			"lines": split(policy, "\n"),
+		},
+		"context": {"location": {"row": 5, "col": 1}},
+	}}
+	items := provider.items with input as regal_module
+
+	items == {{
+		"label": "import",
+		"detail": "import <path>",
+		"kind": 14,
+		"textEdit": {
+			"newText": "import ",
+			"range": {
+				"start": {"character": 0, "line": 4},
+				"end": {"character": 0, "line": 4},
+			},
+		},
+	}}
+}
+
+test_import_completion_on_typing if {
+	policy := `package policy
+
+import rego.v1
+
+imp`
+
+	regal_module := {"regal": {
+		"file": {
+			"name": "p.rego",
+			"lines": split(policy, "\n"),
+		},
+		"context": {"location": {"row": 5, "col": 3}},
+	}}
+	items := provider.items with input as regal_module
+
+	items == {{
+		"label": "import",
+		"detail": "import <path>",
+		"kind": 14,
+		"textEdit": {
+			"newText": "import ",
+			"range": {
+				"start": {"character": 0, "line": 4},
+				"end": {"character": 3, "line": 4},
+			},
+		},
+	}}
+}
diff --git a/bundle/regal/lsp/completion/providers/locals/locals_test.rego b/bundle/regal/lsp/completion/providers/locals/locals_test.rego
@@ -2,7 +2,8 @@ package regal.lsp.completion.providers.locals_test
 
 import rego.v1
 
-import data.regal.lsp.completion.providers.locals
+import data.regal.lsp.completion.providers.locals as provider
+import data.regal.lsp.completion.providers.utils_test as utils
 
 test_no_locals_in_completion_items if {
 	workspace := {"file:///p.rego": `package policy
@@ -26,7 +27,7 @@ bar if {
 			"col": 9,
 		}},
 	}}
-	items := locals.items with input as regal_module with data.workspace.parsed as parsed_modules(workspace)
+	items := provider.items with input as regal_module with data.workspace.parsed as utils.parsed_modules(workspace)
 
 	count(items) == 0
 }
@@ -56,11 +57,11 @@ function(bar) if {
 		}},
 	}}
 
-	items := locals.items with input as regal_module with data.workspace.parsed as parsed_modules(workspace)
+	items := provider.items with input as regal_module with data.workspace.parsed as utils.parsed_modules(workspace)
 
 	count(items) == 2
-	expect_item(items, "bar", {"end": {"character": 9, "line": 8}, "start": {"character": 8, "line": 8}})
-	expect_item(items, "baz", {"end": {"character": 9, "line": 8}, "start": {"character": 8, "line": 8}})
+	utils.expect_item(items, "bar", {"end": {"character": 9, "line": 8}, "start": {"character": 8, "line": 8}})
+	utils.expect_item(items, "baz", {"end": {"character": 9, "line": 8}, "start": {"character": 8, "line": 8}})
 }
 
 test_locals_in_completion_items_function_call if {
@@ -88,11 +89,11 @@ function(bar) if {
 		}},
 	}}
 
-	items := locals.items with input as regal_module with data.workspace.parsed as parsed_modules(workspace)
+	items := provider.items with input as regal_module with data.workspace.parsed as utils.parsed_modules(workspace)
 
 	count(items) == 2
-	expect_item(items, "bar", {"end": {"character": 24, "line": 8}, "start": {"character": 23, "line": 8}})
-	expect_item(items, "baz", {"end": {"character": 24, "line": 8}, "start": {"character": 23, "line": 8}})
+	utils.expect_item(items, "bar", {"end": {"character": 24, "line": 8}, "start": {"character": 23, "line": 8}})
+	utils.expect_item(items, "baz", {"end": {"character": 24, "line": 8}, "start": {"character": 23, "line": 8}})
 }
 
 test_locals_in_completion_items_rule_head_assignment if {
@@ -116,27 +117,8 @@ function(bar) := f if {
 			"col": 19,
 		}},
 	}}
-	items := locals.items with input as regal_module with data.workspace.parsed as parsed_modules(workspace)
+	items := provider.items with input as regal_module with data.workspace.parsed as utils.parsed_modules(workspace)
 
 	count(items) == 1
-	expect_item(items, "foo", {"end": {"character": 18, "line": 4}, "start": {"character": 17, "line": 4}})
-}
-
-parsed_modules(workspace) := {file_uri: parsed_module |
-	some file_uri, contents in workspace
-	parsed_module := regal.parse_module(file_uri, contents)
-}
-
-expect_item(items, label, range) if {
-	expected := {"detail": "local variable", "kind": 6}
-
-	item := object.union(expected, {
-		"label": label,
-		"textEdit": {
-			"newText": label,
-			"range": range,
-		},
-	})
-
-	item in items
+	utils.expect_item(items, "foo", {"end": {"character": 18, "line": 4}, "start": {"character": 17, "line": 4}})
 }
diff --git a/bundle/regal/lsp/completion/providers/utils_test.rego b/bundle/regal/lsp/completion/providers/utils_test.rego
@@ -0,0 +1,22 @@
+package regal.lsp.completion.providers.utils_test
+
+import rego.v1
+
+parsed_modules(workspace) := {file_uri: parsed_module |
+	some file_uri, contents in workspace
+	parsed_module := regal.parse_module(file_uri, contents)
+}
+
+expect_item(items, label, range) if {
+	expected := {"detail": "local variable", "kind": 6}
+
+	item := object.union(expected, {
+		"label": label,
+		"textEdit": {
+			"newText": label,
+			"range": range,
+		},
+	})
+
+	item in items
+}
diff --git a/internal/lsp/completions/manager.go b/internal/lsp/completions/manager.go
@@ -32,7 +32,6 @@ func NewDefaultManager(c *cache.Cache, store storage.Store) *Manager {
 
 	m.RegisterProvider(&providers.Package{})
 	m.RegisterProvider(&providers.PackageName{})
-	m.RegisterProvider(&providers.Import{})
 	m.RegisterProvider(&providers.BuiltIns{})
 	m.RegisterProvider(&providers.RegoV1{})
 	m.RegisterProvider(&providers.PackageRefs{})

diff --git a/internal/lsp/completions/manager_test.go b/internal/lsp/completions/manager_test.go
@@ -55,7 +55,7 @@ func TestManagerEarlyExitInsideComment(t *testing.T) {
 
 	fileContents := `package p
 
-import rego.v1 # modern rego i
+# foo := http
 `
 
 	module := ast.MustParseModule(fileContents)
@@ -64,15 +64,15 @@ import rego.v1 # modern rego i
 	c.SetModule(fileURI, module)
 
 	mgr := NewManager(c, &ManagerOptions{})
-	mgr.RegisterProvider(&providers.Import{})
+	mgr.RegisterProvider(&providers.BuiltIns{})
 
 	completionParams := types.CompletionParams{
 		TextDocument: types.TextDocumentIdentifier{
 			URI: fileURI,
 		},
 		Position: types.Position{
 			Line:      2,
-			Character: 30,
+			Character: 13,
 		},
 	}
 

diff --git a/internal/lsp/completions/providers/import.go b/internal/lsp/completions/providers/import.go
diff --git a/internal/lsp/completions/providers/import_test.go b/internal/lsp/completions/providers/import_test.go
diff --git a/internal/lsp/completions/providers/policy_test.go b/internal/lsp/completions/providers/policy_test.go
@@ -13,7 +13,7 @@ import (
 	"github.com/styrainc/regal/internal/parse"
 )
 
-func TestPolicyProvider(t *testing.T) {
+func TestPolicyProvider_Example1(t *testing.T) {
 	t.Parallel()
 
 	policy := `package p
@@ -71,7 +71,7 @@ allow if {
 	}
 }
 
-func TestPolicyProvider_B(t *testing.T) {
+func TestPolicyProvider_Example2(t *testing.T) {
 	t.Parallel()
 
 	file1 := ast.MustParseModule(`package example