Merge pull request #63 from azriel91/feature/lsp-suggest-custom-entity-types

Suggest custom entity types as a completion option when editing `theme_types_styles`.

Author: azriel91

CHANGELOG.md

@@ -67,6 +67,8 @@
 * Add `thing_layout_edges` to affect node ranks without rendering any visible `<path>`s. ([#61][#61])
 * Reverse path calculation computation when `to` node has a smaller node rank than the `from` node, to reduce unnecessary spacer routing. ([#62][#62])
 * Update `EdgeCurvature::Curved` and `EdgeCurvature::Orthogonal` edge path builders to skip thing description spacers when rank dir spacers are at a larger cross-axis coordinate. ([#62][#62])
+* Suggest custom entity types as a completion option when editing `theme_types_styles`. ([#63][#63])
+* Suggest custom entity types as a completion option when editing `entity_types.<entity_name>.<entity_type_list>`. ([#63][#63])
 
 [#42]: https://github.com/azriel91/disposition/pull/42
 [#43]: https://github.com/azriel91/disposition/pull/43
@@ -89,6 +91,7 @@
 [#60]: https://github.com/azriel91/disposition/pull/60
 [#61]: https://github.com/azriel91/disposition/pull/61
 [#62]: https://github.com/azriel91/disposition/pull/62
+[#63]: https://github.com/azriel91/disposition/pull/63
 
 
 ## 0.3.0 (2026-06-07)

crate/lsp/src/code_action.rs

@@ -11,7 +11,9 @@
 
 use async_lsp::lsp_types::{Position, Range, TextEdit};
 
-use crate::completion::yaml_lines::{indent, is_blank_or_comment, is_list_item, line_map_key};
+use crate::completion::yaml_lines::{
+    indent, is_blank_or_comment, is_list_item, line_map_key, split_flow_items,
+};
 
 pub use self::list_conversion::ListConversion;
 
@@ -162,16 +164,6 @@ impl CodeActionEngine {
     }
 }
 
-/// Splits the inside of a flow sequence (`a, b, c`) into trimmed, non-empty
-/// items.
-fn split_flow_items(inner: &str) -> Vec<&str> {
-    inner
-        .split(',')
-        .map(str::trim)
-        .filter(|item| !item.is_empty())
-        .collect()
-}
-
 /// The range covering the whole of `line` at `line_idx` (column 0 to its end).
 fn whole_line_range(line_idx: usize, line: &str) -> Range {
     Range {

crate/lsp/src/completion/completion_engine.rs

@@ -150,6 +150,7 @@ impl CompletionEngine {
             );
         }
 
+        completion_items_dedup(&mut items);
         items
     }
 
@@ -227,6 +228,7 @@ impl CompletionEngine {
             }));
         }
 
+        completion_items_dedup(&mut items);
         items
     }
 
@@ -259,6 +261,17 @@ impl CompletionEngine {
     }
 }
 
+/// Removes items whose `label` already appeared earlier in `items`.
+///
+/// Document-derived and schema-derived suggestions are merged independently
+/// and can coincide -- e.g. a document declares a custom entity type whose
+/// name happens to match a built-in `EntityType` const -- so this dedups the
+/// combined list by label, keeping the first occurrence.
+fn completion_items_dedup(items: &mut Vec<CompletionItem>) {
+    let mut seen = BTreeSet::new();
+    items.retain(|item| seen.insert(item.label.clone()));
+}
+
 /// Builds the `insert_text` for a value completion `label`.
 ///
 /// Returns `None` when the bare `label` can be inserted as-is. Otherwise the

crate/lsp/src/completion/dynamic_completions.rs

@@ -10,7 +10,7 @@ use std::collections::BTreeSet;
 use crate::completion::{
     id_category::IdCategory,
     key_category::KeyCategory,
-    yaml_lines::{indent, is_blank_or_comment, line_map_key},
+    yaml_lines::{indent, is_blank_or_comment, is_list_item, line_map_key, split_flow_items},
 };
 
 /// IDs defined in the document, grouped by category.
@@ -27,6 +27,10 @@ pub struct DynamicCompletions {
     /// `EdgeGroupId`s declared under `thing_dependencies` /
     /// `thing_interactions`.
     edge_group_ids: BTreeSet<String>,
+    /// Custom `type_*` ids declared as list items under any `entity_types`
+    /// entry (e.g. `type_organisation` from `entity_types.t_aws:
+    /// [type_organisation]`).
+    entity_type_ids: BTreeSet<String>,
 }
 
 impl DynamicCompletions {
@@ -49,12 +53,15 @@ impl DynamicCompletions {
         let mut edge_group_ids = collect_block_keys(&lines, "thing_dependencies", true, true);
         edge_group_ids.extend(collect_block_keys(&lines, "thing_interactions", true, true));
 
+        let entity_type_ids = collect_entity_types_values(&lines);
+
         DynamicCompletions {
             thing_ids,
             tag_ids,
             process_ids,
             step_ids,
             edge_group_ids,
+            entity_type_ids,
         }
     }
 
@@ -65,6 +72,7 @@ impl DynamicCompletions {
             IdCategory::Tag => self.tag_ids.iter().map(String::as_str).collect(),
             IdCategory::ProcessStep => self.step_ids.iter().map(String::as_str).collect(),
             IdCategory::EdgeGroup => self.edge_group_ids.iter().map(String::as_str).collect(),
+            IdCategory::EntityType => self.entity_type_ids.iter().map(String::as_str).collect(),
             IdCategory::Any => self
                 .thing_ids
                 .iter()
@@ -81,7 +89,9 @@ impl DynamicCompletions {
     /// These are the document-derived IDs, templated IDs, and known literal
     /// keys. Schema-derived suggestions (the built-in `StyleAlias` /
     /// `EntityType` keys) are *not* included here -- the completion engine adds
-    /// those from the schema, which this type does not have access to.
+    /// those from the schema, which this type does not have access to. Custom
+    /// `EntityType`s declared in `entity_types`, however, *are*
+    /// document-derived and so are included here.
     pub fn key_suggestions(&self, key_category: KeyCategory) -> Vec<String> {
         let owned = |ids: Vec<&str>| ids.into_iter().map(str::to_string).collect::<Vec<String>>();
 
@@ -113,8 +123,9 @@ impl DynamicCompletions {
                 suggestions.extend(owned(self.ids_for(IdCategory::Tag)));
                 suggestions
             }
-            // Built-in entity types come from the schema, added by the engine.
-            KeyCategory::EntityType => Vec::new(),
+            // Custom entity types declared in `entity_types`; built-in
+            // entity types come from the schema, added by the engine.
+            KeyCategory::EntityType => owned(self.ids_for(IdCategory::EntityType)),
         }
     }
 
@@ -207,3 +218,44 @@ fn collect_block_keys(
 
     keys
 }
+
+/// Collects the custom entity type ids (`type_*` list items) declared under
+/// each entry of the top-level `entity_types` block, e.g. `type_organisation`
+/// from `entity_types.t_aws: [type_organisation]` or a block-sequence
+/// `- type_organisation` item.
+fn collect_entity_types_values(lines: &[&str]) -> BTreeSet<String> {
+    let mut values = BTreeSet::new();
+
+    let Some(block_start) = lines.iter().position(|line| {
+        !is_blank_or_comment(line)
+            && indent(line) == 0
+            && line_map_key(line).as_deref() == Some("entity_types")
+    }) else {
+        return values;
+    };
+
+    for line in &lines[block_start + 1..] {
+        if is_blank_or_comment(line) {
+            continue;
+        }
+        if indent(line) == 0 {
+            break;
+        }
+
+        if is_list_item(line) {
+            let value = line.trim_start().trim_start_matches('-').trim();
+            if !value.is_empty() {
+                values.insert(value.to_string());
+            }
+        } else if line_map_key(line).is_some()
+            && let Some(colon_idx) = line.find(':')
+        {
+            let value = line[colon_idx + 1..].trim();
+            if let Some(inner) = value.strip_prefix('[').and_then(|v| v.strip_suffix(']')) {
+                values.extend(split_flow_items(inner).into_iter().map(str::to_string));
+            }
+        }
+    }
+
+    values
+}

crate/lsp/src/completion/id_category.rs

@@ -19,6 +19,9 @@ pub enum IdCategory {
     EdgeGroup,
     /// Any ID -- the generic `Id` type (e.g. `thing_layouts` keys).
     Any,
+    /// An `EntityType` custom `type_*` id -- declared as a list item under
+    /// any `entity_types` entry.
+    EntityType,
 }
 
 impl IdCategory {
@@ -30,6 +33,7 @@ impl IdCategory {
             "ProcessStepId" => Some(IdCategory::ProcessStep),
             "EdgeGroupId" => Some(IdCategory::EdgeGroup),
             "Id" => Some(IdCategory::Any),
+            "EntityType" => Some(IdCategory::EntityType),
             _ => None,
         }
     }

crate/lsp/src/completion/key_category.rs

@@ -44,7 +44,8 @@ pub enum KeyCategory {
     /// `tag_defaults` plus the tags defined in the document.
     TagFocus,
     /// Keys are `EntityType`s -- `theme_types_styles`. Offers the built-in
-    /// entity types.
+    /// entity types plus any custom `type_*` ids declared under
+    /// `entity_types` in the document.
     EntityType,
 }
 

crate/lsp/src/completion/yaml_lines.rs

@@ -21,6 +21,16 @@ pub fn is_list_item(line: &str) -> bool {
     trimmed == "-" || trimmed.starts_with("- ")
 }
 
+/// Splits the inside of a flow sequence (`a, b, c`) into trimmed, non-empty
+/// items.
+pub fn split_flow_items(inner: &str) -> Vec<&str> {
+    inner
+        .split(',')
+        .map(str::trim)
+        .filter(|item| !item.is_empty())
+        .collect()
+}
+
 /// Returns the map key declared on `line` (`key:` or `key: value`), or `None`
 /// for blank / comment / list-item / non-key lines.
 pub fn line_map_key(line: &str) -> Option<String> {

workspace_tests/src/lsp/completion_engine.rs

@@ -341,6 +341,31 @@ fn builtin_entity_types_offered_in_theme_types_styles() {
     );
 }
 
+#[test]
+fn custom_entity_types_offered_in_theme_types_styles() {
+    let text = "entity_types:\n  t_aws: [type_organisation]\n\
+        theme_types_styles:\n  ";
+    let labels = labels(text, 3, 2);
+
+    for expected in ["type_organisation", "type_thing_default"] {
+        assert!(
+            labels.iter().any(|label| label == expected),
+            "expected entity type `{expected}` in {labels:?}"
+        );
+    }
+}
+
+#[test]
+fn custom_entity_types_offered_in_entity_types_values() {
+    let text = "entity_types:\n  t_a: [type_organisation]\n  t_b:\n    - ";
+    let labels = labels(text, 3, 6);
+
+    assert!(
+        labels.iter().any(|label| label == "type_organisation"),
+        "expected custom entity type `type_organisation` in {labels:?}"
+    );
+}
+
 #[test]
 fn already_declared_keys_filtered_from_suggestions() {
     // `t_a` is already a `thing_names` key, so only `t_b` should remain.

workspace_tests/src/lsp/dynamic_completions.rs

@@ -181,3 +181,26 @@ fn key_suggestions_style_alias_offers_custom_placeholder() {
         dynamic_completions.key_suggestions(KeyCategory::StyleAlias)
     );
 }
+
+#[test]
+fn collects_entity_type_ids_from_flow_and_block_entity_types_lists() {
+    let text = "entity_types:\n  t_aws: [type_organisation]\n  t_github:\n    \
+        - type_organisation\n    - type_open_source\n";
+    let dynamic_completions = DynamicCompletions::from_text(text);
+
+    assert_eq!(
+        vec!["type_open_source", "type_organisation"],
+        dynamic_completions.ids_for(IdCategory::EntityType)
+    );
+}
+
+#[test]
+fn key_suggestions_entity_type_includes_custom_types_from_entity_types() {
+    let text = "entity_types:\n  t_aws: [type_organisation]\n";
+    let dynamic_completions = DynamicCompletions::from_text(text);
+
+    assert_eq!(
+        vec!["type_organisation".to_string()],
+        dynamic_completions.key_suggestions(KeyCategory::EntityType)
+    );
+}