avrae-ls 0.6.4__py3-none-any.whl → 0.7.0__py3-none-any.whl

avrae_ls/server.py

@@ -17,12 +17,12 @@ from .context import ContextBuilder
 from .diagnostics import DiagnosticProvider
 from .runtime import MockExecutor
 from .alias_preview import render_alias_command, simulate_command
-from .parser import find_draconic_blocks
+from .parser import is_alias_module_path
+from .source_context import build_source_context, block_for_line
 from .signature_help import load_signatures, signature_help_for_code
 from .completions import gather_suggestions, completion_items_for_position, hover_for_position
 from .code_actions import code_actions_for_document
 from .symbols import build_symbol_table, document_symbols, find_definition_range, find_references, range_for_word
-from .argument_parsing import apply_argument_parsing
 
 # Prefer package metadata so the server version matches the installed wheel.
 try:
@@ -134,14 +134,14 @@ async def did_change_config(server: AvraeLanguageServer, params: types.DidChange
 @ls.feature(types.TEXT_DOCUMENT_DOCUMENT_SYMBOL)
 def on_document_symbol(server: AvraeLanguageServer, params: types.DocumentSymbolParams):
     doc = server.workspace.get_text_document(params.text_document.uri)
-    symbols = document_symbols(doc.source)
+    symbols = document_symbols(doc.source, treat_as_module=_is_alias_module_document(doc))
     return symbols
 
 
 @ls.feature(types.TEXT_DOCUMENT_DEFINITION)
 def on_definition(server: AvraeLanguageServer, params: types.DefinitionParams):
     doc = server.workspace.get_text_document(params.text_document.uri)
-    table = build_symbol_table(doc.source)
+    table = build_symbol_table(doc.source, treat_as_module=_is_alias_module_document(doc))
     word = doc.word_at_position(params.position)
     rng = find_definition_range(table, word)
     if rng is None:
@@ -152,19 +152,26 @@ def on_definition(server: AvraeLanguageServer, params: types.DefinitionParams):
 @ls.feature(types.TEXT_DOCUMENT_REFERENCES)
 def on_references(server: AvraeLanguageServer, params: types.ReferenceParams):
     doc = server.workspace.get_text_document(params.text_document.uri)
-    table = build_symbol_table(doc.source)
+    is_module = _is_alias_module_document(doc)
+    table = build_symbol_table(doc.source, treat_as_module=is_module)
     word = doc.word_at_position(params.position)
     if not word or not table.lookup(word):
         return []
 
-    ranges = find_references(table, doc.source, word, include_declaration=params.context.include_declaration)
+    ranges = find_references(
+        table,
+        doc.source,
+        word,
+        include_declaration=params.context.include_declaration,
+        treat_as_module=is_module,
+    )
     return [types.Location(uri=params.text_document.uri, range=rng) for rng in ranges]
 
 
 @ls.feature(types.TEXT_DOCUMENT_PREPARE_RENAME)
 def on_prepare_rename(server: AvraeLanguageServer, params: types.PrepareRenameParams):
     doc = server.workspace.get_text_document(params.text_document.uri)
-    table = build_symbol_table(doc.source)
+    table = build_symbol_table(doc.source, treat_as_module=_is_alias_module_document(doc))
     word = doc.word_at_position(params.position)
     if not word or not table.lookup(word):
         return None
@@ -174,12 +181,13 @@ def on_prepare_rename(server: AvraeLanguageServer, params: types.PrepareRenamePa
 @ls.feature(types.TEXT_DOCUMENT_RENAME)
 def on_rename(server: AvraeLanguageServer, params: types.RenameParams):
     doc = server.workspace.get_text_document(params.text_document.uri)
-    table = build_symbol_table(doc.source)
+    is_module = _is_alias_module_document(doc)
+    table = build_symbol_table(doc.source, treat_as_module=is_module)
     word = doc.word_at_position(params.position)
     if not word or not table.lookup(word) or not params.new_name:
         return None
 
-    ranges = find_references(table, doc.source, word, include_declaration=True)
+    ranges = find_references(table, doc.source, word, include_declaration=True, treat_as_module=is_module)
     if not ranges:
         return None
     edits = [types.TextEdit(range=rng, new_text=params.new_name) for rng in ranges]
@@ -191,7 +199,7 @@ def on_workspace_symbol(server: AvraeLanguageServer, params: types.WorkspaceSymb
     symbols: list[types.SymbolInformation] = []
     query = (params.query or "").lower()
     for uri, doc in server.workspace.text_documents.items():
-        table = build_symbol_table(doc.source)
+        table = build_symbol_table(doc.source, treat_as_module=_is_alias_module_document(doc))
         for entry in table.entries:
             if query and query not in entry.name.lower():
                 continue
@@ -208,20 +216,18 @@ def on_workspace_symbol(server: AvraeLanguageServer, params: types.WorkspaceSymb
 @ls.feature(types.TEXT_DOCUMENT_SIGNATURE_HELP)
 def on_signature_help(server: AvraeLanguageServer, params: types.SignatureHelpParams):
     doc = server.workspace.get_text_document(params.text_document.uri)
-    source = apply_argument_parsing(doc.source)
-    blocks = find_draconic_blocks(source)
+    is_module = _is_alias_module_document(doc)
+    source_ctx = build_source_context(doc.source, is_module)
     pos = params.position
-    if not blocks:
-        return signature_help_for_code(source, pos.line, pos.character, server._signatures)
-
-    for block in blocks:
-        start = block.line_offset
-        end = block.line_offset + block.line_count
-        if start <= pos.line <= end:
-            rel_line = pos.line - start
-            help_ = signature_help_for_code(block.code, rel_line, pos.character, server._signatures)
-            if help_:
-                return help_
+    if not source_ctx.blocks:
+        return signature_help_for_code(source_ctx.prepared, pos.line, pos.character, server._signatures)
+
+    block = block_for_line(source_ctx.blocks, pos.line)
+    if block:
+        rel_line = pos.line - block.line_offset
+        help_ = signature_help_for_code(block.code, rel_line, pos.character, server._signatures)
+        if help_:
+            return help_
     return None
 
 
@@ -233,18 +239,16 @@ def on_completion(server: AvraeLanguageServer, params: types.CompletionParams):
     doc = server.workspace.get_text_document(params.text_document.uri)
     ctx_data = server.state.context_builder.build()
     suggestions = gather_suggestions(ctx_data, server.state.context_builder.gvar_resolver, server._signatures)
-    source = apply_argument_parsing(doc.source)
-    blocks = find_draconic_blocks(source)
+    is_module = _is_alias_module_document(doc)
+    source_ctx = build_source_context(doc.source, is_module)
     pos = params.position
-    if not blocks:
-        return completion_items_for_position(source, pos.line, pos.character, suggestions)
-
-    for block in blocks:
-        start = block.line_offset
-        end = block.line_offset + block.line_count
-        if start <= pos.line <= end:
-            rel_line = pos.line - start
-            return completion_items_for_position(block.code, rel_line, pos.character, suggestions)
+    if not source_ctx.blocks:
+        return completion_items_for_position(source_ctx.prepared, pos.line, pos.character, suggestions)
+
+    block = block_for_line(source_ctx.blocks, pos.line)
+    if block:
+        rel_line = pos.line - block.line_offset
+        return completion_items_for_position(block.code, rel_line, pos.character, suggestions)
     return []
 
 
@@ -253,24 +257,41 @@ def on_hover(server: AvraeLanguageServer, params: types.HoverParams):
     doc = server.workspace.get_text_document(params.text_document.uri)
     ctx_data = server.state.context_builder.build()
     pos = params.position
-    source = apply_argument_parsing(doc.source)
-    blocks = find_draconic_blocks(source)
-    if not blocks:
-        return hover_for_position(source, pos.line, pos.character, server._signatures, ctx_data, server.state.context_builder.gvar_resolver)
-
-    for block in blocks:
-        start = block.line_offset
-        end = block.line_offset + block.line_count
-        if start <= pos.line <= end:
-            rel_line = pos.line - start
-            return hover_for_position(block.code, rel_line, pos.character, server._signatures, ctx_data, server.state.context_builder.gvar_resolver)
+    is_module = _is_alias_module_document(doc)
+    source_ctx = build_source_context(doc.source, is_module)
+    if not source_ctx.blocks:
+        return hover_for_position(
+            source_ctx.prepared,
+            pos.line,
+            pos.character,
+            server._signatures,
+            ctx_data,
+            server.state.context_builder.gvar_resolver,
+        )
+
+    block = block_for_line(source_ctx.blocks, pos.line)
+    if block:
+        rel_line = pos.line - block.line_offset
+        return hover_for_position(
+            block.code,
+            rel_line,
+            pos.character,
+            server._signatures,
+            ctx_data,
+            server.state.context_builder.gvar_resolver,
+        )
     return None
 
 
 @ls.feature(types.TEXT_DOCUMENT_CODE_ACTION)
 def on_code_action(server: AvraeLanguageServer, params: types.CodeActionParams):
     doc = server.workspace.get_text_document(params.text_document.uri)
-    return code_actions_for_document(doc.source, params, server.workspace_root)
+    return code_actions_for_document(
+        doc.source,
+        params,
+        server.workspace_root,
+        treat_as_module=_is_alias_module_document(doc),
+    )
 
 
 @ls.command(RUN_ALIAS_COMMAND)
@@ -372,7 +393,10 @@ async def _publish_diagnostics(
     doc = server.workspace.get_text_document(uri)
     ctx_data = server.state.context_builder.build(profile)
     diags = await server.state.diagnostics.analyze(
-        doc.source, ctx_data, server.state.context_builder.gvar_resolver
+        doc.source,
+        ctx_data,
+        server.state.context_builder.gvar_resolver,
+        treat_as_module=_is_alias_module_document(doc),
     )
     if extra:
         diags.extend(extra)
@@ -443,5 +467,19 @@ def _find_using_range(source: str, module: str | None) -> types.Range | None:
     return None
 
 
+def _is_alias_module_document(doc: Any) -> bool:
+    language_id = getattr(doc, "language_id", None)
+    if language_id == "avrae-module":
+        return True
+    uri = getattr(doc, "uri", None)
+    if not isinstance(uri, str):
+        return False
+    try:
+        path = Path(uris.to_fs_path(uri))
+    except Exception:
+        return uri.endswith(".alias-module")
+    return is_alias_module_path(path)
+
+
 def create_server() -> AvraeLanguageServer:
     return ls

avrae_ls/source_context.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Sequence
+
+from .argument_parsing import apply_argument_parsing
+from .parser import DraconicBlock, find_draconic_blocks
+
+
+@dataclass(frozen=True)
+class SourceContext:
+    source: str
+    prepared: str
+    blocks: list[DraconicBlock]
+    treat_as_module: bool
+
+
+def build_source_context(source: str, treat_as_module: bool, *, apply_args: bool = True) -> SourceContext:
+    prepared = apply_argument_parsing(source) if apply_args and not treat_as_module else source
+    blocks = find_draconic_blocks(prepared, treat_as_module=treat_as_module)
+    return SourceContext(source=source, prepared=prepared, blocks=blocks, treat_as_module=treat_as_module)
+
+
+def block_for_line(blocks: Sequence[DraconicBlock], line: int) -> DraconicBlock | None:
+    for block in blocks:
+        start = block.line_offset
+        end = block.line_offset + block.line_count
+        if start <= line <= end:
+            return block
+    return None

avrae_ls/symbols.py

@@ -8,8 +8,9 @@ from typing import Dict, Iterable, List, Optional
 import draconic
 from lsprotocol import types
 
-from .argument_parsing import apply_argument_parsing
-from .parser import find_draconic_blocks
+from .source_context import build_source_context
+from .parser import wrap_draconic
+from .lsp_utils import range_from_positions, shift_range
 
 log = logging.getLogger(__name__)
 
@@ -35,20 +36,19 @@ class SymbolTable:
         return self._index.get(name)
 
 
-def build_symbol_table(source: str) -> SymbolTable:
+def build_symbol_table(source: str, *, treat_as_module: bool = False) -> SymbolTable:
     entries: list[SymbolEntry] = []
-    parsed_source = apply_argument_parsing(source)
-    blocks = find_draconic_blocks(parsed_source)
-    if not blocks:
-        entries.extend(_symbols_from_code(parsed_source, 0, 0))
+    source_ctx = build_source_context(source, treat_as_module)
+    if not source_ctx.blocks:
+        entries.extend(_symbols_from_code(source_ctx.prepared, 0, 0))
     else:
-        for block in blocks:
+        for block in source_ctx.blocks:
             entries.extend(_symbols_from_code(block.code, block.line_offset, block.char_offset))
     return SymbolTable(entries)
 
 
-def document_symbols(source: str) -> List[types.DocumentSymbol]:
-    table = build_symbol_table(source)
+def document_symbols(source: str, *, treat_as_module: bool = False) -> List[types.DocumentSymbol]:
+    table = build_symbol_table(source, treat_as_module=treat_as_module)
     return [
         types.DocumentSymbol(
             name=entry.name,
@@ -68,9 +68,13 @@ def find_definition_range(table: SymbolTable, name: str) -> types.Range | None:
 
 
 def find_references(
-    table: SymbolTable, source: str, name: str, include_declaration: bool = True
+    table: SymbolTable,
+    source: str,
+    name: str,
+    include_declaration: bool = True,
+    *,
+    treat_as_module: bool = False,
 ) -> List[types.Range]:
-    parsed_source = apply_argument_parsing(source)
     ranges: list[types.Range] = []
     entry = table.lookup(name)
     include_stores = include_declaration and entry is None
@@ -78,11 +82,11 @@ def find_references(
         if entry:
             ranges.append(entry.selection_range)
 
-    blocks = find_draconic_blocks(parsed_source)
-    if not blocks:
-        ranges.extend(_references_from_code(parsed_source, name, 0, 0, include_stores))
+    source_ctx = build_source_context(source, treat_as_module)
+    if not source_ctx.blocks:
+        ranges.extend(_references_from_code(source_ctx.prepared, name, 0, 0, include_stores))
     else:
-        for block in blocks:
+        for block in source_ctx.blocks:
             ranges.extend(
                 _references_from_code(block.code, name, block.line_offset, block.char_offset, include_stores)
             )
@@ -158,13 +162,14 @@ def _entry_from_node(node: ast.AST, line_offset: int = 0, char_offset: int = 0)
     else:
         return None
 
-    rng = _range_from_positions(
+    rng = range_from_positions(
         getattr(node, "lineno", 1),
         getattr(node, "col_offset", 0),
         getattr(node, "end_lineno", None),
         getattr(node, "end_col_offset", None),
+        ensure_nonempty=True,
     )
-    rng = _shift_range(rng, line_offset, char_offset)
+    rng = shift_range(rng, line_offset, char_offset)
     return SymbolEntry(name=name, kind=kind, range=rng, selection_range=rng)
 
 
@@ -179,11 +184,12 @@ class _ReferenceCollector(ast.NodeVisitor):
         if node.id == self._target:
             if isinstance(node.ctx, ast.Store) and not self._include_stores:
                 return
-            rng = _range_from_positions(
+            rng = range_from_positions(
                 getattr(node, "lineno", 1),
                 getattr(node, "col_offset", 0),
                 getattr(node, "end_lineno", None),
                 getattr(node, "end_col_offset", None),
+                ensure_nonempty=True,
             )
             self.ranges.append(rng)
         self.generic_visit(node)
@@ -205,7 +211,7 @@ def _references_from_code(
         collector.visit(node)
 
     local_offset = line_offset + offset_adjust
-    return [_shift_range(rng, local_offset, char_offset) for rng in collector.ranges]
+    return [shift_range(rng, local_offset, char_offset) for rng in collector.ranges]
 
 
 def _parse_draconic(code: str) -> tuple[list[ast.AST], int]:
@@ -213,7 +219,7 @@ def _parse_draconic(code: str) -> tuple[list[ast.AST], int]:
     try:
         return parser.parse(code), 0
     except draconic.DraconicSyntaxError:
-        wrapped, added = _wrap_draconic(code)
+        wrapped, added = wrap_draconic(code)
         try:
             return parser.parse(wrapped), -added
         except draconic.DraconicSyntaxError:
@@ -223,45 +229,6 @@ def _parse_draconic(code: str) -> tuple[list[ast.AST], int]:
         return [], 0
 
 
-def _range_from_positions(
-    lineno: int | None,
-    col_offset: int | None,
-    end_lineno: int | None,
-    end_col_offset: int | None,
-) -> types.Range:
-    start_line = max((lineno or 1) - 1, 0)
-    start_char = max(col_offset or 0, 0)
-    end_line = max(((end_lineno or lineno or 1) - 1), 0)
-    raw_end_char = end_col_offset if end_col_offset is not None else col_offset
-    end_char = max(raw_end_char or start_char, start_char)
-    if end_char <= start_char:
-        end_char = start_char + 1
-    return types.Range(
-        start=types.Position(line=start_line, character=start_char),
-        end=types.Position(line=end_line, character=end_char),
-    )
-
-
-def _shift_range(rng: types.Range, line_offset: int, char_offset: int = 0) -> types.Range:
-    start_char = rng.start.character + (char_offset if rng.start.line == 0 else 0)
-    end_char = rng.end.character + (char_offset if rng.end.line == 0 else 0)
-    if line_offset == 0:
-        return types.Range(
-            start=types.Position(line=rng.start.line, character=start_char),
-            end=types.Position(line=rng.end.line, character=end_char),
-        )
-    return types.Range(
-        start=types.Position(line=max(rng.start.line + line_offset, 0), character=max(start_char, 0)),
-        end=types.Position(line=max(rng.end.line + line_offset, 0), character=max(end_char, 0)),
-    )
-
-
-def _wrap_draconic(code: str) -> tuple[str, int]:
-    indented = "\n".join(f"    {line}" for line in code.splitlines())
-    wrapped = f"def __alias_main__():\n{indented}\n__alias_main__()"
-    return wrapped, 1
-
-
 def _dedupe_ranges(ranges: Iterable[types.Range]) -> List[types.Range]:
     seen = set()
     unique: list[types.Range] = []