diff --git a/.eslintrc b/.eslintrc new file mode 100644 index 00000000..ac291699 --- /dev/null +++ b/.eslintrc @@ -0,0 +1,19 @@ +{ + "parser": "@typescript-eslint/parser", + "parserOptions": { + "ecmaVersion": 12, + "sourceType": "module" + }, + "plugins": ["@typescript-eslint"], + "extends": ["eslint:recommended", "plugin:@typescript-eslint/recommended", "prettier"], + + "rules": { + "@typescript-eslint/no-unused-vars": "error", + "@typescript-eslint/consistent-type-definitions": ["error", "type"] + }, + + "env": { + "browser": true, + "es2021": true + } +} \ No newline at end of file diff --git a/.eslintrc.json b/.eslintrc.json new file mode 100644 index 00000000..bb8a0015 --- /dev/null +++ b/.eslintrc.json @@ -0,0 +1,24 @@ +{ + "root": true, + "parser": "@typescript-eslint/parser", + "parserOptions": { + "ecmaVersion": 6, + "sourceType": "module" + }, + "plugins": ["@typescript-eslint"], + "rules": { + "@typescript-eslint/naming-convention": [ + "warn", + { + "selector": "import", + "format": ["camelCase", "PascalCase"] + } + ], + "@typescript-eslint/semi": "warn", + "curly": "warn", + "eqeqeq": "warn", + "no-throw-literal": "warn", + "semi": "off" + }, + "ignorePatterns": ["out", "dist", "**/*.d.ts"] +} diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 00000000..e869ba4a --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,33 @@ +name: VSCode Extension CI + +on: [push, pull_request] + +jobs: + build-and-test: + runs-on: ubuntu-latest + + steps: + - name: Checkout Repository + uses: actions/checkout@v2 + + - name: Install Node.js + uses: actions/setup-node@v2 + with: + node-version: '18.x' + + - name: Install dependencies + run: npm install + + - name: Build VSCode Extension + run: npm run compile + + - name: Run VSCode Extension Formatter + run: npm run format + + - name: Run VSCode Extension Linter + run: npm run lint + + - name: Run headless test + uses: coactions/setup-xvfb@v1 + with: + run: npm test diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..f394924c --- /dev/null +++ b/.gitignore @@ -0,0 +1,40 @@ +.DS_Store +Thumbs.db +.vscode-test/ +.idea/ +.eslintcache +lib-cov +*.log +*.log* +pids + +# Node +.npm/ +node_modules/ +package-lock.json +npm-debug.log + +# Build outputs +dist/ +out/ +build/ +*.tsbuildinfo +.history/ + +# env +.env +.env* + +# Output of 'npm pack' +*.tgz + +# From vscode-python-tools-extension-template +*.vsix +.venv/ +.vs/ +.nox/ +bundled/libs/ +**/__pycache__ +**/.pytest_cache +**/.vs + diff --git a/.prettierrc b/.prettierrc new file mode 100644 index 00000000..fbe6ca4f --- /dev/null +++ b/.prettierrc @@ -0,0 +1,11 @@ +{ + "printWidth": 100, + "tabWidth": 2, + "useTabs": false, + "semi": true, + "proseWrap": "preserve", + "singleQuote": true, + "arrowParens": "avoid", + "trailingComma": "es5", + "bracketSpacing": true +} diff --git a/.pylintrc b/.pylintrc new file mode 100644 index 00000000..8749d63a --- /dev/null +++ b/.pylintrc @@ -0,0 +1,7 @@ +[MESSAGES CONTROL] +disable= + C0103, # Name doesn't conform to naming style + C0415, # Import outside toplevel + W0613, # Unused argument + R0801, # Similar lines in multiple files + R0903, # Too few public methods \ No newline at end of file diff --git a/.vscode-test.mjs b/.vscode-test.mjs new file mode 100644 index 00000000..b62ba25f --- /dev/null +++ b/.vscode-test.mjs @@ -0,0 +1,5 @@ +import { defineConfig } from '@vscode/test-cli'; + +export default defineConfig({ + files: 'out/test/**/*.test.js', +}); diff --git a/.vscode/extensions.json b/.vscode/extensions.json new file mode 100644 index 00000000..9e5e0b73 --- /dev/null +++ b/.vscode/extensions.json @@ -0,0 +1,9 @@ +{ + // See http://go.microsoft.com/fwlink/?LinkId=827846 + // for the documentation about the extensions.json format + "recommendations": [ + "dbaeumer.vscode-eslint", + "amodio.tsl-problem-matcher", + "ms-vscode.extension-test-runner" + ] +} diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 00000000..fe54ee0f --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,18 @@ +// A launch configuration that compiles the extension and then opens it inside a new window +// Use IntelliSense to learn about possible attributes. +// Hover to view descriptions of existing attributes. +// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 +{ + "version": "0.2.0", + "configurations": [ + { + "name": "Run Extension", + "type": "extensionHost", + "request": "launch", + "args": ["--extensionDevelopmentPath=${workspaceFolder}"], + "outFiles": ["${workspaceFolder}/dist/**/*.js"], + "preLaunchTask": "${defaultBuildTask}", + "timeout": 20000 + } + ] +} diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 00000000..0e65fd08 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,19 @@ +// Place your settings in this file to overwrite default and user settings. +{ + "python.defaultInterpreterPath": "", + "files.exclude": { + "out": false, // set this to true to hide the "out" folder with the compiled JS files + "dist": false // set this to true to hide the "dist" folder with the compiled JS files + }, + "search.exclude": { + "out": true, // set this to false to include "out" folder in search results + "dist": true // set this to false to include "dist" folder in search results + }, + // Turn off tsc task auto detection since we have the necessary tasks as npm scripts + "typescript.tsc.autoDetect": "off", + "python.testing.pytestArgs": ["src/test/python_tests"], + "python.testing.unittestEnabled": false, + "python.testing.pytestEnabled": true, + "python.testing.cwd": "${workspaceFolder}", + "python.analysis.extraPaths": ["bundled/libs", "bundled/tool"] +} diff --git a/.vscode/tasks.json b/.vscode/tasks.json new file mode 100644 index 00000000..59cb4ae2 --- /dev/null +++ b/.vscode/tasks.json @@ -0,0 +1,45 @@ +// See https://go.microsoft.com/fwlink/?LinkId=733558 +// for the documentation about the tasks.json format +{ + "version": "2.0.0", + "tasks": [ + { + "type": "npm", + "script": "watch", + "problemMatcher": "$ts-webpack-watch", + "isBackground": true, + "presentation": { + "reveal": "never", + "group": "watchers" + }, + "group": { + "kind": "build", + "isDefault": true + } + }, + { + "type": "npm", + "script": "watch-tests", + "problemMatcher": "$tsc-watch", + "isBackground": true, + "presentation": { + "reveal": "never", + "group": "watchers" + }, + "group": "build" + }, + { + "label": "tasks: watch-tests", + "dependsOn": ["npm: watch", "npm: watch-tests"], + "problemMatcher": [] + }, + { + "type": "npm", + "script": "compile", + "group": "build", + "problemMatcher": [], + "label": "npm: compile", + "detail": "webpack" + } + ] +} diff --git a/.vscodeignore b/.vscodeignore new file mode 100644 index 00000000..fa9763b0 --- /dev/null +++ b/.vscodeignore @@ -0,0 +1,15 @@ +.vscode/** +.vscode-test/** +out/** +node_modules/** +src/** +.gitignore +.yarnrc +webpack.config.js +vsc-extension-quickstart.md +**/tsconfig.json +**/.eslintrc.json +**/*.map +**/*.js.map +**/*.ts +**/.vscode-test.* \ No newline at end of file diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 00000000..eeaf678a --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,100 @@ +# Contributing to ZenML VSCode Extension + +We appreciate your interest in contributing to the ZenML VSCode extension! This guide will help you get started with setting up your development environment, making changes, and proposing those changes back to the project. By following these guidelines, you'll ensure a smooth and efficient contribution process. + +## Setting Up Your Development Environment + +1. **Fork and Clone**: Fork the [zenml-io/vscode-zenml repository](https://github.com/zenml-io/vscode-zenml) and clone it to your local machine. + +```bash +git clone https://github.com/YOUR_USERNAME/vscode-zenml.git +git checkout develop +``` + +2. **Install Dependencies**: Navigate to the cloned repository directory and install the required dependencies. + +```bash +cd vscode-zenml +npm install +``` + +3. **Compile the Project**: Build the TypeScript source code into JavaScript. + +```bash +npm run compile +``` + +### Python Environment Setup + +The extension's Python functionality requires setting up a corresponding Python environment. + +1. Create and activate a Python virtual environment using Python 3.8 or greater. (e.g., `python -m venv .venv` on Windows, or `python3 -m venv .venv` on Unix-based systems). +2. Install `nox` for environment management. + +```bash +python -m pip install nox +``` + +3. Use `nox` to set up the environment. + +```bash +nox --session setup +``` + +4. Install any Python dependencies specified in `requirements.txt`. + +```bash +pip install -r requirements.txt +``` + +## Development Workflow + +- **Running the Extension**: Press `F5` to open a new VSCode window with the + extension running, or click the `Start Debugging` button in the VSCode menubar + under the `Run` menu. +- **Making Changes**: Edit the source code. The TypeScript code is in the `src` directory, and Python logic for the LSP server is in `bundled/tool`. + +### Testing + +- **Writing Tests**: Add tests in the `src/test` directory. Follow the naming convention `*.test.ts` for test files. +- **Running Tests**: Use the provided npm scripts to compile and run tests. + +```bash +npm run test +``` + +### Debugging + +- **VSCode Debug Console**: Utilize the debug console in the VSCode development window for troubleshooting and inspecting values. +- **Extension Host Logs**: Review the extension host logs for runtime errors or unexpected behavior. + +## Contributing Changes + +1. **Create a Branch**: Make your changes in a new git branch based on the `develop` branch. + +```bash +git checkout -b feature/your-feature-name +``` + +2. **Commit Your Changes**: Write clear, concise commit messages following the [conventional commit](https://www.conventionalcommits.org/en/v1.0.0/) guidelines. +3. **Push to Your Fork**: Push your branch to your fork on GitHub. + +```bash +git push origin feature/your-feature-name +``` + +4. **Open a Pull Request**: Go to the original `zenml-io/vscode-zenml` repository and create a pull request from your feature branch. Please follow our [contribution guidelines](https://github.com/zenml-io/zenml/blob/develop/CONTRIBUTING.md) for more details on proposing pull requests. + +## Troubleshooting Common Issues + +- Ensure all dependencies are up to date and compatible. +- Rebuild the project (`npm run compile`) after making changes. +- Reset your development environment if encountering persistent issues by re-running `nox` setup commands and reinstalling dependencies. +- You can also run the `scripts/clear_and_compile.sh` script, which will delete the cache, `dist` folder, and recompile automatically. +- Check the [ZenML documentation](https://docs.zenml.io) and [GitHub issues](https://github.com/zenml-io/zenml/issues) for common problems and solutions. + +### Additional Resources + +- [ZenML VSCode Extension Repository](https://github.com/zenml-io/vscode-zenml) +- [ZenML Documentation](https://docs.zenml.io) +- [ZenML Slack Community](https://zenml.io/slack-invite) diff --git a/README.md b/README.md index ae7c0768..b9592a49 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,54 @@ -# vscode-zenml -VSCode extension for ZenML +# ZenML Extension for Visual Studio Code + +![](https://img.shields.io/github/license/zenml-io/vscode-zenml) + +![](resources/zenml-extension.gif) + +The ZenML VSCode extension seamlessly integrates with [ZenML](https://github.com/zenml-io/zenml) to enhance your MLOps workflow within VSCode. It is designed to accurately mirror the current state of your ZenML environment within your IDE, ensuring a smooth and integrated experience. + +## Features + +- **Server, Stacks, and Pipeline Runs Views**: Interact directly with ML stacks, pipeline runs, and server configurations from 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 Indicators**: Display the current stack name and connection status. + +## Getting Started + +Note that you'll need to have [ZenML](https://github.com/zenml-io/zenml) installed in your Python environment to use +this extension and your Python version needs to be 3.8 or greater. + +1. **Install the Extension**: Search for "ZenML" in the VSCode Extensions view (`Ctrl+Shift+X`) and install it. +2. **Connect to ZenML Server**: Use the `ZenML: Connect` command to connect to your ZenML server. +3. **Explore ZenML Views**: Navigate to the ZenML activity bar to access the Server, Stacks, and Pipeline Runs views. + +## Using ZenML in VSCode + +- **Manage Server Connections**: Connect or disconnect from ZenML servers and refresh server status. +- **Stack Operations**: View stack details, rename, copy, or set active stacks directly from VSCode. +- **Pipeline Runs**: Monitor and manage pipeline runs, including deleting runs from the system. +- **Environment Information**: Get detailed snapshots of the development environment, aiding troubleshooting. + +## 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**: To ensure full functionality and compatibility, make sure you have ZenML version 0.55.0 or newer. +- **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 + +Your feedback and contributions are welcome! Please refer to our [contribution +guidelines](https://github.com/zenml-io/vscode-zenml/CONTRIBUTING.md) for more +information. + +For any further questions or issues, please reach out to us in our [Slack +Community](https://zenml.io/slack-invite). To learn more about ZenML, +please visit our [website](https://zenml.io/) and read [our documentation](https://docs.zenml.io). + +## License + +Apache-2.0 + +--- + +ZenML © 2024, ZenML. Released under the [Apache-2.0 License](LICENSE). diff --git a/bundled/tool/__init__.py b/bundled/tool/__init__.py new file mode 100644 index 00000000..d83bd1ae --- /dev/null +++ b/bundled/tool/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. diff --git a/bundled/tool/_debug_server.py b/bundled/tool/_debug_server.py new file mode 100644 index 00000000..5b9cb755 --- /dev/null +++ b/bundled/tool/_debug_server.py @@ -0,0 +1,52 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Debugging support for LSP.""" + +import os +import pathlib +import runpy +import sys + + +def update_sys_path(path_to_add: str) -> None: + """Add given path to `sys.path`.""" + if path_to_add not in sys.path and os.path.isdir(path_to_add): + sys.path.append(path_to_add) + + +# Ensure debugger is loaded before we load anything else, to debug initialization. +debugger_path = os.getenv("DEBUGPY_PATH", None) +if debugger_path: + if debugger_path.endswith("debugpy"): + debugger_path = os.fspath(pathlib.Path(debugger_path).parent) + + # pylint: disable=wrong-import-position,import-error + import debugpy + + update_sys_path(debugger_path) + + # pylint: disable=wrong-import-position,import-error + + # 5678 is the default port, If you need to change it update it here + # and in launch.json. + debugpy.connect(5678) + + # This will ensure that execution is paused as soon as the debugger + # connects to VS Code. If you don't want to pause here comment this + # line and set breakpoints as appropriate. + debugpy.breakpoint() + +SERVER_PATH = os.fspath(pathlib.Path(__file__).parent / "lsp_server.py") +# NOTE: Set breakpoint in `lsp_server.py` before continuing. +runpy.run_path(SERVER_PATH, run_name="__main__") diff --git a/bundled/tool/constants.py b/bundled/tool/constants.py new file mode 100644 index 00000000..a3c87a7e --- /dev/null +++ b/bundled/tool/constants.py @@ -0,0 +1,26 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Constants for ZenML Tool""" + +TOOL_MODULE_NAME = "zenml-python" +TOOL_DISPLAY_NAME = "ZenML" +MIN_ZENML_VERSION = "0.55.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_REQUIREMENTS_NOT_MET = "zenml/requirementsNotMet" diff --git a/bundled/tool/lazy_import.py b/bundled/tool/lazy_import.py new file mode 100644 index 00000000..2cce2c4b --- /dev/null +++ b/bundled/tool/lazy_import.py @@ -0,0 +1,87 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +""" +Utilities for lazy importing and temporary suppression of stdout and logging. +Reduces noise when integrating with logging-heavy systems like LSP. +Useful for quieter module loads and command executions. +""" + +import importlib +import logging +import os +import sys +from contextlib import contextmanager + + +@contextmanager +def suppress_logging_temporarily(level=logging.ERROR): + """ + Temporarily elevates logging level and suppresses stdout to + minimize console output during imports. + + Parameters: + level (int): Temporary logging level (default: ERROR). + + Yields: + None: While suppressing stdout. + """ + original_level = logging.root.level + original_stdout = sys.stdout + logging.root.setLevel(level) + with open(os.devnull, "w", encoding="utf-8") as fnull: + sys.stdout = fnull + try: + yield + finally: + sys.stdout = original_stdout + logging.root.setLevel(original_level) + + +@contextmanager +def suppress_stdout_temporarily(): + """ + This context manager suppresses stdout for LSP commands, + silencing unnecessary or unwanted output during execution. + + Yields: + None: While suppressing stdout. + """ + with open(os.devnull, "w", encoding="utf-8") as fnull: + original_stdout = sys.stdout + original_stderr = sys.stderr + sys.stdout = fnull + sys.stderr = fnull + try: + yield + finally: + sys.stdout = original_stdout + sys.stderr = original_stderr + + +def lazy_import(module_name, class_name=None): + """ + Lazily imports a module or class, suppressing ZenML log output + to minimize initialization time and noise. + + Args: + module_name (str): The name of the module to import. + class_name (str, optional): The class name within the module. Defaults to None. + + Returns: + The imported module or class. + """ + with suppress_logging_temporarily(): + module = importlib.import_module(module_name) + if class_name: + return getattr(module, class_name) + return module diff --git a/bundled/tool/lsp_jsonrpc.py b/bundled/tool/lsp_jsonrpc.py new file mode 100644 index 00000000..7eb79472 --- /dev/null +++ b/bundled/tool/lsp_jsonrpc.py @@ -0,0 +1,267 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Light-weight JSON-RPC over standard IO.""" + + +import atexit +import io +import json +import pathlib +import subprocess +import threading +import uuid +from concurrent.futures import ThreadPoolExecutor +from typing import BinaryIO, Dict, Optional, Sequence, Union + +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 + + +class StreamClosedException(Exception): + """JSON RPC stream is closed.""" + + pass # pylint: disable=unnecessary-pass + + +class JsonWriter: + """Manages writing JSON-RPC messages to the writer stream.""" + + def __init__(self, writer: io.TextIOWrapper): + self._writer = writer + self._lock = threading.Lock() + + def close(self): + """Closes the underlying writer stream.""" + with self._lock: + if not self._writer.closed: + self._writer.close() + + def write(self, data): + """Writes given data to stream in JSON-RPC format.""" + if self._writer.closed: + raise StreamClosedException() + + with self._lock: + content = json.dumps(data) + length = len(content.encode("utf-8")) + self._writer.write(f"{CONTENT_LENGTH}{length}\r\n\r\n{content}".encode("utf-8")) + self._writer.flush() + + +class JsonReader: + """Manages reading JSON-RPC messages from stream.""" + + def __init__(self, reader: io.TextIOWrapper): + self._reader = reader + + def close(self): + """Closes the underlying reader stream.""" + if not self._reader.closed: + self._reader.close() + + def read(self): + """Reads data from the stream in JSON-RPC format.""" + if self._reader.closed: + raise StreamClosedException + length = None + while not length: + line = to_str(self._readline()) + if line.startswith(CONTENT_LENGTH): + length = int(line[len(CONTENT_LENGTH) :]) + + line = to_str(self._readline()).strip() + while line: + line = to_str(self._readline()).strip() + + content = to_str(self._reader.read(length)) + return json.loads(content) + + def _readline(self): + line = self._reader.readline() + if not line: + raise EOFError + return line + + +class JsonRpc: + """Manages sending and receiving data over JSON-RPC.""" + + def __init__(self, reader: io.TextIOWrapper, writer: io.TextIOWrapper): + self._reader = JsonReader(reader) + self._writer = JsonWriter(writer) + + def close(self): + """Closes the underlying streams.""" + try: + self._reader.close() + except: # pylint: disable=bare-except + pass + try: + self._writer.close() + except: # pylint: disable=bare-except + pass + + def send_data(self, data): + """Send given data in JSON-RPC format.""" + self._writer.write(data) + + def receive_data(self): + """Receive data in JSON-RPC format.""" + return self._reader.read() + + +def create_json_rpc(readable: BinaryIO, writable: BinaryIO) -> JsonRpc: + """Creates JSON-RPC wrapper for the readable and writable streams.""" + return JsonRpc(readable, writable) + + +class ProcessManager: + """Manages sub-processes launched for running tools.""" + + def __init__(self): + self._args: Dict[str, Sequence[str]] = {} + self._processes: Dict[str, subprocess.Popen] = {} + self._rpc: Dict[str, JsonRpc] = {} + self._lock = threading.Lock() + self._thread_pool = ThreadPoolExecutor(10) + + def stop_all_processes(self): + """Send exit command to all processes and shutdown transport.""" + for i in self._rpc.values(): + try: + i.send_data({"id": str(uuid.uuid4()), "method": "exit"}) + except: # pylint: disable=bare-except + pass + self._thread_pool.shutdown(wait=False) + + def start_process(self, workspace: str, args: Sequence[str], cwd: str) -> None: + """Starts a process and establishes JSON-RPC communication over stdio.""" + # pylint: disable=consider-using-with + proc = subprocess.Popen( + args, + cwd=cwd, + stdout=subprocess.PIPE, + stdin=subprocess.PIPE, + ) + self._processes[workspace] = proc + self._rpc[workspace] = create_json_rpc(proc.stdout, proc.stdin) + + def _monitor_process(): + proc.wait() + with self._lock: + try: + del self._processes[workspace] + rpc = self._rpc.pop(workspace) + rpc.close() + except: # pylint: disable=bare-except + pass + + self._thread_pool.submit(_monitor_process) + + def get_json_rpc(self, workspace: str) -> JsonRpc: + """Gets the JSON-RPC wrapper for the a given id.""" + with self._lock: + if workspace in self._rpc: + return self._rpc[workspace] + raise StreamClosedException() + + +_process_manager = ProcessManager() +atexit.register(_process_manager.stop_all_processes) + + +def _get_json_rpc(workspace: str) -> Union[JsonRpc, None]: + try: + return _process_manager.get_json_rpc(workspace) + except StreamClosedException: + return None + except KeyError: + return None + + +def get_or_start_json_rpc( + workspace: str, interpreter: Sequence[str], cwd: str +) -> Union[JsonRpc, None]: + """Gets an existing JSON-RPC connection or starts one and return it.""" + res = _get_json_rpc(workspace) + if not res: + args = [*interpreter, RUNNER_SCRIPT] + _process_manager.start_process(workspace, args, cwd) + res = _get_json_rpc(workspace) + return res + + +class RpcRunResult: + """Object to hold result from running tool over RPC.""" + + def __init__(self, stdout: str, stderr: str, exception: Optional[str] = None): + self.stdout: str = stdout + self.stderr: str = stderr + self.exception: Optional[str] = exception + + +# pylint: disable=too-many-arguments +def run_over_json_rpc( + workspace: str, + interpreter: Sequence[str], + module: str, + argv: Sequence[str], + use_stdin: bool, + cwd: str, + source: str = None, +) -> RpcRunResult: + """Uses JSON-RPC to execute a command.""" + rpc: Union[JsonRpc, None] = get_or_start_json_rpc(workspace, interpreter, cwd) + if not rpc: + # pylint: disable=broad-exception-raised + raise Exception("Failed to run over JSON-RPC.") + + msg_id = str(uuid.uuid4()) + msg = { + "id": msg_id, + "method": "run", + "module": module, + "argv": argv, + "useStdin": use_stdin, + "cwd": cwd, + } + if source: + msg["source"] = source + + rpc.send_data(msg) + + data = rpc.receive_data() + + if data["id"] != msg_id: + return RpcRunResult("", f"Invalid result for request: {json.dumps(msg, indent=4)}") + + result = data["result"] if "result" in data else "" + if "error" in data: + error = data["error"] + + if data.get("exception", False): + return RpcRunResult(result, "", error) + return RpcRunResult(result, error) + + return RpcRunResult(result, "") + + +def shutdown_json_rpc(): + """Shutdown all JSON-RPC processes.""" + _process_manager.stop_all_processes() diff --git a/bundled/tool/lsp_runner.py b/bundled/tool/lsp_runner.py new file mode 100644 index 00000000..b379ad2d --- /dev/null +++ b/bundled/tool/lsp_runner.py @@ -0,0 +1,88 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +""" +Runner to use when running under a different interpreter. +""" + +import os +import pathlib +import sys +import traceback + + +# ********************************************************** +# Update sys.path before importing any bundled libraries. +# ********************************************************** +def update_sys_path(path_to_add: str, strategy: str) -> None: + """Add given path to `sys.path`.""" + if path_to_add not in sys.path and os.path.isdir(path_to_add): + if strategy == "useBundled": + sys.path.insert(0, path_to_add) + elif strategy == "fromEnvironment": + sys.path.append(path_to_add) + + +# Ensure that we can import LSP libraries, and other bundled libraries. +update_sys_path( + os.fspath(pathlib.Path(__file__).parent.parent / "libs"), + os.getenv("LS_IMPORT_STRATEGY", "useBundled"), +) + + +# pylint: disable=wrong-import-position,import-error +import lsp_jsonrpc as jsonrpc +import lsp_utils as utils + +RPC = jsonrpc.create_json_rpc(sys.stdin.buffer, sys.stdout.buffer) + +EXIT_NOW = False +while not EXIT_NOW: + msg = RPC.receive_data() + + method = msg["method"] + if method == "exit": + EXIT_NOW = True + continue + + if method == "run": + is_exception = False + # This is needed to preserve sys.path, pylint modifies + # sys.path and that might not work for this scenario + # next time around. + with utils.substitute_attr(sys, "path", sys.path[:]): + try: + # TODO: `utils.run_module` is equivalent to running `python -m `. + # If your tool supports a programmatic API then replace the function below + # with code for your tool. You can also use `utils.run_api` helper, which + # handles changing working directories, managing io streams, etc. + # Also update `_run_tool_on_document` and `_run_tool` functions in `lsp_server.py`. + result = utils.run_module( + module=msg["module"], + argv=msg["argv"], + use_stdin=msg["useStdin"], + cwd=msg["cwd"], + source=msg["source"] if "source" in msg else None, + ) + except Exception: # pylint: disable=broad-except + result = utils.RunResult("", traceback.format_exc(chain=True)) + is_exception = True + + response = {"id": msg["id"]} + if result.stderr: + response["error"] = result.stderr + response["exception"] = is_exception + elif result.stdout: + response["result"] = result.stdout + + RPC.send_data(response) diff --git a/bundled/tool/lsp_server.py b/bundled/tool/lsp_server.py new file mode 100644 index 00000000..97f56d8f --- /dev/null +++ b/bundled/tool/lsp_server.py @@ -0,0 +1,211 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Implementation of tool support over LSP.""" +from __future__ import annotations + +import json +import os +import pathlib +import sys +from typing import Any, Dict, Optional, Tuple + +from constants import TOOL_DISPLAY_NAME, TOOL_MODULE_NAME, ZENML_CLIENT_INITIALIZED + + +# ********************************************************** +# Update sys.path before importing any bundled libraries. +# ********************************************************** +def update_sys_path(path_to_add: str, strategy: str) -> None: + """Add given path to `sys.path`.""" + if path_to_add not in sys.path and os.path.isdir(path_to_add): + if strategy == "useBundled": + sys.path.insert(0, path_to_add) + elif strategy == "fromEnvironment": + sys.path.append(path_to_add) + + +# Ensure that we can import LSP libraries, and other bundled libraries. +update_sys_path( + os.fspath(pathlib.Path(__file__).parent.parent / "libs"), + os.getenv("LS_IMPORT_STRATEGY", "useBundled"), +) + + +# ********************************************************** +# Imports needed for the language server goes below this. +# ********************************************************** +# pylint: disable=wrong-import-position,import-error +import lsp_jsonrpc as jsonrpc +import lsprotocol.types as lsp +from lsp_zenml import ZenLanguageServer +from pygls import uris, workspace + +WORKSPACE_SETTINGS = {} +GLOBAL_SETTINGS = {} +RUNNER = pathlib.Path(__file__).parent / "lsp_runner.py" + +MAX_WORKERS = 5 + +LSP_SERVER = ZenLanguageServer( + name="zen-language-server", version="0.0.1", max_workers=MAX_WORKERS +) + +# ********************************************************** +# Tool specific code goes below this. +# ********************************************************** +TOOL_MODULE = TOOL_MODULE_NAME +TOOL_DISPLAY = TOOL_DISPLAY_NAME +# Default arguments always passed to zenml. +TOOL_ARGS = [] +# Versions of zenml found by workspace +VERSION_LOOKUP: Dict[str, Tuple[int, int, int]] = {} + + +# ********************************************************** +# Required Language Server Initialization and Exit handlers. +# ********************************************************** +@LSP_SERVER.feature(lsp.INITIALIZE) +async def initialize(params: lsp.InitializeParams) -> None: + """LSP handler for initialize request.""" + # pylint: disable=global-statement + log_to_output(f"CWD Server: {os.getcwd()}") + + paths = "\r\n ".join(sys.path) + log_to_output(f"sys.path used to run Server:\r\n {paths}") + + GLOBAL_SETTINGS.update(**params.initialization_options.get("globalSettings", {})) + + settings = params.initialization_options["settings"] + _update_workspace_settings(settings) + log_to_output( + f"Settings used to run Server:\r\n{json.dumps(settings, indent=4, ensure_ascii=False)}\r\n" + ) + log_to_output( + f"Global settings:\r\n{json.dumps(GLOBAL_SETTINGS, indent=4, ensure_ascii=False)}\r\n" + ) + log_to_output( + f"Workspace settings:\r\n{json.dumps(WORKSPACE_SETTINGS, indent=4, ensure_ascii=False)}\r\n" + ) + + log_to_output("ZenML LSP is initializing.") + LSP_SERVER.send_custom_notification("sanityCheck", "ZenML LSP is initializing.") + + # Below is not needed as the interpreter path gets automatically updated when changed in vscode. + # interpreter_path = WORKSPACE_SETTINGS[os.getcwd()]["interpreter"][0] + # LSP_SERVER.update_python_interpreter(interpreter_path) + + # Check install status and initialize ZenML client if ZenML is installed. + await LSP_SERVER.initialize_zenml_client() + + # Wait for 5 secondsto allow the language client to setup and settle down client side. + ready_status = {"ready": True} if LSP_SERVER.zenml_client else {"ready": False} + LSP_SERVER.send_custom_notification(ZENML_CLIENT_INITIALIZED, ready_status) + + +@LSP_SERVER.feature(lsp.EXIT) +def on_exit(_params: Optional[Any] = None) -> None: + """Handle clean up on exit.""" + jsonrpc.shutdown_json_rpc() + + +@LSP_SERVER.feature(lsp.SHUTDOWN) +def on_shutdown(_params: Optional[Any] = None) -> None: + """Handle clean up on shutdown.""" + jsonrpc.shutdown_json_rpc() + + +# ***************************************************** +# Internal functional and settings management APIs. +# ***************************************************** +def _get_global_defaults(): + return { + "path": GLOBAL_SETTINGS.get("path", []), + "interpreter": GLOBAL_SETTINGS.get("interpreter", [sys.executable]), + "args": GLOBAL_SETTINGS.get("args", []), + "importStrategy": GLOBAL_SETTINGS.get("importStrategy", "useBundled"), + "showNotifications": GLOBAL_SETTINGS.get("showNotifications", "off"), + } + + +def _update_workspace_settings(settings): + if not settings: + key = os.getcwd() + WORKSPACE_SETTINGS[key] = { + "cwd": key, + "workspaceFS": key, + "workspace": uris.from_fs_path(key), + **_get_global_defaults(), + } + return + + for setting in settings: + key = uris.to_fs_path(setting["workspace"]) + WORKSPACE_SETTINGS[key] = { + "cwd": key, + **setting, + "workspaceFS": key, + } + + +# ***************************************************** +# Internal execution APIs. +# ***************************************************** +def get_cwd(settings: Dict[str, Any], document: Optional[workspace.Document]) -> str: + """Returns cwd for the given settings and document.""" + if settings["cwd"] == "${workspaceFolder}": + return settings["workspaceFS"] + + if settings["cwd"] == "${fileDirname}": + if document is not None: + return os.fspath(pathlib.Path(document.path).parent) + return settings["workspaceFS"] + + return settings["cwd"] + + +# ***************************************************** +# Logging and notification. +# ***************************************************** +def log_to_output( + message: str, msg_type: lsp.MessageType = lsp.MessageType.Log +) -> None: + """Log to output.""" + LSP_SERVER.show_message_log(message, msg_type) + + +def log_error(message: str) -> None: + """Log error.""" + LSP_SERVER.show_message_log(message, lsp.MessageType.Error) + if os.getenv("LS_SHOW_NOTIFICATION", "off") in ["onError", "onWarning", "always"]: + LSP_SERVER.show_message(message, lsp.MessageType.Error) + + +def log_warning(message: str) -> None: + """Log warning.""" + LSP_SERVER.show_message_log(message, lsp.MessageType.Warning) + if os.getenv("LS_SHOW_NOTIFICATION", "off") in ["onWarning", "always"]: + LSP_SERVER.show_message(message, lsp.MessageType.Warning) + + +def log_always(message: str) -> None: + """Log message.""" + LSP_SERVER.show_message_log(message, lsp.MessageType.Info) + if os.getenv("LS_SHOW_NOTIFICATION", "off") in ["always"]: + LSP_SERVER.show_message(message, lsp.MessageType.Info) + + +# ***************************************************** +# Start the server. +# ***************************************************** +if __name__ == "__main__": + LSP_SERVER.start_io() diff --git a/bundled/tool/lsp_utils.py b/bundled/tool/lsp_utils.py new file mode 100644 index 00000000..716f79d3 --- /dev/null +++ b/bundled/tool/lsp_utils.py @@ -0,0 +1,218 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +"""Utility functions and classes for use with running tools over LSP.""" +from __future__ import annotations + +import contextlib +import io +import os +import os.path +import runpy +import site +import subprocess +import sys +import threading +from typing import Any, Callable, List, Sequence, Tuple, Union + +# Save the working directory used when loading this module +SERVER_CWD = os.getcwd() +CWD_LOCK = threading.Lock() + + +def as_list(content: Union[Any, List[Any], Tuple[Any]]) -> Union[List[Any], Tuple[Any]]: + """Ensures we always get a list""" + if isinstance(content, (list, tuple)): + return content + return [content] + + +# pylint: disable-next=consider-using-generator +_site_paths = tuple( + [ + os.path.normcase(os.path.normpath(p)) + for p in (as_list(site.getsitepackages()) + as_list(site.getusersitepackages())) + ] +) + + +def is_same_path(file_path1, file_path2) -> bool: + """Returns true if two paths are the same.""" + return os.path.normcase(os.path.normpath(file_path1)) == os.path.normcase( + os.path.normpath(file_path2) + ) + + +def is_current_interpreter(executable) -> bool: + """Returns true if the executable path is same as the current interpreter.""" + return is_same_path(executable, sys.executable) + + +def is_stdlib_file(file_path) -> bool: + """Return True if the file belongs to standard library.""" + return os.path.normcase(os.path.normpath(file_path)).startswith(_site_paths) + + +# pylint: disable-next=too-few-public-methods +class RunResult: + """Object to hold result from running tool.""" + + def __init__(self, stdout: str, stderr: str): + self.stdout: str = stdout + self.stderr: str = stderr + + +class CustomIO(io.TextIOWrapper): + """Custom stream object to replace stdio.""" + + name = None + + def __init__(self, name, encoding="utf-8", newline=None): + self._buffer = io.BytesIO() + self._buffer.name = name + super().__init__(self._buffer, encoding=encoding, newline=newline) + + def close(self): + """Provide this close method which is used by some tools.""" + # This is intentionally empty. + + def get_value(self) -> str: + """Returns value from the buffer as string.""" + self.seek(0) + return self.read() + + +@contextlib.contextmanager +def substitute_attr(obj: Any, attribute: str, new_value: Any): + """Manage object attributes context when using runpy.run_module().""" + old_value = getattr(obj, attribute) + setattr(obj, attribute, new_value) + yield + setattr(obj, attribute, old_value) + + +@contextlib.contextmanager +def redirect_io(stream: str, new_stream): + """Redirect stdio streams to a custom stream.""" + old_stream = getattr(sys, stream) + setattr(sys, stream, new_stream) + yield + setattr(sys, stream, old_stream) + + +@contextlib.contextmanager +def change_cwd(new_cwd): + """Change working directory before running code.""" + os.chdir(new_cwd) + yield + os.chdir(SERVER_CWD) + + +def _run_module(module: str, argv: Sequence[str], use_stdin: bool, source: str = None) -> RunResult: + """Runs as a module.""" + str_output = CustomIO("", encoding="utf-8") + str_error = CustomIO("", encoding="utf-8") + + try: + with substitute_attr(sys, "argv", argv): + with redirect_io("stdout", str_output): + with redirect_io("stderr", str_error): + if use_stdin and source is not None: + str_input = CustomIO("", encoding="utf-8", newline="\n") + with redirect_io("stdin", str_input): + str_input.write(source) + str_input.seek(0) + runpy.run_module(module, run_name="__main__") + else: + runpy.run_module(module, run_name="__main__") + except SystemExit: + pass + + return RunResult(str_output.get_value(), str_error.get_value()) + + +def run_module( + module: str, argv: Sequence[str], use_stdin: bool, cwd: str, source: str = None +) -> RunResult: + """Runs as a module.""" + with CWD_LOCK: + if is_same_path(os.getcwd(), cwd): + return _run_module(module, argv, use_stdin, source) + with change_cwd(cwd): + return _run_module(module, argv, use_stdin, source) + + +def run_path(argv: Sequence[str], use_stdin: bool, cwd: str, source: str = None) -> RunResult: + """Runs as an executable.""" + if use_stdin: + with subprocess.Popen( + argv, + encoding="utf-8", + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + stdin=subprocess.PIPE, + cwd=cwd, + ) as process: + return RunResult(*process.communicate(input=source)) + else: + result = subprocess.run( + argv, + encoding="utf-8", + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + cwd=cwd, + ) + return RunResult(result.stdout, result.stderr) + + +def run_api( + callback: Callable[[Sequence[str], CustomIO, CustomIO, CustomIO | None], None], + argv: Sequence[str], + use_stdin: bool, + cwd: str, + source: str = None, +) -> RunResult: + """Run a API.""" + with CWD_LOCK: + if is_same_path(os.getcwd(), cwd): + return _run_api(callback, argv, use_stdin, source) + with change_cwd(cwd): + return _run_api(callback, argv, use_stdin, source) + + +def _run_api( + callback: Callable[[Sequence[str], CustomIO, CustomIO, CustomIO | None], None], + argv: Sequence[str], + use_stdin: bool, + source: str = None, +) -> RunResult: + str_output = CustomIO("", encoding="utf-8") + str_error = CustomIO("", encoding="utf-8") + + try: + with substitute_attr(sys, "argv", argv): + with redirect_io("stdout", str_output): + with redirect_io("stderr", str_error): + if use_stdin and source is not None: + str_input = CustomIO("", encoding="utf-8", newline="\n") + with redirect_io("stdin", str_input): + str_input.write(source) + str_input.seek(0) + callback(argv, str_output, str_error, str_input) + else: + callback(argv, str_output, str_error) + except SystemExit: + pass + + return RunResult(str_output.get_value(), str_error.get_value()) diff --git a/bundled/tool/lsp_zenml.py b/bundled/tool/lsp_zenml.py new file mode 100644 index 00000000..095946c1 --- /dev/null +++ b/bundled/tool/lsp_zenml.py @@ -0,0 +1,292 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +""" +Extends the main Language Server Protocol (LSP) server for ZenML +by adding custom functionalities. It acts as a wrapper around the core LSP +server implementation (`lsp_server.py`), providing ZenML-specific features +such as checking ZenML installation, verifying version compatibility, and +updating Python interpreter paths. +""" + + +import asyncio +import subprocess +import sys +from functools import wraps + +import lsprotocol.types as lsp +from constants import MIN_ZENML_VERSION, TOOL_MODULE_NAME, IS_ZENML_INSTALLED +from lazy_import import suppress_stdout_temporarily +from packaging.version import parse as parse_version +from pygls.server import LanguageServer +from zen_watcher import ZenConfigWatcher +from zenml_client import ZenMLClient + +zenml_init_error = { + "error": "ZenML is not initialized. Please check ZenML version requirements." +} + + +class ZenLanguageServer(LanguageServer): + """ZenML Language Server implementation.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.python_interpreter = sys.executable + self.zenml_client = None + # self.register_commands() + + async def is_zenml_installed(self) -> bool: + """Asynchronously checks if ZenML is installed.""" + try: + process = await asyncio.create_subprocess_exec( + self.python_interpreter, + "-c", + "import zenml", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + await process.wait() + if process.returncode == 0: + self.show_message_log("✅ ZenML installation check: Successful.") + return True + self.show_message_log( + "❌ ZenML installation check failed.", lsp.MessageType.Error + ) + return False + except Exception as e: + self.show_message_log( + f"Error checking ZenML installation: {str(e)}", lsp.MessageType.Error + ) + return False + + async def initialize_zenml_client(self): + """Initializes the ZenML client.""" + self.send_custom_notification("zenml/client", {"status": "pending"}) + if self.zenml_client is not None: + # Client is already initialized. + self.notify_user("⭐️ ZenML Client Already Initialized ⭐️") + return + + if not await self.is_zenml_installed(): + self.send_custom_notification(IS_ZENML_INSTALLED, {"is_installed": False}) + self.notify_user("❗ ZenML not detected.", lsp.MessageType.Warning) + return + + zenml_version = self.get_zenml_version() + self.send_custom_notification( + IS_ZENML_INSTALLED, {"is_installed": True, "version": zenml_version} + ) + # Initializing ZenML client after successful installation check. + self.log_to_output("🚀 Initializing ZenML client...") + try: + self.zenml_client = ZenMLClient() + self.notify_user("✅ ZenML client initialized successfully.") + # register pytool module commands + self.register_commands() + # initialize watcher + self.initialize_global_config_watcher() + except Exception as e: + self.notify_user( + f"Failed to initialize ZenML client: {str(e)}", lsp.MessageType.Error + ) + + def initialize_global_config_watcher(self): + """Sets up and starts the Global Configuration Watcher.""" + try: + watcher = ZenConfigWatcher(self) + watcher.watch_zenml_config_yaml() + self.log_to_output("👀 Watching ZenML configuration for changes.") + except Exception as e: + self.notify_user( + f"Error setting up the Global Configuration Watcher: {e}", + msg_type=lsp.MessageType.Error, + ) + + def zenml_command(self, wrapper_name=None): + """ + Decorator for executing commands with ZenMLClient or its specified wrapper. + + This decorator ensures that commands are only executed if ZenMLClient is properly + initialized. If a `wrapper_name` is provided, the command targets a specific + wrapper within ZenMLClient; otherwise, it targets ZenMLClient directly. + + Args: + wrapper_name (str, optional): The specific wrapper within ZenMLClient to target. + Defaults to None, targeting the ZenMLClient itself. + """ + + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + client = self.zenml_client + if not client: + self.log_to_output("ZenML client not found in ZenLanguageServer.") + return zenml_init_error + self.log_to_output(f"Executing command with wrapper: {wrapper_name}") + if not client.initialized: + return zenml_init_error + + with suppress_stdout_temporarily(): + if wrapper_name: + wrapper_instance = getattr( + self.zenml_client, wrapper_name, None + ) + if not wrapper_instance: + return {"error": f"Wrapper '{wrapper_name}' not found."} + return func(wrapper_instance, *args, **kwargs) + return func(self.zenml_client, *args, **kwargs) + + return wrapper + + return decorator + + def get_zenml_version(self) -> str: + """Gets the ZenML version.""" + command = [ + self.python_interpreter, + "-c", + "import zenml; print(zenml.__version__)", + ] + result = subprocess.run(command, capture_output=True, text=True, check=True) + return result.stdout.strip() + + def check_zenml_version(self) -> dict: + """Checks if the installed ZenML version meets the minimum requirement.""" + version_str = self.get_zenml_version() + installed_version = parse_version(version_str) + if installed_version < parse_version(MIN_ZENML_VERSION): + return self._construct_version_validation_response(False, version_str) + + return self._construct_version_validation_response(True, version_str) + + def _construct_version_validation_response(self, meets_requirement, version_str): + """Constructs a version validation response.""" + if meets_requirement: + message = "ZenML version requirement is met." + status = {"message": message, "version": version_str, "is_valid": True} + else: + message = f"Supported versions >= {MIN_ZENML_VERSION}. Found version {version_str}." + status = {"message": message, "version": version_str, "is_valid": False} + + self.send_custom_notification("zenml/version", status) + self.notify_user(message) + return status + + def send_custom_notification(self, method: str, args: dict): + """Sends a custom notification to the LSP client.""" + self.show_message_log( + f"Sending custom notification: {method} with args: {args}" + ) + self.send_notification(method, args) + + def update_python_interpreter(self, interpreter_path): + """Updates the Python interpreter path and handles errors.""" + try: + self.python_interpreter = interpreter_path + self.show_message_log( + f"LSP_Python_Interpreter Updated: {self.python_interpreter}" + ) + # pylint: disable=broad-exception-caught + except Exception as e: + self.show_message_log( + f"Failed to update Python interpreter: {str(e)}", lsp.MessageType.Error + ) + + def notify_user( + self, message: str, msg_type: lsp.MessageType = lsp.MessageType.Info + ): + """Logs a message and also notifies the user.""" + self.show_message(message, msg_type) + + def log_to_output( + self, message: str, msg_type: lsp.MessageType = lsp.MessageType.Log + ) -> None: + """Log to output.""" + self.show_message_log(message, msg_type) + + def register_commands(self): + """Registers ZenML Python Tool commands.""" + + @self.command(f"{TOOL_MODULE_NAME}.getGlobalConfig") + @self.zenml_command(wrapper_name="config_wrapper") + def get_global_configuration(wrapper_instance, *args, **kwargs) -> dict: + """Fetches global ZenML configuration settings.""" + return wrapper_instance.get_global_configuration() + + @self.command(f"{TOOL_MODULE_NAME}.getGlobalConfigFilePath") + @self.zenml_command(wrapper_name="config_wrapper") + def get_global_config_file_path(wrapper_instance, *args, **kwargs): + """Retrieves the file path of the global ZenML configuration.""" + return wrapper_instance.get_global_config_file_path() + + @self.command(f"{TOOL_MODULE_NAME}.serverInfo") + @self.zenml_command(wrapper_name="zen_server_wrapper") + def get_server_info(wrapper_instance, *args, **kwargs): + """Gets information about the ZenML server.""" + return wrapper_instance.get_server_info() + + @self.command(f"{TOOL_MODULE_NAME}.connect") + @self.zenml_command(wrapper_name="zen_server_wrapper") + def connect(wrapper_instance, args): + """Connects to a ZenML server with specified arguments.""" + return wrapper_instance.connect(args) + + @self.command(f"{TOOL_MODULE_NAME}.disconnect") + @self.zenml_command(wrapper_name="zen_server_wrapper") + def disconnect(wrapper_instance, *args, **kwargs): + """Disconnects from the current ZenML server.""" + return wrapper_instance.disconnect(*args, **kwargs) + + @self.command(f"{TOOL_MODULE_NAME}.fetchStacks") + @self.zenml_command(wrapper_name="stacks_wrapper") + def fetch_stacks(wrapper_instance, args): + """Fetches a list of all ZenML stacks.""" + return wrapper_instance.fetch_stacks(args) + + @self.command(f"{TOOL_MODULE_NAME}.getActiveStack") + @self.zenml_command(wrapper_name="stacks_wrapper") + def get_active_stack(wrapper_instance, *args, **kwargs): + """Gets the currently active ZenML stack.""" + return wrapper_instance.get_active_stack() + + @self.command(f"{TOOL_MODULE_NAME}.switchActiveStack") + @self.zenml_command(wrapper_name="stacks_wrapper") + def set_active_stack(wrapper_instance, args): + """Sets the active ZenML stack to the specified stack.""" + print(f"args received: {args}") + return wrapper_instance.set_active_stack(args) + + @self.command(f"{TOOL_MODULE_NAME}.renameStack") + @self.zenml_command(wrapper_name="stacks_wrapper") + def rename_stack(wrapper_instance, args): + """Renames a specified ZenML stack.""" + return wrapper_instance.rename_stack(args) + + @self.command(f"{TOOL_MODULE_NAME}.copyStack") + @self.zenml_command(wrapper_name="stacks_wrapper") + def copy_stack(wrapper_instance, args): + """Copies a specified ZenML stack to a new stack.""" + return wrapper_instance.copy_stack(args) + + @self.command(f"{TOOL_MODULE_NAME}.getPipelineRuns") + @self.zenml_command(wrapper_name="pipeline_runs_wrapper") + def fetch_pipeline_runs(wrapper_instance, args): + """Fetches all ZenML pipeline runs.""" + return wrapper_instance.fetch_pipeline_runs(args) + + @self.command(f"{TOOL_MODULE_NAME}.deletePipelineRun") + @self.zenml_command(wrapper_name="pipeline_runs_wrapper") + def delete_pipeline_run(wrapper_instance, args): + """Deletes a specified ZenML pipeline run.""" + return wrapper_instance.delete_pipeline_run(args) diff --git a/bundled/tool/zen_watcher.py b/bundled/tool/zen_watcher.py new file mode 100644 index 00000000..6795fbc6 --- /dev/null +++ b/bundled/tool/zen_watcher.py @@ -0,0 +1,139 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +""" +ZenML Global Configuration Watcher. + +This module contains ZenConfigWatcher, a class that watches for changes +in the ZenML global configuration file and triggers notifications accordingly. +""" +import os +from threading import Timer +from typing import Any, Optional + +import yaml +from constants import ZENML_SERVER_CHANGED, ZENML_STACK_CHANGED +from lazy_import import suppress_stdout_temporarily +from watchdog.events import FileSystemEventHandler +from watchdog.observers import Observer + + +class ZenConfigWatcher(FileSystemEventHandler): + """ + Watches for changes in the ZenML global configuration file. + + Upon modification of the global configuration file, it triggers notifications + to update config details accordingly. + """ + + def __init__(self, lsp_server): + super().__init__() + self.LSP_SERVER = lsp_server + self.observer: Optional[Any] = None + self.debounce_interval: float = 2.0 + self._timer: Optional[Timer] = None + self.last_known_url: str = "" + self.last_known_stack_id: str = "" + self.show_notification: bool = os.getenv("LS_SHOW_NOTIFICATION", "off") in [ + "onError", + "onWarning", + "always", + ] + + def process_config_change(self, config_file_path: str): + """Process the configuration file change.""" + with suppress_stdout_temporarily(): + try: + with open(config_file_path, "r") as f: + config = yaml.safe_load(f) + + new_url = config.get("store", {}).get("url", "") + new_stack_id = config.get("active_stack_id", "") + + url_changed = new_url != self.last_known_url + stack_id_changed = new_stack_id != self.last_known_stack_id + # Send ZENML_SERVER_CHANGED if url changed + if url_changed: + server_details = { + "url": new_url, + "api_token": config.get("store", {}).get("api_token", ""), + "store_type": config.get("store", {}).get("type", ""), + } + self.LSP_SERVER.send_custom_notification( + ZENML_SERVER_CHANGED, + 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 + except (FileNotFoundError, PermissionError) as e: + self.log_error(f"Configuration file access error: {e} - {config_file_path}") + except yaml.YAMLError as e: + self.log_error(f"YAML parsing error in configuration: {e} - {config_file_path}") + except Exception as e: + self.log_error(f"Unexpected error while monitoring configuration: {e}") + + def on_modified(self, event): + """ + Handles the modification event triggered when the global configuration file is changed. + """ + if self._timer is not None: + self._timer.cancel() + self._timer = Timer(self.debounce_interval, self.process_event, [event]) + self._timer.start() + + def process_event(self, event): + """ + Processes the event with a debounce mechanism. + """ + with suppress_stdout_temporarily(): + config_wrapper_instance = self.LSP_SERVER.zenml_client.config_wrapper + config_file_path = config_wrapper_instance.get_global_config_file_path() + if event.src_path == str(config_file_path): + self.process_config_change(config_file_path) + + def watch_zenml_config_yaml(self): + """ + Initializes and starts a file watcher on the ZenML global configuration directory. + Upon detecting a change, it triggers handlers to process these changes. + """ + config_wrapper_instance = self.LSP_SERVER.zenml_client.config_wrapper + config_dir_path = config_wrapper_instance.get_global_config_directory_path() + + # Check if config_dir_path is valid and readable + if os.path.isdir(config_dir_path) and os.access(config_dir_path, os.R_OK): + try: + self.observer = Observer() + self.observer.schedule(self, config_dir_path, recursive=False) + self.observer.start() + self.LSP_SERVER.log_to_output(f"Started watching {config_dir_path} for changes.") + except Exception as e: + self.log_error(f"Failed to start file watcher: {e}") + else: + self.log_error("Config directory path invalid or missing.") + + def stop_watching(self): + """ + Stops the file watcher gracefully. + """ + if self.observer is not None: + self.observer.stop() + self.observer.join() # waits for observer to fully stop + self.LSP_SERVER.log_to_output("Stopped watching config directory for changes.") + + def log_error(self, message: str): + """Log error.""" + self.LSP_SERVER.show_message_log(message, 1) + if self.show_notification: + self.LSP_SERVER.show_message(message, 1) diff --git a/bundled/tool/zenml_client.py b/bundled/tool/zenml_client.py new file mode 100644 index 00000000..9c234745 --- /dev/null +++ b/bundled/tool/zenml_client.py @@ -0,0 +1,39 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either expressc +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""ZenML client class. Initializes all wrappers.""" + + +class ZenMLClient: + """Provides a high-level interface to ZenML functionalities by wrapping core components.""" + + def __init__(self): + """ + Initializes the ZenMLClient with wrappers for managing configurations, + server interactions, stacks, and pipeline runs. + """ + # pylint: disable=wrong-import-position,import-error + from lazy_import import lazy_import + from zenml_wrappers import ( + GlobalConfigWrapper, + PipelineRunsWrapper, + StacksWrapper, + ZenServerWrapper, + ) + + 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.initialized = True diff --git a/bundled/tool/zenml_wrappers.py b/bundled/tool/zenml_wrappers.py new file mode 100644 index 00000000..12946c54 --- /dev/null +++ b/bundled/tool/zenml_wrappers.py @@ -0,0 +1,526 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""This module provides wrappers for ZenML configuration and operations.""" + +import json +import pathlib +from typing import Any + + +class GlobalConfigWrapper: + """Wrapper class for global configuration management.""" + + def __init__(self): + # pylint: disable=wrong-import-position,import-error + from lazy_import import lazy_import + + self.lazy_import = lazy_import + """Initializes the GlobalConfigWrapper instance.""" + self._gc = lazy_import("zenml.config.global_config", "GlobalConfiguration")() + + @property + def gc(self): + """Returns the global configuration instance.""" + return self._gc + + @property + def fileio(self): + """Provides access to file I/O operations.""" + return self.lazy_import("zenml.io", "fileio") + + @property + def get_global_config_directory(self): + """Returns the function to get the global configuration directory.""" + return self.lazy_import("zenml.utils.io_utils", "get_global_config_directory") + + @property + def RestZenStoreConfiguration(self): + """Returns the RestZenStoreConfiguration class for store configuration.""" + # pylint: disable=not-callable + return self.lazy_import( + "zenml.zen_stores.rest_zen_store", "RestZenStoreConfiguration" + ) + + def get_global_config_directory_path(self) -> str: + """Get the global configuration directory path. + + Returns: + str: Path to the global configuration directory. + """ + # pylint: disable=not-callable + config_dir = pathlib.Path(self.get_global_config_directory()) + if self.fileio.exists(str(config_dir)): + return str(config_dir) + return "Configuration directory does not exist." + + def get_global_config_file_path(self) -> str: + """Get the global configuration file path. + + Returns: + str: Path to the global configuration file. + """ + # pylint: disable=not-callable + config_dir = pathlib.Path(self.get_global_config_directory()) + config_path = config_dir / "config.yaml" + if self.fileio.exists(str(config_path)): + return str(config_path) + return "Configuration file does not exist." + + def set_store_configuration(self, remote_url: str, access_token: str): + """Set the store configuration. + + Args: + remote_url (str): Remote URL. + access_token (str): Access token. + """ + # pylint: disable=not-callable + new_store_config = self.RestZenStoreConfiguration( + type="rest", url=remote_url, api_token=access_token, verify_ssl=True + ) + + # Method name changed in 0.55.4 - 0.56.1 + if hasattr(self.gc, "set_store_configuration"): + self.gc.set_store_configuration(new_store_config) + elif hasattr(self.gc, "set_store"): + self.gc.set_store(new_store_config) + else: + raise AttributeError( + "GlobalConfiguration object does not have a method to set store configuration." + ) + self.gc.set_store(new_store_config) + + def get_global_configuration(self) -> dict: + """Get the global configuration. + + Returns: + dict: Global configuration. + """ + gc_dict = json.loads(self.gc.json(indent=2)) + user_id = gc_dict.get("user_id", "") + + if user_id and user_id.startswith("UUID('") and user_id.endswith("')"): + gc_dict["user_id"] = user_id[6:-2] + + return gc_dict + + +class ZenServerWrapper: + """Wrapper class for Zen Server management.""" + + def __init__(self, config_wrapper): + """Initializes ZenServerWrapper with a configuration wrapper.""" + # pylint: disable=wrong-import-position,import-error + from lazy_import import lazy_import + + self.lazy_import = lazy_import + self._config_wrapper = config_wrapper + + @property + def gc(self): + """Returns the global configuration via the config wrapper.""" + return self._config_wrapper.gc + + @property + def web_login(self): + """Provides access to the ZenML web login function.""" + return self.lazy_import("zenml.cli", "web_login") + + @property + def ServerDeploymentNotFoundError(self): + """Returns the ZenML ServerDeploymentNotFoundError class.""" + return self.lazy_import( + "zenml.zen_server.deploy.exceptions", "ServerDeploymentNotFoundError" + ) + + @property + def AuthorizationException(self): + """Returns the ZenML AuthorizationException class.""" + return self.lazy_import("zenml.exceptions", "AuthorizationException") + + @property + def StoreType(self): + """Returns the ZenML StoreType enum.""" + return self.lazy_import("zenml.enums", "StoreType") + + @property + def BaseZenStore(self): + """Returns the BaseZenStore class for ZenML store operations.""" + return self.lazy_import("zenml.zen_stores.base_zen_store", "BaseZenStore") + + @property + def ServerDeployer(self): + """Provides access to the ZenML server deployment utilities.""" + return self.lazy_import("zenml.zen_server.deploy.deployer", "ServerDeployer") + + @property + def get_active_deployment(self): + """Returns the function to get the active ZenML server deployment.""" + return self.lazy_import("zenml.zen_server.utils", "get_active_deployment") + + def get_server_info(self) -> dict: + """Fetches the ZenML server info. + + Returns: + dict: Dictionary containing server info. + """ + store_info = json.loads(self.gc.zen_store.get_store_info().json(indent=2)) + # Handle both 'store' and 'store_configuration' depending on version + store_attr_name = ( + "store_configuration" + if hasattr(self.gc, "store_configuration") + else "store" + ) + store_config = json.loads(getattr(self.gc, store_attr_name).json(indent=2)) + return {"storeInfo": store_info, "storeConfig": store_config} + + def connect(self, args, **kwargs) -> dict: + """Connects to a ZenML server. + + Args: + args (list): List of arguments. + Returns: + dict: Dictionary containing the result of the operation. + """ + url = args[0] + verify_ssl = args[1] if len(args) > 1 else True + + if not url: + return {"error": "Server URL is required."} + + try: + # pylint: disable=not-callable + access_token = self.web_login(url=url, verify_ssl=verify_ssl) + self._config_wrapper.set_store_configuration( + remote_url=url, access_token=access_token + ) + return {"message": "Connected successfully.", "access_token": access_token} + except self.AuthorizationException as e: + return {"error": f"Authorization failed: {str(e)}"} + + def disconnect(self, args) -> dict: + """Disconnects from a ZenML server. + + Args: + args (list): List of arguments. + Returns: + dict: Dictionary containing the result of the operation. + """ + try: + # Adjust for changes from 'store' to 'store_configuration' + store_attr_name = ( + "store_configuration" + if hasattr(self.gc, "store_configuration") + else "store" + ) + url = getattr(self.gc, store_attr_name).url + store_type = self.BaseZenStore.get_store_type(url) + + # pylint: disable=not-callable + server = self.get_active_deployment(local=True) + deployer = self.ServerDeployer() + + messages = [] + + if server: + deployer.remove_server(server.config.name) + messages.append("Shut down the local ZenML server.") + else: + messages.append("No local ZenML server was found running.") + + if store_type == self.StoreType.REST: + deployer.disconnect_from_server() + messages.append("Disconnected from the remote ZenML REST server.") + + self.gc.set_default_store() + + return {"message": " ".join(messages)} + except self.ServerDeploymentNotFoundError as e: + return {"error": f"Failed to disconnect: {str(e)}"} + + +class PipelineRunsWrapper: + """Wrapper for interacting with ZenML pipeline runs.""" + + def __init__(self, client): + """Initializes PipelineRunsWrapper with a ZenML client.""" + # pylint: disable=wrong-import-position,import-error + from lazy_import import lazy_import + + self.lazy_import = lazy_import + self.client = client + + @property + def ValidationError(self): + """Returns the ZenML ZenMLBaseException class.""" + return self.lazy_import("zenml.exceptions", "ValidationError") + + @property + def ZenMLBaseException(self): + """Returns the ZenML ZenMLBaseException class.""" + return self.lazy_import("zenml.exceptions", "ZenMLBaseException") + + def fetch_pipeline_runs(self, args): + """Fetches all ZenML pipeline runs. + + Returns: + list: List of dictionaries containing pipeline run data. + """ + page = args[0] + max_size = args[1] + try: + runs_page = self.client.list_pipeline_runs( + sort_by="desc:updated", page=page, size=max_size, hydrate=True + ) + runs_data = [ + { + "id": str(run.id), + "name": run.body.pipeline.name, + "status": run.body.status, + "version": run.body.pipeline.body.version, + "stackName": run.body.stack.name, + "startTime": ( + run.metadata.start_time.isoformat() + if run.metadata.start_time + else None + ), + "endTime": ( + run.metadata.end_time.isoformat() + if run.metadata.end_time + else None + ), + "os": run.metadata.client_environment.get("os", "Unknown OS"), + "osVersion": run.metadata.client_environment.get( + "os_version", + run.metadata.client_environment.get( + "mac_version", "Unknown Version" + ), + ), + "pythonVersion": run.metadata.client_environment.get( + "python_version", "Unknown" + ), + } + for run in runs_page.items + ] + + return { + "runs": runs_data, + "total": runs_page.total, + "total_pages": runs_page.total_pages, + "current_page": page, + "items_per_page": max_size, + } + except self.ValidationError as e: + return {"error": "ValidationError", "message": str(e)} + except self.ZenMLBaseException as e: + return [{"error": f"Failed to retrieve pipeline runs: {str(e)}"}] + + def delete_pipeline_run(self, args) -> dict: + """Deletes a ZenML pipeline run. + + Args: + args (list): List of arguments. + Returns: + dict: Dictionary containing the result of the operation. + """ + try: + run_id = args[0] + self.client.delete_pipeline_run(run_id) + return {"message": f"Pipeline run `{run_id}` deleted successfully."} + except self.ZenMLBaseException as e: + return {"error": f"Failed to delete pipeline run: {str(e)}"} + + +class StacksWrapper: + """Wrapper class for Stacks management.""" + + def __init__(self, client): + """Initializes StacksWrapper with a ZenML client.""" + # pylint: disable=wrong-import-position,import-error + from lazy_import import lazy_import + + self.lazy_import = lazy_import + self.client = client + + @property + def ZenMLBaseException(self): + """Returns the ZenML ZenMLBaseException class.""" + return self.lazy_import("zenml.exceptions", "ZenMLBaseException") + + @property + def ValidationError(self): + """Returns the ZenML ZenMLBaseException class.""" + return self.lazy_import("zenml.exceptions", "ValidationError") + + @property + def IllegalOperationError(self) -> Any: + """Returns the IllegalOperationError class.""" + return self.lazy_import("zenml.exceptions", "IllegalOperationError") + + @property + def StackComponentValidationError(self): + """Returns the ZenML StackComponentValidationError class.""" + return self.lazy_import("zenml.exceptions", "StackComponentValidationError") + + @property + def ZenKeyError(self) -> Any: + """Returns the ZenKeyError class.""" + return self.lazy_import("zenml.exceptions", "ZenKeyError") + + def fetch_stacks(self, args): + """Fetches all ZenML stacks and components with pagination.""" + page = args[0] + max_size = args[1] + try: + stacks_page = self.client.list_stacks( + page=page, size=max_size, hydrate=True + ) + stacks_data = self.process_stacks(stacks_page.items) + + return { + "stacks": stacks_data, + "total": stacks_page.total, + "total_pages": stacks_page.total_pages, + "current_page": page, + "items_per_page": max_size, + } + except self.ValidationError as e: + return {"error": "ValidationError", "message": str(e)} + except self.ZenMLBaseException as e: + return [{"error": f"Failed to retrieve stacks: {str(e)}"}] + + def process_stacks(self, stacks): + """Process stacks to the desired format.""" + return [ + { + "id": str(stack.id), + "name": stack.name, + "components": { + component_type: [ + { + "id": str(component.id), + "name": component.name, + "flavor": component.flavor, + "type": component.type, + } + for component in components + ] + for component_type, components in stack.components.items() + }, + } + for stack in stacks + ] + + def get_active_stack(self) -> dict: + """Fetches the active ZenML stack. + + Returns: + dict: Dictionary containing active stack data. + """ + try: + active_stack = self.client.active_stack_model + return { + "id": str(active_stack.id), + "name": active_stack.name, + } + except self.ZenMLBaseException as e: + return {"error": f"Failed to retrieve active stack: {str(e)}"} + + def set_active_stack(self, args) -> dict: + """Sets the active ZenML stack. + + Args: + args (list): List containing the stack name or id. + Returns: + dict: Dictionary containing the active stack data. + """ + stack_name_or_id = args[0] + + if not stack_name_or_id: + return {"error": "Missing stack_name_or_id"} + + try: + self.client.activate_stack(stack_name_id_or_prefix=stack_name_or_id) + active_stack = self.client.active_stack_model + return { + "message": f"Active stack set to: {active_stack.name}", + "id": str(active_stack.id), + "name": active_stack.name, + } + except KeyError as err: + return {"error": str(err)} + + def rename_stack(self, args) -> dict: + """Renames a specified ZenML stack. + + Args: + args (list): List containing the stack name or id and the new stack name. + Returns: + dict: Dictionary containing the renamed stack data. + """ + stack_name_or_id = args[0] + new_stack_name = args[1] + + if not stack_name_or_id or not new_stack_name: + return {"error": "Missing stack_name_or_id or new_stack_name"} + + try: + self.client.update_stack( + name_id_or_prefix=stack_name_or_id, + name=new_stack_name, + ) + return { + "message": f"Stack `{stack_name_or_id}` successfully renamed to `{new_stack_name}`!" + } + except (KeyError, self.IllegalOperationError) as err: + return {"error": str(err)} + + def copy_stack(self, args) -> dict: + """Copies a specified ZenML stack to a new stack. + + Args: + args (list): List containing the source stack name or id and the target stack name. + Returns: + dict: Dictionary containing the copied stack data. + """ + source_stack_name_or_id = args[0] + target_stack_name = args[1] + + if not source_stack_name_or_id or not target_stack_name: + return { + "error": "Both source stack name/id and target stack name are required" + } + + try: + stack_to_copy = self.client.get_stack( + name_id_or_prefix=source_stack_name_or_id + ) + component_mapping = { + c_type: [c.id for c in components][0] + for c_type, components in stack_to_copy.components.items() + if components + } + + self.client.create_stack( + name=target_stack_name, components=component_mapping + ) + return { + "message": ( + f"Stack `{source_stack_name_or_id}` successfully copied " + f"to `{target_stack_name}`!" + ) + } + except ( + self.ZenKeyError, + self.StackComponentValidationError, + ) as e: + return {"error": str(e)} diff --git a/noxfile.py b/noxfile.py new file mode 100644 index 00000000..ed856c66 --- /dev/null +++ b/noxfile.py @@ -0,0 +1,179 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""All the action we need during build""" + +import json +import os +import pathlib +import urllib.request as url_lib +from typing import List + +import nox # pylint: disable=import-error + + +def _install_bundle(session: nox.Session) -> None: + session.install( + "-t", + "./bundled/libs", + "--no-cache-dir", + "--implementation", + "py", + "--no-deps", + "--upgrade", + "-r", + "./requirements.txt", + ) + + +def _check_files(names: List[str]) -> None: + root_dir = pathlib.Path(__file__).parent + for name in names: + file_path = root_dir / name + lines: List[str] = file_path.read_text().splitlines() + if any(line for line in lines if line.startswith("# TODO:")): + # pylint: disable=broad-exception-raised + raise Exception(f"Please update {os.fspath(file_path)}.") + + +def _update_pip_packages(session: nox.Session) -> None: + session.run( + "pip-compile", + "--generate-hashes", + "--resolver=backtracking", + "--upgrade", + "./requirements.in", + ) + session.run( + "pip-compile", + "--generate-hashes", + "--resolver=backtracking", + "--upgrade", + "./src/test/python_tests/requirements.in", + ) + + +def _get_package_data(package): + json_uri = f"https://registry.npmjs.org/{package}" + with url_lib.urlopen(json_uri) as response: + return json.loads(response.read()) + + +def _update_npm_packages(session: nox.Session) -> None: + pinned = { + "vscode-languageclient", + "@types/vscode", + "@types/node", + } + package_json_path = pathlib.Path(__file__).parent / "package.json" + package_json = json.loads(package_json_path.read_text(encoding="utf-8")) + + for package in package_json["dependencies"]: + if package not in pinned: + data = _get_package_data(package) + latest = "^" + data["dist-tags"]["latest"] + package_json["dependencies"][package] = latest + + for package in package_json["devDependencies"]: + if package not in pinned: + data = _get_package_data(package) + latest = "^" + data["dist-tags"]["latest"] + package_json["devDependencies"][package] = latest + + # Ensure engine matches the package + if ( + package_json["engines"]["vscode"] + != package_json["devDependencies"]["@types/vscode"] + ): + print( + "Please check VS Code engine version and @types/vscode version in package.json." + ) + + new_package_json = json.dumps(package_json, indent=4) + # JSON dumps uses \n for line ending on all platforms by default + if not new_package_json.endswith("\n"): + new_package_json += "\n" + package_json_path.write_text(new_package_json, encoding="utf-8") + session.run("npm", "install", external=True) + + +def _setup_template_environment(session: nox.Session) -> None: + session.install("wheel", "pip-tools") + session.run( + "pip-compile", + "--generate-hashes", + "--resolver=backtracking", + "--upgrade", + "./requirements.in", + ) + session.run( + "pip-compile", + "--generate-hashes", + "--resolver=backtracking", + "--upgrade", + "./src/test/python_tests/requirements.in", + ) + _install_bundle(session) + + +@nox.session() +def setup(session: nox.Session) -> None: + """Sets up the template for development.""" + _setup_template_environment(session) + print(f"DEBUG – Virtual Environment Interpreter: {session.bin}/python") + + +@nox.session() +def tests(session: nox.Session) -> None: + """Runs all the tests for the extension.""" + session.install("-r", "src/test/python_tests/requirements.txt") + session.run("pytest", "src/test/python_tests") + + +@nox.session() +def lint(session: nox.Session) -> None: + """Runs linter and formatter checks on python files.""" + session.install("-r", "./requirements.txt") + session.install("-r", "src/test/python_tests/requirements.txt") + + session.install("pylint") + session.run("pylint", "-d", "W0511", "./bundled/tool") + session.run( + "pylint", + "-d", + "W0511", + "--ignore=./src/test/python_tests/test_data", + "./src/test/python_tests", + ) + session.run("pylint", "-d", "W0511", "noxfile.py") + + # check formatting using black + session.install("black") + session.run("black", "--check", "./bundled/tool") + session.run("black", "--check", "./src/test/python_tests") + session.run("black", "--check", "noxfile.py") + + # check import sorting using isort + session.install("isort") + session.run("isort", "--check", "./bundled/tool") + session.run("isort", "--check", "./src/test/python_tests") + session.run("isort", "--check", "noxfile.py") + + # check typescript code + session.run("npm", "run", "lint", external=True) + + +@nox.session() +def build_package(session: nox.Session) -> None: + """Builds VSIX package for publishing.""" + _check_files(["README.md", "LICENSE", "SECURITY.md", "SUPPORT.md"]) + _setup_template_environment(session) + session.run("npm", "install", external=True) + session.run("npm", "run", "vsce-package", external=True) + + +@nox.session() +def update_packages(session: nox.Session) -> None: + """Update pip and npm packages.""" + session.install("wheel", "pip-tools") + _update_pip_packages(session) + _update_npm_packages(session) diff --git a/package.json b/package.json new file mode 100644 index 00000000..6774adb1 --- /dev/null +++ b/package.json @@ -0,0 +1,376 @@ +{ + "name": "zenml", + "publisher": "ZenML", + "displayName": "ZenML", + "description": "Integrates ZenML directly into VS Code, enhancing machine learning workflow with support for pipelines, stacks, and server management.", + "version": "0.0.1", + "license": "Apache-2.0", + "categories": [ + "Machine Learning", + "Visualization" + ], + "repository": { + "type": "git", + "url": "https://github.com/zenml-io/vscode-zenml" + }, + "keywords": [ + "zenml", + "ml", + "machine learning", + "mlops", + "stack management", + "pipeline management", + "development tools" + ], + "engines": { + "vscode": "^1.86.0" + }, + "activationEvents": [ + "onStartupFinished", + "onLanguage:python", + "workspaceContains:*.py" + ], + "extensionDependencies": [ + "ms-python.python" + ], + "capabilities": { + "virtualWorkspaces": { + "supported": false, + "description": "Virtual Workspaces are not supported with ." + } + }, + "main": "./dist/extension.js", + "serverInfo": { + "name": "ZenML", + "module": "zenml-python" + }, + "scripts": { + "vscode:prepublish": "npm run package", + "compile": "webpack", + "watch": "webpack --watch", + "package": "webpack --mode production --devtool source-map --config ./webpack.config.js", + "compile-tests": "tsc -p . --outDir out", + "watch-tests": "tsc -p . -w --outDir out", + "pretest": "npm run compile-tests && npm run compile && npm run lint", + "format": "prettier --ignore-path .gitignore --write \"**/*.+(ts|json)\"", + "lint": "eslint src --ext ts", + "test": "vscode-test", + "vsce-package": "vsce package -o zenml.vsix" + }, + "contributes": { + "configuration": { + "title": "ZenML", + "properties": { + "zenml-python.args": { + "default": [], + "description": "Arguments passed in. Each argument is a separate item in the array.", + "items": { + "type": "string" + }, + "scope": "resource", + "type": "array" + }, + "zenml-python.path": { + "default": [], + "scope": "resource", + "items": { + "type": "string" + }, + "type": "array" + }, + "zenml-python.importStrategy": { + "default": "useBundled", + "enum": [ + "useBundled", + "fromEnvironment" + ], + "enumDescriptions": [ + "Always use the bundled version of ``.", + "Use `` from environment, fallback to bundled version only if `` not available in the environment." + ], + "scope": "window", + "type": "string" + }, + "zenml-python.interpreter": { + "default": [], + "description": "When set to a path to python executable, extension will use that to launch the server and any subprocess.", + "scope": "resource", + "items": { + "type": "string" + }, + "type": "array" + }, + "zenml-python.showNotifications": { + "default": "off", + "description": "Controls when notifications are shown by this extension.", + "enum": [ + "off", + "onError", + "onWarning", + "always" + ], + "enumDescriptions": [ + "All notifications are turned off, any errors or warning are still available in the logs.", + "Notifications are shown only in the case of an error.", + "Notifications are shown for errors and warnings.", + "Notifications are show for anything that the server chooses to show." + ], + "scope": "machine", + "type": "string" + }, + "zenml.serverUrl": { + "type": "string", + "default": "", + "description": "ZenML Server URL" + }, + "zenml.accessToken": { + "type": "string", + "default": "", + "description": "Access token for the ZenML server" + }, + "zenml.activeStackId": { + "type": "string", + "default": "", + "description": "Active stack id for the ZenML server" + } + } + }, + "commands": [ + { + "command": "zenml.promptForInterpreter", + "title": "Select Python Interpreter", + "category": "ZenML" + }, + { + "command": "zenml-python.restart", + "title": "Restart LSP Server", + "category": "ZenML" + }, + { + "command": "zenml.connectServer", + "title": "Connect", + "category": "ZenML Server" + }, + { + "command": "zenml.disconnectServer", + "title": "Disconnect", + "category": "ZenML Server" + }, + { + "command": "zenml.refreshServerStatus", + "title": "Refresh Server Status", + "icon": "$(refresh)", + "category": "ZenML Server" + }, + { + "command": "zenml.setStackItemsPerPage", + "title": "Set Stacks Per Page", + "icon": "$(layers)", + "category": "ZenML Stacks" + }, + { + "command": "zenml.refreshStackView", + "title": "Refresh Stack View", + "icon": "$(refresh)", + "category": "ZenML Stacks" + }, + { + "command": "zenml.getActiveStack", + "title": "Get Active Stack", + "category": "ZenML Stacks" + }, + { + "command": "zenml.renameStack", + "title": "Rename Stack", + "icon": "$(edit)", + "category": "ZenML Stacks" + }, + { + "command": "zenml.setActiveStack", + "title": "Set Active Stack", + "icon": "$(check)", + "category": "ZenML Stacks" + }, + { + "command": "zenml.copyStack", + "title": "Copy Stack", + "icon": "$(copy)", + "category": "ZenML" + }, + { + "command": "zenml.setPipelineRunsPerPage", + "title": "Set Runs Per Page", + "icon": "$(layers)", + "category": "ZenML Pipeline Runs" + }, + { + "command": "zenml.refreshPipelineView", + "title": "Refresh Pipeline View", + "icon": "$(refresh)", + "category": "ZenML Pipeline Runs" + }, + { + "command": "zenml.deletePipelineRun", + "title": "Delete Pipeline Run", + "icon": "$(trash)", + "category": "ZenML Pipeline Runs" + }, + { + "command": "zenml.refreshEnvironmentView", + "title": "Refresh Environment View", + "icon": "$(refresh)", + "category": "ZenML Environment" + }, + { + "command": "zenml.setPythonInterpreter", + "title": "Switch Python Interpreter", + "icon": "$(arrow-swap)", + "category": "ZenML Environment" + }, + { + "command": "zenml.restartLspServer", + "title": "Restart LSP Server", + "icon": "$(debug-restart)", + "category": "ZenML Environment" + } + ], + "viewsContainers": { + "activitybar": [ + { + "id": "zenml", + "title": "ZenML", + "icon": "resources/logo.png" + } + ] + }, + "views": { + "zenml": [ + { + "id": "zenmlServerView", + "name": "Server", + "icon": "$(vm)" + }, + { + "id": "zenmlStackView", + "name": "Stacks", + "icon": "$(layers)" + }, + { + "id": "zenmlPipelineView", + "name": "Pipeline Runs", + "icon": "$(beaker)" + }, + { + "id": "zenmlEnvironmentView", + "name": "Environment", + "icon": "$(server-environment)" + } + ] + }, + "menus": { + "view/title": [ + { + "when": "serverCommandsRegistered && view == zenmlServerView", + "command": "zenml.connectServer", + "group": "navigation" + }, + { + "when": "serverCommandsRegistered && view == zenmlServerView", + "command": "zenml.disconnectServer", + "group": "navigation" + }, + { + "when": "serverCommandsRegistered && view == zenmlServerView", + "command": "zenml.refreshServerStatus", + "group": "navigation" + }, + { + "when": "stackCommandsRegistered && view == zenmlStackView", + "command": "zenml.setStackItemsPerPage", + "group": "navigation@1" + }, + { + "when": "stackCommandsRegistered && view == zenmlStackView", + "command": "zenml.refreshStackView", + "group": "navigation@2" + }, + { + "when": "pipelineCommandsRegistered && view == zenmlPipelineView", + "command": "zenml.setPipelineRunsPerPage", + "group": "navigation@1" + }, + { + "when": "pipelineCommandsRegistered && view == zenmlPipelineView", + "command": "zenml.refreshPipelineView", + "group": "navigation@2" + }, + { + "when": "environmentCommandsRegistered && view == zenmlEnvironmentView", + "command": "zenml.restartLspServer", + "group": "navigation" + } + ], + "view/item/context": [ + { + "when": "stackCommandsRegistered && view == zenmlStackView && viewItem == stack", + "command": "zenml.setActiveStack", + "group": "inline@1" + }, + { + "when": "stackCommandsRegistered && view == zenmlStackView && viewItem == stack", + "command": "zenml.renameStack", + "group": "inline@2" + }, + { + "when": "stackCommandsRegistered && view == zenmlStackView && viewItem == stack", + "command": "zenml.copyStack", + "group": "inline@3" + }, + { + "when": "pipelineCommandsRegistered && view == zenmlPipelineView && viewItem == pipelineRun", + "command": "zenml.deletePipelineRun", + "group": "inline" + }, + { + "when": "environmentCommandsRegistered && view == zenmlEnvironmentView && viewItem == interpreter", + "command": "zenml.setPythonInterpreter", + "group": "inline" + } + ] + } + }, + "devDependencies": { + "@types/fs-extra": "^11.0.4", + "@types/mocha": "^10.0.6", + "@types/node": "^18.19.18", + "@types/proxyquire": "^1.3.31", + "@types/sinon": "^17.0.3", + "@types/vscode": "^1.86.0", + "@types/webpack": "^5.28.5", + "@typescript-eslint/eslint-plugin": "^6.19.1", + "@typescript-eslint/parser": "^6.19.1", + "@vscode/test-cli": "^0.0.4", + "@vscode/test-electron": "^2.3.9", + "@vscode/vsce": "^2.24.0", + "eslint": "^8.56.0", + "eslint-config-prettier": "^9.1.0", + "prettier": "^3.2.5", + "proxyquire": "^2.1.3", + "sinon": "^17.0.1", + "ts-loader": "^9.5.1", + "ts-mockito": "^2.6.1", + "ts-node": "^10.9.2", + "tsconfig-paths": "^4.2.0", + "tsconfig-paths-webpack-plugin": "^4.1.0", + "typemoq": "^2.1.0", + "typescript": "^5.3.3", + "webpack": "^5.90.0", + "webpack-cli": "^5.1.4" + }, + "dependencies": { + "@vscode/python-extension": "^1.0.5", + "axios": "^1.6.7", + "fs-extra": "^11.2.0", + "vscode-languageclient": "^9.0.1" + } +} diff --git a/package.nls.json b/package.nls.json new file mode 100644 index 00000000..f2e046db --- /dev/null +++ b/package.nls.json @@ -0,0 +1,35 @@ +{ + "extension.description": "Integrates ZenML directly into VS Code, enhancing machine learning workflow with support for pipelines, stacks, and server management.", + "zenml.promptForInterpreter": "Prompt to select a Python interpreter from the command palette in VSCode.", + "command.connectServer": "Establishes a connection to a specified ZenML server.", + "command.disconnectServer": "Disconnects from the currently connected ZenML server.", + "command.refreshServerStatus": "Refreshes the status of the ZenML server to reflect the current state.", + "command.refreshStackView": "Refreshes the stack view to display the latest information about ZenML stacks.", + "command.setStackItemsPerPage": "Sets the number of stacks to display per page in the stack view.", + "command.nextStackPage": "Navigates to the next page of stacks in the stack view.", + "command.previousStackPage": "Navigates to the previous page of stacks in the stack view.", + "command.getActiveStack": "Retrieves the currently active stack in ZenML.", + "command.renameStack": "Renames a specified ZenML stack.", + "command.setActiveStack": "Sets a specified stack as the active ZenML stack.", + "command.copyStack": "Creates a copy of a specified ZenML stack.", + "command.refreshPipelineView": "Refreshes the pipeline view to display the latest information on ZenML pipeline runs.", + "command.setPipelineRunsPerPage": "Sets the number of pipeline runs to display per page in the pipeline view.", + "command.nextPipelineRunsPage": "Navigates to the next page of pipeline runs in the pipeline view.", + "command.previousPipelineRunsPage": "Navigates to the previous page of pipeline runs in the pipeline view.", + "command.deletePipelineRun": "Deletes a specified ZenML pipeline run.", + "command.setPythonInterpreter": "Sets the Python interpreter for ZenML-related tasks within the VS Code environment.", + "command.refreshEnvironmentView": "Updates the Environment View with the latest ZenML configuration and system information.", + "command.restartLspServer": "Restarts the Language Server to ensure the latest configurations are used.", + "settings.args.description": "Specify arguments to pass to the ZenML CLI. Provide each argument as a separate item in the array.", + "settings.path.description": "Defines the path to the ZenML CLI. If left as an empty array, the default system path will be used.", + "settings.importStrategy.description": "Determines which ZenML CLI to use. Options include using a bundled version (`useBundled`) or attempting to use the CLI installed in the current Python environment (`fromEnvironment`).", + "settings.showNotifications.description": "Controls when notifications are shown by this extension.", + "settings.showNotifications.off.description": "All notifications are turned off, any errors or warnings when formatting Python files are still available in the logs.", + "settings.showNotifications.onError.description": "Notifications are shown only in the case of an error when formatting Python files.", + "settings.showNotifications.onWarning.description": "Notifications are shown for any errors and warnings when formatting Python files.", + "settings.showNotifications.always.description": "Notifications are show for anything that the server chooses to show when formatting Python files.", + "settings.interpreter.description": "Sets the path to the Python interpreter used by ZenML. This is used for launching the server and subprocesses. Leave empty to use the default interpreter.", + "settings.serverUrl.description": "URL to connect to the ZenML server.", + "settings.accessToken.description": "Access token required for authenticating with the ZenML server (not necessary currently).", + "settings.activeStackId.description": "Identifier for the currently active stack in ZenML. This is used to specify which stack is being used for operations." +} diff --git a/requirements.in b/requirements.in new file mode 100644 index 00000000..35b2285a --- /dev/null +++ b/requirements.in @@ -0,0 +1,16 @@ +# This file is used to generate requirements.txt. +# NOTE: +# Use Python 3.8 or greater which ever is the minimum version of the python +# you plan on supporting when creating the environment or using pip-tools. +# Only run the commands below to manully upgrade packages in requirements.txt: +# 1) python -m pip install pip-tools +# 2) pip-compile --generate-hashes --resolver=backtracking --upgrade ./requirements.in +# If you are using nox commands to setup or build package you don't need to +# run the above commands manually. + +# Required packages +pygls +packaging +# Tool-specific packages for ZenML extension +watchdog +PyYAML diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 00000000..47a77428 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,114 @@ +# +# This file is autogenerated by pip-compile with Python 3.11 +# by the following command: +# +# pip-compile --generate-hashes ./requirements.in +# +attrs==23.2.0 \ + --hash=sha256:935dc3b529c262f6cf76e50877d35a4bd3c1de194fd41f47a2b7ae8f19971f30 \ + --hash=sha256:99b87a485a5820b23b879f04c2305b44b951b502fd64be915879d77a7e8fc6f1 + # via + # cattrs + # lsprotocol +cattrs==23.2.3 \ + --hash=sha256:0341994d94971052e9ee70662542699a3162ea1e0c62f7ce1b4a57f563685108 \ + --hash=sha256:a934090d95abaa9e911dac357e3a8699e0b4b14f8529bcc7d2b1ad9d51672b9f + # via + # lsprotocol + # pygls +lsprotocol==2023.0.1 \ + --hash=sha256:c75223c9e4af2f24272b14c6375787438279369236cd568f596d4951052a60f2 \ + --hash=sha256:cc5c15130d2403c18b734304339e51242d3018a05c4f7d0f198ad6e0cd21861d + # via pygls +packaging==24.0 \ + --hash=sha256:2ddfb553fdf02fb784c234c7ba6ccc288296ceabec964ad2eae3777778130bc5 \ + --hash=sha256:eb82c5e3e56209074766e6885bb04b8c38a0c015d0a30036ebe7ece34c9989e9 + # via -r ./requirements.in +pygls==1.3.0 \ + --hash=sha256:1b44ace89c9382437a717534f490eadc6fda7c0c6c16ac1eaaf5568e345e4fb8 \ + --hash=sha256:d4a01414b6ed4e34e7e8fd29b77d3e88c29615df7d0bbff49bf019e15ec04b8f + # via -r ./requirements.in +pyyaml==6.0.1 \ + --hash=sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5 \ + --hash=sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc \ + --hash=sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df \ + --hash=sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741 \ + --hash=sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206 \ + --hash=sha256:18aeb1bf9a78867dc38b259769503436b7c72f7a1f1f4c93ff9a17de54319b27 \ + --hash=sha256:1d4c7e777c441b20e32f52bd377e0c409713e8bb1386e1099c2415f26e479595 \ + --hash=sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62 \ + --hash=sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98 \ + --hash=sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696 \ + --hash=sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290 \ + --hash=sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9 \ + --hash=sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d \ + --hash=sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6 \ + --hash=sha256:4fb147e7a67ef577a588a0e2c17b6db51dda102c71de36f8549b6816a96e1867 \ + --hash=sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47 \ + --hash=sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486 \ + --hash=sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6 \ + --hash=sha256:596106435fa6ad000c2991a98fa58eeb8656ef2325d7e158344fb33864ed87e3 \ + --hash=sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007 \ + --hash=sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938 \ + --hash=sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0 \ + --hash=sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c \ + --hash=sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735 \ + --hash=sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d \ + --hash=sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28 \ + --hash=sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4 \ + --hash=sha256:9046c58c4395dff28dd494285c82ba00b546adfc7ef001486fbf0324bc174fba \ + --hash=sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8 \ + --hash=sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef \ + --hash=sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5 \ + --hash=sha256:afd7e57eddb1a54f0f1a974bc4391af8bcce0b444685d936840f125cf046d5bd \ + --hash=sha256:b1275ad35a5d18c62a7220633c913e1b42d44b46ee12554e5fd39c70a243d6a3 \ + --hash=sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0 \ + --hash=sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515 \ + --hash=sha256:baa90d3f661d43131ca170712d903e6295d1f7a0f595074f151c0aed377c9b9c \ + --hash=sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c \ + --hash=sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924 \ + --hash=sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34 \ + --hash=sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43 \ + --hash=sha256:c8098ddcc2a85b61647b2590f825f3db38891662cfc2fc776415143f599bb859 \ + --hash=sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673 \ + --hash=sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54 \ + --hash=sha256:d858aa552c999bc8a8d57426ed01e40bef403cd8ccdd0fc5f6f04a00414cac2a \ + --hash=sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b \ + --hash=sha256:f003ed9ad21d6a4713f0a9b5a7a0a79e08dd0f221aff4525a2be4c346ee60aab \ + --hash=sha256:f22ac1c3cac4dbc50079e965eba2c1058622631e526bd9afd45fedd49ba781fa \ + --hash=sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c \ + --hash=sha256:fca0e3a251908a499833aa292323f32437106001d436eca0e6e7833256674585 \ + --hash=sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d \ + --hash=sha256:fd66fc5d0da6d9815ba2cebeb4205f95818ff4b79c3ebe268e75d961704af52f + # via -r ./requirements.in +watchdog==4.0.0 \ + --hash=sha256:11e12fafb13372e18ca1bbf12d50f593e7280646687463dd47730fd4f4d5d257 \ + --hash=sha256:2895bf0518361a9728773083908801a376743bcc37dfa252b801af8fd281b1ca \ + --hash=sha256:39cb34b1f1afbf23e9562501673e7146777efe95da24fab5707b88f7fb11649b \ + --hash=sha256:45cc09cc4c3b43fb10b59ef4d07318d9a3ecdbff03abd2e36e77b6dd9f9a5c85 \ + --hash=sha256:4986db5e8880b0e6b7cd52ba36255d4793bf5cdc95bd6264806c233173b1ec0b \ + --hash=sha256:5369136a6474678e02426bd984466343924d1df8e2fd94a9b443cb7e3aa20d19 \ + --hash=sha256:557ba04c816d23ce98a06e70af6abaa0485f6d94994ec78a42b05d1c03dcbd50 \ + --hash=sha256:6a4db54edea37d1058b08947c789a2354ee02972ed5d1e0dca9b0b820f4c7f92 \ + --hash=sha256:6a80d5cae8c265842c7419c560b9961561556c4361b297b4c431903f8c33b269 \ + --hash=sha256:6a9c71a0b02985b4b0b6d14b875a6c86ddea2fdbebd0c9a720a806a8bbffc69f \ + --hash=sha256:6c47bdd680009b11c9ac382163e05ca43baf4127954c5f6d0250e7d772d2b80c \ + --hash=sha256:6e949a8a94186bced05b6508faa61b7adacc911115664ccb1923b9ad1f1ccf7b \ + --hash=sha256:73c7a935e62033bd5e8f0da33a4dcb763da2361921a69a5a95aaf6c93aa03a87 \ + --hash=sha256:76ad8484379695f3fe46228962017a7e1337e9acadafed67eb20aabb175df98b \ + --hash=sha256:8350d4055505412a426b6ad8c521bc7d367d1637a762c70fdd93a3a0d595990b \ + --hash=sha256:87e9df830022488e235dd601478c15ad73a0389628588ba0b028cb74eb72fed8 \ + --hash=sha256:8f9a542c979df62098ae9c58b19e03ad3df1c9d8c6895d96c0d51da17b243b1c \ + --hash=sha256:8fec441f5adcf81dd240a5fe78e3d83767999771630b5ddfc5867827a34fa3d3 \ + --hash=sha256:9a03e16e55465177d416699331b0f3564138f1807ecc5f2de9d55d8f188d08c7 \ + --hash=sha256:ba30a896166f0fee83183cec913298151b73164160d965af2e93a20bbd2ab605 \ + --hash=sha256:c17d98799f32e3f55f181f19dd2021d762eb38fdd381b4a748b9f5a36738e935 \ + --hash=sha256:c522392acc5e962bcac3b22b9592493ffd06d1fc5d755954e6be9f4990de932b \ + --hash=sha256:d0f9bd1fd919134d459d8abf954f63886745f4660ef66480b9d753a7c9d40927 \ + --hash=sha256:d18d7f18a47de6863cd480734613502904611730f8def45fc52a5d97503e5101 \ + --hash=sha256:d31481ccf4694a8416b681544c23bd271f5a123162ab603c7d7d2dd7dd901a07 \ + --hash=sha256:e3e7065cbdabe6183ab82199d7a4f6b3ba0a438c5a512a68559846ccb76a78ec \ + --hash=sha256:eed82cdf79cd7f0232e2fdc1ad05b06a5e102a43e331f7d041e5f0e0a34a51c4 \ + --hash=sha256:f970663fa4f7e80401a7b0cbeec00fa801bf0287d93d48368fc3e6fa32716245 \ + --hash=sha256:f9b2fdca47dc855516b2d66eef3c39f2672cbf7e7a42e7e67ad2cbfcd6ba107d + # via -r ./requirements.in diff --git a/resources/logo.png b/resources/logo.png new file mode 100644 index 00000000..9a4c55c6 Binary files /dev/null and b/resources/logo.png differ diff --git a/resources/python.png b/resources/python.png new file mode 100644 index 00000000..517bd26a Binary files /dev/null and b/resources/python.png differ diff --git a/resources/zenml-extension.gif b/resources/zenml-extension.gif new file mode 100644 index 00000000..f3c30739 Binary files /dev/null and b/resources/zenml-extension.gif differ diff --git a/resources/zenml_logo.png b/resources/zenml_logo.png new file mode 100644 index 00000000..484b67e4 Binary files /dev/null and b/resources/zenml_logo.png differ diff --git a/runtime.txt b/runtime.txt new file mode 100644 index 00000000..7739b5ac --- /dev/null +++ b/runtime.txt @@ -0,0 +1 @@ +python-3.8.18 \ No newline at end of file diff --git a/scripts/clear_and_compile.sh b/scripts/clear_and_compile.sh new file mode 100755 index 00000000..e45d6f6f --- /dev/null +++ b/scripts/clear_and_compile.sh @@ -0,0 +1,21 @@ +#!/bin/bash + +echo "Removing bundled/tool/__pycache__..." +if [ -d "bundled/tool/__pycache__" ]; then + rm -rf bundled/tool/__pycache__ +fi + +echo "Removing dist directory..." +if [ -d "dist" ]; then + rm -rf dist +fi + +echo "Recompiling with npm..." +if ! command -v npm &> /dev/null +then + echo "npm could not be found. Please install npm to proceed." + exit 1 +fi +npm run compile + +echo "Operation completed." diff --git a/scripts/format.sh b/scripts/format.sh new file mode 100755 index 00000000..20e0aef8 --- /dev/null +++ b/scripts/format.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +set -euxo pipefail + +# Source directory for Python tool +python_src="bundled/tool" + +# Process arguments +for arg in "$@"; do + if [ "$arg" = "--dry-run" ]; then + echo "Dry run mode enabled" + dry_run=true + fi +done + +echo "Formatting Python files in bundled/tool..." +find $python_src -name "*.py" -print0 | xargs -0 autoflake --in-place --remove-all-unused-imports +find $python_src -name "*.py" -print0 | xargs -0 isort -- +find $python_src -name "*.py" -print0 | xargs -0 black --line-length 100 -- + +# For now, Typescript formatting is separate from Python +# echo "Formatting TypeScript files..." +# npx prettier --ignore-path .gitignore --write "**/*.+(ts|json)" + +echo "Formatting complete." diff --git a/scripts/lint.sh b/scripts/lint.sh new file mode 100755 index 00000000..e8acc6d6 --- /dev/null +++ b/scripts/lint.sh @@ -0,0 +1,10 @@ +#!/bin/bash + +cd "$(dirname "$0")/.." || exit + +PYTHONPATH="$PYTHONPATH:$(pwd)/bundled/tool" +export PYTHONPATH + +nox --session lint + +unset PYTHONPATH diff --git a/src/commands/environment/cmds.ts b/src/commands/environment/cmds.ts new file mode 100644 index 00000000..9ac987b5 --- /dev/null +++ b/src/commands/environment/cmds.ts @@ -0,0 +1,95 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. + +import { ProgressLocation, commands, window } from 'vscode'; +import { getInterpreterFromWorkspaceSettings } from '../../common/settings'; +import { EnvironmentDataProvider } from '../../views/activityBar/environmentView/EnvironmentDataProvider'; +import { LSP_ZENML_CLIENT_INITIALIZED, PYTOOL_MODULE } from '../../utils/constants'; +import { LSClient } from '../../services/LSClient'; +import { EventBus } from '../../services/EventBus'; +import { REFRESH_ENVIRONMENT_VIEW } from '../../utils/constants'; + +/** + * Set the Python interpreter for the current workspace. + * + * @returns {Promise} Resolves after refreshing the view. + */ +const setPythonInterpreter = async (): Promise => { + await window.withProgress( + { + location: ProgressLocation.Window, + title: 'Refreshing server status...', + }, + async progress => { + progress.report({ increment: 10 }); + const currentInterpreter = await getInterpreterFromWorkspaceSettings(); + + await commands.executeCommand('python.setInterpreter'); + + const newInterpreter = await getInterpreterFromWorkspaceSettings(); + + if (newInterpreter === currentInterpreter) { + console.log('Interpreter selection unchanged or cancelled. No server restart required.'); + window.showInformationMessage('Interpreter selection unchanged. Restart not required.'); + return; + } + progress.report({ increment: 90 }); + console.log('Interpreter selection completed.'); + + window.showInformationMessage( + 'ZenML server will restart to apply the new interpreter settings.' + ); + // The onDidChangePythonInterpreter event will trigger the server restart. + progress.report({ increment: 100 }); + } + ); +}; + +const refreshEnvironmentView = async (): Promise => { + window.withProgress( + { + location: ProgressLocation.Notification, + title: 'Refreshing Environment View...', + cancellable: false, + }, + async () => { + EnvironmentDataProvider.getInstance().refresh(); + } + ); +}; + +const restartLSPServer = async (): Promise => { + await window.withProgress( + { + location: ProgressLocation.Window, + title: 'Restarting LSP Server...', + }, + async progress => { + progress.report({ increment: 10 }); + const lsClient = LSClient.getInstance(); + lsClient.isZenMLReady = false; + lsClient.localZenML = { is_installed: false, version: '' }; + const eventBus = EventBus.getInstance(); + eventBus.emit(REFRESH_ENVIRONMENT_VIEW); + eventBus.emit(LSP_ZENML_CLIENT_INITIALIZED, false); + await commands.executeCommand(`${PYTOOL_MODULE}.restart`); + progress.report({ increment: 100 }); + } + ); +}; + +export const environmentCommands = { + setPythonInterpreter, + refreshEnvironmentView, + restartLSPServer, +}; diff --git a/src/commands/environment/registry.ts b/src/commands/environment/registry.ts new file mode 100644 index 00000000..1153ab31 --- /dev/null +++ b/src/commands/environment/registry.ts @@ -0,0 +1,50 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { registerCommand } from '../../common/vscodeapi'; +import { environmentCommands } from './cmds'; +import { ZenExtension } from '../../services/ZenExtension'; +import { ExtensionContext, commands } from 'vscode'; + +/** + * Registers pipeline-related commands for the extension. + * + * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. + */ +export const registerEnvironmentCommands = (context: ExtensionContext) => { + try { + const registeredCommands = [ + registerCommand( + 'zenml.setPythonInterpreter', + async () => await environmentCommands.setPythonInterpreter() + ), + registerCommand( + 'zenml.refreshEnvironmentView', + async () => await environmentCommands.refreshEnvironmentView() + ), + registerCommand( + 'zenml.restartLspServer', + async () => await environmentCommands.restartLSPServer() + ), + ]; + + registeredCommands.forEach(cmd => { + context.subscriptions.push(cmd); + ZenExtension.commandDisposables.push(cmd); + }); + + commands.executeCommand('setContext', 'environmentCommandsRegistered', true); + } catch (error) { + console.error('Error registering environment commands:', error); + commands.executeCommand('setContext', 'environmentCommandsRegistered', false); + } +}; diff --git a/src/commands/pipelines/cmds.ts b/src/commands/pipelines/cmds.ts new file mode 100644 index 00000000..92e3c4ba --- /dev/null +++ b/src/commands/pipelines/cmds.ts @@ -0,0 +1,78 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { LSClient } from '../../services/LSClient'; +import { showErrorMessage, showInformationMessage } from '../../utils/notifications'; +import { PipelineTreeItem } from '../../views/activityBar'; +import { PipelineDataProvider } from '../../views/activityBar/pipelineView/PipelineDataProvider'; +import * as vscode from 'vscode'; + +/** + * Triggers a refresh of the pipeline view within the UI components. + * + * @returns {Promise} Resolves after refreshing the view. + */ +const refreshPipelineView = async (): Promise => { + await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Window, + title: 'Refreshing server status...', + }, + async () => { + await PipelineDataProvider.getInstance().refresh(); + } + ); +}; + +/** + * Deletes a pipeline run. + * + * @param {PipelineTreeItem} node The pipeline run to delete. + * @returns {Promise} Resolves after deleting the pipeline run. + */ +const deletePipelineRun = async (node: PipelineTreeItem): Promise => { + const userConfirmation = await vscode.window.showWarningMessage( + 'Are you sure you want to delete this pipeline run?', + { modal: true }, + 'Yes', + 'No' + ); + + if (userConfirmation === 'Yes') { + await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Window, + title: 'Deleting pipeline run...', + }, + async () => { + const runId = node.id; + try { + const lsClient = LSClient.getInstance(); + const result = await lsClient.sendLsClientRequest('deletePipelineRun', [runId]); + if (result && 'error' in result) { + throw new Error(result.error); + } + showInformationMessage('Pipeline run deleted successfully.'); + await refreshPipelineView(); + } catch (error: any) { + console.error(`Error deleting pipeline run: ${error}`); + showErrorMessage(`Failed to delete pipeline run: ${error.message}`); + } + } + ); + } +}; + +export const pipelineCommands = { + refreshPipelineView, + deletePipelineRun, +}; diff --git a/src/commands/pipelines/registry.ts b/src/commands/pipelines/registry.ts new file mode 100644 index 00000000..88b18159 --- /dev/null +++ b/src/commands/pipelines/registry.ts @@ -0,0 +1,52 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { pipelineCommands } from './cmds'; +import { registerCommand } from '../../common/vscodeapi'; +import { ZenExtension } from '../../services/ZenExtension'; +import { PipelineDataProvider, PipelineTreeItem } from '../../views/activityBar'; +import { ExtensionContext, commands } from 'vscode'; + +/** + * Registers pipeline-related commands for the extension. + * + * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. + */ +export const registerPipelineCommands = (context: ExtensionContext) => { + const pipelineDataProvider = PipelineDataProvider.getInstance(); + + try { + const registeredCommands = [ + registerCommand( + 'zenml.refreshPipelineView', + async () => await pipelineCommands.refreshPipelineView() + ), + registerCommand( + 'zenml.deletePipelineRun', + async (node: PipelineTreeItem) => await pipelineCommands.deletePipelineRun(node) + ), + registerCommand('zenml.nextPipelineRunsPage', async () => pipelineDataProvider.goToNextPage()), + registerCommand('zenml.previousPipelineRunsPage', async () => pipelineDataProvider.goToPreviousPage()), + registerCommand("zenml.setPipelineRunsPerPage", async () => await pipelineDataProvider.updateItemsPerPage()), + ]; + + registeredCommands.forEach(cmd => { + context.subscriptions.push(cmd); + ZenExtension.commandDisposables.push(cmd); + }); + + commands.executeCommand('setContext', 'pipelineCommandsRegistered', true); + } catch (error) { + console.error('Error registering pipeline commands:', error); + commands.executeCommand('setContext', 'pipelineCommandsRegistered', false); + } +}; diff --git a/src/commands/server/cmds.ts b/src/commands/server/cmds.ts new file mode 100644 index 00000000..8224828c --- /dev/null +++ b/src/commands/server/cmds.ts @@ -0,0 +1,126 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as vscode from 'vscode'; +import { LSClient } from '../../services/LSClient'; +import { + ConnectServerResponse, + GenericLSClientResponse, + RestServerConnectionResponse, +} from '../../types/LSClientResponseTypes'; +import { updateServerUrlAndToken } from '../../utils/global'; +import { showInformationMessage } from '../../utils/notifications'; +import { refreshUtils } from '../../utils/refresh'; +import { ServerDataProvider } from '../../views/activityBar'; +import { promptAndStoreServerUrl } from './utils'; + +/** + * Initiates a connection to the ZenML server using a Flask service for OAuth2 authentication. + * The service handles user authentication, device authorization, and updates the global configuration upon success. + * + * @returns {Promise} Resolves after attempting to connect to the server. + */ +const connectServer = async (): Promise => { + const url = await promptAndStoreServerUrl(); + + if (!url) { + return false; + } + + return new Promise(resolve => { + vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: 'Connecting to ZenML server...', + cancellable: false, + }, + async progress => { + try { + const lsClient = LSClient.getInstance(); + const result = await lsClient.sendLsClientRequest('connect', [ + url, + ]); + + if (result && 'error' in result) { + throw new Error(result.error); + } + + const accessToken = (result as RestServerConnectionResponse).access_token; + await updateServerUrlAndToken(url, accessToken); + await refreshUtils.refreshUIComponents(); + resolve(true); + } catch (error) { + console.error('Failed to connect to ZenML server:', error); + vscode.window.showErrorMessage( + `Failed to connect to ZenML server: ${(error as Error).message}` + ); + resolve(false); + } + } + ); + }); +}; + +/** + * Disconnects from the ZenML server and clears related configuration and state in the application. + * + * @returns {Promise} Resolves after successfully disconnecting from the server. + */ +const disconnectServer = async (): Promise => { + await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: 'Disconnecting from ZenML server...', + cancellable: false, + }, + async progress => { + try { + const lsClient = LSClient.getInstance(); + const result = await lsClient.sendLsClientRequest('disconnect'); + if (result && 'error' in result) { + throw result; + } + await refreshUtils.refreshUIComponents(); + } catch (error: any) { + console.error('Failed to disconnect from ZenML server:', error); + vscode.window.showErrorMessage( + 'Failed to disconnect from ZenML server: ' + error.message || error + ); + } + } + ); +}; + +/** + * Triggers a refresh of the server status within the UI components. + * + * @returns {Promise} Resolves after refreshing the server status. + */ +const refreshServerStatus = async (): Promise => { + await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Window, + title: 'Refreshing server status...', + cancellable: false, + }, + async () => { + await ServerDataProvider.getInstance().refresh(); + showInformationMessage('Server status refreshed.'); + } + ); +}; + +export const serverCommands = { + connectServer, + disconnectServer, + refreshServerStatus, +}; diff --git a/src/commands/server/registry.ts b/src/commands/server/registry.ts new file mode 100644 index 00000000..d84b1e16 --- /dev/null +++ b/src/commands/server/registry.ts @@ -0,0 +1,47 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { serverCommands } from './cmds'; +import { registerCommand } from '../../common/vscodeapi'; +import { ZenExtension } from '../../services/ZenExtension'; +import { ExtensionContext, commands } from 'vscode'; + +/** + * Registers server-related commands for the extension. + * + * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. + */ +export const registerServerCommands = (context: ExtensionContext) => { + try { + const registeredCommands = [ + registerCommand('zenml.connectServer', async () => await serverCommands.connectServer()), + registerCommand( + 'zenml.disconnectServer', + async () => await serverCommands.disconnectServer() + ), + registerCommand( + 'zenml.refreshServerStatus', + async () => await serverCommands.refreshServerStatus() + ), + ]; + + registeredCommands.forEach(cmd => { + context.subscriptions.push(cmd); + ZenExtension.commandDisposables.push(cmd); + }); + + commands.executeCommand('setContext', 'serverCommandsRegistered', true); + } catch (error) { + console.error('Error registering server commands:', error); + commands.executeCommand('setContext', 'serverCommandsRegistered', false); + } +}; diff --git a/src/commands/server/utils.ts b/src/commands/server/utils.ts new file mode 100644 index 00000000..236e729c --- /dev/null +++ b/src/commands/server/utils.ts @@ -0,0 +1,95 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as vscode from 'vscode'; +import { ServerStatus, ZenServerDetails } from '../../types/ServerInfoTypes'; +import { LSClient } from '../../services/LSClient'; +import { INITIAL_ZENML_SERVER_STATUS, PYTOOL_MODULE } from '../../utils/constants'; +import { ServerStatusInfoResponse } from '../../types/LSClientResponseTypes'; +import { ErrorTreeItem, createErrorItem } from '../../views/activityBar/common/ErrorTreeItem'; + +/** + * Prompts the user to enter the ZenML server URL and stores it in the global configuration. + */ +export async function promptAndStoreServerUrl(): Promise { + let serverUrl = await vscode.window.showInputBox({ + prompt: 'Enter the ZenML server URL', + placeHolder: 'https://', + }); + + serverUrl = serverUrl?.trim(); + + if (serverUrl) { + serverUrl = serverUrl.replace(/\/$/, ''); + // Validate the server URL format before storing + if (!/^https?:\/\/[^\s$.?#].[^\s]*$/.test(serverUrl)) { + vscode.window.showErrorMessage('Invalid server URL format.'); + return; + } + const config = vscode.workspace.getConfiguration('zenml'); + await config.update('serverUrl', serverUrl, vscode.ConfigurationTarget.Global); + } + + return serverUrl; +} + +/** + * Retrieves the server status from the language server or the provided server details. + * + * @returns {Promise} A promise that resolves with the server status, parsed from server details. + */ +export async function checkServerStatus(): Promise { + const lsClient = LSClient.getInstance(); + // For debugging + if (!lsClient.clientReady) { + return INITIAL_ZENML_SERVER_STATUS; + } + + try { + const result = await lsClient.sendLsClientRequest('serverInfo'); + if (!result || 'error' in result) { + if ('clientVersion' in result && 'serverVersion' in result) { + return createErrorItem(result); + } + } else if (isZenServerDetails(result)) { + return createServerStatusFromDetails(result); + } + } catch (error) { + console.error('Failed to fetch server information:', error); + } + return INITIAL_ZENML_SERVER_STATUS; +} + +function isZenServerDetails(response: any): response is ZenServerDetails { + return response && 'storeInfo' in response && 'storeConfig' in response; +} + +function createServerStatusFromDetails(details: ZenServerDetails): ServerStatus { + const { storeInfo, storeConfig } = details; + return { + ...storeInfo, + isConnected: storeConfig?.type === 'rest', + url: storeConfig?.url ?? 'unknown', + store_type: storeConfig?.type ?? 'unknown', + }; +} + +export function isServerStatus(obj: any): obj is ServerStatus { + return 'isConnected' in obj && 'url' in obj; +} + +export const serverUtils = { + promptAndStoreServerUrl, + checkServerStatus, + isZenServerDetails, + createServerStatusFromDetails, +}; diff --git a/src/commands/stack/cmds.ts b/src/commands/stack/cmds.ts new file mode 100644 index 00000000..b76b11dd --- /dev/null +++ b/src/commands/stack/cmds.ts @@ -0,0 +1,166 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as vscode from 'vscode'; +import { StackDataProvider, StackTreeItem } from '../../views/activityBar'; +import ZenMLStatusBar from '../../views/statusBar'; +import { switchActiveStack } from './utils'; +import { LSClient } from '../../services/LSClient'; +import { showInformationMessage } from '../../utils/notifications'; + +/** + * Refreshes the stack view. + */ +const refreshStackView = async () => { + vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: 'Refreshing Stack View...', + cancellable: false, + }, + async progress => { + await StackDataProvider.getInstance().refresh(); + } + ); +}; + +/** + * Refreshes the active stack. + */ +const refreshActiveStack = async () => { + const statusBar = ZenMLStatusBar.getInstance(); + + vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: 'Refreshing Active Stack...', + cancellable: false, + }, + async progress => { + await statusBar.refreshActiveStack(); + } + ); +}; + +/** + * Renames the selected stack to a new name. + * + * @param node The stack to rename. + * @returns {Promise} Resolves after renaming the stack. + */ +const renameStack = async (node: StackTreeItem): Promise => { + const newStackName = await vscode.window.showInputBox({ prompt: 'Enter new stack name' }); + if (!newStackName) { + return; + } + vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: 'Renaming Stack...', + cancellable: false, + }, + async () => { + try { + const { label, id } = node; + const lsClient = LSClient.getInstance(); + const result = await lsClient.sendLsClientRequest('renameStack', [id, newStackName]); + if (result && 'error' in result) { + throw new Error(result.error); + } + showInformationMessage(`Stack ${label} successfully renamed to ${newStackName}.`); + await StackDataProvider.getInstance().refresh(); + } catch (error: any) { + if (error.response) { + vscode.window.showErrorMessage(`Failed to rename stack: ${error.response.data.message}`); + } else { + console.error('Failed to rename stack:', error); + vscode.window.showErrorMessage('Failed to rename stack'); + } + } + } + ); +}; + +/** + * Copies the selected stack to a new stack with a specified name. + * + * @param {StackTreeItem} node The stack to copy. + */ +const copyStack = async (node: StackTreeItem) => { + const newStackName = await vscode.window.showInputBox({ + prompt: 'Enter the name for the copied stack', + }); + if (!newStackName) { + return; + } + vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: 'Copying Stack...', + cancellable: false, + }, + async progress => { + try { + const lsClient = LSClient.getInstance(); + const result = await lsClient.sendLsClientRequest('copyStack', [node.id, newStackName]); + if ('error' in result && result.error) { + throw new Error(result.error); + } + showInformationMessage('Stack copied successfully.'); + await StackDataProvider.getInstance().refresh(); + } catch (error: any) { + if (error.response && error.response.data && error.response.data.message) { + vscode.window.showErrorMessage(`Failed to copy stack: ${error.response.data.message}`); + } else { + console.error('Failed to copy stack:', error); + vscode.window.showErrorMessage(`Failed to copy stack: ${error.message || error}`); + } + } + } + ); +}; + +/** + * Sets the selected stack as the active stack and stores it in the global context. + * + * @param {StackTreeItem} node The stack to activate. + * @returns {Promise} Resolves after setting the active stack. + */ +const setActiveStack = async (node: StackTreeItem): Promise => { + vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: 'Setting Active Stack...', + cancellable: false, + }, + async () => { + try { + const result = await switchActiveStack(node.id); + if (result) { + const { id, name } = result; + showInformationMessage(`Active stack set to: ${name}`); + } + } catch (error) { + console.log(error); + vscode.window.showErrorMessage(`Failed to set active stack: ${error}`); + } + } + ); +}; + +export const stackCommands = { + refreshStackView, + refreshActiveStack, + renameStack, + copyStack, + setActiveStack, +}; diff --git a/src/commands/stack/registry.ts b/src/commands/stack/registry.ts new file mode 100644 index 00000000..eba47594 --- /dev/null +++ b/src/commands/stack/registry.ts @@ -0,0 +1,60 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { StackDataProvider, StackTreeItem } from '../../views/activityBar'; +import { stackCommands } from './cmds'; +import { registerCommand } from '../../common/vscodeapi'; +import { ZenExtension } from '../../services/ZenExtension'; +import { ExtensionContext, commands, window } from 'vscode'; + +/** + * Registers stack-related commands for the extension. + * + * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. + */ +export const registerStackCommands = (context: ExtensionContext) => { + const stackDataProvider = StackDataProvider.getInstance(); + try { + const registeredCommands = [ + registerCommand("zenml.setStackItemsPerPage", async () => await stackDataProvider.updateItemsPerPage()), + registerCommand('zenml.refreshStackView', async () => await stackCommands.refreshStackView()), + registerCommand( + 'zenml.refreshActiveStack', + async () => await stackCommands.refreshActiveStack() + ), + registerCommand( + 'zenml.renameStack', + async (node: StackTreeItem) => await stackCommands.renameStack(node) + ), + registerCommand( + 'zenml.setActiveStack', + async (node: StackTreeItem) => await stackCommands.setActiveStack(node) + ), + registerCommand( + 'zenml.copyStack', + async (node: StackTreeItem) => await stackCommands.copyStack(node) + ), + registerCommand('zenml.nextStackPage', async () => stackDataProvider.goToNextPage()), + registerCommand('zenml.previousStackPage', async () => stackDataProvider.goToPreviousPage()), + ]; + + registeredCommands.forEach(cmd => { + context.subscriptions.push(cmd); + ZenExtension.commandDisposables.push(cmd); + }); + + commands.executeCommand('setContext', 'stackCommandsRegistered', true); + } catch (error) { + console.error('Error registering stack commands:', error); + commands.executeCommand('setContext', 'stackCommandsRegistered', false); + } +}; diff --git a/src/commands/stack/utils.ts b/src/commands/stack/utils.ts new file mode 100644 index 00000000..5ec2321f --- /dev/null +++ b/src/commands/stack/utils.ts @@ -0,0 +1,91 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as vscode from 'vscode'; +import { LSClient } from '../../services/LSClient'; +import { GetActiveStackResponse, SetActiveStackResponse } from '../../types/LSClientResponseTypes'; +import { showErrorMessage } from '../../utils/notifications'; + +/** + * Switches the active ZenML stack to the specified stack name. + * + * @param {string} stackNameOrId - The id or name of the ZenML stack to be activated. + * @returns {Promise<{id: string, name: string}>} A promise that resolves with the id and name of the newly activated stack, or undefined on error. + */ +export const switchActiveStack = async ( + stackNameOrId: string +): Promise<{ id: string; name: string } | undefined> => { + try { + const lsClient = LSClient.getInstance(); + const result = await lsClient.sendLsClientRequest('switchActiveStack', [ + stackNameOrId, + ]); + if (result && 'error' in result) { + console.log('Error in switchZenMLStack result', result); + throw new Error(result.error); + } + const { id, name } = result; + await storeActiveStack(id); + return { id, name }; + } catch (error: any) { + console.error(`Error setting active stack: ${error}`); + showErrorMessage(`Failed to set active stack: ${error.message}`); + } +}; + +/** + * Gets the id and name of the active ZenML stack. + * + * @returns {Promise<{id: string, name: string}>} A promise that resolves with the id and name of the active stack, or undefined on error; + */ +export const getActiveStack = async (): Promise<{ id: string; name: string } | undefined> => { + const lsClient = LSClient.getInstance(); + if (!lsClient.clientReady) { + return; + } + + try { + const result = await lsClient.sendLsClientRequest('getActiveStack'); + if (result && 'error' in result) { + throw new Error(result.error); + } + return result; + } catch (error: any) { + console.error(`Error getting active stack information: ${error}`); + showErrorMessage(`Failed to get active stack information: ${error.message}`); + return undefined; + } +}; + +/** + * Stores the specified ZenML stack id in the global configuration. + * + * @param {string} id - The id of the ZenML stack to be stored. + * @returns {Promise} A promise that resolves when the stack information has been successfully stored. + */ +export const storeActiveStack = async (id: string): Promise => { + const config = vscode.workspace.getConfiguration('zenml'); + await config.update('activeStackId', id, vscode.ConfigurationTarget.Global); +}; + +export const getActiveStackIdFromConfig = (): string | undefined => { + const config = vscode.workspace.getConfiguration('zenml'); + return config.get('activeStackId'); +}; + +const stackUtils = { + switchActiveStack, + getActiveStack, + storeActiveStack, +}; + +export default stackUtils; diff --git a/src/common/constants.ts b/src/common/constants.ts new file mode 100644 index 00000000..e50ed77f --- /dev/null +++ b/src/common/constants.ts @@ -0,0 +1,25 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as path from 'path'; + +export const EXTENSION_ID = 'ZenML.zenml'; +const folderName = path.basename(__dirname); +export const EXTENSION_ROOT_DIR = + folderName === 'common' ? path.dirname(path.dirname(__dirname)) : path.dirname(__dirname); +export const BUNDLED_PYTHON_SCRIPTS_DIR = path.join(EXTENSION_ROOT_DIR, 'bundled'); +export const SERVER_SCRIPT_PATH = path.join(BUNDLED_PYTHON_SCRIPTS_DIR, 'tool', `lsp_server.py`); +export const DEBUG_SERVER_SCRIPT_PATH = path.join( + BUNDLED_PYTHON_SCRIPTS_DIR, + 'tool', + `_debug_server.py` +); diff --git a/src/common/log/logging.ts b/src/common/log/logging.ts new file mode 100644 index 00000000..cbd1e152 --- /dev/null +++ b/src/common/log/logging.ts @@ -0,0 +1,70 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. + +import * as util from 'util'; +import { Disposable, LogOutputChannel } from 'vscode'; + +type Arguments = unknown[]; +class OutputChannelLogger { + constructor(private readonly channel: LogOutputChannel) {} + + public traceLog(...data: Arguments): void { + this.channel.appendLine(util.format(...data)); + } + + public traceError(...data: Arguments): void { + this.channel.error(util.format(...data)); + } + + public traceWarn(...data: Arguments): void { + this.channel.warn(util.format(...data)); + } + + public traceInfo(...data: Arguments): void { + this.channel.info(util.format(...data)); + } + + public traceVerbose(...data: Arguments): void { + this.channel.debug(util.format(...data)); + } +} + +let channel: OutputChannelLogger | undefined; +export function registerLogger(logChannel: LogOutputChannel): Disposable { + channel = new OutputChannelLogger(logChannel); + return { + dispose: () => { + channel = undefined; + }, + }; +} + +export function traceLog(...args: Arguments): void { + channel?.traceLog(...args); +} + +export function traceError(...args: Arguments): void { + channel?.traceError(...args); +} + +export function traceWarn(...args: Arguments): void { + channel?.traceWarn(...args); +} + +export function traceInfo(...args: Arguments): void { + channel?.traceInfo(...args); +} + +export function traceVerbose(...args: Arguments): void { + channel?.traceVerbose(...args); +} diff --git a/src/common/python.ts b/src/common/python.ts new file mode 100644 index 00000000..621e9f74 --- /dev/null +++ b/src/common/python.ts @@ -0,0 +1,107 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. + +/* eslint-disable @typescript-eslint/naming-convention */ +import { commands, Disposable, Event, EventEmitter, Uri } from 'vscode'; +import { traceError, traceLog } from './log/logging'; +import { PythonExtension, ResolvedEnvironment } from '@vscode/python-extension'; +export interface IInterpreterDetails { + path?: string[]; + resource?: Uri; +} + +const onDidChangePythonInterpreterEvent = new EventEmitter(); +export const onDidChangePythonInterpreter: Event = + onDidChangePythonInterpreterEvent.event; + +let _api: PythonExtension | undefined; +async function getPythonExtensionAPI(): Promise { + if (_api) { + return _api; + } + _api = await PythonExtension.api(); + return _api; +} + +export async function initializePython(disposables: Disposable[]): Promise { + try { + const api = await getPythonExtensionAPI(); + + if (api) { + disposables.push( + api.environments.onDidChangeActiveEnvironmentPath(e => { + onDidChangePythonInterpreterEvent.fire({ path: [e.path], resource: e.resource?.uri }); + }) + ); + + traceLog('Waiting for interpreter from python extension.'); + onDidChangePythonInterpreterEvent.fire(await getInterpreterDetails()); + } + } catch (error) { + traceError('Error initializing python: ', error); + } +} + +export async function resolveInterpreter( + interpreter: string[] +): Promise { + const api = await getPythonExtensionAPI(); + return api?.environments.resolveEnvironment(interpreter[0]); +} + +export async function getInterpreterDetails(resource?: Uri): Promise { + const api = await getPythonExtensionAPI(); + const environment = await api?.environments.resolveEnvironment( + api?.environments.getActiveEnvironmentPath(resource) + ); + if (environment?.executable.uri && checkVersion(environment)) { + return { path: [environment?.executable.uri.fsPath], resource }; + } + return { path: undefined, resource }; +} + +export async function getDebuggerPath(): Promise { + const api = await getPythonExtensionAPI(); + return api?.debug.getDebuggerPackagePath(); +} + +export async function runPythonExtensionCommand(command: string, ...rest: any[]) { + await getPythonExtensionAPI(); + return await commands.executeCommand(command, ...rest); +} + +export function checkVersion(resolved: ResolvedEnvironment | undefined): boolean { + const version = resolved?.version; + if (version?.major === 3 && version?.minor >= 8) { + traceLog(`Python version ${version?.major}.${version?.minor}.${version?.micro} is supported.`); + return true; + } + traceError(`Python version ${version?.major}.${version?.minor} is not supported.`); + traceError(`Selected python path: ${resolved?.executable.uri?.fsPath}`); + traceError('Supported versions are 3.8 and above.'); + return false; +} + +export function isPythonVersonSupported(resolvedEnv: ResolvedEnvironment | undefined): { + isSupported: boolean; + message?: string; +} { + const version = resolvedEnv?.version; + + if (version?.major === 3 && version?.minor >= 8) { + return { isSupported: true }; + } + + const errorMessage = `Unsupported Python ${version?.major}.${version?.minor}; requires >= 3.8.`; + return { isSupported: false, message: errorMessage }; +} diff --git a/src/common/server.ts b/src/common/server.ts new file mode 100644 index 00000000..e0e0dfe8 --- /dev/null +++ b/src/common/server.ts @@ -0,0 +1,178 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as fsapi from 'fs-extra'; +import * as vscode from 'vscode'; +import { Disposable, env, l10n, LanguageStatusSeverity, LogOutputChannel } from 'vscode'; +import { State } from 'vscode-languageclient'; +import { + LanguageClient, + LanguageClientOptions, + RevealOutputChannelOn, + ServerOptions, +} from 'vscode-languageclient/node'; +import { EventBus } from '../services/EventBus'; +import { LSClient } from '../services/LSClient'; +import { ZenExtension } from '../services/ZenExtension'; +import { LSCLIENT_STATE_CHANGED } from '../utils/constants'; +import { toggleCommands } from '../utils/global'; +import { DEBUG_SERVER_SCRIPT_PATH, SERVER_SCRIPT_PATH } from './constants'; +import { traceError, traceInfo, traceVerbose } from './log/logging'; +import { getDebuggerPath } from './python'; +import { + getExtensionSettings, + getGlobalSettings, + getWorkspaceSettings, + ISettings, +} from './settings'; +import { updateStatus } from './status'; +import { getLSClientTraceLevel, getProjectRoot } from './utilities'; +import { isVirtualWorkspace } from './vscodeapi'; + +export type IInitOptions = { settings: ISettings[]; globalSettings: ISettings }; + +async function createServer( + settings: ISettings, + serverId: string, + serverName: string, + outputChannel: LogOutputChannel, + initializationOptions: IInitOptions +): Promise { + const command = settings.interpreter[0]; + // console.log('command is', command); + + const cwd = settings.cwd; + + // Set debugger path needed for debugging python code. + const newEnv = { ...process.env }; + const debuggerPath = await getDebuggerPath(); + const isDebugScript = await fsapi.pathExists(DEBUG_SERVER_SCRIPT_PATH); + if (newEnv.USE_DEBUGPY && debuggerPath) { + newEnv.DEBUGPY_PATH = debuggerPath; + } else { + newEnv.USE_DEBUGPY = 'False'; + } + + // Set import strategy + newEnv.LS_IMPORT_STRATEGY = settings.importStrategy; + + // Set notification type + newEnv.LS_SHOW_NOTIFICATION = settings.showNotifications; + + const args = + newEnv.USE_DEBUGPY === 'False' || !isDebugScript + ? settings.interpreter.slice(1).concat([SERVER_SCRIPT_PATH]) + : settings.interpreter.slice(1).concat([DEBUG_SERVER_SCRIPT_PATH]); + traceInfo(`Server run command: ${[command, ...args].join(' ')}`); + + const serverOptions: ServerOptions = { + command, + args, + options: { cwd, env: newEnv }, + }; + + // Options to control the language client + const clientOptions: LanguageClientOptions = { + // Register the server for python documents + documentSelector: isVirtualWorkspace() + ? [{ language: 'python' }] + : [ + { scheme: 'file', language: 'python' }, + { scheme: 'untitled', language: 'python' }, + { scheme: 'vscode-notebook', language: 'python' }, + { scheme: 'vscode-notebook-cell', language: 'python' }, + ], + outputChannel: outputChannel, + traceOutputChannel: outputChannel, + revealOutputChannelOn: RevealOutputChannelOn.Never, + initializationOptions: { + ...initializationOptions, + interpreter: settings.interpreter[0], + }, + }; + + return new LanguageClient(serverId, serverName, serverOptions, clientOptions); +} + +let _disposables: Disposable[] = []; +export async function restartServer( + workspaceSetting: ISettings +): Promise { + const lsClientInstance = LSClient.getInstance(); + const lsClient = lsClientInstance.getLanguageClient(); + if (lsClient) { + traceInfo(`Server: Stop requested`); + await lsClient.stop(); + _disposables.forEach(d => d.dispose()); + _disposables = []; + } + updateStatus(undefined, LanguageStatusSeverity.Information, true); + + const newLSClient = await createServer( + workspaceSetting, + ZenExtension.serverId, + ZenExtension.serverName, + ZenExtension.outputChannel, + { + settings: await getExtensionSettings(ZenExtension.serverId, true), + globalSettings: await getGlobalSettings(ZenExtension.serverId, true), + } + ); + + lsClientInstance.updateClient(newLSClient); + + traceInfo(`Server: Start requested.`); + _disposables.push( + newLSClient.onDidChangeState(e => { + EventBus.getInstance().emit(LSCLIENT_STATE_CHANGED, e.newState); + switch (e.newState) { + case State.Stopped: + traceVerbose(`Server State: Stopped`); + break; + case State.Starting: + traceVerbose(`Server State: Starting`); + break; + case State.Running: + traceVerbose(`Server State: Running`); + updateStatus(undefined, LanguageStatusSeverity.Information, false); + break; + } + }) + ); + try { + await ZenExtension.lsClient.startLanguageClient(); + } catch (ex) { + updateStatus(l10n.t('Server failed to start.'), LanguageStatusSeverity.Error); + traceError(`Server: Start failed: ${ex}`); + } + await newLSClient.setTrace( + getLSClientTraceLevel(ZenExtension.outputChannel.logLevel, env.logLevel) + ); + return newLSClient; +} + +export async function runServer() { + await toggleCommands(false); + + const projectRoot = await getProjectRoot(); + const workspaceSetting = await getWorkspaceSettings(ZenExtension.serverId, projectRoot, true); + if (workspaceSetting.interpreter.length === 0) { + updateStatus( + vscode.l10n.t('Please select a Python interpreter.'), + vscode.LanguageStatusSeverity.Error + ); + traceError('Python interpreter missing. Please use Python 3.8 or greater.'); + return; + } + + await restartServer(workspaceSetting); +} diff --git a/src/common/settings.ts b/src/common/settings.ts new file mode 100644 index 00000000..5978ba3c --- /dev/null +++ b/src/common/settings.ts @@ -0,0 +1,177 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { + ConfigurationChangeEvent, + ConfigurationScope, + ConfigurationTarget, + WorkspaceConfiguration, + WorkspaceFolder, + workspace, +} from 'vscode'; +import { getInterpreterDetails } from './python'; +import { getConfiguration, getWorkspaceFolders } from './vscodeapi'; +import path from 'path'; +import * as fs from 'fs'; +import { PYTOOL_MODULE } from '../utils/constants'; +import { getProjectRoot } from './utilities'; + +export interface ISettings { + cwd: string; + workspace: string; + args: string[]; + path: string[]; + interpreter: string[]; + importStrategy: string; + showNotifications: string; +} + +export function getExtensionSettings( + namespace: string, + includeInterpreter?: boolean +): Promise { + return Promise.all( + getWorkspaceFolders().map(w => getWorkspaceSettings(namespace, w, includeInterpreter)) + ); +} + +function resolveVariables( + value: (string | { path: string })[], + workspace?: WorkspaceFolder +): string[] { + const substitutions = new Map(); + const home = process.env.HOME || process.env.USERPROFILE; + if (home) { + substitutions.set('${userHome}', home); + } + if (workspace) { + substitutions.set('${workspaceFolder}', workspace.uri.fsPath); + } + substitutions.set('${cwd}', process.cwd()); + getWorkspaceFolders().forEach(w => { + substitutions.set('${workspaceFolder:' + w.name + '}', w.uri.fsPath); + }); + + return value.map(item => { + // Check if item is an object and has a path property + if (typeof item === 'object' && 'path' in item) { + let path = item.path; + for (const [key, value] of substitutions) { + path = path.replace(key, value); + } + return path; + } else if (typeof item === 'string') { + // Item is a string, proceed as before + for (const [key, value] of substitutions) { + item = item.replace(key, value); + } + return item; + } else { + // Item is not a string or does not match the expected structure, log a warning or handle as needed + console.warn('Item does not match expected format:', item); + return ''; // or return a sensible default + } + }); +} + +export function getInterpreterFromSetting(namespace: string, scope?: ConfigurationScope) { + const config = getConfiguration(namespace, scope); + return config.get('interpreter'); +} + +export async function getWorkspaceSettings( + namespace: string, + workspace: WorkspaceFolder, + includeInterpreter?: boolean +): Promise { + const config = getConfiguration(namespace, workspace.uri); + + let interpreter: string[] = []; + if (includeInterpreter) { + interpreter = getInterpreterFromSetting(namespace, workspace) ?? []; + if (interpreter.length === 0) { + interpreter = (await getInterpreterDetails(workspace.uri)).path ?? []; + } + } + + const workspaceSetting = { + cwd: workspace.uri.fsPath, + workspace: workspace.uri.toString(), + args: resolveVariables(config.get(`args`) ?? [], workspace), + path: resolveVariables(config.get(`path`) ?? [], workspace), + interpreter: resolveVariables(interpreter, workspace), + importStrategy: config.get(`importStrategy`) ?? 'useBundled', + showNotifications: config.get(`showNotifications`) ?? 'off', + }; + + // console.log("WORKSPACE SETTINGS: ", workspaceSetting); + + return workspaceSetting; +} + +function getGlobalValue(config: WorkspaceConfiguration, key: string, defaultValue: T): T { + const inspect = config.inspect(key); + return inspect?.globalValue ?? inspect?.defaultValue ?? defaultValue; +} + +export async function getGlobalSettings( + namespace: string, + includeInterpreter?: boolean +): Promise { + const config = getConfiguration(namespace); + + let interpreter: string[] = []; + if (includeInterpreter) { + interpreter = getGlobalValue(config, 'interpreter', []); + if (interpreter === undefined || interpreter.length === 0) { + interpreter = (await getInterpreterDetails()).path ?? []; + } + } + + // const debugInterpreter = (await getInterpreterDetails()).path ?? []; + // console.log('Global Interpreter: ', debugInterpreter); + + const setting = { + cwd: process.cwd(), + workspace: process.cwd(), + args: getGlobalValue(config, 'args', []), + path: getGlobalValue(config, 'path', []), + interpreter: interpreter, + importStrategy: getGlobalValue(config, 'importStrategy', 'useBundled'), + showNotifications: getGlobalValue(config, 'showNotifications', 'off'), + }; + + // console.log("GLOBAL SETTINGS: ", setting); + + return setting; +} + +export function checkIfConfigurationChanged( + e: ConfigurationChangeEvent, + namespace: string +): boolean { + const settings = [ + `${namespace}.args`, + `${namespace}.path`, + `${namespace}.interpreter`, + `${namespace}.importStrategy`, + `${namespace}.showNotifications`, + ]; + const changed = settings.map(s => e.affectsConfiguration(s)); + return changed.includes(true); +} + +export async function getInterpreterFromWorkspaceSettings(): Promise { + const projectRoot = await getProjectRoot(); + const workspaceSettings = await getWorkspaceSettings(PYTOOL_MODULE, projectRoot, true); + return workspaceSettings.interpreter[0]; +} diff --git a/src/common/status.ts b/src/common/status.ts new file mode 100644 index 00000000..642ff8ab --- /dev/null +++ b/src/common/status.ts @@ -0,0 +1,45 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { LanguageStatusItem, Disposable, l10n, LanguageStatusSeverity } from 'vscode'; +import { createLanguageStatusItem } from './vscodeapi'; +import { Command } from 'vscode-languageclient'; +import { getDocumentSelector } from './utilities'; + +let _status: LanguageStatusItem | undefined; +export function registerLanguageStatusItem(id: string, name: string, command: string): Disposable { + _status = createLanguageStatusItem(id, getDocumentSelector()); + _status.name = name; + _status.text = name; + _status.command = Command.create(l10n.t('Open logs'), command); + + return { + dispose: () => { + _status?.dispose(); + _status = undefined; + }, + }; +} + +export function updateStatus( + status: string | undefined, + severity: LanguageStatusSeverity, + busy?: boolean, + detail?: string +): void { + if (_status) { + _status.text = status && status.length > 0 ? `${_status.name}: ${status}` : `${_status.name}`; + _status.severity = severity; + _status.busy = busy ?? false; + _status.detail = detail; + } +} diff --git a/src/common/utilities.ts b/src/common/utilities.ts new file mode 100644 index 00000000..5b365442 --- /dev/null +++ b/src/common/utilities.ts @@ -0,0 +1,89 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as fs from 'fs-extra'; +import * as path from 'path'; +import { DocumentSelector, LogLevel, Uri, WorkspaceFolder } from 'vscode'; +import { Trace } from 'vscode-jsonrpc/node'; +import { getWorkspaceFolders, isVirtualWorkspace } from './vscodeapi'; + +function logLevelToTrace(logLevel: LogLevel): Trace { + switch (logLevel) { + case LogLevel.Error: + case LogLevel.Warning: + case LogLevel.Info: + return Trace.Messages; + + case LogLevel.Debug: + case LogLevel.Trace: + return Trace.Verbose; + + case LogLevel.Off: + default: + return Trace.Off; + } +} + +export function getLSClientTraceLevel(channelLogLevel: LogLevel, globalLogLevel: LogLevel): Trace { + if (channelLogLevel === LogLevel.Off) { + return logLevelToTrace(globalLogLevel); + } + if (globalLogLevel === LogLevel.Off) { + return logLevelToTrace(channelLogLevel); + } + const level = logLevelToTrace( + channelLogLevel <= globalLogLevel ? channelLogLevel : globalLogLevel + ); + return level; +} + +export async function getProjectRoot(): Promise { + const workspaces: readonly WorkspaceFolder[] = getWorkspaceFolders(); + if (workspaces.length === 0) { + return { + uri: Uri.file(process.cwd()), + name: path.basename(process.cwd()), + index: 0, + }; + } else if (workspaces.length === 1) { + return workspaces[0]; + } else { + let rootWorkspace = workspaces[0]; + let root = undefined; + for (const w of workspaces) { + if (await fs.pathExists(w.uri.fsPath)) { + root = w.uri.fsPath; + rootWorkspace = w; + break; + } + } + + for (const w of workspaces) { + if (root && root.length > w.uri.fsPath.length && (await fs.pathExists(w.uri.fsPath))) { + root = w.uri.fsPath; + rootWorkspace = w; + } + } + return rootWorkspace; + } +} + +export function getDocumentSelector(): DocumentSelector { + return isVirtualWorkspace() + ? [{ language: 'python' }] + : [ + { scheme: 'file', language: 'python' }, + { scheme: 'untitled', language: 'python' }, + { scheme: 'vscode-notebook', language: 'python' }, + { scheme: 'vscode-notebook-cell', language: 'python' }, + ]; +} diff --git a/src/common/vscodeapi.ts b/src/common/vscodeapi.ts new file mode 100644 index 00000000..ff2463bd --- /dev/null +++ b/src/common/vscodeapi.ts @@ -0,0 +1,71 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. + +/* eslint-disable @typescript-eslint/explicit-module-boundary-types */ +/* eslint-disable @typescript-eslint/no-explicit-any */ +import { + commands, + ConfigurationScope, + Disposable, + DocumentSelector, + languages, + LanguageStatusItem, + LogOutputChannel, + Uri, + window, + workspace, + WorkspaceConfiguration, + WorkspaceFolder, +} from 'vscode'; + +export function createOutputChannel(name: string): LogOutputChannel { + return window.createOutputChannel(name, { log: true }); +} + +export function getConfiguration( + config: string, + scope?: ConfigurationScope +): WorkspaceConfiguration { + return workspace.getConfiguration(config, scope); +} + +export function registerCommand( + command: string, + callback: (...args: any[]) => any, + thisArg?: any +): Disposable { + return commands.registerCommand(command, callback, thisArg); +} + +export const { onDidChangeConfiguration } = workspace; + +export function isVirtualWorkspace(): boolean { + const isVirtual = + workspace.workspaceFolders && workspace.workspaceFolders.every(f => f.uri.scheme !== 'file'); + return !!isVirtual; +} + +export function getWorkspaceFolders(): readonly WorkspaceFolder[] { + return workspace.workspaceFolders ?? []; +} + +export function getWorkspaceFolder(uri: Uri): WorkspaceFolder | undefined { + return workspace.getWorkspaceFolder(uri); +} + +export function createLanguageStatusItem( + id: string, + selector: DocumentSelector +): LanguageStatusItem { + return languages.createLanguageStatusItem(id, selector); +} diff --git a/src/extension.ts b/src/extension.ts new file mode 100644 index 00000000..fa6e2026 --- /dev/null +++ b/src/extension.ts @@ -0,0 +1,63 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as vscode from 'vscode'; +import { EventBus } from './services/EventBus'; +import { LSClient } from './services/LSClient'; +import { ZenExtension } from './services/ZenExtension'; +import { refreshUIComponents } from './utils/refresh'; +import { EnvironmentDataProvider } from './views/activityBar/environmentView/EnvironmentDataProvider'; +import { registerEnvironmentCommands } from './commands/environment/registry'; +import { LSP_ZENML_CLIENT_INITIALIZED } from './utils/constants'; +import { toggleCommands } from './utils/global'; + +export async function activate(context: vscode.ExtensionContext) { + const eventBus = EventBus.getInstance(); + const lsClient = LSClient.getInstance(); + + const handleZenMLClientInitialized = async (isInitialized: boolean) => { + console.log('ZenML client initialized: ', isInitialized); + if (isInitialized) { + await toggleCommands(true); + await refreshUIComponents(); + } + }; + + eventBus.on(LSP_ZENML_CLIENT_INITIALIZED, handleZenMLClientInitialized); + + vscode.window.createTreeView('zenmlEnvironmentView', { + treeDataProvider: EnvironmentDataProvider.getInstance(), + }); + registerEnvironmentCommands(context); + + await ZenExtension.activate(context, lsClient); + + context.subscriptions.push( + new vscode.Disposable(() => { + eventBus.off(LSP_ZENML_CLIENT_INITIALIZED, handleZenMLClientInitialized); + }) + ); +} + +/** + * Deactivates the ZenML extension. + * + * @returns {Promise} A promise that resolves to void. + */ +export async function deactivate(): Promise { + const lsClient = LSClient.getInstance().getLanguageClient(); + + if (lsClient) { + await lsClient.stop(); + EventBus.getInstance().emit('lsClientReady', false); + } +} diff --git a/src/services/EventBus.ts b/src/services/EventBus.ts new file mode 100644 index 00000000..8abbe1f8 --- /dev/null +++ b/src/services/EventBus.ts @@ -0,0 +1,33 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { EventEmitter } from 'events'; + +export class EventBus extends EventEmitter { + private static instance: EventBus; + + constructor() { + super(); + } + + /** + * Retrieves the singleton instance of EventBus. + * + * @returns {EventBus} The singleton instance. + */ + public static getInstance(): EventBus { + if (!EventBus.instance) { + EventBus.instance = new EventBus(); + } + return EventBus.instance; + } +} diff --git a/src/services/LSClient.ts b/src/services/LSClient.ts new file mode 100644 index 00000000..ccd52939 --- /dev/null +++ b/src/services/LSClient.ts @@ -0,0 +1,260 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { ProgressLocation, commands, window } from 'vscode'; +import { LanguageClient } from 'vscode-languageclient/node'; +import { storeActiveStack } from '../commands/stack/utils'; +import { GenericLSClientResponse, VersionMismatchError } from '../types/LSClientResponseTypes'; +import { LSNotificationIsZenMLInstalled } from '../types/LSNotificationTypes'; +import { ConfigUpdateDetails } from '../types/ServerInfoTypes'; +import { + LSCLIENT_READY, + LSP_IS_ZENML_INSTALLED, + LSP_ZENML_CLIENT_INITIALIZED, + LSP_ZENML_SERVER_CHANGED, + LSP_ZENML_STACK_CHANGED, + PYTOOL_MODULE, + REFRESH_ENVIRONMENT_VIEW, +} from '../utils/constants'; +import { getZenMLServerUrl, updateServerUrlAndToken } from '../utils/global'; +import { debounce } from '../utils/refresh'; +import { EventBus } from './EventBus'; +import { ServerDataProvider } from '../views/activityBar'; + +export class LSClient { + private static instance: LSClient | null = null; + private client: LanguageClient | null = null; + private eventBus: EventBus = EventBus.getInstance(); + public clientReady: boolean = false; + public isZenMLReady = false; + public localZenML: LSNotificationIsZenMLInstalled = { + is_installed: false, + version: '', + }; + + public restartLSPServerDebounced = debounce(async () => { + await commands.executeCommand(`${PYTOOL_MODULE}.restart`); + // await refreshUIComponents(); + }, 500); + + /** + * Sets up notification listeners for the language client. + * + * @returns void + */ + public setupNotificationListeners(): void { + if (this.client) { + this.client.onNotification(LSP_ZENML_SERVER_CHANGED, this.handleServerChanged.bind(this)); + this.client.onNotification(LSP_ZENML_STACK_CHANGED, this.handleStackChanged.bind(this)); + this.client.onNotification(LSP_IS_ZENML_INSTALLED, this.handleZenMLInstalled.bind(this)); + this.client.onNotification(LSP_ZENML_CLIENT_INITIALIZED, this.handleZenMLReady.bind(this)); + } + } + + /** + * Starts the language client. + * + * @returns A promise resolving to void. + */ + public async startLanguageClient(): Promise { + try { + if (this.client) { + await this.client.start(); + this.clientReady = true; + this.eventBus.emit(LSCLIENT_READY, true); + console.log('Language client started successfully.'); + } + } catch (error) { + console.error('Failed to start the language client:', error); + } + } + + /** + * Handles the zenml/isInstalled notification. + * + * @param params The installation status of ZenML. + */ + public handleZenMLInstalled(params: { is_installed: boolean; version?: string }): void { + console.log(`Received ${LSP_IS_ZENML_INSTALLED} notification: `, params.is_installed); + this.localZenML = params; + this.eventBus.emit(REFRESH_ENVIRONMENT_VIEW); + } + + /** + * Handles the zenml/ready notification. + * + * @param params The ready status of ZenML. + * @returns A promise resolving to void. + */ + public async handleZenMLReady(params: { ready: boolean }): Promise { + console.log(`Received ${LSP_ZENML_CLIENT_INITIALIZED} notification: `, params.ready); + if (!params.ready) { + this.eventBus.emit(LSP_ZENML_CLIENT_INITIALIZED, false); + await commands.executeCommand('zenml.promptForInterpreter'); + } else { + this.eventBus.emit(LSP_ZENML_CLIENT_INITIALIZED, true); + } + this.isZenMLReady = params.ready; + this.eventBus.emit(REFRESH_ENVIRONMENT_VIEW); + } + + /** + * Handles the zenml/serverChanged notification. + * + * @param details The details of the server update. + */ + public async handleServerChanged(details: ConfigUpdateDetails): Promise { + if (this.isZenMLReady) { + console.log(`Received ${LSP_ZENML_SERVER_CHANGED} notification`); + + const currentServerUrl = getZenMLServerUrl(); + const { url, api_token } = details; + if (currentServerUrl !== url) { + window.withProgress( + { + location: ProgressLocation.Notification, + title: 'ZenML config change detected', + cancellable: false, + }, + async progress => { + await this.stopLanguageClient(); + await updateServerUrlAndToken(url, api_token); + this.restartLSPServerDebounced(); + } + ); + } + } + } + + /** + * Stops the language client. + * + * @returns A promise resolving to void. + */ + public async stopLanguageClient(): Promise { + this.clientReady = false; + try { + if (this.client) { + await this.client.stop(); + this.eventBus.emit(LSCLIENT_READY, false); + console.log('Language client stopped successfully.'); + } + } catch (error) { + console.error('Failed to stop the language client:', error); + } + } + + /** + * Handles the zenml/stackChanged notification. + * + * @param activeStackId The ID of the active stack. + */ + public async handleStackChanged(activeStackId: string): Promise { + console.log(`Received ${LSP_ZENML_STACK_CHANGED} notification:`, activeStackId); + await storeActiveStack(activeStackId); + this.eventBus.emit(LSP_ZENML_STACK_CHANGED, activeStackId); + } + + /** + * Sends a request to the language server. + * + * @param {string} command The command to send to the language server. + * @param {any[]} [args] The arguments to send with the command. + * @returns {Promise} A promise resolving to the response from the language server. + */ + public async sendLsClientRequest( + command: string, + args?: any[] + ): Promise { + if (!this.client || !this.clientReady) { + console.error(`${command}: LSClient is not ready yet.`); + return { error: 'LSClient is not ready yet.' } as T; + } + if (!this.isZenMLReady) { + console.error(`${command}: ZenML Client is not initialized yet.`); + return { error: 'ZenML Client is not initialized.' } as T; + } + try { + const result = await this.client.sendRequest('workspace/executeCommand', { + command: `${PYTOOL_MODULE}.${command}`, + arguments: args, + }); + return result as T; + } catch (error: any) { + const errorMessage = error.message; + console.error(`Failed to execute command ${command}:`, errorMessage || error); + if (errorMessage.includes('ValidationError') || errorMessage.includes('RuntimeError')) { + return this.handleKnownErrors(error); + } + return { error: errorMessage } as T; + } + } + + private handleKnownErrors(error: any): T { + let errorType = 'Error'; + let serverVersion = 'N/A'; + let errorMessage = error.message; + let newErrorMessage = ''; + const versionRegex = /\b\d+\.\d+\.\d+\b/; + + if (errorMessage.includes('ValidationError')) { + errorType = 'ValidationError'; + } else if (errorMessage.includes('RuntimeError')) { + errorType = 'RuntimeError'; + if (errorMessage.includes('revision identified by')) { + const matches = errorMessage.match(versionRegex); + if (matches) { + serverVersion = matches[0]; + newErrorMessage = `Can't locate revision identified by ${serverVersion}`; + } + } + } + + return { + error: errorType, + message: newErrorMessage || errorMessage, + clientVersion: this.localZenML.version || 'N/A', + serverVersion, + } as T; + } + + /** + * Updates the language client. + * + * @param {LanguageClient} updatedCLient The new language client. + */ + public updateClient(updatedCLient: LanguageClient): void { + this.client = updatedCLient; + this.setupNotificationListeners(); + } + + /** + * Gets the language client. + * + * @returns {LanguageClient | null} The language client. + */ + public getLanguageClient(): LanguageClient | null { + return this.client; + } + + /** + * Retrieves the singleton instance of LSClient. + * + * @returns {LSClient} The singleton instance. + */ + public static getInstance(): LSClient { + if (!this.instance) { + this.instance = new LSClient(); + } + return this.instance; + } +} diff --git a/src/services/ZenExtension.ts b/src/services/ZenExtension.ts new file mode 100644 index 00000000..7d1f1f48 --- /dev/null +++ b/src/services/ZenExtension.ts @@ -0,0 +1,258 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as fs from 'fs-extra'; +import * as path from 'path'; +import * as vscode from 'vscode'; + +import { registerPipelineCommands } from '../commands/pipelines/registry'; +import { registerServerCommands } from '../commands/server/registry'; +import { registerStackCommands } from '../commands/stack/registry'; +import { EXTENSION_ROOT_DIR } from '../common/constants'; +import { registerLogger, traceLog, traceVerbose } from '../common/log/logging'; +import { + IInterpreterDetails, + initializePython, + isPythonVersonSupported, + onDidChangePythonInterpreter, + resolveInterpreter, +} from '../common/python'; +import { runServer } from '../common/server'; +import { checkIfConfigurationChanged, getInterpreterFromSetting } from '../common/settings'; +import { registerLanguageStatusItem } from '../common/status'; +import { getLSClientTraceLevel } from '../common/utilities'; +import { + createOutputChannel, + onDidChangeConfiguration, + registerCommand, +} from '../common/vscodeapi'; +import { refreshUIComponents } from '../utils/refresh'; +import { PipelineDataProvider, ServerDataProvider, StackDataProvider } from '../views/activityBar'; +import ZenMLStatusBar from '../views/statusBar'; +import { LSClient } from './LSClient'; +import { toggleCommands } from '../utils/global'; + +export interface IServerInfo { + name: string; + module: string; +} + +export class ZenExtension { + private static context: vscode.ExtensionContext; + static commandDisposables: vscode.Disposable[] = []; + static viewDisposables: vscode.Disposable[] = []; + public static lsClient: LSClient; + public static outputChannel: vscode.LogOutputChannel; + public static serverId: string; + public static serverName: string; + private static viewsAndCommandsSetup = false; + public static interpreterCheckInProgress = false; + + private static dataProviders = new Map>([ + ['zenmlServerView', ServerDataProvider.getInstance()], + ['zenmlStackView', StackDataProvider.getInstance()], + ['zenmlPipelineView', PipelineDataProvider.getInstance()], + ]); + + private static registries = [ + registerServerCommands, + registerStackCommands, + registerPipelineCommands, + ]; + + /** + * Initializes the extension services and saves the context for reuse. + * + * @param context The extension context provided by VS Code on activation. + */ + static async activate(context: vscode.ExtensionContext, lsClient: LSClient): Promise { + this.context = context; + this.lsClient = lsClient; + const serverDefaults = this.loadServerDefaults(); + this.serverName = serverDefaults.name; + this.serverId = serverDefaults.module; + + this.setupLoggingAndTrace(); + this.subscribeToCoreEvents(); + this.deferredInitialize(); + } + + /** + * Deferred initialization tasks to be run after initializing other tasks. + */ + static deferredInitialize(): void { + setImmediate(async () => { + const interpreter = getInterpreterFromSetting(this.serverId); + if (interpreter === undefined || interpreter.length === 0) { + traceLog(`Python extension loading`); + await initializePython(this.context.subscriptions); + traceLog(`Python extension loaded`); + } else { + await runServer(); + } + await this.setupViewsAndCommands(); + }); + } + + /** + * Sets up the views and commands for the ZenML extension. + */ + static async setupViewsAndCommands(): Promise { + if (this.viewsAndCommandsSetup) { + console.log('Views and commands have already been set up. Refreshing views...'); + await toggleCommands(true); + return; + } + + ZenMLStatusBar.getInstance(); + this.dataProviders.forEach((provider, viewId) => { + const view = vscode.window.createTreeView(viewId, { treeDataProvider: provider }); + this.viewDisposables.push(view); + }); + this.registries.forEach(register => register(this.context)); + await toggleCommands(true); + this.viewsAndCommandsSetup = true; + } + + /** + * Registers command and configuration event handlers to the extension context. + */ + private static subscribeToCoreEvents(): void { + this.context.subscriptions.push( + onDidChangePythonInterpreter(async (interpreterDetails: IInterpreterDetails) => { + this.interpreterCheckInProgress = true; + if (interpreterDetails.path) { + const resolvedEnv = await resolveInterpreter(interpreterDetails.path); + const { isSupported, message } = isPythonVersonSupported(resolvedEnv); + if (!isSupported) { + vscode.window.showErrorMessage(`Interpreter not supported: ${message}`); + this.interpreterCheckInProgress = false; + return; + } + await runServer(); + if (!this.lsClient.isZenMLReady) { + console.log('ZenML Client is not initialized yet.'); + await this.promptForPythonInterpreter(); + } else { + vscode.window.showInformationMessage('🚀 ZenML installation found. Ready to use.'); + await refreshUIComponents(); + } + } + this.interpreterCheckInProgress = false; + }), + registerCommand(`${this.serverId}.showLogs`, async () => { + this.outputChannel.show(); + }), + onDidChangeConfiguration(async (e: vscode.ConfigurationChangeEvent) => { + if (checkIfConfigurationChanged(e, this.serverId)) { + console.log('Configuration changed, restarting LSP server...', e); + await runServer(); + } + }), + registerCommand(`${this.serverId}.restart`, async () => { + await runServer(); + }), + registerCommand(`zenml.promptForInterpreter`, async () => { + if (!this.interpreterCheckInProgress && !this.lsClient.isZenMLReady) { + await this.promptForPythonInterpreter(); + } + }), + registerLanguageStatusItem(this.serverId, this.serverName, `${this.serverId}.showLogs`) + ); + } + + /** + * Prompts the user to select a Python interpreter. + * + * @returns {Promise} A promise that resolves to void. + */ + static async promptForPythonInterpreter(): Promise { + if (this.interpreterCheckInProgress) { + console.log('Interpreter check already in progress. Skipping prompt.'); + return; + } + if (this.lsClient.isZenMLReady) { + console.log('ZenML is already installed, no need to prompt for interpreter.'); + return; + } + try { + const selected = await vscode.window.showInformationMessage( + 'ZenML not found with the current Python interpreter. Would you like to select a different interpreter?', + 'Select Interpreter', + 'Cancel' + ); + if (selected === 'Select Interpreter') { + await vscode.commands.executeCommand('python.setInterpreter'); + console.log('Interpreter selection completed.'); + } else { + console.log('Interpreter selection cancelled.'); + } + } catch (err) { + console.error('Error selecting Python interpreter:', err); + } + } + + /** + * Initializes the outputChannel and logging for the ZenML extension. + */ + private static setupLoggingAndTrace(): void { + this.outputChannel = createOutputChannel(this.serverName); + + this.context.subscriptions.push(this.outputChannel, registerLogger(this.outputChannel)); + const changeLogLevel = async (c: vscode.LogLevel, g: vscode.LogLevel) => { + const level = getLSClientTraceLevel(c, g); + const lsClient = LSClient.getInstance().getLanguageClient(); + await lsClient?.setTrace(level); + }; + + this.context.subscriptions.push( + this.outputChannel.onDidChangeLogLevel( + async e => await changeLogLevel(e, vscode.env.logLevel) + ), + vscode.env.onDidChangeLogLevel( + async e => await changeLogLevel(this.outputChannel.logLevel, e) + ) + ); + + traceLog(`Name: ${this.serverName}`); + traceLog(`Module: ${this.serverId}`); + traceVerbose( + `Full Server Info: ${JSON.stringify({ name: this.serverName, module: this.serverId })}` + ); + } + + /** + * Loads the server defaults from the package.json file. + * + * @returns {IServerInfo} The server defaults. + */ + private static loadServerDefaults(): IServerInfo { + const packageJson = path.join(EXTENSION_ROOT_DIR, 'package.json'); + const content = fs.readFileSync(packageJson).toString(); + const config = JSON.parse(content); + return config.serverInfo as IServerInfo; + } + + /** + * Deactivates ZenML features when requirements not met. + * + * @returns {Promise} A promise that resolves to void. + */ + static async deactivateFeatures(): Promise { + this.commandDisposables.forEach(disposable => disposable.dispose()); + this.commandDisposables = []; + + this.viewDisposables.forEach(disposable => disposable.dispose()); + this.viewDisposables = []; + console.log('Features deactivated due to unmet requirements.'); + } +} diff --git a/src/test/python_tests/__init__.py b/src/test/python_tests/__init__.py new file mode 100644 index 00000000..d83bd1ae --- /dev/null +++ b/src/test/python_tests/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. diff --git a/src/test/python_tests/lsp_test_client/__init__.py b/src/test/python_tests/lsp_test_client/__init__.py new file mode 100644 index 00000000..d83bd1ae --- /dev/null +++ b/src/test/python_tests/lsp_test_client/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. diff --git a/src/test/python_tests/lsp_test_client/constants.py b/src/test/python_tests/lsp_test_client/constants.py new file mode 100644 index 00000000..a2d76fae --- /dev/null +++ b/src/test/python_tests/lsp_test_client/constants.py @@ -0,0 +1,21 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +""" +Constants for use with tests. +""" +import pathlib + +TEST_ROOT = pathlib.Path(__file__).parent.parent +PROJECT_ROOT = TEST_ROOT.parent.parent.parent +TEST_DATA = TEST_ROOT / "test_data" diff --git a/src/test/python_tests/lsp_test_client/defaults.py b/src/test/python_tests/lsp_test_client/defaults.py new file mode 100644 index 00000000..c242decf --- /dev/null +++ b/src/test/python_tests/lsp_test_client/defaults.py @@ -0,0 +1,230 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +""" +Default initialize request params. +""" + +import os + +from .constants import PROJECT_ROOT +from .utils import as_uri, get_initialization_options + +VSCODE_DEFAULT_INITIALIZE = { + "processId": os.getpid(), + "clientInfo": {"name": "vscode", "version": "1.45.0"}, + "rootPath": str(PROJECT_ROOT), + "rootUri": as_uri(str(PROJECT_ROOT)), + "capabilities": { + "workspace": { + "applyEdit": True, + "workspaceEdit": { + "documentChanges": True, + "resourceOperations": ["create", "rename", "delete"], + "failureHandling": "textOnlyTransactional", + }, + "didChangeConfiguration": {"dynamicRegistration": True}, + "didChangeWatchedFiles": {"dynamicRegistration": True}, + "symbol": { + "dynamicRegistration": True, + "symbolKind": { + "valueSet": [ + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + ] + }, + "tagSupport": {"valueSet": [1]}, + }, + "executeCommand": {"dynamicRegistration": True}, + "configuration": True, + "workspaceFolders": True, + }, + "textDocument": { + "publishDiagnostics": { + "relatedInformation": True, + "versionSupport": False, + "tagSupport": {"valueSet": [1, 2]}, + "complexDiagnosticCodeSupport": True, + }, + "synchronization": { + "dynamicRegistration": True, + "willSave": True, + "willSaveWaitUntil": True, + "didSave": True, + }, + "completion": { + "dynamicRegistration": True, + "contextSupport": True, + "completionItem": { + "snippetSupport": True, + "commitCharactersSupport": True, + "documentationFormat": ["markdown", "plaintext"], + "deprecatedSupport": True, + "preselectSupport": True, + "tagSupport": {"valueSet": [1]}, + "insertReplaceSupport": True, + }, + "completionItemKind": { + "valueSet": [ + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + ] + }, + }, + "hover": { + "dynamicRegistration": True, + "contentFormat": ["markdown", "plaintext"], + }, + "signatureHelp": { + "dynamicRegistration": True, + "signatureInformation": { + "documentationFormat": ["markdown", "plaintext"], + "parameterInformation": {"labelOffsetSupport": True}, + }, + "contextSupport": True, + }, + "definition": {"dynamicRegistration": True, "linkSupport": True}, + "references": {"dynamicRegistration": True}, + "documentHighlight": {"dynamicRegistration": True}, + "documentSymbol": { + "dynamicRegistration": True, + "symbolKind": { + "valueSet": [ + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + ] + }, + "hierarchicalDocumentSymbolSupport": True, + "tagSupport": {"valueSet": [1]}, + }, + "codeAction": { + "dynamicRegistration": True, + "isPreferredSupport": True, + "codeActionLiteralSupport": { + "codeActionKind": { + "valueSet": [ + "", + "quickfix", + "refactor", + "refactor.extract", + "refactor.inline", + "refactor.rewrite", + "source", + "source.organizeImports", + ] + } + }, + }, + "codeLens": {"dynamicRegistration": True}, + "formatting": {"dynamicRegistration": True}, + "rangeFormatting": {"dynamicRegistration": True}, + "onTypeFormatting": {"dynamicRegistration": True}, + "rename": {"dynamicRegistration": True, "prepareSupport": True}, + "documentLink": { + "dynamicRegistration": True, + "tooltipSupport": True, + }, + "typeDefinition": { + "dynamicRegistration": True, + "linkSupport": True, + }, + "implementation": { + "dynamicRegistration": True, + "linkSupport": True, + }, + "colorProvider": {"dynamicRegistration": True}, + "foldingRange": { + "dynamicRegistration": True, + "rangeLimit": 5000, + "lineFoldingOnly": True, + }, + "declaration": {"dynamicRegistration": True, "linkSupport": True}, + "selectionRange": {"dynamicRegistration": True}, + }, + "window": {"workDoneProgress": True}, + }, + "trace": "verbose", + "workspaceFolders": [{"uri": as_uri(str(PROJECT_ROOT)), "name": "my_project"}], + "initializationOptions": get_initialization_options(), +} diff --git a/src/test/python_tests/lsp_test_client/session.py b/src/test/python_tests/lsp_test_client/session.py new file mode 100644 index 00000000..c895772c --- /dev/null +++ b/src/test/python_tests/lsp_test_client/session.py @@ -0,0 +1,224 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +""" +LSP session client for testing. +""" + +import os +import subprocess +import sys +from concurrent.futures import Future, ThreadPoolExecutor +from threading import Event + +from pyls_jsonrpc.dispatchers import MethodDispatcher +from pyls_jsonrpc.endpoint import Endpoint +from pyls_jsonrpc.streams import JsonRpcStreamReader, JsonRpcStreamWriter + +from .constants import PROJECT_ROOT +from .defaults import VSCODE_DEFAULT_INITIALIZE + +LSP_EXIT_TIMEOUT = 5000 + + +PUBLISH_DIAGNOSTICS = "textDocument/publishDiagnostics" +WINDOW_LOG_MESSAGE = "window/logMessage" +WINDOW_SHOW_MESSAGE = "window/showMessage" + + +# pylint: disable=too-many-instance-attributes +class LspSession(MethodDispatcher): + """Send and Receive messages over LSP as a test LS Client.""" + + def __init__(self, cwd=None, script=None): + self.cwd = cwd if cwd else os.getcwd() + # pylint: disable=consider-using-with + self._thread_pool = ThreadPoolExecutor() + self._sub = None + self._writer = None + self._reader = None + self._endpoint = None + self._notification_callbacks = {} + self.script = ( + script if script else (PROJECT_ROOT / "bundled" / "tool" / "lsp_server.py") + ) + + def __enter__(self): + """Context manager entrypoint. + + shell=True needed for pytest-cov to work in subprocess. + """ + # pylint: disable=consider-using-with + self._sub = subprocess.Popen( + [sys.executable, str(self.script)], + stdout=subprocess.PIPE, + stdin=subprocess.PIPE, + bufsize=0, + cwd=self.cwd, + env=os.environ, + shell="WITH_COVERAGE" in os.environ, + ) + + self._writer = JsonRpcStreamWriter(os.fdopen(self._sub.stdin.fileno(), "wb")) + self._reader = JsonRpcStreamReader(os.fdopen(self._sub.stdout.fileno(), "rb")) + + dispatcher = { + PUBLISH_DIAGNOSTICS: self._publish_diagnostics, + WINDOW_SHOW_MESSAGE: self._window_show_message, + WINDOW_LOG_MESSAGE: self._window_log_message, + } + self._endpoint = Endpoint(dispatcher, self._writer.write) + self._thread_pool.submit(self._reader.listen, self._endpoint.consume) + return self + + def __exit__(self, typ, value, _tb): + self.shutdown(True) + try: + self._sub.terminate() + except Exception: # pylint:disable=broad-except + pass + self._endpoint.shutdown() + self._thread_pool.shutdown() + + def initialize( + self, + initialize_params=None, + process_server_capabilities=None, + ): + """Sends the initialize request to LSP server.""" + if initialize_params is None: + initialize_params = VSCODE_DEFAULT_INITIALIZE + server_initialized = Event() + + def _after_initialize(fut): + if process_server_capabilities: + process_server_capabilities(fut.result()) + self.initialized() + server_initialized.set() + + self._send_request( + "initialize", + params=( + initialize_params + if initialize_params is not None + else VSCODE_DEFAULT_INITIALIZE + ), + handle_response=_after_initialize, + ) + + server_initialized.wait() + + def initialized(self, initialized_params=None): + """Sends the initialized notification to LSP server.""" + self._endpoint.notify("initialized", initialized_params or {}) + + def shutdown(self, should_exit, exit_timeout=LSP_EXIT_TIMEOUT): + """Sends the shutdown request to LSP server.""" + + def _after_shutdown(_): + if should_exit: + self.exit_lsp(exit_timeout) + + self._send_request("shutdown", handle_response=_after_shutdown) + + def exit_lsp(self, exit_timeout=LSP_EXIT_TIMEOUT): + """Handles LSP server process exit.""" + self._endpoint.notify("exit") + assert self._sub.wait(exit_timeout) == 0 + + def notify_did_change(self, did_change_params): + """Sends did change notification to LSP Server.""" + self._send_notification("textDocument/didChange", params=did_change_params) + + def notify_did_save(self, did_save_params): + """Sends did save notification to LSP Server.""" + self._send_notification("textDocument/didSave", params=did_save_params) + + def notify_did_open(self, did_open_params): + """Sends did open notification to LSP Server.""" + self._send_notification("textDocument/didOpen", params=did_open_params) + + def notify_did_close(self, did_close_params): + """Sends did close notification to LSP Server.""" + self._send_notification("textDocument/didClose", params=did_close_params) + + def text_document_formatting(self, formatting_params): + """Sends text document references request to LSP server.""" + fut = self._send_request("textDocument/formatting", params=formatting_params) + return fut.result() + + def text_document_code_action(self, code_action_params): + """Sends text document code actions request to LSP server.""" + fut = self._send_request("textDocument/codeAction", params=code_action_params) + return fut.result() + + def code_action_resolve(self, code_action_resolve_params): + """Sends text document code actions resolve request to LSP server.""" + fut = self._send_request( + "codeAction/resolve", params=code_action_resolve_params + ) + return fut.result() + + def set_notification_callback(self, notification_name, callback): + """Set custom LS notification handler.""" + self._notification_callbacks[notification_name] = callback + + def get_notification_callback(self, notification_name): + """Gets callback if set or default callback for a given LS + notification.""" + try: + return self._notification_callbacks[notification_name] + except KeyError: + + def _default_handler(_params): + """Default notification handler.""" + + return _default_handler + + def _publish_diagnostics(self, publish_diagnostics_params): + """Internal handler for text document publish diagnostics.""" + return self._handle_notification( + PUBLISH_DIAGNOSTICS, publish_diagnostics_params + ) + + def _window_log_message(self, window_log_message_params): + """Internal handler for window log message.""" + return self._handle_notification(WINDOW_LOG_MESSAGE, window_log_message_params) + + def _window_show_message(self, window_show_message_params): + """Internal handler for window show message.""" + return self._handle_notification( + WINDOW_SHOW_MESSAGE, window_show_message_params + ) + + def _handle_notification(self, notification_name, params): + """Internal handler for notifications.""" + fut = Future() + + def _handler(): + callback = self.get_notification_callback(notification_name) + callback(params) + fut.set_result(None) + + self._thread_pool.submit(_handler) + return fut + + def _send_request(self, name, params=None, handle_response=lambda f: f.done()): + """Sends {name} request to the LSP server.""" + fut = self._endpoint.request(name, params) + fut.add_done_callback(handle_response) + return fut + + def _send_notification(self, name, params=None): + """Sends {name} notification to the LSP server.""" + self._endpoint.notify(name, params) diff --git a/src/test/python_tests/lsp_test_client/utils.py b/src/test/python_tests/lsp_test_client/utils.py new file mode 100644 index 00000000..509b63de --- /dev/null +++ b/src/test/python_tests/lsp_test_client/utils.py @@ -0,0 +1,84 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +""" +Utility functions for use with tests. +""" +import json +import os +import pathlib +import platform +from random import choice + +from .constants import PROJECT_ROOT + + +def normalizecase(path: str) -> str: + """Fixes 'file' uri or path case for easier testing in windows.""" + if platform.system() == "Windows": + return path.lower() + return path + + +def as_uri(path: str) -> str: + """Return 'file' uri as string.""" + return normalizecase(pathlib.Path(path).as_uri()) + + +class PythonFile: + """Create python file on demand for testing.""" + + def __init__(self, contents, root): + self.contents = contents + self.basename = "".join( + choice("abcdefghijklmnopqrstuvwxyz") if i < 8 else ".py" for i in range(9) + ) + self.fullpath = os.path.join(root, self.basename) + + def __enter__(self): + """Creates a python file for testing.""" + with open(self.fullpath, "w", encoding="utf8") as py_file: + py_file.write(self.contents) + return self + + def __exit__(self, typ, value, _tb): + """Cleans up and deletes the python file.""" + os.unlink(self.fullpath) + + +def get_server_info_defaults(): + """Returns server info from package.json""" + package_json_path = PROJECT_ROOT / "package.json" + package_json = json.loads(package_json_path.read_text()) + return package_json["serverInfo"] + + +def get_initialization_options(): + """Returns initialization options from package.json""" + package_json_path = PROJECT_ROOT / "package.json" + package_json = json.loads(package_json_path.read_text()) + + server_info = package_json["serverInfo"] + server_id = server_info["module"] + + properties = package_json["contributes"]["configuration"]["properties"] + setting = {} + for prop in properties: + name = prop[len(server_id) + 1 :] + value = properties[prop]["default"] + setting[name] = value + + setting["workspace"] = as_uri(str(PROJECT_ROOT)) + setting["interpreter"] = [] + + return {"settings": [setting]} diff --git a/src/test/python_tests/requirements.in b/src/test/python_tests/requirements.in new file mode 100644 index 00000000..15bbd40d --- /dev/null +++ b/src/test/python_tests/requirements.in @@ -0,0 +1,14 @@ +# This file is used to generate ./src/test/python_tests/requirements.txt. +# NOTE: +# Use Python 3.8 or greater which ever is the minimum version of the python +# you plan on supporting when creating the environment or using pip-tools. +# Only run the commands below to manully upgrade packages in requirements.txt: +# 1) python -m pip install pip-tools +# 2) pip-compile --generate-hashes --upgrade ./src/test/python_tests/requirements.in +# If you are using nox commands to setup or build package you don't need to +# run the above commands manually. + +# Packages needed by the testing framework. +pytest +PyHamcrest +python-jsonrpc-server diff --git a/src/test/python_tests/requirements.txt b/src/test/python_tests/requirements.txt new file mode 100644 index 00000000..537379cc --- /dev/null +++ b/src/test/python_tests/requirements.txt @@ -0,0 +1,97 @@ +# +# This file is autogenerated by pip-compile with Python 3.11 +# by the following command: +# +# pip-compile --generate-hashes ./src/test/python_tests/requirements.in +# +iniconfig==2.0.0 \ + --hash=sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3 \ + --hash=sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374 + # via pytest +packaging==24.0 \ + --hash=sha256:2ddfb553fdf02fb784c234c7ba6ccc288296ceabec964ad2eae3777778130bc5 \ + --hash=sha256:eb82c5e3e56209074766e6885bb04b8c38a0c015d0a30036ebe7ece34c9989e9 + # via pytest +pluggy==1.4.0 \ + --hash=sha256:7db9f7b503d67d1c5b95f59773ebb58a8c1c288129a88665838012cfb07b8981 \ + --hash=sha256:8c85c2876142a764e5b7548e7d9a0e0ddb46f5185161049a79b7e974454223be + # via pytest +pyhamcrest==2.1.0 \ + --hash=sha256:c6acbec0923d0cb7e72c22af1926f3e7c97b8e8d69fc7498eabacaf7c975bd9c \ + --hash=sha256:f6913d2f392e30e0375b3ecbd7aee79e5d1faa25d345c8f4ff597665dcac2587 + # via -r ./src/test/python_tests/requirements.in +pytest==8.1.1 \ + --hash=sha256:2a8386cfc11fa9d2c50ee7b2a57e7d898ef90470a7a34c4b949ff59662bb78b7 \ + --hash=sha256:ac978141a75948948817d360297b7aae0fcb9d6ff6bc9ec6d514b85d5a65c044 + # via -r ./src/test/python_tests/requirements.in +python-jsonrpc-server==0.4.0 \ + --hash=sha256:62c543e541f101ec5b57dc654efc212d2c2e3ea47ff6f54b2e7dcb36ecf20595 \ + --hash=sha256:e5a908ff182e620aac07db5f57887eeb0afe33993008f57dc1b85b594cea250c + # via -r ./src/test/python_tests/requirements.in +ujson==5.9.0 \ + --hash=sha256:07e0cfdde5fd91f54cd2d7ffb3482c8ff1bf558abf32a8b953a5d169575ae1cd \ + --hash=sha256:0b159efece9ab5c01f70b9d10bbb77241ce111a45bc8d21a44c219a2aec8ddfd \ + --hash=sha256:0c4d6adb2c7bb9eb7c71ad6f6f612e13b264942e841f8cc3314a21a289a76c4e \ + --hash=sha256:10ca3c41e80509fd9805f7c149068fa8dbee18872bbdc03d7cca928926a358d5 \ + --hash=sha256:20509a8c9f775b3a511e308bbe0b72897ba6b800767a7c90c5cca59d20d7c42c \ + --hash=sha256:25fa46e4ff0a2deecbcf7100af3a5d70090b461906f2299506485ff31d9ec437 \ + --hash=sha256:2a8ea0f55a1396708e564595aaa6696c0d8af532340f477162ff6927ecc46e21 \ + --hash=sha256:2fbb90aa5c23cb3d4b803c12aa220d26778c31b6e4b7a13a1f49971f6c7d088e \ + --hash=sha256:323279e68c195110ef85cbe5edce885219e3d4a48705448720ad925d88c9f851 \ + --hash=sha256:32bba5870c8fa2a97f4a68f6401038d3f1922e66c34280d710af00b14a3ca562 \ + --hash=sha256:3382a3ce0ccc0558b1c1668950008cece9bf463ebb17463ebf6a8bfc060dae34 \ + --hash=sha256:37ef92e42535a81bf72179d0e252c9af42a4ed966dc6be6967ebfb929a87bc60 \ + --hash=sha256:3b23bbb46334ce51ddb5dded60c662fbf7bb74a37b8f87221c5b0fec1ec6454b \ + --hash=sha256:473fb8dff1d58f49912323d7cb0859df5585cfc932e4b9c053bf8cf7f2d7c5c4 \ + --hash=sha256:4a566e465cb2fcfdf040c2447b7dd9718799d0d90134b37a20dff1e27c0e9096 \ + --hash=sha256:4e35d7885ed612feb6b3dd1b7de28e89baaba4011ecdf995e88be9ac614765e9 \ + --hash=sha256:506a45e5fcbb2d46f1a51fead991c39529fc3737c0f5d47c9b4a1d762578fc30 \ + --hash=sha256:5635b78b636a54a86fdbf6f027e461aa6c6b948363bdf8d4fbb56a42b7388320 \ + --hash=sha256:5ca35f484622fd208f55041b042d9d94f3b2c9c5add4e9af5ee9946d2d30db01 \ + --hash=sha256:60718f1720a61560618eff3b56fd517d107518d3c0160ca7a5a66ac949c6cf1c \ + --hash=sha256:63fb2e6599d96fdffdb553af0ed3f76b85fda63281063f1cb5b1141a6fcd0617 \ + --hash=sha256:6974b3a7c17bbf829e6c3bfdc5823c67922e44ff169851a755eab79a3dd31ec0 \ + --hash=sha256:6adef377ed583477cf005b58c3025051b5faa6b8cc25876e594afbb772578f21 \ + --hash=sha256:6bbd91a151a8f3358c29355a491e915eb203f607267a25e6ab10531b3b157c5e \ + --hash=sha256:6eecbd09b316cea1fd929b1e25f70382917542ab11b692cb46ec9b0a26c7427f \ + --hash=sha256:70e06849dfeb2548be48fdd3ceb53300640bc8100c379d6e19d78045e9c26120 \ + --hash=sha256:7309d063cd392811acc49b5016728a5e1b46ab9907d321ebbe1c2156bc3c0b99 \ + --hash=sha256:779a2a88c53039bebfbccca934430dabb5c62cc179e09a9c27a322023f363e0d \ + --hash=sha256:7a365eac66f5aa7a7fdf57e5066ada6226700884fc7dce2ba5483538bc16c8c5 \ + --hash=sha256:7b1c0991c4fe256f5fdb19758f7eac7f47caac29a6c57d0de16a19048eb86bad \ + --hash=sha256:7cc7e605d2aa6ae6b7321c3ae250d2e050f06082e71ab1a4200b4ae64d25863c \ + --hash=sha256:829a69d451a49c0de14a9fecb2a2d544a9b2c884c2b542adb243b683a6f15908 \ + --hash=sha256:829b824953ebad76d46e4ae709e940bb229e8999e40881338b3cc94c771b876c \ + --hash=sha256:82b5a56609f1235d72835ee109163c7041b30920d70fe7dac9176c64df87c164 \ + --hash=sha256:89cc92e73d5501b8a7f48575eeb14ad27156ad092c2e9fc7e3cf949f07e75532 \ + --hash=sha256:8ba7cac47dd65ff88571eceeff48bf30ed5eb9c67b34b88cb22869b7aa19600d \ + --hash=sha256:8fc2aa18b13d97b3c8ccecdf1a3c405f411a6e96adeee94233058c44ff92617d \ + --hash=sha256:9ac92d86ff34296f881e12aa955f7014d276895e0e4e868ba7fddebbde38e378 \ + --hash=sha256:9d302bd17989b6bd90d49bade66943c78f9e3670407dbc53ebcf61271cadc399 \ + --hash=sha256:9f21315f51e0db8ee245e33a649dd2d9dce0594522de6f278d62f15f998e050e \ + --hash=sha256:a6d3f10eb8ccba4316a6b5465b705ed70a06011c6f82418b59278fbc919bef6f \ + --hash=sha256:a807ae73c46ad5db161a7e883eec0fbe1bebc6a54890152ccc63072c4884823b \ + --hash=sha256:ab71bf27b002eaf7d047c54a68e60230fbd5cd9da60de7ca0aa87d0bccead8fa \ + --hash=sha256:b048aa93eace8571eedbd67b3766623e7f0acbf08ee291bef7d8106210432427 \ + --hash=sha256:b28407cfe315bd1b34f1ebe65d3bd735d6b36d409b334100be8cdffae2177b2f \ + --hash=sha256:b5964ea916edfe24af1f4cc68488448fbb1ec27a3ddcddc2b236da575c12c8ae \ + --hash=sha256:b68a0caab33f359b4cbbc10065c88e3758c9f73a11a65a91f024b2e7a1257106 \ + --hash=sha256:ba0823cb70866f0d6a4ad48d998dd338dce7314598721bc1b7986d054d782dfd \ + --hash=sha256:bd4ea86c2afd41429751d22a3ccd03311c067bd6aeee2d054f83f97e41e11d8f \ + --hash=sha256:bdf7fc21a03bafe4ba208dafa84ae38e04e5d36c0e1c746726edf5392e9f9f36 \ + --hash=sha256:c4eec2ddc046360d087cf35659c7ba0cbd101f32035e19047013162274e71fcf \ + --hash=sha256:cdcb02cabcb1e44381221840a7af04433c1dc3297af76fde924a50c3054c708c \ + --hash=sha256:d0fd2eba664a22447102062814bd13e63c6130540222c0aa620701dd01f4be81 \ + --hash=sha256:d581db9db9e41d8ea0b2705c90518ba623cbdc74f8d644d7eb0d107be0d85d9c \ + --hash=sha256:dc80f0f5abf33bd7099f7ac94ab1206730a3c0a2d17549911ed2cb6b7aa36d2d \ + --hash=sha256:e015122b337858dba5a3dc3533af2a8fc0410ee9e2374092f6a5b88b182e9fcc \ + --hash=sha256:e208d3bf02c6963e6ef7324dadf1d73239fb7008491fdf523208f60be6437402 \ + --hash=sha256:e2f909bc08ce01f122fd9c24bc6f9876aa087188dfaf3c4116fe6e4daf7e194f \ + --hash=sha256:f0cb4a7814940ddd6619bdce6be637a4b37a8c4760de9373bac54bb7b229698b \ + --hash=sha256:f4b3917296630a075e04d3d07601ce2a176479c23af838b6cf90a2d6b39b0d95 \ + --hash=sha256:f69f16b8f1c69da00e38dc5f2d08a86b0e781d0ad3e4cc6a13ea033a439c4844 \ + --hash=sha256:f833c529e922577226a05bc25b6a8b3eb6c4fb155b72dd88d33de99d53113124 \ + --hash=sha256:f91719c6abafe429c1a144cfe27883eace9fb1c09a9c5ef1bcb3ae80a3076a4e \ + --hash=sha256:ff741a5b4be2d08fceaab681c9d4bc89abf3c9db600ab435e20b9b6d4dfef12e \ + --hash=sha256:ffdfebd819f492e48e4f31c97cb593b9c1a8251933d8f8972e81697f00326ff1 + # via python-jsonrpc-server diff --git a/src/test/python_tests/test_data/sample1/sample.py b/src/test/python_tests/test_data/sample1/sample.py new file mode 100644 index 00000000..16ae5f39 --- /dev/null +++ b/src/test/python_tests/test_data/sample1/sample.py @@ -0,0 +1,3 @@ +import sys + +print(x) diff --git a/src/test/python_tests/test_data/sample1/sample.unformatted b/src/test/python_tests/test_data/sample1/sample.unformatted new file mode 100644 index 00000000..6c8771b9 --- /dev/null +++ b/src/test/python_tests/test_data/sample1/sample.unformatted @@ -0,0 +1 @@ +import sys;print(x) \ No newline at end of file diff --git a/src/test/python_tests/test_server.py b/src/test/python_tests/test_server.py new file mode 100644 index 00000000..fe4505c6 --- /dev/null +++ b/src/test/python_tests/test_server.py @@ -0,0 +1,148 @@ +# Copyright (c) ZenML GmbH 2024. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +""" +Test for linting over LSP. +""" + +from threading import Event + +from hamcrest import assert_that, is_ + +from .lsp_test_client import constants, defaults, session, utils + +TEST_FILE_PATH = constants.TEST_DATA / "sample1" / "sample.py" +TEST_FILE_URI = utils.as_uri(str(TEST_FILE_PATH)) +SERVER_INFO = utils.get_server_info_defaults() +TIMEOUT = 10 # 10 seconds + + +def test_linting_example(): + """Test to linting on file open.""" + contents = TEST_FILE_PATH.read_text() + + actual = [] + with session.LspSession() as ls_session: + ls_session.initialize(defaults.VSCODE_DEFAULT_INITIALIZE) + + done = Event() + + def _handler(params): + nonlocal actual + actual = params + done.set() + + ls_session.set_notification_callback(session.PUBLISH_DIAGNOSTICS, _handler) + + ls_session.notify_did_open( + { + "textDocument": { + "uri": TEST_FILE_URI, + "languageId": "python", + "version": 1, + "text": contents, + } + } + ) + + # wait for some time to receive all notifications + done.wait(TIMEOUT) + + # TODO: Add your linter specific diagnostic result here + expected = { + "uri": TEST_FILE_URI, + "diagnostics": [ + { + # "range": { + # "start": {"line": 0, "character": 0}, + # "end": {"line": 0, "character": 0}, + # }, + # "message": "Missing module docstring", + # "severity": 3, + # "code": "C0114:missing-module-docstring", + "source": SERVER_INFO["name"], + }, + { + # "range": { + # "start": {"line": 2, "character": 6}, + # "end": { + # "line": 2, + # "character": 7, + # }, + # }, + # "message": "Undefined variable 'x'", + # "severity": 1, + # "code": "E0602:undefined-variable", + "source": SERVER_INFO["name"], + }, + { + # "range": { + # "start": {"line": 0, "character": 0}, + # "end": { + # "line": 0, + # "character": 10, + # }, + # }, + # "message": "Unused import sys", + # "severity": 2, + # "code": "W0611:unused-import", + "source": SERVER_INFO["name"], + }, + ], + } + + assert_that(actual, is_(expected)) + + +def test_formatting_example(): + """Test formatting a python file.""" + FORMATTED_TEST_FILE_PATH = constants.TEST_DATA / "sample1" / "sample.py" + UNFORMATTED_TEST_FILE_PATH = constants.TEST_DATA / "sample1" / "sample.unformatted" + + contents = UNFORMATTED_TEST_FILE_PATH.read_text() + lines = contents.splitlines(keepends=False) + + actual = [] + with utils.PythonFile(contents, UNFORMATTED_TEST_FILE_PATH.parent) as pf: + uri = utils.as_uri(str(pf.fullpath)) + + with session.LspSession() as ls_session: + ls_session.initialize() + ls_session.notify_did_open( + { + "textDocument": { + "uri": uri, + "languageId": "python", + "version": 1, + "text": contents, + } + } + ) + actual = ls_session.text_document_formatting( + { + "textDocument": {"uri": uri}, + # `options` is not used by black + "options": {"tabSize": 4, "insertSpaces": True}, + } + ) + + expected = [ + { + "range": { + "start": {"line": 0, "character": 0}, + "end": {"line": len(lines), "character": 0}, + }, + "newText": FORMATTED_TEST_FILE_PATH.read_text(), + } + ] + + assert_that(actual, is_(expected)) diff --git a/src/test/ts_tests/__mocks__/MockEventBus.ts b/src/test/ts_tests/__mocks__/MockEventBus.ts new file mode 100644 index 00000000..d32515ea --- /dev/null +++ b/src/test/ts_tests/__mocks__/MockEventBus.ts @@ -0,0 +1,55 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied.See the License for the specific language governing +import { EventEmitter } from 'stream'; + +export class MockEventBus extends EventEmitter { + public lsClientReady: boolean = false; + private static instance: MockEventBus; + public zenmlReady: boolean = false; + + constructor() { + super(); + this.on('lsClientReady', (isReady: boolean) => { + this.lsClientReady = isReady; + }); + } + + /** + * Retrieves the singleton instance of EventBus. + * + * @returns {MockEventBus} The singleton instance. + */ + public static getInstance(): MockEventBus { + if (!MockEventBus.instance) { + MockEventBus.instance = new MockEventBus(); + } + return MockEventBus.instance; + } + + /** + * Clears all event handlers. + */ + public clearAllHandlers() { + this.removeAllListeners(); + } + + /** + * Simulates setting the LS Client readiness state. + * + * @param isReady A boolean indicating whether the LS Client is ready. + * @returns void + */ + public setLsClientReady(isReady: boolean): void { + this.lsClientReady = isReady; + this.emit('lsClientReady', isReady); + } +} diff --git a/src/test/ts_tests/__mocks__/MockLSClient.ts b/src/test/ts_tests/__mocks__/MockLSClient.ts new file mode 100644 index 00000000..affddf62 --- /dev/null +++ b/src/test/ts_tests/__mocks__/MockLSClient.ts @@ -0,0 +1,151 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied.See the License for the specific language governing +import { MockEventBus } from './MockEventBus'; +import { MOCK_ACCESS_TOKEN, MOCK_REST_SERVER_DETAILS, MOCK_REST_SERVER_URL } from './constants'; + +interface MockLanguageClient { + start: () => Promise; + onNotification: (type: string, handler: (params: any) => void) => void; + sendRequest: (command: string, args?: any) => Promise; +} + +export class MockLSClient { + notificationHandlers: Map void> = new Map(); + mockLanguageClient: MockLanguageClient; + eventBus: MockEventBus; + private static instance: MockLSClient; + public clientReady: boolean = true; + + constructor(eventBus: MockEventBus) { + this.eventBus = eventBus; + this.mockLanguageClient = { + start: async () => {}, + onNotification: (type: string, handler: (params: any) => void) => { + this.notificationHandlers.set(type, handler); + }, + sendRequest: async (command: string, args?: any) => { + if (command === 'workspace/executeCommand') { + return this.sendLsClientRequest(args); + } else { + throw new Error(`Unmocked command: ${command}`); + } + }, + }; + } + + /** + * Retrieves the singleton instance of EventBus. + * + * @returns {MockLSClient} The singleton instance. + */ + public static getInstance(mockEventBus: MockEventBus): MockLSClient { + if (!MockLSClient.instance) { + MockLSClient.instance = new MockLSClient(mockEventBus); + } + return MockLSClient.instance; + } + + /** + * Starts the language client. + */ + public startLanguageClient(): Promise { + return this.mockLanguageClient.start(); + } + + /** + * Gets the mocked language client. + * + * @returns {MockLanguageClient} The mocked language client. + */ + public getLanguageClient(): MockLanguageClient { + return this.mockLanguageClient; + } + + /** + * Mocks sending a request to the language server. + * + * @param command The command to send to the language server. + * @param args The arguments to send with the command. + * @returns A promise resolving to a mocked response from the language server. + */ + async sendLsClientRequest(command: string, args: any[] = []): Promise { + switch (command) { + case 'connect': + if (args[0] === MOCK_REST_SERVER_URL) { + return Promise.resolve({ + message: 'Connected successfully', + access_token: MOCK_ACCESS_TOKEN, + }); + } else { + return Promise.reject(new Error('Failed to connect with incorrect URL')); + } + case 'disconnect': + return Promise.resolve({ message: 'Disconnected successfully' }); + case `serverInfo`: + return Promise.resolve(MOCK_REST_SERVER_DETAILS); + + case `renameStack`: + const [renameStackId, newStackName] = args; + if (renameStackId && newStackName) { + return Promise.resolve({ + message: `Stack ${renameStackId} successfully renamed to ${newStackName}.`, + }); + } else { + return Promise.resolve({ error: 'Failed to rename stack' }); + } + + case `copyStack`: + const [copyStackId, copyNewStackName] = args; + if (copyStackId && copyNewStackName) { + return Promise.resolve({ + message: `Stack ${copyStackId} successfully copied to ${copyNewStackName}.`, + }); + } else { + return Promise.resolve({ error: 'Failed to copy stack' }); + } + + case `switchActiveStack`: + const [stackNameOrId] = args; + if (stackNameOrId) { + return Promise.resolve({ message: `Active stack set to: ${stackNameOrId}` }); + } else { + return Promise.resolve({ error: 'Failed to set active stack' }); + } + + default: + return Promise.reject(new Error(`Unmocked command: ${command}`)); + } + } + + /** + * Triggers a notification with the given type and parameters. + * + * @param type The type of the notification. + * @param params The parameters of the notification. + * @returns void + */ + public triggerNotification(type: string, params: any): void { + const handler = this.notificationHandlers.get(type); + if (handler) { + handler(params); + if (type === 'zenml/serverChanged') { + this.eventBus.emit('zenml/serverChanged', { + updatedServerConfig: params, + }); + } else if (type === 'zenml/requirementNotMet') { + this.eventBus.emit('lsClientReady', false); + } else if (type === 'zenml/ready') { + this.eventBus.emit('lsClientReady', true); + } + } + } +} diff --git a/src/test/ts_tests/__mocks__/MockViewProviders.ts b/src/test/ts_tests/__mocks__/MockViewProviders.ts new file mode 100644 index 00000000..08cf1333 --- /dev/null +++ b/src/test/ts_tests/__mocks__/MockViewProviders.ts @@ -0,0 +1,48 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { StackDataProvider } from '../../../views/activityBar'; +import { ServerDataProvider } from '../../../views/activityBar'; +import { ServerStatus, ZenServerDetails } from '../../../types/ServerInfoTypes'; +import { INITIAL_ZENML_SERVER_STATUS } from '../../../utils/constants'; +import ZenMLStatusBar from '../../../views/statusBar'; +import sinon from 'sinon'; + +export class MockZenMLStatusBar extends ZenMLStatusBar { + public refreshActiveStack = sinon.stub().resolves(); +} + +export class MockStackDataProvider extends StackDataProvider { + public refresh = sinon.stub().resolves(); +} + +export class MockServerDataProvider extends ServerDataProvider { + public refreshCalled: boolean = false; + public currentServerStatus: ServerStatus = INITIAL_ZENML_SERVER_STATUS; + + public async refresh(updatedServerConfig?: ZenServerDetails): Promise { + this.refreshCalled = true; + if (updatedServerConfig) { + this.currentServerStatus = { + ...updatedServerConfig.storeInfo, + isConnected: updatedServerConfig.storeConfig.type === 'rest', + url: updatedServerConfig.storeConfig.url, + store_type: updatedServerConfig.storeConfig.type, + }; + } + } + + public resetMock(): void { + this.refreshCalled = false; + this.currentServerStatus = INITIAL_ZENML_SERVER_STATUS; + } +} diff --git a/src/test/ts_tests/__mocks__/constants.ts b/src/test/ts_tests/__mocks__/constants.ts new file mode 100644 index 00000000..de54177d --- /dev/null +++ b/src/test/ts_tests/__mocks__/constants.ts @@ -0,0 +1,117 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as sinon from 'sinon'; +import * as vscode from 'vscode'; +import { ServerStatus, ZenServerDetails } from '../../../types/ServerInfoTypes'; + +export const MOCK_REST_SERVER_URL = 'https://zenml.example.com'; +export const MOCK_SQL_SERVER_URL = 'sqlite:///path/to/sqlite.db'; +export const MOCK_SERVER_ID = 'test-server'; +export const MOCK_AUTH_SCHEME = 'OAUTH2_PASSWORD_BEARER'; +export const MOCK_ZENML_VERSION = '0.55.5'; +export const MOCK_ACCESS_TOKEN = 'valid_token'; + +export const MOCK_CONTEXT = { + subscriptions: [], + extensionUri: vscode.Uri.parse('file:///extension/path'), + storagePath: '/path/to/storage', + globalStoragePath: '/path/to/global/storage', + workspaceState: { get: sinon.stub(), update: sinon.stub() }, + globalState: { get: sinon.stub(), update: sinon.stub(), setKeysForSync: sinon.stub() }, + logPath: '/path/to/log', + asAbsolutePath: sinon.stub(), +} as any; + +export const MOCK_REST_SERVER_STATUS: ServerStatus = { + isConnected: true, + id: MOCK_SERVER_ID, + store_type: 'rest', + url: MOCK_REST_SERVER_URL, + version: MOCK_ZENML_VERSION, + debug: false, + deployment_type: 'kubernetes', + database_type: 'sqlite', + secrets_store_type: 'sql', + auth_scheme: MOCK_AUTH_SCHEME, +}; + +export const MOCK_REST_SERVER_DETAILS: ZenServerDetails = { + storeInfo: { + id: MOCK_SERVER_ID, + version: MOCK_ZENML_VERSION, + debug: false, + deployment_type: 'kubernetes', + database_type: 'sqlite', + secrets_store_type: 'sql', + auth_scheme: MOCK_AUTH_SCHEME, + }, + storeConfig: { + type: 'rest', + url: MOCK_REST_SERVER_URL, + secrets_store: null, + backup_secrets_store: null, + username: null, + password: null, + api_key: 'api_key', + verify_ssl: true, + pool_pre_ping: true, + http_timeout: 30, + }, +}; + +export const MOCK_SQL_SERVER_STATUS: ServerStatus = { + isConnected: false, + id: MOCK_SERVER_ID, + store_type: 'sql', + url: MOCK_SQL_SERVER_URL, + version: MOCK_ZENML_VERSION, + debug: false, + deployment_type: 'local', + database_type: 'sqlite', + secrets_store_type: 'sql', + auth_scheme: MOCK_AUTH_SCHEME, +}; + +export const MOCK_SQL_SERVER_DETAILS: ZenServerDetails = { + storeInfo: { + id: MOCK_SERVER_ID, + version: MOCK_ZENML_VERSION, + debug: false, + deployment_type: 'local', + database_type: 'sqlite', + secrets_store_type: 'sql', + auth_scheme: MOCK_AUTH_SCHEME, + }, + storeConfig: { + type: 'sql', + url: MOCK_SQL_SERVER_URL, + secrets_store: null, + backup_secrets_store: null, + username: null, + password: null, + verify_ssl: false, + pool_pre_ping: true, + http_timeout: 30, + driver: '', + database: '', + ssl_ca: '', + ssl_key: '', + ssl_verify_server_cert: false, + ssl_cert: '', + pool_size: 0, + max_overflow: 0, + backup_strategy: '', + backup_directory: '', + backup_database: '', + }, +}; diff --git a/src/test/ts_tests/commands/serverCommands.test.ts b/src/test/ts_tests/commands/serverCommands.test.ts new file mode 100644 index 00000000..2079530c --- /dev/null +++ b/src/test/ts_tests/commands/serverCommands.test.ts @@ -0,0 +1,130 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as assert from 'assert'; +import * as sinon from 'sinon'; +import * as vscode from 'vscode'; +import { serverCommands } from '../../../commands/server/cmds'; +import { EventBus } from '../../../services/EventBus'; +import { LSClient } from '../../../services/LSClient'; +import { refreshUtils } from '../../../utils/refresh'; +import { ServerDataProvider } from '../../../views/activityBar'; +import { MockEventBus } from '../__mocks__/MockEventBus'; +import { MockLSClient } from '../__mocks__/MockLSClient'; +import { MOCK_ACCESS_TOKEN, MOCK_REST_SERVER_URL } from '../__mocks__/constants'; + +suite('Server Commands Tests', () => { + let sandbox: sinon.SinonSandbox; + let showErrorMessageStub: sinon.SinonStub; + let mockLSClient: any; + let mockEventBus = new MockEventBus(); + let emitSpy: sinon.SinonSpy; + let configurationMock: any; + let showInputBoxStub: sinon.SinonStub; + let refreshUIComponentsStub: sinon.SinonStub; + + setup(() => { + sandbox = sinon.createSandbox(); + mockLSClient = new MockLSClient(mockEventBus); + emitSpy = sandbox.spy(mockEventBus, 'emit'); + sandbox.stub(LSClient, 'getInstance').returns(mockLSClient); + sandbox.stub(EventBus, 'getInstance').returns(mockEventBus); + showInputBoxStub = sandbox.stub(vscode.window, 'showInputBox'); + showErrorMessageStub = sandbox.stub(vscode.window, 'showErrorMessage'); + + configurationMock = { + get: sandbox.stub().withArgs('serverUrl').returns(MOCK_REST_SERVER_URL), + update: sandbox.stub().resolves(), + has: sandbox.stub().returns(false), + inspect: sandbox.stub().returns({ globalValue: undefined }), + }; + sandbox.stub(vscode.workspace, 'getConfiguration').returns(configurationMock); + sandbox.stub(vscode.window, 'withProgress').callsFake(async (options, task) => { + const mockProgress = { + report: sandbox.stub(), + }; + const mockCancellationToken = new vscode.CancellationTokenSource(); + await task(mockProgress, mockCancellationToken.token); + }); + + refreshUIComponentsStub = sandbox + .stub(refreshUtils, 'refreshUIComponents') + .callsFake(async () => { + console.log('Stubbed refreshUIComponents called'); + }); + }); + + teardown(() => { + sandbox.restore(); + }); + + test('connectServer successfully connects to the server', async () => { + showInputBoxStub.resolves(MOCK_REST_SERVER_URL); + sandbox + .stub(mockLSClient, 'sendLsClientRequest') + .withArgs('connect', [MOCK_REST_SERVER_URL]) + .resolves({ + message: 'Connected successfully', + access_token: MOCK_ACCESS_TOKEN, + }); + + const result = await serverCommands.connectServer(); + + assert.strictEqual(result, true, 'Should successfully connect to the server'); + sinon.assert.calledOnce(showInputBoxStub); + sinon.assert.calledWith( + configurationMock.update, + 'serverUrl', + MOCK_REST_SERVER_URL, + vscode.ConfigurationTarget.Global + ); + sinon.assert.calledWith( + configurationMock.update, + 'accessToken', + MOCK_ACCESS_TOKEN, + vscode.ConfigurationTarget.Global + ); + }); + + test('disconnectServer successfully disconnects from the server', async () => { + sandbox + .stub(mockLSClient, 'sendLsClientRequest') + .withArgs('disconnect') + .resolves({ message: 'Disconnected successfully' }); + + await serverCommands.disconnectServer(); + + sinon.assert.calledOnce(refreshUIComponentsStub); + }); + + test('connectServer fails with incorrect URL', async () => { + showInputBoxStub.resolves('invalid.url'); + sandbox + .stub(mockLSClient, 'sendLsClientRequest') + .withArgs('connect', ['invalid.url']) + .rejects(new Error('Failed to connect')); + + const result = await serverCommands.connectServer(); + assert.strictEqual(result, false, 'Should fail to connect to the server with incorrect URL'); + sinon.assert.calledOnce(showErrorMessageStub); + }); + + test('refreshServerStatus refreshes the server status', async () => { + const serverDataProviderRefreshStub = sandbox + .stub(ServerDataProvider.getInstance(), 'refresh') + .resolves(); + + await serverCommands.refreshServerStatus(); + + sinon.assert.calledOnce(serverDataProviderRefreshStub); + }); +}); diff --git a/src/test/ts_tests/commands/stackCommands.test.ts b/src/test/ts_tests/commands/stackCommands.test.ts new file mode 100644 index 00000000..1da40c6b --- /dev/null +++ b/src/test/ts_tests/commands/stackCommands.test.ts @@ -0,0 +1,130 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as assert from 'assert'; +import sinon from 'sinon'; +import * as vscode from 'vscode'; +import { stackCommands } from '../../../commands/stack/cmds'; +import ZenMLStatusBar from '../../../views/statusBar'; +import { StackDataProvider, StackTreeItem } from '../../../views/activityBar'; +import { MockEventBus } from '../__mocks__/MockEventBus'; +import { MockLSClient } from '../__mocks__/MockLSClient'; +import { LSClient } from '../../../services/LSClient'; +import { EventBus } from '../../../services/EventBus'; +import stackUtils from '../../../commands/stack/utils'; +import { MockStackDataProvider, MockZenMLStatusBar } from '../__mocks__/MockViewProviders'; + +suite('Stack Commands Test Suite', () => { + let sandbox: sinon.SinonSandbox; + let showInputBoxStub: sinon.SinonStub; + let showInformationMessageStub: sinon.SinonStub; + let mockLSClient: any; + let mockEventBus: any; + let mockStackDataProvider: MockStackDataProvider; + let mockStatusBar: MockZenMLStatusBar; + let switchActiveStackStub: sinon.SinonStub; + let setActiveStackStub: sinon.SinonStub; + + setup(() => { + sandbox = sinon.createSandbox(); + mockEventBus = new MockEventBus(); + mockLSClient = new MockLSClient(mockEventBus); + mockStackDataProvider = new MockStackDataProvider(); + mockStatusBar = new MockZenMLStatusBar(); + + // Stub classes to return mock instances + sandbox.stub(StackDataProvider, 'getInstance').returns(mockStackDataProvider); + sandbox.stub(ZenMLStatusBar, 'getInstance').returns(mockStatusBar); + sandbox.stub(LSClient, 'getInstance').returns(mockLSClient); + sandbox.stub(EventBus, 'getInstance').returns(mockEventBus); + sandbox.stub(stackUtils, 'storeActiveStack').resolves(); + + showInputBoxStub = sandbox.stub(vscode.window, 'showInputBox'); + showInformationMessageStub = sandbox.stub(vscode.window, 'showInformationMessage'); + + switchActiveStackStub = sandbox + .stub(stackUtils, 'switchActiveStack') + .callsFake(async (stackNameOrId: string) => { + console.log('switchActiveStack stub called with', stackNameOrId); + return Promise.resolve({ id: stackNameOrId, name: `MockStackName` }); + }); + + setActiveStackStub = sandbox + .stub(stackCommands, 'setActiveStack') + .callsFake(async (node: StackTreeItem) => { + await switchActiveStackStub(node.id); + showInformationMessageStub(`Active stack set to: ${node.label}`); + await mockStatusBar.refreshActiveStack(); + await mockStackDataProvider.refresh(); + }); + + sandbox.stub(vscode.window, 'withProgress').callsFake(async (options, task) => { + const mockProgress = { + report: sandbox.stub(), + }; + const mockCancellationToken = new vscode.CancellationTokenSource(); + await task(mockProgress, mockCancellationToken.token); + }); + }); + + teardown(() => { + sandbox.restore(); + mockEventBus.clearAllHandlers(); + }); + + test('renameStack successfully renames a stack', async () => { + const stackId = 'stack-id-123'; + const newStackName = 'New Stack Name'; + + showInputBoxStub.resolves(newStackName); + await stackCommands.renameStack({ label: 'Old Stack', id: stackId } as any); + + assert.strictEqual(showInputBoxStub.calledOnce, true); + }); + + test('copyStack successfully copies a stack', async () => { + const sourceStackId = 'stack-id-789'; + const targetStackName = 'Copied Stack'; + + showInputBoxStub.resolves(targetStackName); + await stackCommands.copyStack({ label: 'Source Stack', id: sourceStackId } as any); + + sinon.assert.calledOnce(showInputBoxStub); + assert.strictEqual( + showInputBoxStub.calledWithExactly({ prompt: 'Enter the name for the copied stack' }), + true, + 'Input box was not called with the correct prompt' + ); + }); + + test('stackDataProviderMock.refresh can be called directly', async () => { + await mockStackDataProvider.refresh(); + sinon.assert.calledOnce(mockStackDataProvider.refresh); + }); + + test('refreshActiveStack successfully refreshes the active stack', async () => { + await stackCommands.refreshActiveStack(); + sinon.assert.calledOnce(mockStatusBar.refreshActiveStack); + }); + + test('setActiveStack successfully switches to a new stack', async () => { + const fakeStackNode = new StackTreeItem('MockStackName', 'fake-stack-id', [], false); + + await stackCommands.setActiveStack(fakeStackNode); + + sinon.assert.calledOnce(switchActiveStackStub); + sinon.assert.calledOnce(showInformationMessageStub); + sinon.assert.calledWith(showInformationMessageStub, `Active stack set to: MockStackName`); + sinon.assert.calledOnce(mockStackDataProvider.refresh); + sinon.assert.calledOnce(mockStatusBar.refreshActiveStack); + }); +}); diff --git a/src/test/ts_tests/extension.test.ts b/src/test/ts_tests/extension.test.ts new file mode 100644 index 00000000..54e4fbe2 --- /dev/null +++ b/src/test/ts_tests/extension.test.ts @@ -0,0 +1,50 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as assert from 'assert'; +import sinon from 'sinon'; +import * as vscode from 'vscode'; +import * as extension from '../../extension'; +import { EventBus } from '../../services/EventBus'; +import { LSClient } from '../../services/LSClient'; +import { ZenExtension } from '../../services/ZenExtension'; +import { MockEventBus } from './__mocks__/MockEventBus'; + +suite('Extension Activation Test Suite', () => { + let sandbox: sinon.SinonSandbox; + let contextMock: any; + let initializeSpy: sinon.SinonSpy; + let mockEventBus = new MockEventBus(); + let lsClient: LSClient; + + setup(() => { + sandbox = sinon.createSandbox(); + contextMock = { subscriptions: [] }; + initializeSpy = sinon.spy(ZenExtension, 'activate'); + lsClient = LSClient.getInstance(); + sandbox.stub(EventBus, 'getInstance').returns(mockEventBus); + }); + + teardown(() => { + sandbox.restore(); + initializeSpy.restore(); + }); + + test('ZenML Extension should be present', () => { + assert.ok(vscode.extensions.getExtension('ZenML.zenml')); + }); + + test('activate function behaves as expected', async () => { + await extension.activate(contextMock); + sinon.assert.calledOnceWithExactly(initializeSpy, contextMock, lsClient); + }); +}); diff --git a/src/test/ts_tests/integration/serverConfigUpdate.test.ts b/src/test/ts_tests/integration/serverConfigUpdate.test.ts new file mode 100644 index 00000000..913f1d65 --- /dev/null +++ b/src/test/ts_tests/integration/serverConfigUpdate.test.ts @@ -0,0 +1,85 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as sinon from 'sinon'; +import * as vscode from 'vscode'; +import assert from 'assert'; +import { EventBus } from '../../../services/EventBus'; +import { ZenServerDetails } from '../../../types/ServerInfoTypes'; +import { MOCK_REST_SERVER_DETAILS } from '../__mocks__/constants'; +import { MockLSClient } from '../__mocks__/MockLSClient'; +import { MockEventBus } from '../__mocks__/MockEventBus'; +import { LSCLIENT_READY, LSP_ZENML_SERVER_CHANGED } from '../../../utils/constants'; + +suite('Server Configuration Update Flow Tests', () => { + let sandbox: sinon.SinonSandbox; + let mockEventBus = new MockEventBus(); + let mockLSClientInstance: MockLSClient; + let mockLSClient: any; + let refreshUIComponentsStub: sinon.SinonStub; + + setup(() => { + sandbox = sinon.createSandbox(); + sandbox.stub(vscode.window, 'showInformationMessage'); + refreshUIComponentsStub = sandbox.stub().resolves(); + + // Mock LSClient + mockLSClientInstance = new MockLSClient(mockEventBus); + mockLSClient = mockLSClientInstance.getLanguageClient(); + sandbox.stub(mockLSClientInstance, 'startLanguageClient').resolves(); + + // Mock EventBus + sandbox.stub(EventBus, 'getInstance').returns(mockEventBus); + mockEventBus.on(LSCLIENT_READY, async (isReady: boolean) => { + if (isReady) { + await refreshUIComponentsStub(); + } + }); + mockEventBus.on(LSP_ZENML_SERVER_CHANGED, async (updatedServerConfig: ZenServerDetails) => { + await refreshUIComponentsStub(updatedServerConfig); + }); + }); + + teardown(() => { + sandbox.restore(); + mockEventBus.clearAllHandlers(); + }); + + test('LSClientReady event triggers UI refresh', async () => { + mockEventBus.setLsClientReady(true); + sinon.assert.calledOnce(refreshUIComponentsStub); + }); + + test('MockLSClient triggerNotification works as expected', async () => { + const mockNotificationType = 'testNotification'; + const mockData = { key: 'value' }; + mockLSClientInstance.mockLanguageClient.onNotification(mockNotificationType, (data: any) => { + assert.deepStrictEqual(data, mockData); + }); + mockLSClientInstance.triggerNotification(mockNotificationType, mockData); + }); + + test('zenml/serverChanged event updates global configuration and refreshes UI', async () => { + mockLSClientInstance.mockLanguageClient.onNotification( + LSP_ZENML_SERVER_CHANGED, + (data: ZenServerDetails) => { + assert.deepStrictEqual(data, MOCK_REST_SERVER_DETAILS); + } + ); + + mockLSClientInstance.triggerNotification(LSP_ZENML_SERVER_CHANGED, MOCK_REST_SERVER_DETAILS); + + await new Promise(resolve => setTimeout(resolve, 0)); + + sinon.assert.calledOnce(refreshUIComponentsStub); + }); +}); diff --git a/src/test/ts_tests/unit/ServerDataProvider.test.ts b/src/test/ts_tests/unit/ServerDataProvider.test.ts new file mode 100644 index 00000000..101b0684 --- /dev/null +++ b/src/test/ts_tests/unit/ServerDataProvider.test.ts @@ -0,0 +1,95 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as assert from 'assert'; +import * as sinon from 'sinon'; +import { serverUtils } from '../../../commands/server/utils'; +import { EventBus } from '../../../services/EventBus'; +import { LSClient } from '../../../services/LSClient'; +import { ServerDataProvider } from '../../../views/activityBar'; +import { MockLSClient } from '../__mocks__/MockLSClient'; +import { MOCK_REST_SERVER_STATUS, MOCK_SQL_SERVER_STATUS } from '../__mocks__/constants'; +import { MockEventBus } from '../__mocks__/MockEventBus'; +import { ServerStatus } from '../../../types/ServerInfoTypes'; +import { LOADING_TREE_ITEMS } from '../../../views/activityBar/common/LoadingTreeItem'; + +suite('ServerDataProvider Tests', () => { + let sandbox: sinon.SinonSandbox; + let mockEventBus: MockEventBus; + let serverDataProvider: ServerDataProvider; + let mockLSClientInstance: any; + let mockLSClient: any; + + setup(() => { + sandbox = sinon.createSandbox(); + serverDataProvider = ServerDataProvider.getInstance(); + mockEventBus = new MockEventBus(); + mockLSClientInstance = MockLSClient.getInstance(mockEventBus); + mockLSClient = mockLSClientInstance.getLanguageClient(); + sandbox.stub(LSClient, 'getInstance').returns(mockLSClientInstance); + sandbox.stub(EventBus, 'getInstance').returns(mockEventBus); + sandbox.stub(mockLSClientInstance, 'startLanguageClient').resolves(); + serverDataProvider['zenmlClientReady'] = true; + }); + + teardown(() => { + sandbox.restore(); + }); + + test('ServerDataProvider initializes correctly', async () => { + assert.ok(serverDataProvider); + }); + + test('ServerDataProvider should update server status correctly', async () => { + sandbox.stub(serverUtils, 'checkServerStatus').callsFake(async () => { + return Promise.resolve(MOCK_REST_SERVER_STATUS); + }); + + await serverDataProvider.refresh(); + const serverStatus = serverDataProvider.getCurrentStatus() as ServerStatus; + + assert.strictEqual( + serverStatus.isConnected, + true, + 'Server should be reported as connected for REST config' + ); + }); + + test('ServerDataProvider should update server status to disconnected for non-REST type', async () => { + sandbox.restore(); + + sandbox.stub(serverUtils, 'checkServerStatus').callsFake(async () => { + return Promise.resolve(MOCK_SQL_SERVER_STATUS); + }); + + await serverDataProvider.refresh(); + const serverStatus = serverDataProvider.getCurrentStatus() as ServerStatus; + + assert.strictEqual( + serverStatus.isConnected, + false, + 'Server should be reported as disconnected for SQL config' + ); + }); + test('ServerDataProvider should handle zenmlClient not ready state', async () => { + serverDataProvider['zenmlClientReady'] = false; + + await serverDataProvider.refresh(); + assert.deepStrictEqual( + serverDataProvider.getCurrentStatus(), + [LOADING_TREE_ITEMS.get('zenmlClient')!], + 'ServerDataProvider should show loading state for ZenML client not ready' + ); + + serverDataProvider['zenmlClientReady'] = true; + }); +}); diff --git a/src/test/ts_tests/unit/eventBus.test.ts b/src/test/ts_tests/unit/eventBus.test.ts new file mode 100644 index 00000000..deb70edc --- /dev/null +++ b/src/test/ts_tests/unit/eventBus.test.ts @@ -0,0 +1,50 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import assert from 'assert'; +import sinon from 'sinon'; +import { MockEventBus } from '../__mocks__/MockEventBus'; +import { LSCLIENT_READY } from '../../../utils/constants'; + +suite('MockEventBus and Event Handling', () => { + let eventBus: MockEventBus; + let spy: sinon.SinonSpy; + + setup(() => { + eventBus = new MockEventBus(); + spy = sinon.spy(); + }); + + test('handles lsClientReady event correctly with mock', () => { + eventBus.on(LSCLIENT_READY, spy); + eventBus.emit(LSCLIENT_READY, true); + assert.ok( + spy.calledWith(true), + 'lsClientReady event handler was not called with expected argument' + ); + }); + + test('can clear all event handlers and not trigger events', () => { + eventBus.on(LSCLIENT_READY, spy); + eventBus.clearAllHandlers(); + + // Try emitting the event after clearing all handlers + eventBus.emit(LSCLIENT_READY, true); + + // Verify the spy was not called since all handlers were cleared + assert.strictEqual( + spy.called, + false, + 'lsClientReady event handler was called despite clearing all handlers' + ); + }); +}); diff --git a/src/types/HydratedTypes.ts b/src/types/HydratedTypes.ts new file mode 100644 index 00000000..525a0577 --- /dev/null +++ b/src/types/HydratedTypes.ts @@ -0,0 +1,141 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. + +/************************************************************************************************ + * Hydrated User types from the ZenML Client. + ************************************************************************************************/ +interface User { + body: UserBody; + metadata?: UserMetadata; + resources?: null; + id: string; + permission_denied: boolean; + name: string; +} + +interface UserMetadata { + email?: string | null; + hub_token?: string | null; + external_user_id?: string | null; +} + +interface UserBody { + created: string; + updated: string; + active: boolean; + activation_token?: null; + full_name?: string; + email_opted_in?: null; + is_service_account: boolean; +} + +/************************************************************************************************ + * Hydrated Workspace types from the ZenML Client. + ************************************************************************************************/ +interface Workspace { + body: WorkspaceBody; + metadata?: { + description?: string; + }; + resources?: any | null; + id: string; + permission_denied: boolean; + name: string; +} + +interface WorkspaceBody { + created: string; + updated: string; +} + +interface Workspace { + body: WorkspaceBody; + metadata?: { + description?: string; + }; + resources?: any | null; + id: string; + permission_denied: boolean; + name: string; +} + +interface WorkspaceBody { + created: string; + updated: string; +} + +/************************************************************************************************ + * Hydrated Stack / Components types from the ZenML Client. + ************************************************************************************************/ +interface HydratedStack { + id: string; + name: string; + permission_denied: boolean; + body: { + created: string; + updated: string; + user?: User | null; + }; + metadata: StackMetadata; + resources?: null; +} + +interface HydratedStackComponent { + body: StackComponentBody; + metadata: ComponentMetadata; + resources?: null; + id: string; + permission_denied: boolean; + name: string; +} + +interface StackComponentBody { + created: string; + updated: string; + user?: User | null; + type: string; + flavor: string; +} + +interface ComponentMetadata { + workspace: Workspace; + configuration?: any; + labels?: null; + component_spec_path?: null; + connector_resource_id?: null; + connector?: null; +} + +interface StackMetadata { + workspace: Workspace; + components: HydratedComponents; + description: string; + stack_spec_path?: null; +} + +interface HydratedComponents { + orchestrator?: HydratedStackComponent[]; + artifact_store?: HydratedStackComponent[]; + container_registry?: HydratedStackComponent[]; + model_registry?: HydratedStackComponent[]; + step_operator?: HydratedStackComponent[]; + feature_store?: HydratedStackComponent[]; + model_deployer?: HydratedStackComponent[]; + experiment_tracker?: HydratedStackComponent[]; + alerter?: HydratedStackComponent[]; + annotator?: HydratedStackComponent[]; + data_validator?: HydratedStackComponent[]; + image_builder?: HydratedStackComponent[]; +} + +export { Workspace, WorkspaceBody, User, UserBody, UserMetadata }; diff --git a/src/types/LSClientResponseTypes.ts b/src/types/LSClientResponseTypes.ts new file mode 100644 index 00000000..f86f1f80 --- /dev/null +++ b/src/types/LSClientResponseTypes.ts @@ -0,0 +1,54 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. + +import { ZenServerDetails } from './ServerInfoTypes'; + +/***** Generic Response Types *****/ +export interface SuccessMessageResponse { + message: string; +} + +export interface ErrorMessageResponse { + error: string; + message: string; +} + +export interface VersionMismatchError { + error: string; + message: string; + clientVersion: string; + serverVersion: string; +} + +export type GenericLSClientResponse = SuccessMessageResponse | ErrorMessageResponse; + +/***** Server Response Types *****/ +export interface RestServerConnectionResponse { + message: string; + access_token: string; +} + +export type ServerStatusInfoResponse = + | ZenServerDetails + | VersionMismatchError + | ErrorMessageResponse; +export type ConnectServerResponse = RestServerConnectionResponse | ErrorMessageResponse; + +/***** Stack Response Types *****/ +export interface ActiveStackResponse { + id: string; + name: string; +} + +export type SetActiveStackResponse = ActiveStackResponse | ErrorMessageResponse; +export type GetActiveStackResponse = ActiveStackResponse | ErrorMessageResponse; diff --git a/src/types/LSNotificationTypes.ts b/src/types/LSNotificationTypes.ts new file mode 100644 index 00000000..944782a4 --- /dev/null +++ b/src/types/LSNotificationTypes.ts @@ -0,0 +1,23 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. + +// create type for LSP server notification that returns {is_installed: boolean, version: string} + +export interface LSNotificationIsZenMLInstalled { + is_installed: boolean; + version?: string; +} + +export interface LSNotification { + is_ready: boolean; +} diff --git a/src/types/PipelineTypes.ts b/src/types/PipelineTypes.ts new file mode 100644 index 00000000..96ca4930 --- /dev/null +++ b/src/types/PipelineTypes.ts @@ -0,0 +1,38 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. + +import { ErrorMessageResponse, VersionMismatchError } from './LSClientResponseTypes'; + + +interface PipelineRunsData { + runs: PipelineRun[]; + total: number; + total_pages: number; + current_page: number; + items_per_page: number; +} + +export interface PipelineRun { + id: string; + name: string; + status: string; + version: string; + stackName: string; + startTime: string; + endTime: string; + os: string; + osVersion: string; + pythonVersion: string; +} + +export type PipelineRunsResponse = PipelineRunsData | ErrorMessageResponse | VersionMismatchError; diff --git a/src/types/ServerInfoTypes.ts b/src/types/ServerInfoTypes.ts new file mode 100644 index 00000000..60d02700 --- /dev/null +++ b/src/types/ServerInfoTypes.ts @@ -0,0 +1,102 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +export interface ServerStatus { + isConnected: boolean; + url: string; + version: string; + store_type: string; + deployment_type: string; + database_type: string; + secrets_store_type: string; + database?: string; + backup_directory?: string; + backup_strategy?: string; + auth_scheme?: string; + debug?: boolean; + id?: string; + username?: string | null; +} + +/************************************************************************************************ + * This is the object returned by the @LSP_SERVER.command(zenml.serverInfo") + ************************************************************************************************/ +export interface ZenServerDetails { + storeInfo: ZenServerStoreInfo; + storeConfig: ZenServerStoreConfig; +} + +export interface ConfigUpdateDetails { + url: string; + api_token: string; + store_type: string; +} + +/************************************************************************************************ + * This is the response from the `zen_store.get_store_info()` method in the ZenML Client. + ************************************************************************************************/ +export interface ZenServerStoreInfo { + id: string; + version: string; + debug: boolean; + deployment_type: string; + database_type: string; + secrets_store_type: string; + auth_scheme: string; + base_url?: string; + metadata?: any; +} + +/************************************************************************************************ + * This is the response from the `zen_store.get_store_config()` method in the ZenML Client. + ************************************************************************************************/ +export type ZenServerStoreConfig = RestZenServerStoreConfig | SQLZenServerStoreConfig; + +/************************************************************************************************ + * REST Zen Server Store Config (type === 'rest') + ************************************************************************************************/ +export interface RestZenServerStoreConfig { + type: string; + url: string; + secrets_store: any; + backup_secrets_store: any; + username: string | null; + password: string | null; + api_key: any; + api_token?: string; + verify_ssl: boolean; + http_timeout: number; +} + +/************************************************************************************************ + * SQL Zen Server Store Config (type === 'sql') + ************************************************************************************************/ +export interface SQLZenServerStoreConfig { + type: string; + url: string; + secrets_store: any; + backup_secrets_store: any; + driver: string; + database: string; + username: string | null; + password: string | null; + ssl_ca: string | null; + ssl_cert: string | null; + ssl_key: string | null; + ssl_verify_server_cert: boolean; + pool_size: number; + max_overflow: number; + pool_pre_ping: boolean; + backup_strategy: string; + backup_directory: string; + backup_database: string | null; +} diff --git a/src/types/StackTypes.ts b/src/types/StackTypes.ts new file mode 100644 index 00000000..ca487dd5 --- /dev/null +++ b/src/types/StackTypes.ts @@ -0,0 +1,47 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. + +import { ErrorMessageResponse, VersionMismatchError } from './LSClientResponseTypes'; + +/************************************************************************************************ + * LSClient parses the JSON response from the ZenML Client, and returns the following types. + * Hydrated types are in the HydratedTypes.ts file. + ************************************************************************************************/ +interface StacksData { + stacks: Stack[]; + total: number; + total_pages: number; + current_page: number; + items_per_page: number; +} + +interface Stack { + id: string; + name: string; + components: Components; +} + +interface Components { + [componentType: string]: StackComponent[]; +} + +interface StackComponent { + id: string; + name: string; + flavor: string; + type: string; +} + +export type StacksReponse = StacksData | ErrorMessageResponse | VersionMismatchError; + +export { Stack, Components, StackComponent, StacksData }; diff --git a/src/utils/constants.ts b/src/utils/constants.ts new file mode 100644 index 00000000..6d83407c --- /dev/null +++ b/src/utils/constants.ts @@ -0,0 +1,58 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { ServerStatus } from '../types/ServerInfoTypes'; + +export const PYTOOL_MODULE = 'zenml-python'; +export const PYTOOL_DISPLAY_NAME = 'ZenML'; +export const LANGUAGE_SERVER_NAME = 'zen-language-server'; +export const MIN_ZENML_VERSION = '0.55.0'; +export const ZENML_EMOJI = '⛩️'; + +export const ZENML_PYPI_URL = 'https://pypi.org/pypi/zenml/json'; +export const DEFAULT_LOCAL_ZENML_SERVER_URL = 'http://127.0.0.1:8237'; + +// LSP server notifications +export const LSP_IS_ZENML_INSTALLED = 'zenml/isInstalled'; +export const LSP_ZENML_CLIENT_INITIALIZED = 'zenml/clientInitialized'; +export const LSP_ZENML_SERVER_CHANGED = 'zenml/serverChanged'; +export const LSP_ZENML_STACK_CHANGED = 'zenml/stackChanged'; +export const LSP_ZENML_REQUIREMENTS_NOT_MET = 'zenml/requirementsNotMet'; + +// EventBus emitted events +export const LSCLIENT_READY = 'lsClientReady'; +export const LSCLIENT_STATE_CHANGED = 'lsClientStateChanged'; +export const REFRESH_ENVIRONMENT_VIEW = 'refreshEnvironmentView'; + +export const REFRESH_SERVER_STATUS = 'refreshServerStatus'; +export const SERVER_STATUS_UPDATED = 'serverStatusUpdated'; +export const ITEMS_PER_PAGE_OPTIONS = ['5', '10', '15', '20', '25', '30', '35', '40', '45', '50']; + +export const INITIAL_ZENML_SERVER_STATUS: ServerStatus = { + isConnected: false, + url: '', + store_type: '', + deployment_type: '', + version: '', + debug: false, + database_type: '', + secrets_store_type: '', + username: null, +}; + +export const PIPELINE_RUN_STATUS_ICONS: Record = { + initializing: 'loading~spin', + failed: 'error', + completed: 'check', + running: 'clock', + cached: 'history', +}; diff --git a/src/utils/global.ts b/src/utils/global.ts new file mode 100644 index 00000000..8645a153 --- /dev/null +++ b/src/utils/global.ts @@ -0,0 +1,129 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as vscode from 'vscode'; +import { PYTOOL_MODULE } from './constants'; + +/** + * Resets the ZenML Server URL and access token in the VSCode workspace configuration. + */ +export const resetGlobalConfiguration = async () => { + const config = vscode.workspace.getConfiguration('zenml'); + await config.update('serverUrl', '', vscode.ConfigurationTarget.Global); + await config.update('accessToken', '', vscode.ConfigurationTarget.Global); +}; + +/** + * Retrieves the ZenML Server URL from the VSCode workspace configuration. + */ +export const getZenMLServerUrl = (): string => { + const config = vscode.workspace.getConfiguration('zenml'); + return config.get('serverUrl') || ''; +}; + +/** + * Retrieves the ZenML access token from the VSCode workspace configuration. + */ +export const getZenMLAccessToken = (): string => { + const config = vscode.workspace.getConfiguration('zenml'); + return config.get('accessToken') || ''; +}; + +/** + * Updates the ZenML Server URL and access token in the VSCode workspace configuration. + * + * @param {string} url - The new ZenML Server URL to be updated. + * @param {string} token - The new access token to be updated. + * @returns {Promise} A promise that resolves after the configuration has been updated. + */ +export const updateServerUrlAndToken = async (url: string, token: string): Promise => { + try { + const config = vscode.workspace.getConfiguration('zenml'); + await config.update('serverUrl', url, vscode.ConfigurationTarget.Global); + await config.update('accessToken', token, vscode.ConfigurationTarget.Global); + } catch (error: any) { + console.error(`Failed to update ZenML configuration: ${error.message}`); + throw new Error('Failed to update ZenML Server URL and access token.'); + } +}; + +/** + * Updates the ZenML Server URL in the VSCode workspace configuration. + * + * @param {string} serverUrl - The new ZenML Server URL to be updated. Pass an empty string if you want to clear it. + */ +export const updateServerUrl = async (serverUrl: string): Promise => { + const config = vscode.workspace.getConfiguration('zenml'); + try { + await config.update('serverUrl', serverUrl, vscode.ConfigurationTarget.Global); + console.log('ZenML Server URL has been updated successfully.'); + } catch (error: any) { + console.error(`Failed to update ZenML configuration: ${error.message}`); + throw new Error('Failed to update ZenML Server URL.'); + } +}; + +/** + * Updates the ZenML access token in the VSCode workspace configuration. + * + * @param {string} accessToken - The new access token to be updated. Pass an empty string if you want to clear it. + */ +export const updateAccessToken = async (accessToken: string): Promise => { + const config = vscode.workspace.getConfiguration('zenml'); + try { + await config.update('accessToken', accessToken, vscode.ConfigurationTarget.Global); + console.log('ZenML access token has been updated successfully.'); + } catch (error: any) { + console.error(`Failed to update ZenML configuration: ${error.message}`); + throw new Error('Failed to update ZenML access token.'); + } +}; + +/** + * Updates the default Python interpreter path globally. + * @param interpreterPath The new default Python interpreter path. + */ +export async function updateDefaultPythonInterpreterPath(interpreterPath: string): Promise { + const config = vscode.workspace.getConfiguration('python'); + await config.update('defaultInterpreterPath', interpreterPath, vscode.ConfigurationTarget.Global); +} + +/** + * Updates the ZenML Python interpreter setting. + * @param interpreterPath The new path to the python environminterpreterent. + */ +export async function updatePytoolInterpreter(interpreterPath: string): Promise { + const config = vscode.workspace.getConfiguration(PYTOOL_MODULE); + await config.update('interpreter', [interpreterPath], vscode.ConfigurationTarget.Workspace); +} + +/** + * Retrieves the virtual environment path from the VSCode workspace configuration. + * @returns The path to the virtual environment. + */ +export function getDefaultPythonInterpreterPath(): string { + const config = vscode.workspace.getConfiguration('python'); + const defaultInterpreterPath = config.get('defaultInterpreterPath', ''); + return defaultInterpreterPath; +} + +/** + * Toggles the registration of commands for the extension. + * + * @param state The state to set the commands to. + */ +export async function toggleCommands(state: boolean): Promise { + await vscode.commands.executeCommand('setContext', 'stackCommandsRegistered', state); + await vscode.commands.executeCommand('setContext', 'serverCommandsRegistered', state); + await vscode.commands.executeCommand('setContext', 'pipelineCommandsRegistered', state); + await vscode.commands.executeCommand('setContext', 'environmentCommandsRegistered', state); +} diff --git a/src/utils/notifications.ts b/src/utils/notifications.ts new file mode 100644 index 00000000..3ebe6552 --- /dev/null +++ b/src/utils/notifications.ts @@ -0,0 +1,87 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as vscode from 'vscode'; + +/** + * Shows an information message in the status bar for a specified duration. + * + * @param message The message to show. + * @param duration Duration in milliseconds after which the message will disappear. + */ +export const showStatusBarInfoMessage = (message: string, duration: number = 5000): void => { + const disposable = vscode.window.setStatusBarMessage(message); + setTimeout(() => disposable.dispose(), duration); +}; + +/** + * Shows a warning message in the status bar for a specified duration. + * + * @param message The message to show. + * @param duration Duration in milliseconds after which the message will disappear. + */ +export const showStatusBarWarningMessage = (message: string, duration: number = 5000): void => { + const disposable = vscode.window.setStatusBarMessage(`$(alert) ${message}`); + setTimeout(() => disposable.dispose(), duration); +}; + +/** + * Shows an error message in the status bar for a specified duration. + * + * @param message The message to show. + * @param duration Duration in milliseconds after which the message will disappear. + */ +export const showStatusBarErrorMessage = (message: string, duration: number = 5000): void => { + const disposable = vscode.window.setStatusBarMessage(`$(error) ${message}`); + setTimeout(() => disposable.dispose(), duration); +}; + +/** + * Shows a modal pop up information message. + * + * @param message The message to display. + */ +export const showInformationMessage = (message: string): void => { + vscode.window.showInformationMessage(message); +}; + +/** + * Shows a modal pop up error message, + * + * @param message The message to display. + */ +export const showErrorMessage = (message: string): void => { + vscode.window.showErrorMessage(message); +}; + +/** + * Shows a warning message with actions (buttons) for the user to select. + * + * @param message The warning message to display. + * @param actions An array of actions, each action being an object with a title and an action callback. + */ +export async function showWarningMessageWithActions( + message: string, + ...actions: { title: string; action: () => void | Promise }[] +): Promise { + // Map actions to their titles to display as buttons. + const items = actions.map(action => action.title); + + // Show warning message with buttons. + const selection = await vscode.window.showWarningMessage(message, ...items); + + // Find the selected action based on the title and execute its callback. + const selectedAction = actions.find(action => action.title === selection); + if (selectedAction) { + await selectedAction.action(); + } +} diff --git a/src/utils/refresh.ts b/src/utils/refresh.ts new file mode 100644 index 00000000..2f616277 --- /dev/null +++ b/src/utils/refresh.ts @@ -0,0 +1,119 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { EventBus } from '../services/EventBus'; +import { ZenServerDetails } from '../types/ServerInfoTypes'; +import { PipelineDataProvider, ServerDataProvider, StackDataProvider } from '../views/activityBar'; +import { REFRESH_SERVER_STATUS } from './constants'; + +// Type definition for a refresh function that takes a global configuration object +type RefreshFunction = (updatedServerConfig?: ZenServerDetails) => Promise; + +/** + * Debounces a function to prevent it from being called too frequently. + * + * @param func The function to debounce + * @param wait The time to wait before calling the function + * @returns A debounced version of the function + */ +export function debounce Promise>( + func: T, + wait: number +): (...args: Parameters) => Promise { + let timeout: NodeJS.Timeout | null = null; + + return async (...args: Parameters): Promise => { + return new Promise((resolve, reject) => { + const later = () => { + timeout = null; + func(...args) + .then(resolve) + .catch(reject); + }; + if (timeout) { + clearTimeout(timeout); + } + timeout = setTimeout(later, wait); + }); + }; +} + +/** + * Creates a debounced and delayed refresh function. + * + * @param refreshFn The refresh function to debounce and delay. + * @param delayMs The delay in milliseconds before executing the refresh. + * @returns A function that, when called, starts the debounced and delayed execution process. + */ +export function delayRefresh( + refreshFn: RefreshFunction, + delayMs: number = 5000 +): (updatedServerConfig?: ZenServerDetails) => void { + const debouncedRefreshFn = debounce(refreshFn, delayMs); + return (updatedServerConfig?: ZenServerDetails) => { + debouncedRefreshFn(updatedServerConfig); + }; +} + +/** + * Immediately invokes and retries it after , for a specified number of , + * applying debounce to prevent rapid successive calls. + * + * @param refreshFn The refresh function to attempt. + * @param delayMs The time in milliseconds to delay before each attempt. Default is 5s. + * @param attempts The number of attempts to make before giving up. + * @returns A function that, when called, initiates the delayed attempts to refresh. + */ + +export function delayRefreshWithRetry( + refreshFn: RefreshFunction, + delayMs: number = 5000, + attempts: number = 2 +): (updatedServerConfig?: ZenServerDetails) => void { + let refreshCount = 0; + + const executeRefresh = async (updatedServerConfig?: ZenServerDetails) => { + refreshCount++; + // refresh is called immediately + await refreshFn(updatedServerConfig); + if (refreshCount < attempts) { + setTimeout(() => executeRefresh(updatedServerConfig), delayMs); + } + }; + + const debouncedExecuteRefresh = debounce(executeRefresh, delayMs); + + return (updatedServerConfig?: ZenServerDetails) => { + debouncedExecuteRefresh(updatedServerConfig); + }; +} + +/** + * Triggers a refresh of the UI components. + * + * @returns {Promise} A promise that resolves to void. + */ +export async function refreshUIComponents(): Promise { + await ServerDataProvider.getInstance().refresh(); + await StackDataProvider.getInstance().refresh(); + await PipelineDataProvider.getInstance().refresh(); + setTimeout(() => { + EventBus.getInstance().emit(REFRESH_SERVER_STATUS); + }, 500); +} + +export const refreshUtils = { + debounce, + delayRefresh, + delayRefreshWithRetry, + refreshUIComponents, +}; diff --git a/src/views/activityBar/common/ErrorTreeItem.ts b/src/views/activityBar/common/ErrorTreeItem.ts new file mode 100644 index 00000000..980e2858 --- /dev/null +++ b/src/views/activityBar/common/ErrorTreeItem.ts @@ -0,0 +1,47 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { ThemeIcon, TreeItem } from 'vscode'; + +export interface GenericErrorTreeItem { + label: string; + description: string; + message?: string; + icon?: string; +} + +export type ErrorTreeItemType = VersionMismatchTreeItem | ErrorTreeItem; + +export class ErrorTreeItem extends TreeItem { + constructor(label: string, description: string) { + super(label); + this.description = description; + this.iconPath = new ThemeIcon('error'); + } +} + +export class VersionMismatchTreeItem extends ErrorTreeItem { + constructor(clientVersion: string, serverVersion: string) { + super(`Version mismatch detected`, `Client: ${clientVersion} – Server: ${serverVersion}`); + this.iconPath = new ThemeIcon('warning'); + } +} + +export const createErrorItem = (error: any): TreeItem[] => { + const errorItems: TreeItem[] = []; + console.log('Creating error item', error); + if (error.clientVersion || error.serverVersion) { + errorItems.push(new VersionMismatchTreeItem(error.clientVersion, error.serverVersion)); + } + errorItems.push(new ErrorTreeItem(error.errorType || 'Error', error.message)); + return errorItems; +}; diff --git a/src/views/activityBar/common/LoadingTreeItem.ts b/src/views/activityBar/common/LoadingTreeItem.ts new file mode 100644 index 00000000..61317ba0 --- /dev/null +++ b/src/views/activityBar/common/LoadingTreeItem.ts @@ -0,0 +1,31 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { TreeItem, TreeItemCollapsibleState, ThemeIcon } from 'vscode'; + +export class LoadingTreeItem extends TreeItem { + constructor(message: string) { + super(message, TreeItemCollapsibleState.None); + this.description = 'Refreshing...'; + this.iconPath = new ThemeIcon('sync~spin'); + } +} + +// create a MAP of loading tree items and labels +export const LOADING_TREE_ITEMS = new Map([ + ['server', new LoadingTreeItem('Refreshing Server View...')], + ['stacks', new LoadingTreeItem('Refreshing Stacks View...')], + ['pipelineRuns', new LoadingTreeItem('Refreshing Pipeline Runs...')], + ['environment', new LoadingTreeItem('Refreshing Environments...')], + ['lsClient', new LoadingTreeItem('Waiting for Language Server to start...')], + ['zenmlClient', new LoadingTreeItem('Waiting for ZenML Client to initialize...')], +]); diff --git a/src/views/activityBar/common/PaginationTreeItems.ts b/src/views/activityBar/common/PaginationTreeItems.ts new file mode 100644 index 00000000..fd337ab0 --- /dev/null +++ b/src/views/activityBar/common/PaginationTreeItems.ts @@ -0,0 +1,47 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { ThemeIcon, TreeItem, TreeItemCollapsibleState } from 'vscode'; + +/** + * A TreeItem for displaying pagination in the VSCode TreeView. + */ +export class CommandTreeItem extends TreeItem { + constructor( + public readonly label: string, + commandId: string, + commandArguments?: any[], + icon?: string + ) { + super(label); + this.command = { + title: label, + command: commandId, + arguments: commandArguments, + }; + if (icon) { + this.iconPath = new ThemeIcon(icon); + } + } +} + +export class SetItemsPerPageTreeItem extends TreeItem { + constructor() { + super("Set items per page", TreeItemCollapsibleState.None); + this.tooltip = "Click to set the number of items shown per page"; + this.command = { + command: "zenml.setStacksPerPage", + title: "Set Stack Items Per Page", + arguments: [], + }; + } +} diff --git a/src/views/activityBar/environmentView/EnvironmentDataProvider.ts b/src/views/activityBar/environmentView/EnvironmentDataProvider.ts new file mode 100644 index 00000000..d613fff7 --- /dev/null +++ b/src/views/activityBar/environmentView/EnvironmentDataProvider.ts @@ -0,0 +1,106 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { EventEmitter, TreeDataProvider, TreeItem } from 'vscode'; +import { State } from 'vscode-languageclient'; +import { EventBus } from '../../../services/EventBus'; +import { LSCLIENT_STATE_CHANGED, REFRESH_ENVIRONMENT_VIEW } from '../../../utils/constants'; +import { EnvironmentItem } from './EnvironmentItem'; +import { + createInterpreterDetails, + createLSClientItem, + createWorkspaceSettingsItems, + createZenMLStatusItems, +} from './viewHelpers'; + +export class EnvironmentDataProvider implements TreeDataProvider { + private static instance: EnvironmentDataProvider | null = null; + private _onDidChangeTreeData = new EventEmitter(); + readonly onDidChangeTreeData = this._onDidChangeTreeData.event; + + private lsClientStatus: State = State.Stopped; + + private eventBus = EventBus.getInstance(); + + constructor() { + this.eventBus.off(LSCLIENT_STATE_CHANGED, this.handleLsClientStateChangey.bind(this)); + this.eventBus.on(LSCLIENT_STATE_CHANGED, this.handleLsClientStateChangey.bind(this)); + + this.eventBus.off(REFRESH_ENVIRONMENT_VIEW, () => this.refresh()); + this.eventBus.on(REFRESH_ENVIRONMENT_VIEW, () => this.refresh()); + } + + /** + * Retrieves the singleton instance of ServerDataProvider. + * + * @returns {PipelineDataProvider} The singleton instance. + */ + public static getInstance(): EnvironmentDataProvider { + if (!this.instance) { + this.instance = new EnvironmentDataProvider(); + } + return this.instance; + } + + /** + * Handles the change in the LSP client state. + * + * @param {State} status The new LSP client state. + */ + private handleLsClientStateChangey(status: State) { + this.lsClientStatus = status; + this.refresh(); + } + + /** + * Refreshes the "Pipeline Runs" view by fetching the latest pipeline run data and updating the view. + */ + public refresh(): void { + this._onDidChangeTreeData.fire(); + } + + /** + * Retrieves the tree item for a given pipeline run. + * + * @param element The pipeline run item. + * @returns The corresponding VS Code tree item. + */ + getTreeItem(element: EnvironmentItem): TreeItem { + return element; + } + + /** + * Adjusts createRootItems to set each item to not collapsible and directly return items for Interpreter, Workspace, etc. + */ + private async createRootItems(): Promise { + const items: EnvironmentItem[] = []; + const lsClientStatusItem = createLSClientItem(this.lsClientStatus); + items.push(lsClientStatusItem); + items.push(...(await createZenMLStatusItems())); + items.push(...(await createInterpreterDetails())); + items.push(...(await createWorkspaceSettingsItems())); + + return items; + } + + /** + * Simplifies getChildren by always returning root items, as there are no collapsible items now. + */ + async getChildren(element?: EnvironmentItem): Promise { + if (!element) { + return this.createRootItems(); + } else { + // Since there are no collapsible items, no need to fetch children + return []; + } + } +} diff --git a/src/views/activityBar/environmentView/EnvironmentItem.ts b/src/views/activityBar/environmentView/EnvironmentItem.ts new file mode 100644 index 00000000..6ced0300 --- /dev/null +++ b/src/views/activityBar/environmentView/EnvironmentItem.ts @@ -0,0 +1,65 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import path from 'path'; +import { ThemeIcon, TreeItem, TreeItemCollapsibleState } from 'vscode'; + +export class EnvironmentItem extends TreeItem { + constructor( + public readonly label: string, + public readonly description?: string, + public readonly collapsibleState: TreeItemCollapsibleState = TreeItemCollapsibleState.None, + private readonly customIcon?: string, + public readonly contextValue?: string + ) { + super(label, collapsibleState); + this.iconPath = this.determineIcon(label); + this.contextValue = contextValue; + } + + /** + * Determines the icon for the tree item based on the label. + * + * @param label The label of the tree item. + * @returns The icon for the tree item. + */ + private determineIcon(label: string): { light: string; dark: string } | ThemeIcon | undefined { + if (this.customIcon) { + return new ThemeIcon(this.customIcon); + } + switch (label) { + case 'Workspace': + case 'CWD': + case 'File System': + return new ThemeIcon('folder'); + case 'Interpreter': + case 'Name': + case 'Python Version': + case 'Path': + case 'EnvType': + const pythonLogo = path.join(__dirname, '..', 'resources', 'python.png'); + return { + light: pythonLogo, + dark: pythonLogo, + }; + case 'ZenML Local': + case 'ZenML Client': + const zenmlLogo = path.join(__dirname, '..', 'resources', 'logo.png'); + return { + light: zenmlLogo, + dark: zenmlLogo, + }; + default: + return undefined; + } + } +} diff --git a/src/views/activityBar/environmentView/viewHelpers.ts b/src/views/activityBar/environmentView/viewHelpers.ts new file mode 100644 index 00000000..96dd386d --- /dev/null +++ b/src/views/activityBar/environmentView/viewHelpers.ts @@ -0,0 +1,154 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { TreeItemCollapsibleState } from 'vscode'; +import { EnvironmentItem } from './EnvironmentItem'; +import { getInterpreterDetails, resolveInterpreter } from '../../../common/python'; +import { getWorkspaceSettings } from '../../../common/settings'; +import { PYTOOL_MODULE } from '../../../utils/constants'; +import { getProjectRoot } from '../../../common/utilities'; +import { LSClient } from '../../../services/LSClient'; +import { State } from 'vscode-languageclient'; + +/** + * Creates the LSP client item for the environment view. + * + * @returns {EnvironmentItem} The LSP client item. + */ +export function createLSClientItem(lsClientStatus: State): EnvironmentItem { + const statusMappings = { + [State.Running]: { description: 'Running', icon: 'globe' }, + [State.Starting]: { description: 'Initializing…', icon: 'sync~spin' }, + [State.Stopped]: { description: 'Stopped', icon: 'close' }, + }; + + const { description, icon } = statusMappings[lsClientStatus]; + + return new EnvironmentItem( + 'LSP Client', + description, + TreeItemCollapsibleState.None, + icon, + 'lsClient' + ); +} + +/** + * Creates the ZenML status items for the environment view. + * + * @returns {Promise} The ZenML status items. + */ +export async function createZenMLStatusItems(): Promise { + const zenmlReady = LSClient.getInstance().isZenMLReady; + const localZenML = LSClient.getInstance().localZenML; + + const zenMLLocalInstallationItem = new EnvironmentItem( + 'ZenML Local', + localZenML.is_installed ? `${localZenML.version}` : 'Not found', + TreeItemCollapsibleState.None, + localZenML.is_installed ? 'check' : 'warning' + ); + + const zenMLClientStatusItem = new EnvironmentItem( + 'ZenML Client', + !localZenML.is_installed ? '' : zenmlReady ? 'Initialized' : 'Awaiting Initialization', + TreeItemCollapsibleState.None, + !localZenML.is_installed ? 'error' : zenmlReady ? 'check' : 'sync~spin' + ); + + return [zenMLLocalInstallationItem, zenMLClientStatusItem]; +} + +/** + * Creates the workspace settings items for the environment view. + * + * @returns {Promise} The workspace settings items. + */ +export async function createWorkspaceSettingsItems(): Promise { + const settings = await getWorkspaceSettings(PYTOOL_MODULE, await getProjectRoot(), true); + + return [ + new EnvironmentItem('CWD', settings.cwd), + new EnvironmentItem('File System', settings.workspace), + ...(settings.path && settings.path.length + ? [new EnvironmentItem('Path', settings.path.join('; '))] + : []), + ]; +} + +/** + * Creates the interpreter details items for the environment view. + * + * @returns {Promise} The interpreter details items. + */ +export async function createInterpreterDetails(): Promise { + const interpreterDetails = await getInterpreterDetails(); + const interpreterPath = interpreterDetails.path?.[0]; + if (!interpreterPath) { + return []; + } + + const resolvedEnv = await resolveInterpreter([interpreterPath]); + if (!resolvedEnv) { + return [ + new EnvironmentItem( + 'Details', + 'Could not resolve environment details', + TreeItemCollapsibleState.None + ), + ]; + } + const pythonVersion = `${resolvedEnv.version?.major}.${resolvedEnv.version?.minor}.${resolvedEnv.version?.micro}`; + const simplifiedPath = simplifyPath(resolvedEnv.path); + + return [ + new EnvironmentItem( + 'Python Version', + pythonVersion, + TreeItemCollapsibleState.None, + '', + 'interpreter' + ), + new EnvironmentItem( + 'Name', + resolvedEnv?.environment?.name, + TreeItemCollapsibleState.None, + '', + 'interpreter' + ), + new EnvironmentItem( + 'EnvType', + resolvedEnv?.environment?.type, + TreeItemCollapsibleState.None, + '', + 'interpreter' + ), + new EnvironmentItem('Path', simplifiedPath, TreeItemCollapsibleState.None, '', 'interpreter'), + ]; +} + +/** + * Simplifies the path by replacing the home directory with '~'. + * + * @param path The path to simplify. + * @returns {string} The simplified path. + */ +function simplifyPath(path: string): string { + if (!path) { + return ''; + } + const homeDir = process.env.HOME || process.env.USERPROFILE; + if (homeDir) { + return path.replace(homeDir, '~'); + } + return path; +} diff --git a/src/views/activityBar/index.ts b/src/views/activityBar/index.ts new file mode 100644 index 00000000..088d37a3 --- /dev/null +++ b/src/views/activityBar/index.ts @@ -0,0 +1,6 @@ +// /src/views/treeViews/index.ts +export * from './stackView/StackDataProvider'; +export * from './serverView/ServerDataProvider'; +export * from './stackView/StackTreeItems'; +export * from './pipelineView/PipelineDataProvider'; +export * from './pipelineView/PipelineTreeItems'; diff --git a/src/views/activityBar/pipelineView/PipelineDataProvider.ts b/src/views/activityBar/pipelineView/PipelineDataProvider.ts new file mode 100644 index 00000000..66234c31 --- /dev/null +++ b/src/views/activityBar/pipelineView/PipelineDataProvider.ts @@ -0,0 +1,228 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { EventEmitter, TreeDataProvider, TreeItem, window } from 'vscode'; +import { State } from 'vscode-languageclient'; +import { EventBus } from '../../../services/EventBus'; +import { LSClient } from '../../../services/LSClient'; +import { PipelineRun, PipelineRunsResponse } from '../../../types/PipelineTypes'; +import { + ITEMS_PER_PAGE_OPTIONS, + LSCLIENT_STATE_CHANGED, + LSP_ZENML_CLIENT_INITIALIZED, + LSP_ZENML_STACK_CHANGED, +} from '../../../utils/constants'; +import { ErrorTreeItem, createErrorItem } from '../common/ErrorTreeItem'; +import { LOADING_TREE_ITEMS } from '../common/LoadingTreeItem'; +import { PipelineRunTreeItem, PipelineTreeItem } from './PipelineTreeItems'; +import { CommandTreeItem } from '../common/PaginationTreeItems'; + +/** + * Provides data for the pipeline run tree view, displaying detailed information about each pipeline run. + */ +export class PipelineDataProvider implements TreeDataProvider { + private _onDidChangeTreeData = new EventEmitter(); + readonly onDidChangeTreeData = this._onDidChangeTreeData.event; + + private static instance: PipelineDataProvider | null = null; + private eventBus = EventBus.getInstance(); + private zenmlClientReady = false; + private pipelineRuns: PipelineTreeItem[] | TreeItem[] = [LOADING_TREE_ITEMS.get('pipelineRuns')!]; + + private pagination = { + currentPage: 1, + itemsPerPage: 20, + totalItems: 0, + totalPages: 0, + }; + + + constructor() { + this.subscribeToEvents(); + } + + /** + * Subscribes to relevant events to trigger a refresh of the tree view. + */ + public subscribeToEvents(): void { + this.eventBus.on(LSCLIENT_STATE_CHANGED, (newState: State) => { + if (newState === State.Running) { + this.refresh(); + } else { + this.pipelineRuns = [LOADING_TREE_ITEMS.get('lsClient')!]; + this._onDidChangeTreeData.fire(undefined); + } + }); + + this.eventBus.on(LSP_ZENML_CLIENT_INITIALIZED, (isInitialized: boolean) => { + this.zenmlClientReady = isInitialized; + + if (!isInitialized) { + this.pipelineRuns = [LOADING_TREE_ITEMS.get('pipelineRuns')!]; + this._onDidChangeTreeData.fire(undefined); + return; + } + + this.refresh(); + this.eventBus.off(LSP_ZENML_STACK_CHANGED, () => this.refresh()); + this.eventBus.on(LSP_ZENML_STACK_CHANGED, () => this.refresh()); + }); + } + + /** + * Retrieves the singleton instance of ServerDataProvider. + * + * @returns {PipelineDataProvider} The singleton instance. + */ + public static getInstance(): PipelineDataProvider { + if (!this.instance) { + this.instance = new PipelineDataProvider(); + } + return this.instance; + } + + /** + * Refreshes the "Pipeline Runs" view by fetching the latest pipeline run data and updating the view. + * + * @returns A promise resolving to void. + */ + public async refresh(): Promise { + this.pipelineRuns = [LOADING_TREE_ITEMS.get('pipelineRuns')!]; + this._onDidChangeTreeData.fire(undefined); + const page = this.pagination.currentPage; + const itemsPerPage = this.pagination.itemsPerPage; + + try { + const newPipelineData = await this.fetchPipelineRuns(page, itemsPerPage); + this.pipelineRuns = newPipelineData; + } catch (error: any) { + this.pipelineRuns = createErrorItem(error); + } + + this._onDidChangeTreeData.fire(undefined); + } + + /** + * Retrieves the tree item for a given pipeline run. + * + * @param element The pipeline run item. + * @returns The corresponding VS Code tree item. + */ + getTreeItem(element: TreeItem): TreeItem { + return element; + } + + /** + * Retrieves the children for a given tree item. + * + * @param element The parent tree item. If undefined, root pipeline runs are fetched. + * @returns A promise resolving to an array of child tree items or undefined if there are no children. + */ + async getChildren(element?: TreeItem): Promise { + if (!element) { + const runs = await this.fetchPipelineRuns(this.pagination.currentPage, this.pagination.itemsPerPage); + if (this.pagination.currentPage < this.pagination.totalPages) { + runs.push(new CommandTreeItem("Next Page", 'zenml.nextPipelineRunsPage', undefined, 'arrow-circle-right')); + } + if (this.pagination.currentPage > 1) { + runs.unshift(new CommandTreeItem("Previous Page", 'zenml.previousPipelineRunsPage', undefined, 'arrow-circle-left')); + } + return runs; + } else if (element instanceof PipelineTreeItem) { + return element.children; + } + return undefined; + } + /** + * Fetches pipeline runs from the server and maps them to tree items for display. + * + * @returns A promise resolving to an array of PipelineTreeItems representing fetched pipeline runs. + */ + async fetchPipelineRuns(page: number = 1, itemsPerPage: number = 20): Promise { + if (!this.zenmlClientReady) { + return [LOADING_TREE_ITEMS.get('zenmlClient')!]; + } + try { + const lsClient = LSClient.getInstance(); + const result = await lsClient.sendLsClientRequest( + 'getPipelineRuns', + [page, itemsPerPage] + ); + if (!result || 'error' in result) { + if ('clientVersion' in result && 'serverVersion' in result) { + return createErrorItem(result); + } else { + console.error(`Failed to fetch pipeline runs: ${result.error}`); + return []; + } + } + + if ('runs' in result) { + const { runs, total, total_pages, current_page, items_per_page } = result; + + this.pagination = { + currentPage: current_page, + itemsPerPage: items_per_page, + totalItems: total, + totalPages: total_pages, + }; + + return runs.map((run: PipelineRun) => { + const formattedStartTime = new Date(run.startTime).toLocaleString(); + const formattedEndTime = run.endTime ? new Date(run.endTime).toLocaleString() : 'N/A'; + + const children = [ + new PipelineRunTreeItem('run name', run.name), + new PipelineRunTreeItem('stack', run.stackName), + new PipelineRunTreeItem('start time', formattedStartTime), + new PipelineRunTreeItem('end time', formattedEndTime), + new PipelineRunTreeItem('os', `${run.os} ${run.osVersion}`), + new PipelineRunTreeItem('python version', run.pythonVersion), + ]; + + return new PipelineTreeItem(run, run.id, children); + }); + } else { + console.error(`Unexpected response format:`, result); + return []; + } + } catch (error: any) { + console.error(`Failed to fetch stacks: ${error}`); + return [new ErrorTreeItem("Error", `Failed to fetch pipeline runs: ${error.message || error.toString()}`)]; + } + } + + public async goToNextPage() { + if (this.pagination.currentPage < this.pagination.totalPages) { + this.pagination.currentPage++; + await this.refresh(); + } + } + + public async goToPreviousPage() { + if (this.pagination.currentPage > 1) { + this.pagination.currentPage--; + await this.refresh(); + } + } + + public async updateItemsPerPage() { + const selected = await window.showQuickPick(ITEMS_PER_PAGE_OPTIONS, { + placeHolder: "Choose the max number of pipeline runs to display per page", + }); + if (selected) { + this.pagination.itemsPerPage = parseInt(selected, 10); + this.pagination.currentPage = 1; + await this.refresh(); + } + } +} diff --git a/src/views/activityBar/pipelineView/PipelineTreeItems.ts b/src/views/activityBar/pipelineView/PipelineTreeItems.ts new file mode 100644 index 00000000..128c50b8 --- /dev/null +++ b/src/views/activityBar/pipelineView/PipelineTreeItems.ts @@ -0,0 +1,58 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as vscode from 'vscode'; +import { PipelineRun } from '../../../types/PipelineTypes'; +import { PIPELINE_RUN_STATUS_ICONS } from '../../../utils/constants'; + +/** + * Represents a Pipeline Run Tree Item in the VS Code tree view. + * Displays its name, version and status. + */ +export class PipelineTreeItem extends vscode.TreeItem { + public children: PipelineRunTreeItem[] | undefined; + + constructor( + public readonly run: PipelineRun, + public readonly id: string, + children?: PipelineRunTreeItem[] + ) { + super( + run.name, + children === undefined + ? vscode.TreeItemCollapsibleState.None + : vscode.TreeItemCollapsibleState.Collapsed + ); + this.tooltip = `${run.name} - Status: ${run.status}`; + this.description = `version: ${run.version}, status: ${run.status}`; + this.iconPath = new vscode.ThemeIcon(PIPELINE_RUN_STATUS_ICONS[run.status]); + this.children = children; + } + + contextValue = 'pipelineRun'; +} + +/** + * Represents details of a Pipeline Run Tree Item in the VS Code tree view. + * Displays the stack name for the run, its start time, end time, machine details and Python version. + */ +export class PipelineRunTreeItem extends vscode.TreeItem { + constructor( + public readonly label: string, + public readonly description: string + ) { + super(label, vscode.TreeItemCollapsibleState.None); + this.tooltip = `${label}: ${description}`; + } + + contextValue = 'pipelineRunDetail'; +} diff --git a/src/views/activityBar/serverView/ServerDataProvider.ts b/src/views/activityBar/serverView/ServerDataProvider.ts new file mode 100644 index 00000000..1d1d24b4 --- /dev/null +++ b/src/views/activityBar/serverView/ServerDataProvider.ts @@ -0,0 +1,176 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { EventEmitter, ThemeIcon, TreeDataProvider, TreeItem } from 'vscode'; +import { State } from 'vscode-languageclient'; +import { checkServerStatus, isServerStatus } from '../../../commands/server/utils'; +import { EventBus } from '../../../services/EventBus'; +import { ServerStatus } from '../../../types/ServerInfoTypes'; +import { + INITIAL_ZENML_SERVER_STATUS, + LSCLIENT_STATE_CHANGED, + LSP_ZENML_CLIENT_INITIALIZED, + REFRESH_SERVER_STATUS, +} from '../../../utils/constants'; +import { LOADING_TREE_ITEMS } from '../common/LoadingTreeItem'; +import { ServerTreeItem } from './ServerTreeItems'; + +export class ServerDataProvider implements TreeDataProvider { + private _onDidChangeTreeData = new EventEmitter(); + readonly onDidChangeTreeData = this._onDidChangeTreeData.event; + + private static instance: ServerDataProvider | null = null; + private eventBus = EventBus.getInstance(); + private zenmlClientReady = false; + private currentStatus: ServerStatus | TreeItem[] = INITIAL_ZENML_SERVER_STATUS; + + constructor() { + this.subscribeToEvents(); + } + + /** + * Subscribes to relevant events to trigger a refresh of the tree view. + */ + public subscribeToEvents(): void { + this.eventBus.on(LSCLIENT_STATE_CHANGED, (newState: State) => { + if (newState === State.Running) { + this.refresh(); + } else { + this.currentStatus = [LOADING_TREE_ITEMS.get('lsClient')!]; + this._onDidChangeTreeData.fire(undefined); + } + }); + + this.eventBus.on(LSP_ZENML_CLIENT_INITIALIZED, (isInitialized: boolean) => { + this.zenmlClientReady = isInitialized; + if (!isInitialized) { + this.currentStatus = [LOADING_TREE_ITEMS.get('zenmlClient')!]; + this._onDidChangeTreeData.fire(undefined); + return; + } + + this.refresh(); + this.eventBus.off(REFRESH_SERVER_STATUS, async () => await this.refresh()); + this.eventBus.on(REFRESH_SERVER_STATUS, async () => await this.refresh()); + }); + } + + /** + * Retrieves the singleton instance of ServerDataProvider. + * + * @returns {ServerDataProvider} The singleton instance. + */ + public static getInstance(): ServerDataProvider { + if (!this.instance) { + this.instance = new ServerDataProvider(); + } + return this.instance; + } + + /** + * Updates the server status to the provided status (used for tests). + * + * @param {ServerStatus} status The new server status. + */ + public updateStatus(status: ServerStatus): void { + this.currentStatus = status; + } + + /** + * Updates the server status and triggers a UI refresh to reflect the latest server status. + * If the server status has changed, it emits a serverStatusUpdated event. + * + * @returns {Promise} A promise resolving to void. + */ + public async refresh(): Promise { + this.currentStatus = [LOADING_TREE_ITEMS.get('server')!]; + this._onDidChangeTreeData.fire(undefined); + + if (!this.zenmlClientReady) { + this.currentStatus = [LOADING_TREE_ITEMS.get('zenmlClient')!]; + this._onDidChangeTreeData.fire(undefined); + return; + } + + const serverStatus = await checkServerStatus(); + if (isServerStatus(serverStatus)) { + if (JSON.stringify(serverStatus) !== JSON.stringify(this.currentStatus)) { + this.eventBus.emit('serverStatusUpdated', { + isConnected: serverStatus.isConnected, + serverUrl: serverStatus.url, + }); + } + } + + this.currentStatus = serverStatus; + this._onDidChangeTreeData.fire(undefined); + } + + /** + * Gets the current status of the ZenML server. + * + * @returns {ServerStatus} The current status of the ZenML server, including connectivity, host, port, store type, and store URL. + */ + public getCurrentStatus(): ServerStatus | TreeItem[] { + return this.currentStatus; + } + + /** + * Retrieves the tree item for a given element, applying appropriate icons based on the server's connectivity status. + * + * @param element The tree item to retrieve. + * @returns The corresponding VS Code tree item. + */ + getTreeItem(element: TreeItem): TreeItem { + if (element instanceof ServerTreeItem) { + if (element.serverStatus.isConnected) { + element.iconPath = new ThemeIcon('vm-active'); + } else { + element.iconPath = new ThemeIcon('vm-connect'); + } + } + return element; + } + + /** + * Asynchronously fetches the children for a given tree item. + * + * @param element The parent tree item. If undefined, the root server status is fetched. + * @returns A promise resolving to an array of child tree items or undefined if there are no children. + */ + async getChildren(element?: TreeItem): Promise { + if (!element) { + if (isServerStatus(this.currentStatus)) { + console.log(this.currentStatus); + const updatedServerTreeItem = new ServerTreeItem('Server Status', this.currentStatus); + return [updatedServerTreeItem]; + } else if (Array.isArray(this.currentStatus)) { + return this.currentStatus; + } + } else if (element instanceof ServerTreeItem) { + return element.children; + } + return undefined; + } + + /** + * Retrieves the server version. + * + * @returns The server version. + */ + public getServerVersion(): string { + if (isServerStatus(this.currentStatus)) { + return this.currentStatus.version; + } + return 'N/A'; + } +} diff --git a/src/views/activityBar/serverView/ServerTreeItems.ts b/src/views/activityBar/serverView/ServerTreeItems.ts new file mode 100644 index 00000000..1c89dd2a --- /dev/null +++ b/src/views/activityBar/serverView/ServerTreeItems.ts @@ -0,0 +1,102 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as vscode from 'vscode'; +import { ServerStatus } from '../../../types/ServerInfoTypes'; + +/** + * A specialized TreeItem for displaying details about a server's status. + */ +class ServerDetailTreeItem extends vscode.TreeItem { + /** + * Constructs a new ServerDetailTreeItem instance. + * + * @param {string} label The detail label. + * @param {string} detail The detail value. + */ + constructor(label: string, detail: string) { + super(`${label}: ${detail}`, vscode.TreeItemCollapsibleState.None); + // make URL item clickable (if not local db path) + if (detail.startsWith('http://') || detail.startsWith('https://')) { + this.command = { + title: 'Open URL', + command: 'vscode.open', + arguments: [vscode.Uri.parse(detail)], + }; + this.tooltip = `Click to open ${detail}`; + } + } +} + +/** + * TreeItem for representing and visualizing server status in a tree view. Includes details such as connectivity, + * host, and port as children items when connected, or storeType and storeUrl when disconnected. + */ +export class ServerTreeItem extends vscode.TreeItem { + public children: vscode.TreeItem[] | ServerDetailTreeItem[] | undefined; + + constructor( + public readonly label: string, + public readonly serverStatus: ServerStatus + ) { + super( + label, + serverStatus.isConnected + ? vscode.TreeItemCollapsibleState.Expanded + : vscode.TreeItemCollapsibleState.Expanded + ); + + this.description = `${this.serverStatus.isConnected ? 'Connected ✅' : 'Disconnected'}`; + this.children = this.determineChildrenBasedOnStatus(); + } + + private determineChildrenBasedOnStatus(): ServerDetailTreeItem[] { + const children: ServerDetailTreeItem[] = [ + new ServerDetailTreeItem('URL', this.serverStatus.url), + new ServerDetailTreeItem('Version', this.serverStatus.version), + new ServerDetailTreeItem('Store Type', this.serverStatus.store_type || 'N/A'), + new ServerDetailTreeItem('Deployment Type', this.serverStatus.deployment_type), + new ServerDetailTreeItem('Database Type', this.serverStatus.database_type), + new ServerDetailTreeItem('Secrets Store Type', this.serverStatus.secrets_store_type), + ]; + + // Conditional children based on server status type + if (this.serverStatus.id) { + children.push(new ServerDetailTreeItem('ID', this.serverStatus.id)); + } + if (this.serverStatus.username) { + children.push(new ServerDetailTreeItem('Username', this.serverStatus.username)); + } + if (this.serverStatus.debug !== undefined) { + children.push(new ServerDetailTreeItem('Debug', this.serverStatus.debug ? 'true' : 'false')); + } + if (this.serverStatus.auth_scheme) { + children.push(new ServerDetailTreeItem('Auth Scheme', this.serverStatus.auth_scheme)); + } + // Specific to SQL Server Status + if (this.serverStatus.database) { + children.push(new ServerDetailTreeItem('Database', this.serverStatus.database)); + } + if (this.serverStatus.backup_directory) { + children.push( + new ServerDetailTreeItem('Backup Directory', this.serverStatus.backup_directory) + ); + } + if (this.serverStatus.backup_strategy) { + children.push(new ServerDetailTreeItem('Backup Strategy', this.serverStatus.backup_strategy)); + } + + return children; + } + + contextValue = 'server'; +} diff --git a/src/views/activityBar/stackView/StackDataProvider.ts b/src/views/activityBar/stackView/StackDataProvider.ts new file mode 100644 index 00000000..51c1327a --- /dev/null +++ b/src/views/activityBar/stackView/StackDataProvider.ts @@ -0,0 +1,242 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import { Event, EventEmitter, TreeDataProvider, TreeItem, window, workspace } from 'vscode'; +import { State } from 'vscode-languageclient'; +import { EventBus } from '../../../services/EventBus'; +import { LSClient } from '../../../services/LSClient'; +import { Stack, StackComponent, StacksReponse } from '../../../types/StackTypes'; +import { + ITEMS_PER_PAGE_OPTIONS, + LSCLIENT_STATE_CHANGED, + LSP_ZENML_CLIENT_INITIALIZED, + LSP_ZENML_STACK_CHANGED, +} from '../../../utils/constants'; +import { ErrorTreeItem, createErrorItem } from '../common/ErrorTreeItem'; +import { LOADING_TREE_ITEMS } from '../common/LoadingTreeItem'; +import { StackComponentTreeItem, StackTreeItem } from './StackTreeItems'; +import { CommandTreeItem } from '../common/PaginationTreeItems'; + +export class StackDataProvider implements TreeDataProvider { + private _onDidChangeTreeData = new EventEmitter(); + readonly onDidChangeTreeData: Event = + this._onDidChangeTreeData.event; + + private static instance: StackDataProvider | null = null; + private eventBus = EventBus.getInstance(); + private zenmlClientReady = false; + public stacks: Stack[] | TreeItem[] = [LOADING_TREE_ITEMS.get('stacks')!]; + + private pagination = { + currentPage: 1, + itemsPerPage: 20, + totalItems: 0, + totalPages: 0, + }; + + + constructor() { + this.subscribeToEvents(); + } + + /** + * Subscribes to relevant events to trigger a refresh of the tree view. + */ + public subscribeToEvents(): void { + this.eventBus.on(LSCLIENT_STATE_CHANGED, (newState: State) => { + if (newState === State.Running) { + this.refresh(); + } else { + this.stacks = [LOADING_TREE_ITEMS.get('lsClient')!]; + this._onDidChangeTreeData.fire(undefined); + } + }); + + this.eventBus.on(LSP_ZENML_CLIENT_INITIALIZED, (isInitialized: boolean) => { + this.zenmlClientReady = isInitialized; + + if (!isInitialized) { + this.stacks = [LOADING_TREE_ITEMS.get('stacks')!]; + this._onDidChangeTreeData.fire(undefined); + return; + } + this.refresh(); + this.eventBus.off(LSP_ZENML_STACK_CHANGED, () => this.refresh()); + this.eventBus.on(LSP_ZENML_STACK_CHANGED, () => this.refresh()); + }); + } + + /** + * Retrieves the singleton instance of ServerDataProvider. + * + * @returns {StackDataProvider} The singleton instance. + */ + public static getInstance(): StackDataProvider { + if (!this.instance) { + this.instance = new StackDataProvider(); + } + return this.instance; + } + + /** + * Returns the provided tree item. + * + * @param {TreeItem} element The tree item to return. + * @returns The corresponding VS Code tree item. + */ + getTreeItem(element: TreeItem): TreeItem { + return element; + } + + /** + * Refreshes the tree view data by refetching stacks and triggering the onDidChangeTreeData event. + * + * @returns {Promise} A promise that resolves when the tree view data has been refreshed. + */ + public async refresh(): Promise { + this.stacks = [LOADING_TREE_ITEMS.get('stacks')!]; + this._onDidChangeTreeData.fire(undefined); + const page = this.pagination.currentPage; + const itemsPerPage = this.pagination.itemsPerPage; + + try { + const newStacksData = await this.fetchStacksWithComponents(page, itemsPerPage); + this.stacks = newStacksData; + } catch (error: any) { + this.stacks = createErrorItem(error); + } + + this._onDidChangeTreeData.fire(undefined); + } + + /** + * Retrieves detailed stack information, including components, from the server. + * + * @returns {Promise} A promise that resolves with an array of `StackTreeItem` objects. + */ + async fetchStacksWithComponents(page: number = 1, itemsPerPage: number = 20): Promise { + if (!this.zenmlClientReady) { + return [LOADING_TREE_ITEMS.get('zenmlClient')!]; + } + + try { + const lsClient = LSClient.getInstance(); + const result = await lsClient.sendLsClientRequest( + 'fetchStacks', + [page, itemsPerPage] + ); + if (!result || 'error' in result) { + if ('clientVersion' in result && 'serverVersion' in result) { + return createErrorItem(result); + } else { + console.error(`Failed to fetch stacks: ${result.error}`); + return []; + } + } + + if ('stacks' in result) { + const { stacks, total, total_pages, current_page, items_per_page } = result; + + this.pagination = { + currentPage: current_page, + itemsPerPage: items_per_page, + totalItems: total, + totalPages: total_pages, + }; + + return stacks.map((stack: Stack) => + this.convertToStackTreeItem(stack, this.isActiveStack(stack.id)) + ); + } else { + console.error(`Unexpected response format:`, result); + return []; + } + } catch (error: any) { + console.error(`Failed to fetch stacks: ${error}`); + return [new ErrorTreeItem("Error", `Failed to fetch stacks: ${error.message || error.toString()}`)]; + + } + } + + public async goToNextPage() { + if (this.pagination.currentPage < this.pagination.totalPages) { + this.pagination.currentPage++; + await this.refresh(); + } + } + + public async goToPreviousPage() { + if (this.pagination.currentPage > 1) { + this.pagination.currentPage--; + await this.refresh(); + } + } + + public async updateItemsPerPage() { + const selected = await window.showQuickPick(ITEMS_PER_PAGE_OPTIONS, { + placeHolder: "Choose the max number of stacks to display per page", + }); + if (selected) { + this.pagination.itemsPerPage = parseInt(selected, 10); + this.pagination.currentPage = 1; + await this.refresh(); + } + } + + /** + * Retrieves the children of a given tree item. + * + * @param {TreeItem} element The tree item whose children to retrieve. + * @returns A promise resolving to an array of child tree items or undefined if there are no children. + */ + async getChildren(element?: TreeItem): Promise { + if (!element) { + const stacks = await this.fetchStacksWithComponents(this.pagination.currentPage, this.pagination.itemsPerPage); + if (this.pagination.currentPage < this.pagination.totalPages) { + stacks.push(new CommandTreeItem("Next Page", 'zenml.nextStackPage', undefined, 'arrow-circle-right')); + } + if (this.pagination.currentPage > 1) { + stacks.unshift(new CommandTreeItem("Previous Page", 'zenml.previousStackPage', undefined, 'arrow-circle-left')); + } + return stacks; + } else if (element instanceof StackTreeItem) { + return element.children; + } + return undefined; + } + + /** + * Helper method to determine if a stack is the active stack. + * + * @param {string} stackId The ID of the stack. + * @returns {boolean} True if the stack is active; otherwise, false. + */ + private isActiveStack(stackId: string): boolean { + const activeStackId = workspace.getConfiguration('zenml').get('activeStackId'); + return stackId === activeStackId; + } + + /** + * Transforms a stack from the API into a `StackTreeItem` with component sub-items. + * + * @param {any} stack - The stack object fetched from the API. + * @returns {StackTreeItem} A `StackTreeItem` object representing the stack and its components. + */ + private convertToStackTreeItem(stack: Stack, isActive: boolean): StackTreeItem { + const componentTreeItems = Object.entries(stack.components).flatMap(([type, componentsArray]) => + componentsArray.map( + (component: StackComponent) => new StackComponentTreeItem(component, stack.id) + ) + ); + return new StackTreeItem(stack.name, stack.id, componentTreeItems, isActive); + } +} diff --git a/src/views/activityBar/stackView/StackTreeItems.ts b/src/views/activityBar/stackView/StackTreeItems.ts new file mode 100644 index 00000000..909e0405 --- /dev/null +++ b/src/views/activityBar/stackView/StackTreeItems.ts @@ -0,0 +1,56 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as vscode from 'vscode'; +import { StackComponent } from '../../../types/StackTypes'; + +/** + * A TreeItem for displaying a stack in the VSCode TreeView. + * This item can be expanded to show the components of the stack. + */ +export class StackTreeItem extends vscode.TreeItem { + public children: vscode.TreeItem[] | undefined; + public isActive: boolean; + + constructor( + public readonly label: string, + public readonly id: string, + components: StackComponentTreeItem[], + isActive?: boolean + ) { + super(label, vscode.TreeItemCollapsibleState.Collapsed); + this.children = components; + this.contextValue = 'stack'; + this.isActive = isActive || false; + + if (isActive) { + this.label = `${this.label} 🟢`; + } + } +} + +/** + * A TreeItem for displaying a stack component in the VSCode TreeView. + */ +export class StackComponentTreeItem extends vscode.TreeItem { + constructor( + public component: StackComponent, + public stackId: string + ) { + super(component.name, vscode.TreeItemCollapsibleState.None); + + this.tooltip = `Type: ${component.type}, Flavor: ${component.flavor}, ID: ${stackId}`; + this.description = `${component.type} (${component.flavor})`; + this.contextValue = 'stackComponent'; + this.id = `${stackId}-${component.id}`; + } +} diff --git a/src/views/statusBar/index.ts b/src/views/statusBar/index.ts new file mode 100644 index 00000000..da6e9ae9 --- /dev/null +++ b/src/views/statusBar/index.ts @@ -0,0 +1,103 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// 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. +import * as vscode from 'vscode'; +import { stackCommands } from '../../commands/stack/cmds'; +import { getActiveStack } from '../../commands/stack/utils'; +import { EventBus } from '../../services/EventBus'; +import { LSP_ZENML_STACK_CHANGED, SERVER_STATUS_UPDATED } from '../../utils/constants'; + +/** + * Represents the ZenML extension's status bar. + * This class manages two main status indicators: the server status and the active stack name. + */ +export default class ZenMLStatusBar { + private static instance: ZenMLStatusBar; + private serverStatusItem: vscode.StatusBarItem; + private activeStackItem: vscode.StatusBarItem; + private activeStack: string = 'Loading...'; + private eventBus = EventBus.getInstance(); + + /** + * Initializes a new instance of the ZenMLStatusBar class. + * Sets up the status bar items for server status and active stack, subscribes to server status updates, + * and initiates the initial refresh of the status bar state. + */ + constructor() { + this.serverStatusItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Right, 100); + this.activeStackItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Right, 99); + this.subscribeToEvents(); + } + + /** + * Subscribes to relevant events to trigger a refresh of the status bar. + * + * @returns void + */ + private subscribeToEvents(): void { + this.eventBus.on(LSP_ZENML_STACK_CHANGED, async () => { + await this.refreshActiveStack(); + await stackCommands.refreshStackView(); + }); + + this.eventBus.on(SERVER_STATUS_UPDATED, ({ isConnected, serverUrl }) => { + this.updateServerStatusIndicator(isConnected, serverUrl); + }); + } + + /** + * Retrieves or creates an instance of the ZenMLStatusBar. + * This method implements the Singleton pattern to ensure that only one instance of ZenMLStatusBar exists. + * + * @returns {ZenMLStatusBar} The singleton instance of the ZenMLStatusBar. + */ + public static getInstance(): ZenMLStatusBar { + if (!ZenMLStatusBar.instance) { + ZenMLStatusBar.instance = new ZenMLStatusBar(); + } + return ZenMLStatusBar.instance; + } + + /** + * Asynchronously refreshes the active stack display in the status bar. + * Attempts to retrieve the current active stack name and updates the status bar item accordingly. + * Displays an error message in the status bar if unable to fetch the active stack. + */ + public async refreshActiveStack(): Promise { + try { + const activeStack = await getActiveStack(); + this.activeStack = activeStack?.name || 'default'; + this.activeStackItem.text = `${this.activeStack}`; + this.activeStackItem.tooltip = 'Active ZenML stack.'; + this.activeStackItem.show(); + } catch (error) { + console.error('Failed to fetch active ZenML stack:', error); + this.activeStack = 'Error'; + } + } + + /** + * Updates the server status indicator in the status bar. + * Sets the text, color, and tooltip of the server status item based on the connection status. + * + * @param {boolean} isConnected Whether the server is currently connected. + * @param {string} serverAddress The address of the server, used in the tooltip. + */ + public updateServerStatusIndicator(isConnected: boolean, serverAddress: string) { + this.serverStatusItem.text = isConnected ? `$(vm-active)` : `$(vm-connect)`; + this.serverStatusItem.color = isConnected ? 'green' : ''; + this.serverStatusItem.tooltip = isConnected + ? `Server running at ${serverAddress}.` + : 'Server not running. Click to refresh status.'; + this.serverStatusItem.show(); + } +} diff --git a/tsconfig.json b/tsconfig.json new file mode 100644 index 00000000..45e4ce9b --- /dev/null +++ b/tsconfig.json @@ -0,0 +1,11 @@ +{ + "compilerOptions": { + "module": "Node16", + "target": "ES2022", + "lib": ["ES2022"], + "sourceMap": true, + "rootDir": "src", + "strict": true, + "skipLibCheck": true + } +} diff --git a/webpack.config.js b/webpack.config.js new file mode 100644 index 00000000..37d7024f --- /dev/null +++ b/webpack.config.js @@ -0,0 +1,48 @@ +//@ts-check + +'use strict'; + +const path = require('path'); + +//@ts-check +/** @typedef {import('webpack').Configuration} WebpackConfig **/ + +/** @type WebpackConfig */ +const extensionConfig = { + target: 'node', // VS Code extensions run in a Node.js-context 📖 -> https://webpack.js.org/configuration/node/ + mode: 'none', // this leaves the source code as close as possible to the original (when packaging we set this to 'production') + + entry: './src/extension.ts', // the entry point of this extension, 📖 -> https://webpack.js.org/configuration/entry-context/ + output: { + // the bundle is stored in the 'dist' folder (check package.json), 📖 -> https://webpack.js.org/configuration/output/ + path: path.resolve(__dirname, 'dist'), + filename: 'extension.js', + libraryTarget: 'commonjs2' + }, + externals: { + vscode: 'commonjs vscode' // the vscode-module is created on-the-fly and must be excluded. Add other modules that cannot be webpack'ed, 📖 -> https://webpack.js.org/configuration/externals/ + // modules added here also need to be added in the .vscodeignore file + }, + resolve: { + // support reading TypeScript and JavaScript files, 📖 -> https://github.com/TypeStrong/ts-loader + extensions: ['.ts', '.js'] + }, + module: { + rules: [ + { + test: /\.ts$/, + exclude: /node_modules/, + use: [ + { + loader: 'ts-loader' + } + ] + } + ] + }, + devtool: 'nosources-source-map', + infrastructureLogging: { + level: "log", // enables logging required for problem matchers + }, +}; +module.exports = [ extensionConfig ]; \ No newline at end of file