progress support

Author: bagel897

src/codegen/extensions/lsp/progress.py

@@ -0,0 +1,61 @@
+import uuid
+
+from lsprotocol import types
+from lsprotocol.types import ProgressToken
+from pygls.lsp.server import LanguageServer
+
+from codegen.sdk.codebase.progress.progress import Progress
+from codegen.sdk.codebase.progress.stub_task import StubTask
+from codegen.sdk.codebase.progress.task import Task
+
+
+class LSPTask(Task):
+    count: int | None
+
+    def __init__(self, server: LanguageServer, message: str, token: ProgressToken, count: int | None = None, create_token: bool = True) -> None:
+        self.token = token
+        if create_token:
+            server.work_done_progress.begin(self.token, types.WorkDoneProgressBegin(title=message))
+        self.server = server
+        self.message = message
+        self.count = count
+        self.create_token = create_token
+
+    def update(self, message: str, count: int | None = None) -> None:
+        if self.count is not None and count is not None:
+            percent = int(count * 100 / self.count)
+        else:
+            percent = None
+        self.server.work_done_progress.report(self.token, types.WorkDoneProgressReport(message=message, percentage=percent))
+
+    def end(self) -> None:
+        if self.create_token:
+            self.server.work_done_progress.end(self.token, value=types.WorkDoneProgressEnd())
+
+
+class LSPProgress(Progress[LSPTask | StubTask]):
+    initialized = False
+
+    def __init__(self, server: LanguageServer, initial_token: ProgressToken | None = None):
+        self.server = server
+        self.initial_token = initial_token
+        if initial_token is not None:
+            self.server.work_done_progress.begin(initial_token, types.WorkDoneProgressBegin(title="Parsing codebase..."))
+
+    def begin_with_token(self, message: str, token: ProgressToken, *, count: int | None = None) -> LSPTask:
+        return LSPTask(self.server, message, token, count, create_token=False)
+
+    def begin(self, message: str, count: int | None = None) -> LSPTask | StubTask:
+        if self.initialized:
+            token = str(uuid.uuid4())
+            self.server.work_done_progress.create(token).result()
+            return LSPTask(self.server, message, token, count, create_token=False)
+        elif self.initial_token is not None:
+            return self.begin_with_token(message, self.initial_token, count=None)
+        else:
+            return StubTask()
+
+    def finish_initialization(self) -> None:
+        self.initialized = True
+        if self.initial_token is not None:
+            self.server.work_done_progress.end(self.initial_token, value=types.WorkDoneProgressEnd())

src/codegen/extensions/lsp/protocol.py

@@ -2,10 +2,11 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from lsprotocol.types import INITIALIZE, InitializeParams, InitializeResult, WorkDoneProgressBegin, WorkDoneProgressEnd
+from lsprotocol.types import INITIALIZE, InitializeParams, InitializeResult
 from pygls.protocol import LanguageServerProtocol, lsp_method
 
 from codegen.extensions.lsp.io import LSPIO
+from codegen.extensions.lsp.progress import LSPProgress
 from codegen.extensions.lsp.utils import get_path
 from codegen.sdk.codebase.config import CodebaseConfig
 from codegen.sdk.core.codebase import Codebase
@@ -19,6 +20,7 @@ class CodegenLanguageServerProtocol(LanguageServerProtocol):
     _server: "CodegenLanguageServer"
 
     def _init_codebase(self, params: InitializeParams) -> None:
+        progress = LSPProgress(self._server, params.work_done_token)
         if params.root_path:
             root = Path(params.root_path)
         elif params.root_uri:
@@ -27,15 +29,12 @@ def _init_codebase(self, params: InitializeParams) -> None:
             root = os.getcwd()
         config = CodebaseConfig(feature_flags=CodebaseFeatureFlags(full_range_index=True))
         io = LSPIO(self.workspace)
-        self._server.codebase = Codebase(repo_path=str(root), config=config, io=io)
+        self._server.codebase = Codebase(repo_path=str(root), config=config, io=io, progress=progress)
         self._server.io = io
-        if params.work_done_token:
-            self._server.work_done_progress.end(params.work_done_token, WorkDoneProgressEnd(message="Parsing codebase..."))
+        progress.finish_initialization()
 
     @lsp_method(INITIALIZE)
     def lsp_initialize(self, params: InitializeParams) -> InitializeResult:
         ret = super().lsp_initialize(params)
-        if params.work_done_token:
-            self._server.work_done_progress.begin(params.work_done_token, WorkDoneProgressBegin(title="Parsing codebase..."))
         self._init_codebase(params)
         return ret

src/codegen/sdk/codebase/codebase_context.py

@@ -16,6 +16,7 @@
 from codegen.sdk.codebase.diff_lite import ChangeType, DiffLite
 from codegen.sdk.codebase.flagging.flags import Flags
 from codegen.sdk.codebase.io.file_io import FileIO
+from codegen.sdk.codebase.progress.stub_progress import StubProgress
 from codegen.sdk.codebase.transaction_manager import TransactionManager
 from codegen.sdk.codebase.validation import get_edges, post_reset_validation
 from codegen.sdk.core.autocommit import AutoCommit, commiter
@@ -39,6 +40,7 @@
     from codegen.git.repo_operator.repo_operator import RepoOperator
     from codegen.sdk.codebase.io.io import IO
     from codegen.sdk.codebase.node_classes.node_classes import NodeClasses
+    from codegen.sdk.codebase.progress.progress import Progress
     from codegen.sdk.core.dataclasses.usage import Usage
     from codegen.sdk.core.expressions import Expression
     from codegen.sdk.core.external_module import ExternalModule
@@ -111,16 +113,19 @@ class CodebaseContext:
     projects: list[ProjectConfig]
     unapplied_diffs: list[DiffLite]
     io: IO
+    progress: Progress
 
     def __init__(
         self,
         projects: list[ProjectConfig],
         config: CodebaseConfig = DefaultConfig,
         io: IO | None = None,
+        progress: Progress | None = None,
     ) -> None:
         """Initializes codebase graph and TransactionManager"""
         from codegen.sdk.core.parser import Parser
 
+        self.progress = progress or StubProgress()
         self._graph = PyDiGraph()
         self.filepath_idx = {}
         self._ext_module_idx = {}
@@ -371,7 +376,6 @@ def _process_diff_files(self, files_to_sync: Mapping[SyncType, list[Path]], incr
         skip_uncache = incremental and ((len(files_to_sync[SyncType.DELETE]) + len(files_to_sync[SyncType.REPARSE])) == 0)
         if not skip_uncache:
             uncache_all()
-
         # Step 0: Start the dependency manager and language engine if they exist
         # Start the dependency manager. This may or may not run asynchronously, depending on the implementation
         if self.dependency_manager is not None:
@@ -429,24 +433,29 @@ def _process_diff_files(self, files_to_sync: Mapping[SyncType, list[Path]], incr
             file = self.get_file(file_path)
             file.remove_internal_edges()
 
+        task = self.progress.begin("Reparsing updated files", count=len(files_to_sync[SyncType.REPARSE]))
         files_to_resolve = []
         # Step 4: Reparse updated files
-        for file_path in files_to_sync[SyncType.REPARSE]:
+        for idx, file_path in enumerate(files_to_sync[SyncType.REPARSE]):
+            task.update(f"Reparsing {self.to_relative(file_path)}", count=idx)
             file = self.get_file(file_path)
             to_resolve.extend(file.unparse(reparse=True))
             to_resolve = list(filter(lambda node: self.has_node(node.node_id) and node is not None, to_resolve))
             file.sync_with_file_content()
             files_to_resolve.append(file)
-
+        task.end()
         # Step 5: Add new files as nodes to graph (does not yet add edges)
-        for filepath in files_to_sync[SyncType.ADD]:
+        task = self.progress.begin("Adding new files", count=len(files_to_sync[SyncType.ADD]))
+        for idx, filepath in enumerate(files_to_sync[SyncType.ADD]):
+            task.update(f"Adding {self.to_relative(filepath)}", count=idx)
             content = self.io.read_text(filepath)
             # TODO: this is wrong with context changes
             if filepath.suffix in self.extensions:
                 file_cls = self.node_classes.file_cls
                 new_file = file_cls.from_content(filepath, content, self, sync=False, verify_syntax=False)
                 if new_file is not None:
                     files_to_resolve.append(new_file)
+        task.end()
         for file in files_to_resolve:
             to_resolve.append(file)
             to_resolve.extend(file.get_nodes())
@@ -474,27 +483,35 @@ def _process_diff_files(self, files_to_sync: Mapping[SyncType, list[Path]], incr
             self._computing = True
             try:
                 logger.info(f"> Computing import resolution edges for {counter[NodeType.IMPORT]} imports")
+                task = self.progress.begin("Resolving imports", count=counter[NodeType.IMPORT])
                 for node in to_resolve:
                     if node.node_type == NodeType.IMPORT:
+                        task.update(f"Resolving imports in {node.filepath}", count=idx)
                         node._remove_internal_edges(EdgeType.IMPORT_SYMBOL_RESOLUTION)
                         node.add_symbol_resolution_edge()
                         to_resolve.extend(node.symbol_usages)
+                task.end()
                 if counter[NodeType.EXPORT] > 0:
                     logger.info(f"> Computing export dependencies for {counter[NodeType.EXPORT]} exports")
+                    task = self.progress.begin("Computing export dependencies", count=counter[NodeType.EXPORT])
                     for node in to_resolve:
                         if node.node_type == NodeType.EXPORT:
+                            task.update(f"Computing export dependencies for {node.filepath}", count=idx)
                             node._remove_internal_edges(EdgeType.EXPORT)
                             node.compute_export_dependencies()
                             to_resolve.extend(node.symbol_usages)
+                    task.end()
                 if counter[NodeType.SYMBOL] > 0:
                     from codegen.sdk.core.interfaces.inherits import Inherits
 
                     logger.info("> Computing superclass dependencies")
+                    task = self.progress.begin("Computing superclass dependencies", count=counter[NodeType.SYMBOL])
                     for symbol in to_resolve:
                         if isinstance(symbol, Inherits):
+                            task.update(f"Computing superclass dependencies for {symbol.filepath}", count=idx)
                             symbol._remove_internal_edges(EdgeType.SUBCLASS)
                             symbol.compute_superclass_dependencies()
-
+                    task.end()
                 if not skip_uncache:
                     uncache_all()
                 self._compute_dependencies(to_resolve, incremental)
@@ -504,17 +521,20 @@ def _process_diff_files(self, files_to_sync: Mapping[SyncType, list[Path]], incr
     def _compute_dependencies(self, to_update: list[Importable], incremental: bool):
         seen = set()
         while to_update:
+            task = self.progress.begin("Computing dependencies", count=len(to_update))
             step = to_update.copy()
             to_update.clear()
             logger.info(f"> Incrementally computing dependencies for {len(step)} nodes")
-            for current in step:
+            for idx, current in enumerate(step):
+                task.update(f"Computing dependencies for {current.filepath}", count=idx)
                 if current not in seen:
                     seen.add(current)
                     to_update.extend(current.recompute(incremental))
             if not incremental:
                 for node in self._graph.nodes():
                     if node not in seen:
                         to_update.append(node)
+            task.end()
         seen.clear()
 
     def build_subgraph(self, nodes: list[NodeId]) -> PyDiGraph[Importable, Edge]:

src/codegen/sdk/codebase/progress/progress.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+if TYPE_CHECKING:
+    from codegen.sdk.codebase.progress.task import Task
+
+T = TypeVar("T", bound="Task")
+
+
+class Progress(ABC, Generic[T]):
+    @abstractmethod
+    def begin(self, message: str, count: int | None = None) -> T:
+        pass

src/codegen/sdk/codebase/progress/stub_progress.py

@@ -0,0 +1,7 @@
+from codegen.sdk.codebase.progress.progress import Progress
+from codegen.sdk.codebase.progress.stub_task import StubTask
+
+
+class StubProgress(Progress[StubTask]):
+    def begin(self, message: str, count: int | None = None) -> StubTask:
+        return StubTask()

src/codegen/sdk/codebase/progress/stub_task.py

@@ -0,0 +1,9 @@
+from codegen.sdk.codebase.progress.task import Task
+
+
+class StubTask(Task):
+    def update(self, message: str, count: int | None = None) -> None:
+        pass
+
+    def end(self) -> None:
+        pass

src/codegen/sdk/codebase/progress/task.py

@@ -0,0 +1,11 @@
+from abc import ABC, abstractmethod
+
+
+class Task(ABC):
+    @abstractmethod
+    def update(self, message: str, count: int | None = None) -> None:
+        pass
+
+    @abstractmethod
+    def end(self) -> None:
+        pass

src/codegen/sdk/core/codebase.py

@@ -36,6 +36,7 @@
 from codegen.sdk.codebase.flagging.enums import FlagKwargs
 from codegen.sdk.codebase.flagging.group import Group
 from codegen.sdk.codebase.io.io import IO
+from codegen.sdk.codebase.progress.progress import Progress
 from codegen.sdk.codebase.span import Span
 from codegen.sdk.core.assignment import Assignment
 from codegen.sdk.core.class_definition import Class
@@ -132,6 +133,7 @@ def __init__(
         projects: list[ProjectConfig] | ProjectConfig,
         config: CodebaseConfig = DefaultConfig,
         io: IO | None = None,
+        progress: Progress | None = None,
     ) -> None: ...
 
     @overload
@@ -143,6 +145,7 @@ def __init__(
         projects: None = None,
         config: CodebaseConfig = DefaultConfig,
         io: IO | None = None,
+        progress: Progress | None = None,
     ) -> None: ...
 
     def __init__(
@@ -153,6 +156,7 @@ def __init__(
         projects: list[ProjectConfig] | ProjectConfig | None = None,
         config: CodebaseConfig = DefaultConfig,
         io: IO | None = None,
+        progress: Progress | None = None,
     ) -> None:
         # Sanity check inputs
         if repo_path is not None and projects is not None:
@@ -182,7 +186,7 @@ def __init__(
         self._op = main_project.repo_operator
         self.viz = VisualizationManager(op=self._op)
         self.repo_path = Path(self._op.repo_path)
-        self.ctx = CodebaseContext(projects, config=config, io=io)
+        self.ctx = CodebaseContext(projects, config=config, io=io, progress=progress)
         self.console = Console(record=True, soft_wrap=True)
 
     @noapidoc