From 8fde389664ca59773c38dc0ec1a434a98bc2428b Mon Sep 17 00:00:00 2001 From: Martin Bielik Date: Mon, 16 Dec 2024 23:46:17 +0100 Subject: command-type only roles --- py/context.py | 3 ++- py/roles.py | 10 +++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) (limited to 'py') diff --git a/py/context.py b/py/context.py index f1eacd6..87c3a14 100644 --- a/py/context.py +++ b/py/context.py @@ -92,7 +92,8 @@ def load_role_config(role): enhance_roles_with_custom_function(roles) - if not role in roles: + postfixes = ["", ".complete", ".edit", ".chat"] + if not any([f"{role}{postfix}" in roles for postfix in postfixes]): raise Exception(f"Role `{role}` not found") if is_deprecated_role_syntax(roles, role): diff --git a/py/roles.py b/py/roles.py index 37e5b4d..692ac6c 100644 --- a/py/roles.py +++ b/py/roles.py @@ -6,7 +6,7 @@ if "PYTEST_VERSION" in os.environ: roles_py_imported = True -def load_ai_role_names(): +def load_ai_role_names(command_type): roles_config_path = os.path.expanduser(vim.eval("g:vim_ai_roles_config_file")) if not os.path.exists(roles_config_path): raise Exception(f"Role config file does not exist: {roles_config_path}") @@ -16,6 +16,10 @@ def load_ai_role_names(): enhance_roles_with_custom_function(roles) - role_names = [name for name in roles.sections() if not '.' in name] + role_names = set() + for name in roles.sections(): + parts = name.split('.') + if len(parts) == 1 or parts[-1] == command_type: + role_names.add(parts[0]) - return role_names + return list(role_names) -- cgit v1.2.3