ZenML 0.80.0 Migration and Project View Implementation

README.md

@@ -8,12 +8,12 @@ The ZenML VSCode extension seamlessly integrates with [ZenML](https://github.com
 
 ## Features
 
-- **Server, Stacks, and Pipeline Runs Views**: Interact directly with ML stacks, pipeline runs, and server configurations from the Activity Bar.
-- **DAG Visualization for Pipeline Runs**: Explore Directed Acyclic Graphs for each pipeline view directly directly on the Activity Bar.
+- **Server, Projects, Stacks, and Pipeline Runs Views**: Interact directly with ML stacks, pipeline runs, and server configurations from the Activity Bar.
+- **Project Management**: Explore, switch between, and manage your ZenML projects directly from the Activity Bar.
+- **DAG Visualization for Pipeline Runs**: Explore Directed Acyclic Graphs for each pipeline view directly on the Activity Bar.
 - **Python Tool Integration**: Utilizes a Language Server Protocol (LSP) server for real-time synchronization with the ZenML environment.
 - **Real-Time Configuration Monitoring**: Leverages `watchdog` to dynamically update configurations, keeping the extension in sync with your ZenML setup.
-- **Status Bar**: Display the current stack name and connection status. You can
-  also change your active stack from the status bar.
+- **Status Bar**: Display the current active stack, active project, and connection status. You can also change your active stack from the status bar.
 
 ## Getting Started
 
@@ -27,6 +27,7 @@ this extension and your Python version needs to be 3.8 or greater.
 ## Using ZenML in VSCode
 
 - **Manage Server Connections**: Connect or disconnect from ZenML servers and refresh server status.
+- **Project Operations**: Browse available projects, view project details, set active projects, and refresh project information.
 - **Stack Operations**: View stack details, register, update, delete, copy, or set active stacks directly from VSCode.
 - **Stack Component Operations**: View stack component details, register, update, or delete stack components directly from VSCode.
 - **Pipeline Runs**: Monitor and manage pipeline runs, including deleting runs from the system and rendering DAGs.
@@ -48,7 +49,7 @@ this extension and your Python version needs to be 3.8 or greater.
 ## Requirements
 
 - **ZenML Installation:** ZenML needs to be installed in the local Python environment associated with the Python interpreter selected in the current VS Code workspace. This extension interacts directly with your ZenML environment, so ensuring that ZenML is installed and properly configured is essential.
-- **ZenML Version**: This extension is fully compatible with ZenML versions 0.63.0 through 0.75.0. While newer versions may work, they haven't been officially tested. If you encounter any issues with newer versions, please report them.
+- **ZenML Version**: This extension is compatible with ZenML version 0.80.0 or newer. While it may work with different versions, we recommend keeping your ZenML up to date for the best experience.
 - **Python Version**: Python 3.8 or greater is required for the operation of the LSP server, which is a part of this extension.
 
 ## Feedback and Contributions

RELEASE.md

@@ -65,8 +65,8 @@ Once you've tested the release and everything looks good, you can publish it:
 1. **Tag the Release**
 
    ```bash
-   git tag -a v0.0.x -m "Version 0.0.x"
-   git push origin v0.0.x
+   git tag -a 0.0.x -m "Version 0.0.x"
+   git push origin 0.0.x
    ```
 
 2. **Create a GitHub Release**

bundled/tool/constants.py

@@ -15,13 +15,13 @@
 
 TOOL_MODULE_NAME = "zenml-python"
 TOOL_DISPLAY_NAME = "ZenML"
-MIN_ZENML_VERSION = "0.63.0"
-MAX_ZENML_VERSION = "0.75.0"
+MIN_ZENML_VERSION = "0.80.0"
 
 """Constants for ZenML Notifications and Events"""
 
 IS_ZENML_INSTALLED = "zenml/isInstalled"
 ZENML_CLIENT_INITIALIZED = "zenml/clientInitialized"
 ZENML_SERVER_CHANGED = "zenml/serverChanged"
 ZENML_STACK_CHANGED = "zenml/stackChanged"
+ZENML_PROJECT_CHANGED = "zenml/projectChanged"
 ZENML_REQUIREMENTS_NOT_MET = "zenml/requirementsNotMet"

bundled/tool/lsp_jsonrpc.py

@@ -17,19 +17,28 @@
 import io
 import json
 import pathlib
+import re
 import subprocess
 import threading
 import uuid
 from concurrent.futures import ThreadPoolExecutor
 from typing import BinaryIO, Dict, Optional, Sequence, Union, cast
 
+ANSI_ESCAPE_PATTERN = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+
 CONTENT_LENGTH = "Content-Length: "
 RUNNER_SCRIPT = str(pathlib.Path(__file__).parent / "lsp_runner.py")
 
 
 def to_str(text) -> str:
     """Convert bytes to string as needed."""
-    return text.decode("utf-8") if isinstance(text, bytes) else text
+    result = text.decode("utf-8") if isinstance(text, bytes) else text
+    return result
+
+
+def strip_ansi_codes(text: str) -> str:
+    """Remove ANSI escape sequences from text."""
+    return ANSI_ESCAPE_PATTERN.sub("", text)
 
 
 class StreamClosedException(Exception):
@@ -57,13 +66,30 @@ def write(self, data):
             raise StreamClosedException()
 
         with self._lock:
-            content = json.dumps(data)
+            if isinstance(data, dict):
+                # recursively clean string values in dict
+                clean_data = self._clean_ansi_in_data(data)
+            else:
+                clean_data = data
+
+            content = json.dumps(clean_data)
             length = len(content.encode("utf-8"))
             # Make sure we're writing a string, not bytes, to the TextIOWrapper
             message = f"{CONTENT_LENGTH}{length}\r\n\r\n{content}"
             self._writer.write(message)
             self._writer.flush()
 
+    def _clean_ansi_in_data(self, data):
+        """Recursively clean ANSI codes from string values in data structure."""
+        if isinstance(data, dict):
+            return {k: self._clean_ansi_in_data(v) for k, v in data.items()}
+        elif isinstance(data, list):
+            return [self._clean_ansi_in_data(item) for item in data]
+        elif isinstance(data, str):
+            return strip_ansi_codes(data)
+        else:
+            return data
+
 
 class JsonReader:
     """Manages reading JSON-RPC messages from stream."""
@@ -91,13 +117,18 @@ def read(self):
             line = to_str(self._readline()).strip()
 
         content = to_str(self._reader.read(length))
-        return json.loads(content)
+        # clean any ANSI escape sequences that might cause JSON parsing issues
+        clean_content = strip_ansi_codes(content)
+        return json.loads(clean_content)
 
     def _readline(self):
         line = self._reader.readline()
         if not line:
             raise EOFError
-        return line
+        # strip any ANSI color codes that might interfere with header parsing
+        if isinstance(line, bytes):
+            return line
+        return strip_ansi_codes(line)
 
 
 class JsonRpc:

bundled/tool/lsp_zenml.py

@@ -350,3 +350,43 @@ def get_run_artifact(wrapper_instance, args):
         def get_run_dag(wrapper_instance, args):
             """Gets graph data for a specified ZenML pipeline run"""
             return wrapper_instance.get_pipeline_run_graph(args)
+
+        @self.command(f"{TOOL_MODULE_NAME}.listWorkspaces")
+        @self.zenml_command(wrapper_name="workspaces_wrapper")
+        def list_workspaces(wrapper_instance, args):
+            """Lists workspaces from ZenML Pro"""
+            return wrapper_instance.list_workspaces(args)
+
+        @self.command(f"{TOOL_MODULE_NAME}.getActiveWorkspace")
+        @self.zenml_command(wrapper_name="workspaces_wrapper")
+        def get_active_workspace(wrapper_instance, *args, **kwargs):
+            """Gets the active workspace for the current user"""
+            return wrapper_instance.get_active_workspace()
+
+        @self.command(f"{TOOL_MODULE_NAME}.listProjects")
+        @self.zenml_command(wrapper_name="projects_wrapper")
+        def list_projects(wrapper_instance, args):
+            """Lists projects from ZenML"""
+            return wrapper_instance.list_projects(args)
+
+        @self.command(f"{TOOL_MODULE_NAME}.getActiveProject")
+        @self.zenml_command(wrapper_name="projects_wrapper")
+        def get_active_project(wrapper_instance, *args, **kwargs):
+            """Gets the active project for the current user"""
+            return wrapper_instance.get_active_project()
+
+        @self.command(f"{TOOL_MODULE_NAME}.setActiveProject")
+        @self.zenml_command(wrapper_name="projects_wrapper")
+        def set_active_project(wrapper_instance, args):
+            """Sets the active project for the current user"""
+            result = wrapper_instance.set_active_project(args)
+            if not isinstance(result, dict) or "error" not in result:
+                # Only send notification if successful
+                self.send_custom_notification("zenml/projectChanged", result["id"])
+            return result
+
+        @self.command(f"{TOOL_MODULE_NAME}.getProjectByName")
+        @self.zenml_command(wrapper_name="projects_wrapper")
+        def get_project_by_name(wrapper_instance, args):
+            """Gets a project by name"""
+            return wrapper_instance.get_project_by_name(args[0])

bundled/tool/type_hints.py

@@ -10,9 +10,20 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
-from typing import Any, Dict, List, Optional, TypedDict
+from typing import Any, Dict, List, Optional, Set, Tuple, TypedDict, Union
 from uuid import UUID
 
+MetadataType = Union[
+    str,
+    int,
+    float,
+    bool,
+    Dict[Any, Any],
+    List[Any],
+    Set[Any],
+    Tuple[Any, ...],
+]
+
 
 class StepArtifactBody(TypedDict):
     type: str
@@ -143,3 +154,62 @@ class ListFlavorsResponse(TypedDict):
     total_pages: int
     total: int
     items: List[Flavor]
+
+
+class PipelineRun(TypedDict):
+    id: str
+    name: str
+    status: str
+    stackName: str
+    pipelineName: str
+    startTime: Optional[str]
+    endTime: Optional[str]
+    config: Optional[Dict[str, Any]]
+    steps: Optional[Dict[str, Any]]
+
+
+class ListPipelineRunsResponse(TypedDict):
+    runs: List[PipelineRun]
+    total: int
+    total_pages: int
+    current_page: int
+    items_per_page: int
+    project_name: Optional[str]
+
+
+class Workspace(TypedDict):
+    id: str
+    name: str
+    description: Optional[str]
+    display_name: Optional[str]
+    organization_id: str
+    organization_name: str
+    status: str
+    zenml_version: str
+    zenml_server_url: str
+    dashboard_url: str
+    dashboard_organization_url: str
+
+
+class ListWorkspacesResponse(TypedDict):
+    workspaces: List[Workspace]
+    total: int
+    offset: int
+    limit: int
+
+
+class Project(TypedDict):
+    id: str
+    name: str
+    display_name: Optional[str]
+    created: Optional[str]
+    updated: Optional[str]
+    metadata: Optional[MetadataType]
+
+
+class ListProjectsResponse(TypedDict):
+    projects: List[Project]
+    total: int
+    total_pages: int
+    current_page: int
+    items_per_page: int

bundled/tool/zen_watcher.py

@@ -22,7 +22,7 @@
 from typing import Any, Optional
 
 import yaml
-from constants import ZENML_SERVER_CHANGED, ZENML_STACK_CHANGED
+from constants import ZENML_PROJECT_CHANGED, ZENML_SERVER_CHANGED, ZENML_STACK_CHANGED
 from lazy_import import suppress_stdout_temporarily
 from watchdog.events import FileSystemEventHandler
 from watchdog.observers import Observer
@@ -44,6 +44,7 @@ def __init__(self, lsp_server):
         self._timer: Optional[Timer] = None
         self.last_known_url: str = ""
         self.last_known_stack_id: str = ""
+        self.last_known_project_name: str = ""
         self.show_notification: bool = os.getenv("LS_SHOW_NOTIFICATION", "off") in [
             "onError",
             "onWarning",
@@ -66,9 +67,12 @@ def process_config_change(self, config_file_path: str):
 
                     new_url = config.get("store", {}).get("url", "")
                     new_stack_id = config.get("active_stack_id", "")
+                    new_project_name = config.get("active_project_name", "")
 
                     url_changed = new_url != self.last_known_url
                     stack_id_changed = new_stack_id != self.last_known_stack_id
+                    project_name_changed = new_project_name != self.last_known_project_name
+
                     # Send ZENML_SERVER_CHANGED if url changed
                     if url_changed:
                         server_details = {
@@ -81,10 +85,18 @@ def process_config_change(self, config_file_path: str):
                             server_details,
                         )
                         self.last_known_url = new_url
+
                     # Send ZENML_STACK_CHANGED if stack_id changed
                     if stack_id_changed:
                         self.LSP_SERVER.send_custom_notification(ZENML_STACK_CHANGED, new_stack_id)
                         self.last_known_stack_id = new_stack_id
+
+                    # Send ZENML_PROJECT_CHANGED if project_name changed
+                    if project_name_changed and new_project_name:
+                        self.LSP_SERVER.send_custom_notification(
+                            ZENML_PROJECT_CHANGED, new_project_name
+                        )
+                        self.last_known_project_name = new_project_name
             except (FileNotFoundError, PermissionError) as e:
                 self.log_error(f"Configuration file access error: {e} - {config_file_path}")
             except yaml.YAMLError as e:

bundled/tool/zenml_client.py

@@ -22,18 +22,25 @@ def __init__(self):
         server interactions, stacks, and pipeline runs.
         """
         # pylint: disable=wrong-import-position,import-error
-        from lazy_import import lazy_import
+        from lazy_import import lazy_import, suppress_stdout_temporarily
         from zenml_wrappers import (
             GlobalConfigWrapper,
             PipelineRunsWrapper,
+            ProjectsWrapper,
             StacksWrapper,
+            WorkspacesWrapper,
             ZenServerWrapper,
         )
 
-        self.client = lazy_import("zenml.client", "Client")()
+        # Suppress colorful warnings during client initialization
+        with suppress_stdout_temporarily():
+            self.client = lazy_import("zenml.client", "Client")()
+
         # initialize zenml library wrappers
         self.config_wrapper = GlobalConfigWrapper()
-        self.zen_server_wrapper = ZenServerWrapper(self.config_wrapper)
         self.stacks_wrapper = StacksWrapper(self.client)
         self.pipeline_runs_wrapper = PipelineRunsWrapper(self.client)
+        self.workspaces_wrapper = WorkspacesWrapper(self.client, self.config_wrapper)
+        self.projects_wrapper = ProjectsWrapper(self.client)
+        self.zen_server_wrapper = ZenServerWrapper(self.config_wrapper, self.projects_wrapper)
         self.initialized = True