From 0d39ef1ad89d6663c6e3a65da07c6e5f7d3fcb95 Mon Sep 17 00:00:00 2001
From: Alex Carney <alcarneyme@gmail.com>
Date: Tue, 8 Dec 2020 22:24:28 +0000
Subject: [PATCH] Ensure domain directives are discovered

---
 lib/esbonio/changes/31.fix.rst            |  2 ++
 lib/esbonio/esbonio/lsp/completion.py     |  1 -
 lib/esbonio/esbonio/lsp/initialize.py     | 22 +++++++++++++---
 lib/esbonio/tests/lsp/test_completions.py | 32 +++++++++++++++++++++++
 4 files changed, 53 insertions(+), 4 deletions(-)
 create mode 100644 lib/esbonio/changes/31.fix.rst

diff --git a/lib/esbonio/changes/31.fix.rst b/lib/esbonio/changes/31.fix.rst
new file mode 100644
index 000000000..0a1b90cc3
--- /dev/null
+++ b/lib/esbonio/changes/31.fix.rst
@@ -0,0 +1,2 @@
+**Language Server** Directives that are part of the ``std`` or ``py`` Sphinx domains
+will now be included in completion suggestions
\ No newline at end of file
diff --git a/lib/esbonio/esbonio/lsp/completion.py b/lib/esbonio/esbonio/lsp/completion.py
index 1ee8a2ef3..a473c2e42 100644
--- a/lib/esbonio/esbonio/lsp/completion.py
+++ b/lib/esbonio/esbonio/lsp/completion.py
@@ -6,7 +6,6 @@
 from pygls.types import (
     CompletionList,
     CompletionItem,
-    CompletionItemKind,
     CompletionParams,
     Position,
 )
diff --git a/lib/esbonio/esbonio/lsp/initialize.py b/lib/esbonio/esbonio/lsp/initialize.py
index 12def6615..8a9b9f73c 100644
--- a/lib/esbonio/esbonio/lsp/initialize.py
+++ b/lib/esbonio/esbonio/lsp/initialize.py
@@ -41,9 +41,11 @@ def discover_completion_items(rst: RstLanguageServer):
     - roles
     - directives
     """
-    # Lookup the directives and roles that have been registered
-    dirs = {**directives._directive_registry, **directives._directives}
-    rst.directives = {k: completion_from_directive(k, v) for k, v in dirs.items()}
+
+    rst.directives = {
+        k: completion_from_directive(k, v)
+        for k, v in discover_directives(rst.app).items()
+    }
 
     rst.roles = {
         k: completion_from_role(k, v) for k, v in discover_roles(rst.app).items()
@@ -53,6 +55,20 @@ def discover_completion_items(rst: RstLanguageServer):
     rst.logger.debug("Discovered %s roles", len(rst.roles))
 
 
+def discover_directives(app: Sphinx):
+    """Discover directives that we can offer as autocomplete suggestions."""
+    # Lookup the directives and roles that have been registered
+    dirs = {**directives._directive_registry, **directives._directives}
+
+    # Don't forget to include the directives that are stored under Sphinx domains.
+    # TODO: Implement proper domain handling, focus on std + python for now.
+    domains = app.registry.domains
+    std_directives = domains["std"].directives
+    py_directives = domains["py"].directives
+
+    return {**dirs, **std_directives, **py_directives}
+
+
 def discover_roles(app: Sphinx):
     """Discover roles that we can offer as autocomplete suggestions."""
     # Pull out the roles that are available via docutils.
diff --git a/lib/esbonio/tests/lsp/test_completions.py b/lib/esbonio/tests/lsp/test_completions.py
index 758c8a793..a30ae9e5b 100644
--- a/lib/esbonio/tests/lsp/test_completions.py
+++ b/lib/esbonio/tests/lsp/test_completions.py
@@ -25,6 +25,7 @@
 
 from esbonio.lsp.completion import completions
 from esbonio.lsp.initialize import (
+    discover_directives,
     discover_roles,
     discover_target_types,
     discover_targets,
@@ -189,6 +190,37 @@ def test_role_discovery(sphinx, project, expected, unexpected):
         assert name not in roles.keys(), "Unexpected role '{}'".format(name)
 
 
+@py.test.mark.parametrize(
+    "project,expected,unexpected",
+    [
+        (
+            "sphinx-default",
+            [
+                "figure",
+                "function",
+                "glossary",
+                "image",
+                "list-table",
+                "module",
+                "toctree",
+            ],
+            ["testcode", "autoclass", "automodule"],
+        )
+    ],
+)
+def test_directive_discovery(sphinx, project, expected, unexpected):
+    """Ensure that we can discover directives to offer as completion suggestions."""
+
+    app = sphinx(project)
+    directives = discover_directives(app)
+
+    for name in expected:
+        assert name in directives.keys(), "Missing expected directive '{}'".format(name)
+
+    for name in unexpected:
+        assert name not in directives.keys(), "Unexpected directive '{}'".format(name)
+
+
 @py.test.mark.parametrize(
     "project,type,kind,expected",
     [