├── .coderabbit.yml ├── .github ├── actions │ └── lint │ │ └── action.yml ├── dependabot.yml └── workflows │ ├── ci.yml │ ├── release.yml │ ├── spellcheck.yml │ └── test-publish.yml ├── .gitignore ├── .nvmrc ├── .prettierrc ├── .pylintrc ├── .tool-versions ├── .typos.toml ├── .vscode-test.mjs ├── .vscode ├── extensions.json ├── launch.json ├── settings.json └── tasks.json ├── .vscodeignore ├── .zen └── config.yaml ├── CHANGELOG.md ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── RELEASE.md ├── VERSIONS.md ├── ZENML_BREAKING_CHANGES_ANALYSIS.md ├── bundled └── tool │ ├── __init__.py │ ├── _debug_server.py │ ├── constants.py │ ├── lazy_import.py │ ├── lsp_jsonrpc.py │ ├── lsp_runner.py │ ├── lsp_server.py │ ├── lsp_utils.py │ ├── lsp_zenml.py │ ├── type_hints.py │ ├── zen_watcher.py │ ├── zenml_client.py │ ├── zenml_grapher.py │ ├── zenml_serializers.py │ └── zenml_wrappers.py ├── eslint.config.mjs ├── mypy.ini ├── noxfile.py ├── package-lock.json ├── package.json ├── package.nls.json ├── pyproject.toml ├── requirements-dev.txt ├── requirements.in ├── requirements.txt ├── resources ├── components-form │ ├── components.css │ └── components.js ├── dag-view │ ├── dag.css │ ├── dag.js │ └── icons │ │ ├── alert.svg │ │ ├── cached.svg │ │ ├── check.svg │ │ ├── database.svg │ │ ├── dataflow.svg │ │ ├── initializing.svg │ │ └── play.svg ├── extension-logo.png ├── logo.png ├── logo.svg ├── python.png ├── stacks-form │ ├── stacks.css │ └── stacks.js ├── zenml-extension-dag.gif ├── zenml-extension-new.gif ├── zenml-extension.gif └── zenml_logo.png ├── runtime.txt ├── scripts ├── clear_and_compile.sh ├── format.sh ├── lint.sh └── spellcheck.sh ├── src ├── commands │ ├── components │ │ ├── ComponentsForm.ts │ │ ├── cmds.ts │ │ └── registry.ts │ ├── environment │ │ ├── cmds.ts │ │ └── registry.ts │ ├── models │ │ ├── cmds.ts │ │ └── registry.ts │ ├── pipelines │ │ ├── cmds.ts │ │ ├── registry.ts │ │ └── utils.ts │ ├── projects │ │ ├── cmds.ts │ │ ├── registry.ts │ │ └── utils.ts │ ├── server │ │ ├── cmds.ts │ │ ├── registry.ts │ │ └── utils.ts │ └── stack │ │ ├── StackForm.ts │ │ ├── cmds.ts │ │ ├── registry.ts │ │ └── utils.ts ├── common │ ├── WebviewBase.ts │ ├── api.ts │ ├── constants.ts │ ├── log │ │ └── logging.ts │ ├── panels.ts │ ├── python.ts │ ├── server.ts │ ├── settings.ts │ ├── status.ts │ ├── utilities.ts │ └── vscodeapi.ts ├── dag │ ├── DagConfig.ts │ ├── index.ts │ ├── renderer │ │ ├── DagRenderer.ts │ │ ├── HtmlTemplateBuilder.ts │ │ └── SvgRenderer.ts │ └── utils │ │ ├── IconLoader.ts │ │ └── StatusUtils.ts ├── extension.ts ├── services │ ├── EventBus.ts │ ├── LSClient.ts │ └── ZenExtension.ts ├── test │ ├── python_tests │ │ ├── __init__.py │ │ ├── lsp_test_client │ │ │ ├── __init__.py │ │ │ ├── constants.py │ │ │ ├── defaults.py │ │ │ ├── session.py │ │ │ └── utils.py │ │ ├── requirements.in │ │ ├── requirements.txt │ │ ├── test_data │ │ │ └── sample1 │ │ │ │ ├── sample.py │ │ │ │ └── sample.unformatted │ │ └── test_server.py │ └── ts_tests │ │ ├── __mocks__ │ │ ├── MockEventBus.ts │ │ ├── MockLSClient.ts │ │ ├── MockViewProviders.ts │ │ └── constants.ts │ │ ├── commands │ │ ├── dagRender.test.ts │ │ ├── modelCommands.test.ts │ │ ├── pipelineUtils.test.ts │ │ ├── projectCommands.test.ts │ │ ├── serverCommands.test.ts │ │ └── stackCommands.test.ts │ │ ├── common │ │ └── WebviewBase.test.ts │ │ ├── extension.test.ts │ │ ├── integration │ │ └── serverConfigUpdate.test.ts │ │ ├── services │ │ └── ZenExtension.test.ts │ │ ├── unit │ │ ├── ServerDataProvider.test.ts │ │ └── eventBus.test.ts │ │ └── views │ │ └── PipelineDataProvider.test.ts ├── types │ ├── HydratedTypes.ts │ ├── LSClientResponseTypes.ts │ ├── LSNotificationTypes.ts │ ├── ModelTypes.ts │ ├── PipelineTypes.ts │ ├── ProjectTypes.ts │ ├── QuickPickItemTypes.ts │ ├── ServerInfoTypes.ts │ └── StackTypes.ts ├── utils │ ├── componentUtils.ts │ ├── constants.ts │ ├── global.ts │ ├── notifications.ts │ ├── refresh.ts │ └── ui-constants.ts └── views │ ├── activityBar │ ├── common │ │ ├── ErrorTreeItem.ts │ │ ├── LoadingTreeItem.ts │ │ ├── PaginatedDataProvider.ts │ │ ├── PaginationTreeItems.ts │ │ ├── TreeItemWithChildren.ts │ │ └── index.ts │ ├── componentView │ │ ├── ComponentDataProvider.ts │ │ └── ComponentTreeItems.ts │ ├── environmentView │ │ ├── EnvironmentDataProvider.ts │ │ ├── EnvironmentItem.ts │ │ └── viewHelpers.ts │ ├── index.ts │ ├── modelView │ │ ├── ModelDataProvider.ts │ │ └── ModelTreeItems.ts │ ├── pipelineView │ │ ├── PipelineDataProvider.ts │ │ └── PipelineTreeItems.ts │ ├── projectView │ │ ├── ProjectDataProvider.ts │ │ └── ProjectTreeItems.ts │ ├── serverView │ │ ├── ServerDataProvider.ts │ │ └── ServerTreeItems.ts │ └── stackView │ │ ├── StackDataProvider.ts │ │ └── StackTreeItems.ts │ ├── panel │ └── panelView │ │ ├── PanelDataProvider.ts │ │ └── PanelTreeItem.ts │ └── statusBar │ └── index.ts ├── tsconfig.json └── webpack.config.mjs /.coderabbit.yml: -------------------------------------------------------------------------------- 1 | language: "en" 2 | early_access: false 3 | reviews: 4 | request_changes_workflow: false 5 | high_level_summary: true 6 | poem: true 7 | review_status: true 8 | collapse_walkthrough: false 9 | path_filters: 10 | - "!**/.xml" 11 | path_instructions: 12 | - path: "**/*.js" 13 | instructions: "Review the JavaScript code for conformity with the Google JavaScript style guide, highlighting any deviations." 14 | - path: "**/*.ts" 15 | instructions: "Review the Typescript code for conformity with industry standards and best practices, highlighting any deviations." 16 | - path: "**/*.py" 17 | instructions: | 18 | "Review the Python code for conformity with Python best practices and industry standards, highlighting any deviations." 19 | auto_review: 20 | enabled: true 21 | ignore_title_keywords: 22 | - "WIP" 23 | - "DO NOT MERGE" 24 | drafts: false 25 | base_branches: 26 | - "develop" 27 | chat: 28 | auto_reply: true 29 | -------------------------------------------------------------------------------- /.github/actions/lint/action.yml: -------------------------------------------------------------------------------- 1 | --- 2 | name: Lint 3 | description: Lint TypeScript and Python code 4 | runs: 5 | using: composite 6 | steps: 7 | - name: Install Node 8 | uses: actions/setup-node@v3 9 | with: 10 | node-version: 20.x 11 | cache: npm 12 | cache-dependency-path: package-lock.json 13 | - name: Install Node dependencies 14 | run: npm ci 15 | shell: bash 16 | - name: Check TypeScript format 17 | run: npm run format-check 18 | shell: bash 19 | - name: Install Python 20 | uses: actions/setup-python@v4 21 | with: 22 | python-version: '3.8' 23 | - name: Pip cache 24 | uses: actions/cache@v3 25 | with: 26 | path: ~/.cache/pip 27 | key: ${{ runner.os }}-pip-lint-${{ hashFiles('**/requirements.txt') }} 28 | restore-keys: | 29 | ${{ runner.os }}-pip-lint- 30 | - name: Install Python dependencies 31 | run: | 32 | python -m pip install -U pip 33 | pip install -r requirements-dev.txt 34 | shell: bash 35 | - name: Lint Python and YAML code 36 | run: ./scripts/lint.sh 37 | shell: bash 38 | -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | --- 2 | # To get started with Dependabot version updates, you'll need to specify which 3 | # package ecosystems to update and where the package manifests are located. 4 | # Please see the documentation for all configuration options: 5 | # https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file 6 | version: 2 7 | updates: 8 | - package-ecosystem: npm 9 | directory: / # Location of package manifests 10 | schedule: 11 | interval: weekly 12 | groups: 13 | dependencies: 14 | patterns: 15 | - '*' # This will group all dependencies together 16 | update-types: [minor, patch] 17 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | --- 2 | name: VSCode Extension CI 3 | on: [push, pull_request] 4 | jobs: 5 | lint: 6 | name: Lint 7 | runs-on: ubuntu-latest 8 | steps: 9 | - name: Checkout 10 | uses: actions/checkout@v3 11 | - name: Lint 12 | uses: ./.github/actions/lint 13 | build-and-test: 14 | needs: lint 15 | runs-on: ubuntu-latest 16 | strategy: 17 | fail-fast: false 18 | steps: 19 | - name: Checkout Repository 20 | uses: actions/checkout@v3 21 | - name: Install Node.js 22 | uses: actions/setup-node@v3 23 | with: 24 | node-version: 20.x 25 | cache: npm 26 | cache-dependency-path: package-lock.json 27 | - name: Install Node dependencies 28 | run: npm ci 29 | - name: Compile TS tests 30 | run: npm run pretest 31 | - name: Run headless test 32 | uses: coactions/setup-xvfb@v1 33 | with: 34 | run: npm test 35 | -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | --- 2 | name: Publish VSCode Extension 3 | on: 4 | release: 5 | types: [created] 6 | jobs: 7 | release: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - name: Checkout Repository 11 | uses: actions/checkout@v4 12 | - name: Get Extension Version 13 | id: get_version 14 | run: echo "version=$(jq -r .version package.json)" >> $GITHUB_OUTPUT 15 | - name: Install Node.js 16 | uses: actions/setup-node@v4 17 | with: 18 | node-version: 20.x 19 | cache: npm 20 | - name: Install dependencies 21 | run: npm ci 22 | - name: Run headless tests 23 | uses: coactions/setup-xvfb@v1 24 | with: 25 | run: npm test 26 | - name: Build Extension 27 | run: npm run package 28 | - name: Package Extension 29 | run: npm run vsce-package 30 | - name: Verify Package 31 | run: | 32 | echo "=== VSIX Package Verification ===" 33 | # List contents of the VSIX package 34 | echo "Listing package contents:" 35 | npx @vscode/vsce ls 36 | 37 | # Validate the package.json 38 | echo -e "\nValidating package.json:" 39 | if jq -e '.name and .publisher and .version and .engines.vscode' package.json > /dev/null; then 40 | echo "✅ package.json contains all required fields" 41 | else 42 | echo "❌ package.json is missing required fields" 43 | exit 1 44 | fi 45 | - name: Publish Extension 46 | if: success() && startsWith(github.ref, 'refs/tags/') 47 | run: npx @vscode/vsce publish 48 | env: 49 | VSCE_PAT: ${{ secrets.VSCE_PAT }} 50 | - name: Generate Changelog for Release Notes 51 | if: success() && startsWith(github.ref, 'refs/tags/') 52 | id: changelog 53 | run: | 54 | VERSION="${{ steps.get_version.outputs.version }}" 55 | 56 | # Create header for release notes 57 | cat > RELEASE_NOTES.md << EOF 58 | ## Release $VERSION 59 | EOF 60 | 61 | # Extract just the current version section from CHANGELOG.md 62 | sed -n "/## \[${VERSION}\]/,/## \[/p" CHANGELOG.md | sed '$d' >> RELEASE_NOTES.md 63 | echo "=== RELEASE NOTES ===" 64 | cat RELEASE_NOTES.md 65 | echo "====================" 66 | - name: Create GitHub Release Assets 67 | if: success() && startsWith(github.ref, 'refs/tags/') 68 | run: | 69 | # Create a VERSION file that can be included in the release 70 | echo "${{ steps.get_version.outputs.version }}" > VERSION 71 | - name: Update GitHub Release 72 | if: success() && startsWith(github.ref, 'refs/tags/') 73 | uses: ncipollo/release-action@v1 74 | with: 75 | artifacts: zenml.vsix,VERSION 76 | bodyFile: RELEASE_NOTES.md 77 | allowUpdates: true 78 | tag: ${{ github.ref_name }} 79 | token: ${{ secrets.GITHUB_TOKEN }} 80 | -------------------------------------------------------------------------------- /.github/workflows/spellcheck.yml: -------------------------------------------------------------------------------- 1 | --- 2 | name: Spellcheck 3 | on: [pull_request] 4 | jobs: 5 | spellcheck: 6 | runs-on: ubuntu-latest 7 | steps: 8 | - name: Checkout Actions Repository 9 | uses: actions/checkout@v4 10 | - name: Spell Check Repo 11 | uses: crate-ci/typos@v1 12 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | Thumbs.db 3 | .vscode-test/ 4 | .idea/ 5 | .eslintcache 6 | lib-cov 7 | *.log 8 | *.log* 9 | pids 10 | 11 | # Node 12 | .npm/ 13 | node_modules/ 14 | npm-debug.log 15 | 16 | # Build outputs 17 | dist/ 18 | out/ 19 | build/ 20 | *.tsbuildinfo 21 | .history/ 22 | dag-packed.js 23 | dag-packed.js.map 24 | 25 | # env 26 | .env 27 | .env* 28 | 29 | # Output of 'npm pack' 30 | *.tgz 31 | 32 | # mypy 33 | .mypy_cache/ 34 | 35 | # From vscode-python-tools-extension-template 36 | *.vsix 37 | .venv/ 38 | .vs/ 39 | .nox/ 40 | bundled/libs/ 41 | **/__pycache__ 42 | **/.pytest_cache 43 | **/.vs -------------------------------------------------------------------------------- /.nvmrc: -------------------------------------------------------------------------------- 1 | v22.14.0 -------------------------------------------------------------------------------- /.prettierrc: -------------------------------------------------------------------------------- 1 | { 2 | "printWidth": 100, 3 | "tabWidth": 2, 4 | "useTabs": false, 5 | "semi": true, 6 | "proseWrap": "preserve", 7 | "singleQuote": true, 8 | "arrowParens": "avoid", 9 | "trailingComma": "es5", 10 | "bracketSpacing": true 11 | } 12 | -------------------------------------------------------------------------------- /.pylintrc: -------------------------------------------------------------------------------- 1 | [MESSAGES CONTROL] 2 | disable= 3 | C0103, # Name doesn't conform to naming style 4 | C0415, # Import outside toplevel 5 | W0613, # Unused argument 6 | R0801, # Similar lines in multiple files 7 | R0903, # Too few public methods -------------------------------------------------------------------------------- /.tool-versions: -------------------------------------------------------------------------------- 1 | nodejs 20.10.0 2 | -------------------------------------------------------------------------------- /.typos.toml: -------------------------------------------------------------------------------- 1 | [default] 2 | locale = "en-us" 3 | 4 | [default.extend-words] 5 | # Ignore list of words. Usage example: 6 | # teh = "teh" 7 | -------------------------------------------------------------------------------- /.vscode-test.mjs: -------------------------------------------------------------------------------- 1 | import { defineConfig } from '@vscode/test-cli'; 2 | 3 | export default defineConfig({ 4 | files: 'out/test/**/*.test.js', 5 | }); 6 | -------------------------------------------------------------------------------- /.vscode/extensions.json: -------------------------------------------------------------------------------- 1 | { 2 | // See http://go.microsoft.com/fwlink/?LinkId=827846 3 | // for the documentation about the extensions.json format 4 | "recommendations": [ 5 | "dbaeumer.vscode-eslint", 6 | "amodio.tsl-problem-matcher", 7 | "ms-vscode.extension-test-runner" 8 | ] 9 | } 10 | -------------------------------------------------------------------------------- /.vscode/launch.json: -------------------------------------------------------------------------------- 1 | // A launch configuration that compiles the extension and then opens it inside a new window 2 | // Use IntelliSense to learn about possible attributes. 3 | // Hover to view descriptions of existing attributes. 4 | // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 5 | { 6 | "version": "0.2.0", 7 | "configurations": [ 8 | { 9 | "name": "Run Extension", 10 | "type": "extensionHost", 11 | "request": "launch", 12 | "args": ["--extensionDevelopmentPath=${workspaceFolder}"], 13 | "outFiles": ["${workspaceFolder}/dist/**/*.js"], 14 | "preLaunchTask": "${defaultBuildTask}", 15 | "timeout": 20000 16 | } 17 | ] 18 | } 19 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | // Place your settings in this file to overwrite default and user settings. 2 | { 3 | "python.defaultInterpreterPath": "~/.pyenv/versions/3.11.6/envs/vscode-zenml/bin/python", 4 | "files.exclude": { 5 | "out": false, // set this to true to hide the "out" folder with the compiled JS files 6 | "dist": false // set this to true to hide the "dist" folder with the compiled JS files 7 | }, 8 | "search.exclude": { 9 | "out": true, // set this to false to include "out" folder in search results 10 | "dist": true // set this to false to include "dist" folder in search results 11 | }, 12 | // Turn off tsc task auto detection since we have the necessary tasks as npm scripts 13 | "typescript.tsc.autoDetect": "off", 14 | "python.testing.pytestArgs": ["src/test/python_tests"], 15 | "python.testing.unittestEnabled": false, 16 | "python.testing.pytestEnabled": true, 17 | "python.testing.cwd": "${workspaceFolder}", 18 | "python.analysis.extraPaths": ["bundled/libs", "bundled/tool"], 19 | "python.analysis.typeCheckingMode": "off" 20 | } 21 | -------------------------------------------------------------------------------- /.vscode/tasks.json: -------------------------------------------------------------------------------- 1 | // See https://go.microsoft.com/fwlink/?LinkId=733558 2 | // for the documentation about the tasks.json format 3 | { 4 | "version": "2.0.0", 5 | "tasks": [ 6 | { 7 | "type": "npm", 8 | "script": "watch", 9 | "problemMatcher": "$ts-webpack-watch", 10 | "isBackground": true, 11 | "presentation": { 12 | "reveal": "never", 13 | "group": "watchers" 14 | }, 15 | "group": { 16 | "kind": "build", 17 | "isDefault": true 18 | } 19 | }, 20 | { 21 | "type": "npm", 22 | "script": "watch-tests", 23 | "problemMatcher": "$tsc-watch", 24 | "isBackground": true, 25 | "presentation": { 26 | "reveal": "never", 27 | "group": "watchers" 28 | }, 29 | "group": "build" 30 | }, 31 | { 32 | "label": "tasks: watch-tests", 33 | "dependsOn": ["npm: watch", "npm: watch-tests"], 34 | "problemMatcher": [] 35 | }, 36 | { 37 | "type": "npm", 38 | "script": "compile", 39 | "group": "build", 40 | "problemMatcher": [], 41 | "label": "npm: compile", 42 | "detail": "webpack" 43 | } 44 | ] 45 | } 46 | -------------------------------------------------------------------------------- /.vscodeignore: -------------------------------------------------------------------------------- 1 | .vscode/** 2 | .vscode-test/** 3 | out/** 4 | node_modules/** 5 | src/** 6 | .gitignore 7 | .yarnrc 8 | webpack.config.mjs 9 | vsc-extension-quickstart.md 10 | **/tsconfig.json 11 | **/.eslintrc.json 12 | **/*.map 13 | **/*.js.map 14 | **/*.ts 15 | **/.vscode-test.* 16 | -------------------------------------------------------------------------------- /.zen/config.yaml: -------------------------------------------------------------------------------- 1 | active_project_id: 6e23c046-6cc4-411c-8b9f-75f0c8a1a818 2 | active_stack_id: a51bf906-a221-46d7-81b8-c81f8adc5097 3 | -------------------------------------------------------------------------------- /RELEASE.md: -------------------------------------------------------------------------------- 1 | # Release Process Guide 2 | 3 | This document describes the process for testing and releasing new versions of the ZenML VSCode extension. 4 | 5 | ## Testing Releases 6 | 7 | Before publishing a new version of the extension, you should test the release process to ensure everything works correctly. We have a dedicated GitHub Actions workflow for this purpose. 8 | 9 | ### Using the Test Workflow 10 | 11 | 1. **Prepare your changes** 12 | 13 | - Update the version in `package.json` 14 | - Update `CHANGELOG.md` with the new version and changes 15 | - Commit these changes to your branch 16 | 17 | 2. **Run the test workflow** 18 | 19 | - Go to the Actions tab in the GitHub repository 20 | - Select "Test VSCode Extension Release" workflow 21 | - Click "Run workflow" 22 | - Select the branch containing your changes 23 | - Click "Run workflow" again 24 | 25 | 3. **Review the workflow output** 26 | 27 | - Check that all steps completed successfully 28 | - Verify the package content and validation passed 29 | - Review the generated changelog and release notes 30 | 31 | 4. **Test the extension locally** 32 | - Download the VSIX artifact from the workflow run 33 | - Install it in VS Code using: `code --install-extension zenml.vsix` 34 | - Test that all features work as expected 35 | 36 | ### How the Test Workflow Works 37 | 38 | The test workflow (`.github/workflows/test-publish.yml`) performs these steps: 39 | 40 | 1. **Build and Package** 41 | 42 | - Builds the extension source code 43 | - Packages it into a VSIX file 44 | - Runs the test suite 45 | 46 | 2. **Validate** 47 | 48 | - Verifies the package content 49 | - Validates required fields in package.json 50 | - Simulates the publishing process without actually publishing 51 | 52 | 3. **Generate Documentation** 53 | 54 | - Creates a changelog from git history since the last tag 55 | - Extracts the current version's information from CHANGELOG.md 56 | 57 | 4. **Upload Artifacts** 58 | - Makes the VSIX package available for download 59 | - Provides the generated changelog and release notes for review 60 | 61 | ## Publishing a Release 62 | 63 | Once you've tested the release and everything looks good, you can publish it: 64 | 65 | 1. **Tag the Release** 66 | 67 | ```bash 68 | git tag -a 0.0.x -m "Version 0.0.x" 69 | git push origin 0.0.x 70 | ``` 71 | 72 | 2. **Create a GitHub Release** 73 | - Go to the Releases page in the GitHub repository 74 | - Click "Draft a new release" 75 | - Select the tag you just pushed 76 | - Title the release "Version 0.0.x" 77 | - The release workflow will automatically: 78 | - Publish the extension to the VS Code Marketplace 79 | - Attach the VSIX file to the GitHub release 80 | - Use the content from CHANGELOG.md as release notes 81 | 82 | ### Release Workflow Details 83 | 84 | The release workflow (`.github/workflows/release.yml`) is triggered when a GitHub release is created. It: 85 | 86 | 1. Builds and packages the extension 87 | 2. Runs tests to ensure quality 88 | 3. Publishes the extension to the VS Code Marketplace 89 | 4. Updates the GitHub release with artifacts and release notes 90 | 91 | ## Troubleshooting 92 | 93 | If you encounter issues during the release process: 94 | 95 | - **Build failures**: Check the logs for specific error messages 96 | - **Test failures**: Run tests locally to debug before retrying 97 | - **Publication issues**: Verify the VSCE_PAT secret is valid in the repository settings 98 | - **Version conflicts**: Ensure the version in package.json matches your tag 99 | -------------------------------------------------------------------------------- /VERSIONS.md: -------------------------------------------------------------------------------- 1 | # ZenML Extension Version Compatibility 2 | 3 | This file serves as a reference for ZenML extension version compatibility with different ZenML library versions. If you're working with an older version of ZenML, please install the appropriate extension version for best compatibility. 4 | 5 | ## Version Compatibility Table 6 | 7 | | Extension Version | ZenML Library Version | Notes | 8 | | ----------------- | --------------------- | ------------------------------------------------------------- | 9 | | 0.0.21 | 0.83.x | Performance optimizations, DAG loading spinner, TTL caching | 10 | | 0.0.20 | 0.80.x - 0.82.x | Models tree view and stability improvements | 11 | | 0.0.16 | 0.80.x | Includes stability fixes and DAG view improvements | 12 | | 0.0.15 | 0.80.0 | Adds stability fixes and improves project management | 13 | | 0.0.13 – 0.0.14 | 0.80.0 | Includes Project View and full project management support | 14 | | 0.0.12 | 0.75.0 – 0.79.x | Supports stack registration and components visualization | 15 | | 0.0.11 | 0.63.0 – 0.74.0 | Includes DAG Visualizer and pipeline run management | 16 | | 0.0.5 – 0.0.10 | 0.55.0 – 0.63.0 | Early versions with basic functionality | 17 | 18 | ## Installing a Specific Extension Version 19 | 20 | If you need to work with an older ZenML version, you can install a compatible extension version using the VS Code UI or the command line: 21 | 22 | ### Using VS Code UI: 23 | 24 | 1. Go to the Extensions view (Ctrl+Shift+X) 25 | 2. Search for "ZenML" 26 | 3. Click the dropdown next to the Install button 27 | 4. Select "Install Another Version..." 28 | 5. Choose the version that matches your ZenML library version 29 | 30 | ### Using Command Line: 31 | 32 | ```bash 33 | # Example for installing version 0.0.11 34 | code --install-extension zenml.zenml-vscode@0.0.11 35 | ``` 36 | 37 | ## Upgrading Your ZenML Library 38 | 39 | For the best experience, we recommend upgrading your ZenML library to the latest version: 40 | 41 | ```bash 42 | pip install -U zenml 43 | ``` 44 | 45 | This will allow you to take advantage of all the features available in the latest extension version. 46 | -------------------------------------------------------------------------------- /bundled/tool/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) ZenML GmbH 2024. All Rights Reserved. 2 | # Licensed under the Apache License, Version 2.0 (the "License"); 3 | # you may not use this file except in compliance with the License. 4 | # You may obtain a copy of the License at: 5 | # 6 | # http://www.apache.org/licenses/LICENSE-2.0 7 | # 8 | # Unless required by applicable law or agreed to in writing, software 9 | # distributed under the License is distributed on an "AS IS" BASIS, 10 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | # or implied. See the License for the specific language governing 12 | # permissions and limitations under the License. 13 | -------------------------------------------------------------------------------- /bundled/tool/_debug_server.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) ZenML GmbH 2024. All Rights Reserved. 2 | # Licensed under the Apache License, Version 2.0 (the "License"); 3 | # you may not use this file except in compliance with the License. 4 | # You may obtain a copy of the License at: 5 | # 6 | # http://www.apache.org/licenses/LICENSE-2.0 7 | # 8 | # Unless required by applicable law or agreed to in writing, software 9 | # distributed under the License is distributed on an "AS IS" BASIS, 10 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | # or implied. See the License for the specific language governing 12 | # permissions and limitations under the License. 13 | 14 | """Debugging support for LSP.""" 15 | 16 | import os 17 | import pathlib 18 | import runpy 19 | import sys 20 | 21 | 22 | def update_sys_path(path_to_add: str) -> None: 23 | """Add given path to `sys.path`.""" 24 | if path_to_add not in sys.path and os.path.isdir(path_to_add): 25 | sys.path.append(path_to_add) 26 | 27 | 28 | # Ensure debugger is loaded before we load anything else, to debug initialization. 29 | debugger_path = os.getenv("DEBUGPY_PATH", None) 30 | if debugger_path: 31 | if debugger_path.endswith("debugpy"): 32 | debugger_path = os.fspath(pathlib.Path(debugger_path).parent) 33 | 34 | # pylint: disable=wrong-import-position,import-error 35 | import debugpy 36 | 37 | update_sys_path(debugger_path) 38 | 39 | # pylint: disable=wrong-import-position,import-error 40 | 41 | # 5678 is the default port, If you need to change it update it here 42 | # and in launch.json. 43 | debugpy.connect(5678) 44 | 45 | # This will ensure that execution is paused as soon as the debugger 46 | # connects to VS Code. If you don't want to pause here comment this 47 | # line and set breakpoints as appropriate. 48 | debugpy.breakpoint() 49 | 50 | SERVER_PATH = os.fspath(pathlib.Path(__file__).parent / "lsp_server.py") 51 | # NOTE: Set breakpoint in `lsp_server.py` before continuing. 52 | runpy.run_path(SERVER_PATH, run_name="__main__") 53 | -------------------------------------------------------------------------------- /bundled/tool/constants.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) ZenML GmbH 2024. All Rights Reserved. 2 | # Licensed under the Apache License, Version 2.0 (the "License"); 3 | # you may not use this file except in compliance with the License. 4 | # You may obtain a copy of the License at: 5 | # 6 | # http://www.apache.org/licenses/LICENSE-2.0 7 | # 8 | # Unless required by applicable law or agreed to in writing, software 9 | # distributed under the License is distributed on an "AS IS" BASIS, 10 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | # or implied. See the License for the specific language governing 12 | # permissions and limitations under the License. 13 | 14 | """Constants for ZenML Tool""" 15 | 16 | TOOL_MODULE_NAME = "zenml-python" 17 | TOOL_DISPLAY_NAME = "ZenML" 18 | MIN_ZENML_VERSION = "0.80.0" 19 | 20 | """Constants for ZenML Notifications and Events""" 21 | 22 | IS_ZENML_INSTALLED = "zenml/isInstalled" 23 | ZENML_CLIENT_INITIALIZED = "zenml/clientInitialized" 24 | ZENML_SERVER_CHANGED = "zenml/serverChanged" 25 | ZENML_STACK_CHANGED = "zenml/stackChanged" 26 | ZENML_PROJECT_CHANGED = "zenml/projectChanged" 27 | ZENML_REQUIREMENTS_NOT_MET = "zenml/requirementsNotMet" 28 | -------------------------------------------------------------------------------- /bundled/tool/lazy_import.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) ZenML GmbH 2024. All Rights Reserved. 2 | # Licensed under the Apache License, Version 2.0 (the "License"); 3 | # you may not use this file except in compliance with the License. 4 | # You may obtain a copy of the License at: 5 | # 6 | # http://www.apache.org/licenses/LICENSE-2.0 7 | # 8 | # Unless required by applicable law or agreed to in writing, software 9 | # distributed under the License is distributed on an "AS IS" BASIS, 10 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | # or implied. See the License for the specific language governing 12 | # permissions and limitations under the License. 13 | """ 14 | Utilities for lazy importing and temporary suppression of stdout and logging. 15 | Reduces noise when integrating with logging-heavy systems like LSP. 16 | Useful for quieter module loads and command executions. 17 | """ 18 | 19 | import importlib 20 | import logging 21 | import os 22 | import sys 23 | from contextlib import contextmanager 24 | 25 | 26 | @contextmanager 27 | def suppress_logging_temporarily(level=logging.ERROR): 28 | """ 29 | Temporarily elevates logging level and suppresses stdout to 30 | minimize console output during imports. 31 | 32 | Parameters: 33 | level (int): Temporary logging level (default: ERROR). 34 | 35 | Yields: 36 | None: While suppressing stdout. 37 | """ 38 | original_level = logging.root.level 39 | original_stdout = sys.stdout 40 | logging.root.setLevel(level) 41 | with open(os.devnull, "w", encoding="utf-8") as fnull: 42 | sys.stdout = fnull 43 | try: 44 | yield 45 | finally: 46 | sys.stdout = original_stdout 47 | logging.root.setLevel(original_level) 48 | 49 | 50 | @contextmanager 51 | def suppress_stdout_temporarily(): 52 | """ 53 | This context manager suppresses stdout for LSP commands, 54 | silencing unnecessary or unwanted output during execution. 55 | 56 | Yields: 57 | None: While suppressing stdout. 58 | """ 59 | with open(os.devnull, "w", encoding="utf-8") as fnull: 60 | original_stdout = sys.stdout 61 | original_stderr = sys.stderr 62 | sys.stdout = fnull 63 | sys.stderr = fnull 64 | try: 65 | yield 66 | finally: 67 | sys.stdout = original_stdout 68 | sys.stderr = original_stderr 69 | 70 | 71 | def lazy_import(module_name, class_name=None): 72 | """ 73 | Lazily imports a module or class, suppressing ZenML log output 74 | to minimize initialization time and noise. 75 | 76 | Args: 77 | module_name (str): The name of the module to import. 78 | class_name (str, optional): The class name within the module. Defaults to None. 79 | 80 | Returns: 81 | The imported module or class. 82 | """ 83 | with suppress_logging_temporarily(): 84 | module = importlib.import_module(module_name) 85 | if class_name: 86 | return getattr(module, class_name) 87 | return module 88 | -------------------------------------------------------------------------------- /bundled/tool/lsp_runner.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) ZenML GmbH 2024. All Rights Reserved. 2 | # Licensed under the Apache License, Version 2.0 (the "License"); 3 | # you may not use this file except in compliance with the License. 4 | # You may obtain a copy of the License at: 5 | # 6 | # http://www.apache.org/licenses/LICENSE-2.0 7 | # 8 | # Unless required by applicable law or agreed to in writing, software 9 | # distributed under the License is distributed on an "AS IS" BASIS, 10 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | # or implied. See the License for the specific language governing 12 | # permissions and limitations under the License. 13 | 14 | """ 15 | Runner to use when running under a different interpreter. 16 | """ 17 | 18 | import os 19 | import pathlib 20 | import sys 21 | import traceback 22 | 23 | 24 | # ********************************************************** 25 | # Update sys.path before importing any bundled libraries. 26 | # ********************************************************** 27 | def update_sys_path(path_to_add: str, strategy: str) -> None: 28 | """Add given path to `sys.path`.""" 29 | if path_to_add not in sys.path and os.path.isdir(path_to_add): 30 | if strategy == "useBundled": 31 | sys.path.insert(0, path_to_add) 32 | elif strategy == "fromEnvironment": 33 | sys.path.append(path_to_add) 34 | 35 | 36 | # Ensure that we can import LSP libraries, and other bundled libraries. 37 | update_sys_path( 38 | os.fspath(pathlib.Path(__file__).parent.parent / "libs"), 39 | os.getenv("LS_IMPORT_STRATEGY", "useBundled"), 40 | ) 41 | 42 | 43 | # pylint: disable=wrong-import-position,import-error 44 | import lsp_jsonrpc as jsonrpc # noqa: E402 45 | import lsp_utils as utils # noqa: E402 46 | 47 | RPC = jsonrpc.create_json_rpc(sys.stdin.buffer, sys.stdout.buffer) 48 | 49 | EXIT_NOW = False 50 | while not EXIT_NOW: 51 | msg = RPC.receive_data() 52 | 53 | method = msg["method"] 54 | if method == "exit": 55 | EXIT_NOW = True 56 | continue 57 | 58 | if method == "run": 59 | is_exception = False 60 | # This is needed to preserve sys.path, pylint modifies 61 | # sys.path and that might not work for this scenario 62 | # next time around. 63 | with utils.substitute_attr(sys, "path", sys.path[:]): 64 | try: 65 | # TODO: `utils.run_module` is equivalent to running `python -m `. 66 | # If your tool supports a programmatic API then replace the function below 67 | # with code for your tool. You can also use `utils.run_api` helper, which 68 | # handles changing working directories, managing io streams, etc. 69 | # Also update `_run_tool_on_document` and `_run_tool` functions in `lsp_server.py`. 70 | result = utils.run_module( 71 | module=msg["module"], 72 | argv=msg["argv"], 73 | use_stdin=msg["useStdin"], 74 | cwd=msg["cwd"], 75 | source=msg["source"] if "source" in msg else None, 76 | ) 77 | except Exception: # pylint: disable=broad-except 78 | result = utils.RunResult("", traceback.format_exc(chain=True)) 79 | is_exception = True 80 | 81 | response = {"id": msg["id"]} 82 | if result.stderr: 83 | response["error"] = result.stderr 84 | response["exception"] = is_exception 85 | elif result.stdout: 86 | response["result"] = result.stdout 87 | 88 | RPC.send_data(response) 89 | -------------------------------------------------------------------------------- /bundled/tool/zenml_client.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) ZenML GmbH 2024. All Rights Reserved. 2 | # Licensed under the Apache License, Version 2.0 (the "License"); 3 | # you may not use this file except in compliance with the License. 4 | # You may obtain a copy of the License at: 5 | # 6 | # http://www.apache.org/licenses/LICENSE-2.0 7 | # 8 | # Unless required by applicable law or agreed to in writing, software 9 | # distributed under the License is distributed on an "AS IS" BASIS, 10 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either expressc 11 | # or implied. See the License for the specific language governing 12 | # permissions and limitations under the License. 13 | """ZenML client class. Initializes all wrappers.""" 14 | 15 | 16 | class ZenMLClient: 17 | """Provides a high-level interface to ZenML functionalities by wrapping core components.""" 18 | 19 | def __init__(self): 20 | """ 21 | Initializes the ZenMLClient with wrappers for managing configurations, 22 | server interactions, stacks, and pipeline runs. 23 | """ 24 | # pylint: disable=wrong-import-position,import-error 25 | from lazy_import import lazy_import, suppress_stdout_temporarily 26 | from zenml_wrappers import ( 27 | GlobalConfigWrapper, 28 | ModelsWrapper, 29 | PipelineRunsWrapper, 30 | ProjectsWrapper, 31 | StacksWrapper, 32 | WorkspacesWrapper, 33 | ZenServerWrapper, 34 | ) 35 | 36 | # Suppress colorful warnings during client initialization 37 | with suppress_stdout_temporarily(): 38 | self.client = lazy_import("zenml.client", "Client")() 39 | 40 | # initialize zenml library wrappers 41 | self.config_wrapper = GlobalConfigWrapper() 42 | self.stacks_wrapper = StacksWrapper(self.client) 43 | self.pipeline_runs_wrapper = PipelineRunsWrapper(self.client) 44 | self.workspaces_wrapper = WorkspacesWrapper(self.client, self.config_wrapper) 45 | self.projects_wrapper = ProjectsWrapper(self.client) 46 | self.models_wrapper = ModelsWrapper(self.client, self.projects_wrapper) 47 | self.zen_server_wrapper = ZenServerWrapper(self.config_wrapper, self.projects_wrapper) 48 | self.initialized = True 49 | -------------------------------------------------------------------------------- /eslint.config.mjs: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import eslint from '@eslint/js'; 15 | import prettierConfig from 'eslint-config-prettier'; 16 | import prettier from 'eslint-plugin-prettier'; 17 | import tseslint from 'typescript-eslint'; 18 | 19 | export default tseslint.config( 20 | eslint.configs.recommended, 21 | ...tseslint.configs.recommended, 22 | prettierConfig, 23 | { 24 | languageOptions: { 25 | parser: tseslint.parser, 26 | parserOptions: { 27 | ecmaVersion: 'latest', 28 | sourceType: 'module', 29 | }, 30 | }, 31 | files: ['**/*.ts'], 32 | plugins: { 33 | '@typescript-eslint': tseslint.plugin, 34 | prettier: prettier, 35 | }, 36 | rules: { 37 | // ESLint recommended rules 38 | semi: ['warn', 'always'], 39 | curly: 'warn', 40 | eqeqeq: 'warn', 41 | 'no-throw-literal': 'warn', 42 | // TypeScript specific rules 43 | '@typescript-eslint/no-unused-vars': 'warn', 44 | 'prefer-const': 'warn', 45 | '@typescript-eslint/naming-convention': [ 46 | 'warn', 47 | { 48 | selector: 'import', 49 | format: ['camelCase', 'PascalCase'], 50 | }, 51 | ], 52 | // Disable rules 53 | '@typescript-eslint/no-explicit-any': 'off', 54 | '@typescript-eslint/no-unsafe-function-type': 'warn', 55 | '@typescript-eslint/consistent-type-definitions': 'off', 56 | // Prettier rules 57 | 'prettier/prettier': 'warn', 58 | }, 59 | linterOptions: { 60 | reportUnusedDisableDirectives: true, 61 | }, 62 | ignores: ['out/**', 'dist/**', '**/*.d.ts'], 63 | } 64 | ); 65 | -------------------------------------------------------------------------------- /mypy.ini: -------------------------------------------------------------------------------- 1 | [mypy] 2 | python_version = 3.8 3 | ignore_missing_imports = True 4 | check_untyped_defs = True 5 | disallow_untyped_defs = False 6 | warn_unused_ignores = True 7 | 8 | [mypy-bundled.tool.*] 9 | ignore_errors = False 10 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.ruff] 2 | # Same as Black. 3 | line-length = 100 4 | indent-width = 4 5 | 6 | exclude = [ 7 | ".git", 8 | ".mypy_cache", 9 | ".ruff_cache", 10 | "dist", 11 | "node_modules", 12 | "venv", 13 | ".venv", 14 | ] 15 | 16 | [tool.ruff.format] 17 | quote-style = "double" 18 | indent-style = "space" 19 | skip-magic-trailing-comma = false 20 | line-ending = "auto" 21 | 22 | [tool.ruff.lint] 23 | # Enable flake8-bugbear rules 24 | select = ["E", "F", "B", "I"] 25 | ignore = [] 26 | fixable = ["ALL"] # Allow autofix for all enabled rules. 27 | unfixable = [] 28 | 29 | [tool.ruff.lint.isort] 30 | known-first-party = ["zenml"] 31 | 32 | [tool.ruff.lint.per-file-ignores] 33 | "__init__.py" = ["F401"] -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | mypy 2 | ruff 3 | yamlfix 4 | PyYAML 5 | types-PyYAML 6 | debugpy 7 | -------------------------------------------------------------------------------- /requirements.in: -------------------------------------------------------------------------------- 1 | # This file is used to generate requirements.txt. 2 | # NOTE: 3 | # Use Python 3.8 or greater which ever is the minimum version of the python 4 | # you plan on supporting when creating the environment or using pip-tools. 5 | # Only run the commands below to manually upgrade packages in requirements.txt: 6 | # 1) python -m pip install pip-tools 7 | # 2) pip-compile --generate-hashes --resolver=backtracking --upgrade ./requirements.in 8 | # If you are using nox commands to setup or build package you don't need to 9 | # run the above commands manually. 10 | 11 | # Required packages 12 | pygls 13 | packaging 14 | # Tool-specific packages for ZenML extension 15 | watchdog 16 | PyYAML 17 | types-PyYAML 18 | # Conda env fixes: Ensure hash generation for typing-extensions and exceptiongroup on Python < 3.11 19 | typing-extensions>=4.1.0,!=4.6.3; python_version < "3.11" 20 | exceptiongroup==1.2.0; python_version < '3.11' -------------------------------------------------------------------------------- /resources/components-form/components.css: -------------------------------------------------------------------------------- 1 | /* Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | Licensed under the Apache License, Version 2.0(the "License"); 3 | you may not use this file except in compliance with the License. 4 | You may obtain a copy of the License at: 5 | 6 | http://www.apache.org/licenses/LICENSE-2.0 7 | 8 | Unless required by applicable law or agreed to in writing, software 9 | distributed under the License is distributed on an "AS IS" BASIS, 10 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | or implied. See the License for the specific language governing 12 | permissions and limitations under the License. */ 13 | .container { 14 | max-width: 600px; 15 | margin: auto; 16 | } 17 | 18 | .block { 19 | padding: 10px 10px 5px 10px; 20 | border: 2px var(--vscode-editor-foreground) solid; 21 | margin-bottom: 10px; 22 | } 23 | 24 | .logo { 25 | float: right; 26 | max-width: 100px; 27 | } 28 | 29 | .docs { 30 | clear: both; 31 | display: flex; 32 | justify-content: space-around; 33 | padding: 5px 0px; 34 | } 35 | 36 | .button, 37 | button { 38 | padding: 2px 5px; 39 | background-color: var(--vscode-editor-background); 40 | color: var(--vscode-editor-foreground); 41 | border: 2px var(--vscode-editor-foreground) solid; 42 | border-radius: 10px; 43 | } 44 | 45 | .field { 46 | margin-bottom: 10px; 47 | } 48 | 49 | .value { 50 | width: 100%; 51 | box-sizing: border-box; 52 | padding-left: 20px; 53 | } 54 | 55 | .input { 56 | width: 100%; 57 | } 58 | 59 | .center { 60 | display: flex; 61 | align-items: center; 62 | justify-content: center; 63 | } 64 | 65 | .loader { 66 | width: 20px; 67 | height: 20px; 68 | border: 5px solid #fff; 69 | border-bottom-color: #ff3d00; 70 | border-radius: 50%; 71 | display: inline-block; 72 | box-sizing: border-box; 73 | animation: rotation 1s linear infinite; 74 | } 75 | 76 | @keyframes rotation { 77 | 0% { 78 | transform: rotate(0deg); 79 | } 80 | 100% { 81 | transform: rotate(360deg); 82 | } 83 | } 84 | 85 | .hidden { 86 | display: none; 87 | } 88 | -------------------------------------------------------------------------------- /resources/dag-view/icons/alert.svg: -------------------------------------------------------------------------------- 1 | 2 | 5 | -------------------------------------------------------------------------------- /resources/dag-view/icons/cached.svg: -------------------------------------------------------------------------------- 1 | 2 | 4 | -------------------------------------------------------------------------------- /resources/dag-view/icons/check.svg: -------------------------------------------------------------------------------- 1 | 2 | 4 | 5 | -------------------------------------------------------------------------------- /resources/dag-view/icons/database.svg: -------------------------------------------------------------------------------- 1 | 2 | 4 | 5 | -------------------------------------------------------------------------------- /resources/dag-view/icons/initializing.svg: -------------------------------------------------------------------------------- 1 | 2 | 4 | -------------------------------------------------------------------------------- /resources/dag-view/icons/play.svg: -------------------------------------------------------------------------------- 1 | 2 | 4 | 5 | -------------------------------------------------------------------------------- /resources/extension-logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zenml-io/vscode-zenml/a969342b9ec0863d459d77342fb8daa96f0dced0/resources/extension-logo.png -------------------------------------------------------------------------------- /resources/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zenml-io/vscode-zenml/a969342b9ec0863d459d77342fb8daa96f0dced0/resources/logo.png -------------------------------------------------------------------------------- /resources/python.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zenml-io/vscode-zenml/a969342b9ec0863d459d77342fb8daa96f0dced0/resources/python.png -------------------------------------------------------------------------------- /resources/stacks-form/stacks.css: -------------------------------------------------------------------------------- 1 | /* Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | Licensed under the Apache License, Version 2.0(the "License"); 3 | you may not use this file except in compliance with the License. 4 | You may obtain a copy of the License at: 5 | 6 | http://www.apache.org/licenses/LICENSE-2.0 7 | 8 | Unless required by applicable law or agreed to in writing, software 9 | distributed under the License is distributed on an "AS IS" BASIS, 10 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | or implied. See the License for the specific language governing 12 | permissions and limitations under the License. */ 13 | h2 { 14 | text-align: center; 15 | } 16 | 17 | input { 18 | background-color: var(--vscode-editor-background); 19 | color: var(--vscode-editor-foreground); 20 | } 21 | 22 | input[type='radio'] { 23 | appearance: none; 24 | width: 15px; 25 | height: 15px; 26 | border-radius: 50%; 27 | background-clip: content-box; 28 | border: 2px solid var(--vscode-editor-foreground); 29 | background-color: var(--vscode-editor-background); 30 | } 31 | 32 | input[type='radio']:checked { 33 | background-color: var(--vscode-editor-foreground); 34 | padding: 2px; 35 | } 36 | 37 | p { 38 | margin: 0; 39 | padding: 0; 40 | } 41 | 42 | .options { 43 | display: flex; 44 | flex-wrap: nowrap; 45 | overflow-x: auto; 46 | align-items: center; 47 | width: 100%; 48 | border: 2px var(--vscode-editor-foreground) solid; 49 | border-radius: 5px; 50 | scrollbar-color: var(--vscode-editor-foreground) var(--vscode-editor-background); 51 | } 52 | 53 | .single-option { 54 | display: flex; 55 | flex-direction: row; 56 | align-items: start; 57 | justify-content: center; 58 | padding: 5px; 59 | margin: 10px; 60 | flex-shrink: 0; 61 | } 62 | 63 | .single-option input { 64 | margin-right: 5px; 65 | } 66 | 67 | .single-option label { 68 | display: flex; 69 | flex-direction: column; 70 | justify-content: center; 71 | align-items: center; 72 | background-color: #eee; 73 | color: #111; 74 | text-align: center; 75 | padding: 5px; 76 | border: 2px var(--vscode-editor-foreground) solid; 77 | border-radius: 5px; 78 | } 79 | 80 | .single-option img { 81 | width: 50px; 82 | height: 50px; 83 | flex-shrink: 0; 84 | } 85 | 86 | .center { 87 | display: flex; 88 | justify-content: center; 89 | align-items: center; 90 | margin: 10px; 91 | } 92 | 93 | .loader { 94 | width: 20px; 95 | height: 20px; 96 | border: 5px solid #fff; 97 | border-bottom-color: #ff3d00; 98 | border-radius: 50%; 99 | display: inline-block; 100 | box-sizing: border-box; 101 | animation: rotation 1s linear infinite; 102 | } 103 | 104 | @keyframes rotation { 105 | 0% { 106 | transform: rotate(0deg); 107 | } 108 | 100% { 109 | transform: rotate(360deg); 110 | } 111 | } 112 | 113 | .hidden { 114 | display: none; 115 | } 116 | -------------------------------------------------------------------------------- /resources/stacks-form/stacks.js: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | document.querySelector('input[name="orchestrator"]').toggleAttribute('required'); 14 | document.querySelector('input[name="artifact_store"]').toggleAttribute('required'); 15 | 16 | const form = document.querySelector('form'); 17 | const submit = document.querySelector('input[type="submit"]'); 18 | const spinner = document.querySelector('.loader'); 19 | let previousValues = {}; 20 | let id = undefined; 21 | let mode = 'register'; 22 | 23 | form.addEventListener('click', evt => { 24 | const target = evt.target; 25 | let input = null; 26 | 27 | if (target instanceof HTMLLabelElement) { 28 | input = document.getElementById(target.htmlFor); 29 | } else if (target instanceof HTMLInputElement && target.type === 'radio') { 30 | input = target; 31 | } 32 | 33 | if (!input) { 34 | return; 35 | } 36 | 37 | const value = input.value; 38 | const name = input.name; 39 | if (previousValues[name] === value) { 40 | delete previousValues[name]; 41 | input.checked = false; 42 | } else { 43 | previousValues[name] = value; 44 | } 45 | }); 46 | 47 | (() => { 48 | const vscode = acquireVsCodeApi(); 49 | 50 | form.addEventListener('submit', evt => { 51 | evt.preventDefault(); 52 | submit.disabled = true; 53 | spinner.classList.remove('hidden'); 54 | const data = Object.fromEntries(new FormData(evt.target)); 55 | 56 | if (id) { 57 | data.id = id; 58 | } 59 | 60 | vscode.postMessage({ 61 | command: mode, 62 | data, 63 | }); 64 | }); 65 | })(); 66 | 67 | const title = document.querySelector('h2'); 68 | const nameInput = document.querySelector('input[name="name"]'); 69 | 70 | window.addEventListener('message', evt => { 71 | const message = evt.data; 72 | 73 | switch (message.command) { 74 | case 'register': 75 | mode = 'register'; 76 | title.innerText = 'Register Stack'; 77 | id = undefined; 78 | previousValues = {}; 79 | form.reset(); 80 | break; 81 | 82 | case 'update': 83 | mode = 'update'; 84 | title.innerText = 'Update Stack'; 85 | id = message.data.id; 86 | nameInput.value = message.data.name; 87 | previousValues = message.data.components; 88 | Object.entries(message.data.components).forEach(([type, id]) => { 89 | const input = document.querySelector(`[name="${type}"][value="${id}"]`); 90 | input.checked = true; 91 | }); 92 | break; 93 | 94 | case 'fail': 95 | spinner.classList.add('hidden'); 96 | submit.disabled = false; 97 | break; 98 | } 99 | }); 100 | -------------------------------------------------------------------------------- /resources/zenml-extension-dag.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zenml-io/vscode-zenml/a969342b9ec0863d459d77342fb8daa96f0dced0/resources/zenml-extension-dag.gif -------------------------------------------------------------------------------- /resources/zenml-extension-new.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zenml-io/vscode-zenml/a969342b9ec0863d459d77342fb8daa96f0dced0/resources/zenml-extension-new.gif -------------------------------------------------------------------------------- /resources/zenml-extension.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zenml-io/vscode-zenml/a969342b9ec0863d459d77342fb8daa96f0dced0/resources/zenml-extension.gif -------------------------------------------------------------------------------- /resources/zenml_logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/zenml-io/vscode-zenml/a969342b9ec0863d459d77342fb8daa96f0dced0/resources/zenml_logo.png -------------------------------------------------------------------------------- /runtime.txt: -------------------------------------------------------------------------------- 1 | python-3.8.18 -------------------------------------------------------------------------------- /scripts/clear_and_compile.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | echo "Removing bundled/tool/__pycache__..." 4 | if [ -d "bundled/tool/__pycache__" ]; then 5 | rm -rf bundled/tool/__pycache__ 6 | fi 7 | 8 | echo "Removing dist directory..." 9 | if [ -d "dist" ]; then 10 | rm -rf dist 11 | fi 12 | 13 | echo "Recompiling with npm..." 14 | if ! command -v npm &> /dev/null 15 | then 16 | echo "npm could not be found. Please install npm to proceed." 17 | exit 1 18 | fi 19 | npm run compile 20 | 21 | echo "Operation completed." 22 | -------------------------------------------------------------------------------- /scripts/format.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -euxo pipefail 3 | 4 | # source directory for python tool 5 | python_src="bundled/tool" 6 | 7 | echo "Formatting Python files in bundled/tool..." 8 | # autoflake replacement: removes unused imports and variables 9 | ruff check $python_src --select F401,F841 --fix --exclude "__init__.py" --isolated 10 | 11 | # isort replacement: sorts imports 12 | ruff check $python_src --select I --fix --ignore D 13 | 14 | # black replacement: formats code 15 | ruff format $python_src 16 | 17 | echo "Formatting TypeScript files..." 18 | npm run format 19 | 20 | echo "Formatting YAML files..." 21 | find .github -name "*.yml" -print0 | xargs -0 yamlfix -- 22 | 23 | echo "Formatting complete." 24 | -------------------------------------------------------------------------------- /scripts/lint.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -euxo pipefail 3 | 4 | cd "$(dirname "$0")/.." || exit 5 | printf "Working directory: %s\n" "$(pwd)" 6 | # set PYTHONPATH to include bundled/tool 7 | PYTHONPATH="${PYTHONPATH-}:$(pwd)/bundled/tool" 8 | export PYTHONPATH 9 | 10 | PYTHON_SRC="bundled/tool" 11 | 12 | echo "Linting Python files..." 13 | ruff check $PYTHON_SRC 14 | 15 | echo "Checking for unused imports and variables..." 16 | ruff check $PYTHON_SRC --select F401,F841 --exclude "__init__.py" --isolated 17 | 18 | echo "Checking Python formatting..." 19 | ruff format $PYTHON_SRC --check 20 | 21 | echo "Type checking python files with mypy..." 22 | mypy $PYTHON_SRC 23 | 24 | echo "Linting TypeScript files with eslint..." 25 | npm run lint 26 | 27 | echo "Checking yaml files with yamlfix..." 28 | yamlfix .github/workflows/*.yml --check 29 | 30 | unset PYTHONPATH 31 | -------------------------------------------------------------------------------- /scripts/spellcheck.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | set -euo pipefail 3 | 4 | REPO_ROOT="$(git rev-parse --show-toplevel)" 5 | DEFAULT_CONFIG_PATH="${REPO_ROOT}/.typos.toml" 6 | 7 | if [[ "${1:-}" == "-h" ]] || [[ "${1:-}" == "--help" ]]; then 8 | echo "This is a wrapper around typos that enforces repository config." 9 | echo "For available options, see: typos --help" 10 | exit 0 11 | fi 12 | 13 | # Error out if the default config is not found 14 | if [[ ! -f "$DEFAULT_CONFIG_PATH" ]]; then 15 | echo "ERROR: Repository config not found at $DEFAULT_CONFIG_PATH" 16 | exit 1 17 | fi 18 | 19 | # Remove any -c/--config arguments to prevent override attempts 20 | args=() 21 | skip_next=0 22 | for arg in "$@"; do 23 | if ((skip_next)); then 24 | skip_next=0 25 | continue 26 | fi 27 | if [[ "$arg" == "--config" ]] || [[ "$arg" == "-c" ]]; then 28 | echo "WARNING: Custom config path ignored, using repository config" 29 | skip_next=1 30 | continue 31 | fi 32 | if [[ "$arg" == --config=* ]] || [[ "$arg" == -c=* ]]; then 33 | echo "WARNING: Custom config path ignored, using repository config" 34 | continue 35 | fi 36 | args+=("$arg") 37 | done 38 | 39 | exec typos --config "$DEFAULT_CONFIG_PATH" "${args[@]+"${args[@]}"}" 40 | -------------------------------------------------------------------------------- /src/commands/components/registry.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { ExtensionContext, commands } from 'vscode'; 14 | import { registerCommand } from '../../common/vscodeapi'; 15 | import { ZenExtension } from '../../services/ZenExtension'; 16 | import { ComponentDataProvider } from '../../views/activityBar/componentView/ComponentDataProvider'; 17 | import { StackComponentTreeItem } from '../../views/activityBar/componentView/ComponentTreeItems'; 18 | import { componentCommands } from './cmds'; 19 | 20 | /** 21 | * Registers stack component-related commands for the extension. 22 | * 23 | * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. 24 | */ 25 | export const registerComponentCommands = (context: ExtensionContext) => { 26 | const componentDataProvider = ComponentDataProvider.getInstance(); 27 | try { 28 | const registeredCommands = [ 29 | registerCommand('zenml.setComponentItemsPerPage', async () => 30 | componentDataProvider.updateItemsPerPage() 31 | ), 32 | registerCommand('zenml.refreshComponentView', async () => 33 | componentCommands.refreshComponentView() 34 | ), 35 | registerCommand('zenml.registerComponent', async () => componentCommands.registerComponent()), 36 | registerCommand('zenml.updateComponent', async (node: StackComponentTreeItem) => 37 | componentCommands.updateComponent(node) 38 | ), 39 | registerCommand('zenml.deleteComponent', async (node: StackComponentTreeItem) => 40 | componentCommands.deleteComponent(node) 41 | ), 42 | registerCommand('zenml.nextComponentPage', async () => componentDataProvider.goToNextPage()), 43 | registerCommand('zenml.previousComponentPage', async () => 44 | componentDataProvider.goToPreviousPage() 45 | ), 46 | ]; 47 | 48 | registeredCommands.forEach(cmd => { 49 | context.subscriptions.push(cmd); 50 | ZenExtension.commandDisposables.push(cmd); 51 | }); 52 | 53 | commands.executeCommand('setContext', 'componentCommandsRegistered', true); 54 | } catch (e) { 55 | console.error('Error registering component commands:', e); 56 | commands.executeCommand('setContext', 'componentCommandsRegistered', false); 57 | } 58 | }; 59 | -------------------------------------------------------------------------------- /src/commands/environment/cmds.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import { ProgressLocation, commands, window } from 'vscode'; 15 | import { getInterpreterFromWorkspaceSettings } from '../../common/settings'; 16 | import { EventBus } from '../../services/EventBus'; 17 | import { LSClient } from '../../services/LSClient'; 18 | import { 19 | LSP_ZENML_CLIENT_INITIALIZED, 20 | PYTOOL_MODULE, 21 | REFRESH_ENVIRONMENT_VIEW, 22 | } from '../../utils/constants'; 23 | import { EnvironmentDataProvider } from '../../views/activityBar/environmentView/EnvironmentDataProvider'; 24 | 25 | /** 26 | * Set the Python interpreter for the current workspace. 27 | * 28 | * @returns {Promise} Resolves after refreshing the view. 29 | */ 30 | const setPythonInterpreter = async (): Promise => { 31 | await window.withProgress( 32 | { 33 | location: ProgressLocation.Notification, 34 | }, 35 | async progress => { 36 | progress.report({ increment: 10 }); 37 | const currentInterpreter = await getInterpreterFromWorkspaceSettings(); 38 | 39 | await commands.executeCommand('python.setInterpreter'); 40 | 41 | const newInterpreter = await getInterpreterFromWorkspaceSettings(); 42 | 43 | if (newInterpreter === currentInterpreter) { 44 | console.log('Interpreter selection unchanged or cancelled. No server restart required.'); 45 | return; 46 | } 47 | progress.report({ increment: 90 }); 48 | console.log('ZenML server restarting to apply the new interpreter settings.'); 49 | // The onDidChangePythonInterpreter event will trigger the server restart. 50 | progress.report({ increment: 100 }); 51 | } 52 | ); 53 | }; 54 | 55 | const refreshEnvironmentView = async (): Promise => { 56 | window.withProgress( 57 | { 58 | location: ProgressLocation.Notification, 59 | cancellable: false, 60 | }, 61 | async () => { 62 | EnvironmentDataProvider.getInstance().refresh(); 63 | } 64 | ); 65 | }; 66 | 67 | const restartLSPServer = async (): Promise => { 68 | await window.withProgress( 69 | { 70 | location: ProgressLocation.Window, 71 | title: 'Restarting LSP Server...', 72 | }, 73 | async progress => { 74 | progress.report({ increment: 10 }); 75 | const lsClient = LSClient.getInstance(); 76 | lsClient.isZenMLReady = false; 77 | lsClient.localZenML = { is_installed: false, version: '' }; 78 | const eventBus = EventBus.getInstance(); 79 | eventBus.emit(REFRESH_ENVIRONMENT_VIEW); 80 | eventBus.emit(LSP_ZENML_CLIENT_INITIALIZED, false); 81 | await commands.executeCommand(`${PYTOOL_MODULE}.restart`); 82 | progress.report({ increment: 100 }); 83 | } 84 | ); 85 | }; 86 | 87 | export const environmentCommands = { 88 | setPythonInterpreter, 89 | refreshEnvironmentView, 90 | restartLSPServer, 91 | }; 92 | -------------------------------------------------------------------------------- /src/commands/environment/registry.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { ExtensionContext, commands } from 'vscode'; 14 | import { registerCommand } from '../../common/vscodeapi'; 15 | import { ZenExtension } from '../../services/ZenExtension'; 16 | import { environmentCommands } from './cmds'; 17 | 18 | /** 19 | * Registers pipeline-related commands for the extension. 20 | * 21 | * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. 22 | */ 23 | export const registerEnvironmentCommands = (context: ExtensionContext) => { 24 | try { 25 | const registeredCommands = [ 26 | registerCommand('zenml.setPythonInterpreter', async () => 27 | environmentCommands.setPythonInterpreter() 28 | ), 29 | registerCommand('zenml.refreshEnvironmentView', async () => 30 | environmentCommands.refreshEnvironmentView() 31 | ), 32 | registerCommand('zenml.restartLspServer', async () => environmentCommands.restartLSPServer()), 33 | ]; 34 | 35 | registeredCommands.forEach(cmd => { 36 | context.subscriptions.push(cmd); 37 | ZenExtension.commandDisposables.push(cmd); 38 | }); 39 | 40 | commands.executeCommand('setContext', 'environmentCommandsRegistered', true); 41 | } catch (error) { 42 | console.error('Error registering environment commands:', error); 43 | commands.executeCommand('setContext', 'environmentCommandsRegistered', false); 44 | } 45 | }; 46 | -------------------------------------------------------------------------------- /src/commands/models/cmds.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import { ModelDataProvider } from '../../views/activityBar/modelView/ModelDataProvider'; 15 | 16 | /** 17 | * Refreshes the Models view. 18 | */ 19 | async function refreshModelView(): Promise { 20 | const modelProvider = ModelDataProvider.getInstance(); 21 | void modelProvider.refresh(); 22 | } 23 | 24 | export const modelCommands = { 25 | refreshModelView, 26 | }; 27 | -------------------------------------------------------------------------------- /src/commands/models/registry.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { ExtensionContext, commands } from 'vscode'; 14 | import { registerCommand } from '../../common/vscodeapi'; 15 | import { ZenExtension } from '../../services/ZenExtension'; 16 | import { ModelDataProvider } from '../../views/activityBar/modelView/ModelDataProvider'; 17 | import { modelCommands } from './cmds'; 18 | 19 | /** 20 | * Registers model commands with VS Code. 21 | * 22 | * @param {ExtensionContext} context - The extension context. 23 | */ 24 | export function registerModelCommands(context: ExtensionContext): void { 25 | const modelDataProvider = ModelDataProvider.getInstance(); 26 | 27 | try { 28 | const registeredCommands = [ 29 | registerCommand('zenml.refreshModelView', async () => modelCommands.refreshModelView()), 30 | registerCommand('zenml.setModelsPerPage', async () => modelDataProvider.updateItemsPerPage()), 31 | registerCommand('zenml.nextModelPage', async () => modelDataProvider.goToNextPage()), 32 | registerCommand('zenml.previousModelPage', async () => modelDataProvider.goToPreviousPage()), 33 | ]; 34 | 35 | registeredCommands.forEach(cmd => { 36 | context.subscriptions.push(cmd); 37 | ZenExtension.commandDisposables.push(cmd); 38 | }); 39 | 40 | commands.executeCommand('setContext', 'modelCommandsRegistered', true); 41 | } catch (error) { 42 | console.error('Error registering model commands:', error); 43 | commands.executeCommand('setContext', 'modelCommandsRegistered', false); 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /src/commands/pipelines/cmds.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as vscode from 'vscode'; 14 | import DagRenderer from '../../dag/renderer/DagRenderer'; 15 | import { LSClient } from '../../services/LSClient'; 16 | import { PipelineTreeItem } from '../../views/activityBar'; 17 | import { createCommandErrorItem } from '../../views/activityBar/common/ErrorTreeItem'; 18 | import { PipelineDataProvider } from '../../views/activityBar/pipelineView/PipelineDataProvider'; 19 | import { getPipelineRunDashboardUrl } from './utils'; 20 | 21 | /** 22 | * Triggers a refresh of the pipeline view within the UI components. 23 | * 24 | * @returns {Promise} Resolves after refreshing the view. 25 | */ 26 | const refreshPipelineView = async (): Promise => { 27 | await vscode.window.withProgress( 28 | { 29 | location: vscode.ProgressLocation.Window, 30 | }, 31 | async () => { 32 | await PipelineDataProvider.getInstance().refresh(); 33 | } 34 | ); 35 | }; 36 | 37 | /** 38 | * Deletes a pipeline run. 39 | * 40 | * @param {PipelineTreeItem} node The pipeline run to delete. 41 | * @returns {Promise} Resolves after deleting the pipeline run. 42 | */ 43 | const deletePipelineRun = async (node: PipelineTreeItem): Promise => { 44 | const userConfirmation = await vscode.window.showWarningMessage( 45 | 'Are you sure you want to delete this pipeline run?', 46 | { modal: true }, 47 | 'Yes', 48 | 'No' 49 | ); 50 | 51 | if (userConfirmation === 'Yes') { 52 | await vscode.window.withProgress( 53 | { 54 | location: vscode.ProgressLocation.Window, 55 | title: 'Deleting pipeline run...', 56 | }, 57 | async () => { 58 | const runId = node.id; 59 | try { 60 | const lsClient = LSClient.getInstance(); 61 | const result = await lsClient.sendLsClientRequest('deletePipelineRun', [runId]); 62 | if (result && 'error' in result) { 63 | throw new Error(result.error); 64 | } 65 | vscode.window.showInformationMessage('Pipeline run deleted successfully'); 66 | await refreshPipelineView(); 67 | } catch (error: any) { 68 | console.error(`Error deleting pipeline run: ${error}`); 69 | vscode.window.showErrorMessage(`Failed to delete pipeline run: ${error.message}`); 70 | } 71 | } 72 | ); 73 | } 74 | }; 75 | 76 | /** 77 | * Opens the selected pipieline run in the ZenML Dashboard in the browser 78 | * 79 | * @param {PipelineTreeItem} node The pipeline run to open. 80 | */ 81 | const goToPipelineUrl = (node: PipelineTreeItem): void => { 82 | const url = getPipelineRunDashboardUrl(node.id); 83 | 84 | if (url) { 85 | try { 86 | const parsedUrl = vscode.Uri.parse(url); 87 | 88 | vscode.env.openExternal(parsedUrl); 89 | } catch (error) { 90 | console.log(error); 91 | const pipelineProvider = PipelineDataProvider.getInstance(); 92 | pipelineProvider.showCommandError( 93 | createCommandErrorItem('open pipeline URL', `Failed to open pipeline run URL: ${error}`) 94 | ); 95 | } 96 | } 97 | }; 98 | 99 | /** 100 | * Opens the selected pipeline run in a DAG visualizer Webview Panel 101 | * 102 | * @param {PipelineTreeItem} node The pipeline run to render. 103 | */ 104 | const renderDag = (node: PipelineTreeItem): void => { 105 | DagRenderer.getInstance()?.createView(node); 106 | }; 107 | 108 | export const pipelineCommands = { 109 | refreshPipelineView, 110 | deletePipelineRun, 111 | goToPipelineUrl, 112 | renderDag, 113 | }; 114 | -------------------------------------------------------------------------------- /src/commands/pipelines/registry.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { ExtensionContext, commands } from 'vscode'; 14 | import { registerCommand } from '../../common/vscodeapi'; 15 | import { ZenExtension } from '../../services/ZenExtension'; 16 | import { PipelineDataProvider, PipelineTreeItem } from '../../views/activityBar'; 17 | import { pipelineCommands } from './cmds'; 18 | 19 | /** 20 | * Registers pipeline-related commands for the extension. 21 | * 22 | * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. 23 | */ 24 | export const registerPipelineCommands = (context: ExtensionContext) => { 25 | const pipelineDataProvider = PipelineDataProvider.getInstance(); 26 | 27 | try { 28 | const registeredCommands = [ 29 | registerCommand('zenml.refreshPipelineView', async () => 30 | pipelineCommands.refreshPipelineView() 31 | ), 32 | registerCommand('zenml.deletePipelineRun', async (node: PipelineTreeItem) => 33 | pipelineCommands.deletePipelineRun(node) 34 | ), 35 | registerCommand('zenml.goToPipelineUrl', async (node: PipelineTreeItem) => 36 | pipelineCommands.goToPipelineUrl(node) 37 | ), 38 | registerCommand('zenml.renderDag', async (node: PipelineTreeItem) => 39 | pipelineCommands.renderDag(node) 40 | ), 41 | registerCommand('zenml.nextPipelineRunsPage', async () => 42 | pipelineDataProvider.goToNextPage() 43 | ), 44 | registerCommand('zenml.previousPipelineRunsPage', async () => 45 | pipelineDataProvider.goToPreviousPage() 46 | ), 47 | registerCommand('zenml.setPipelineRunsPerPage', async () => 48 | pipelineDataProvider.updateItemsPerPage() 49 | ), 50 | ]; 51 | 52 | registeredCommands.forEach(cmd => { 53 | context.subscriptions.push(cmd); 54 | ZenExtension.commandDisposables.push(cmd); 55 | }); 56 | 57 | commands.executeCommand('setContext', 'pipelineCommandsRegistered', true); 58 | } catch (error) { 59 | console.error('Error registering pipeline commands:', error); 60 | commands.executeCommand('setContext', 'pipelineCommandsRegistered', false); 61 | } 62 | }; 63 | -------------------------------------------------------------------------------- /src/commands/pipelines/utils.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import { ServerDataProvider } from '../../views/activityBar'; 15 | import { buildWorkspaceProjectUrl, getBaseUrl, isServerStatus } from '../server/utils'; 16 | 17 | /** 18 | * Gets the Dashboard URL for the corresponding ZenML pipeline run 19 | * 20 | * @param {string} id - The id of the ZenML pipeline run to be opened 21 | * @returns {string} - The URL corresponding to the pipeline run in the ZenML Dashboard 22 | */ 23 | export const getPipelineRunDashboardUrl = (id: string): string => { 24 | if (!id) { 25 | return ''; 26 | } 27 | 28 | const status = ServerDataProvider.getInstance().getCurrentStatus(); 29 | if (!status || !isServerStatus(status) || status.deployment_type === 'other') { 30 | return ''; 31 | } 32 | 33 | const baseUrl = getBaseUrl(status.dashboard_url); 34 | const suffix = `/runs/${id}?tab=overview`; 35 | 36 | const url = buildWorkspaceProjectUrl(baseUrl, status, suffix); 37 | 38 | return url; 39 | }; 40 | 41 | const pipelineUtils = { 42 | getPipelineRunDashboardUrl, 43 | }; 44 | 45 | export default pipelineUtils; 46 | -------------------------------------------------------------------------------- /src/commands/projects/cmds.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as vscode from 'vscode'; 14 | import { EventBus } from '../../services/EventBus'; 15 | import { LSP_ZENML_PROJECT_CHANGED } from '../../utils/constants'; 16 | import { ProjectDataProvider } from '../../views/activityBar/projectView/ProjectDataProvider'; 17 | import { ProjectTreeItem } from '../../views/activityBar/projectView/ProjectTreeItems'; 18 | import { getProjectDashboardUrl, switchActiveProject } from './utils'; 19 | 20 | /** 21 | * Refreshes the project view. 22 | */ 23 | const refreshProjectView = async () => { 24 | vscode.window.withProgress( 25 | { 26 | location: vscode.ProgressLocation.Notification, 27 | cancellable: false, 28 | }, 29 | async () => { 30 | await ProjectDataProvider.getInstance().refresh(); 31 | } 32 | ); 33 | }; 34 | 35 | /** 36 | * Sets the selected project as the active project and stores it in the global context. 37 | * 38 | * @param {ProjectTreeItem} node The project to activate. 39 | * @returns {Promise} Resolves after setting the active project. 40 | */ 41 | const setActiveProject = async (node: ProjectTreeItem): Promise => { 42 | vscode.window.withProgress( 43 | { 44 | location: vscode.ProgressLocation.Notification, 45 | title: 'Setting Active Project...', 46 | cancellable: false, 47 | }, 48 | async () => { 49 | try { 50 | const result = await switchActiveProject(node.name); 51 | if (result) { 52 | EventBus.getInstance().emit(LSP_ZENML_PROJECT_CHANGED, node.name); 53 | vscode.window.showInformationMessage(`Active project set to: ${node.name}`); 54 | } 55 | } catch (error) { 56 | console.log(error); 57 | const errorMessage = error instanceof Error ? error.message : String(error); 58 | vscode.window.showErrorMessage(`Failed to set active project: ${errorMessage}`); 59 | } 60 | } 61 | ); 62 | }; 63 | 64 | /** 65 | * Opens the selected project in the ZenML Dashboard in the browser 66 | * 67 | * @param {ProjectTreeItem} node The project to open. 68 | */ 69 | const goToProjectUrl = (node: ProjectTreeItem) => { 70 | const url = getProjectDashboardUrl(node.project.name); 71 | 72 | if (url) { 73 | try { 74 | const parsedUrl = vscode.Uri.parse(url); 75 | 76 | vscode.env.openExternal(parsedUrl); 77 | } catch (error) { 78 | console.log(error); 79 | vscode.window.showErrorMessage(`Failed to open project URL: ${error}`); 80 | } 81 | } else { 82 | vscode.window.showErrorMessage(`Could not determine URL for project: ${node.project.name}`); 83 | } 84 | }; 85 | 86 | export const projectCommands = { 87 | refreshProjectView, 88 | setActiveProject, 89 | goToProjectUrl, 90 | }; 91 | -------------------------------------------------------------------------------- /src/commands/projects/registry.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { ExtensionContext, commands } from 'vscode'; 14 | import { registerCommand } from '../../common/vscodeapi'; 15 | import { ZenExtension } from '../../services/ZenExtension'; 16 | import { ProjectDataProvider } from '../../views/activityBar/projectView/ProjectDataProvider'; 17 | import { ProjectTreeItem } from '../../views/activityBar/projectView/ProjectTreeItems'; 18 | import { projectCommands } from './cmds'; 19 | 20 | /** 21 | * Registers project-related commands for the extension. 22 | * 23 | * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. 24 | */ 25 | export const registerProjectCommands = (context: ExtensionContext) => { 26 | const projectDataProvider = ProjectDataProvider.getInstance(); 27 | try { 28 | const registeredCommands = [ 29 | registerCommand('zenml.setProjectItemsPerPage', async () => 30 | projectDataProvider.updateItemsPerPage() 31 | ), 32 | registerCommand('zenml.refreshProjectView', async () => projectCommands.refreshProjectView()), 33 | registerCommand('zenml.setActiveProject', async (node: ProjectTreeItem) => 34 | projectCommands.setActiveProject(node) 35 | ), 36 | registerCommand('zenml.goToProjectUrl', (node: ProjectTreeItem) => 37 | projectCommands.goToProjectUrl(node) 38 | ), 39 | registerCommand('zenml.nextProjectPage', async () => projectDataProvider.goToNextPage()), 40 | registerCommand('zenml.previousProjectPage', async () => 41 | projectDataProvider.goToPreviousPage() 42 | ), 43 | ]; 44 | 45 | registeredCommands.forEach(cmd => { 46 | context.subscriptions.push(cmd); 47 | ZenExtension.commandDisposables.push(cmd); 48 | }); 49 | 50 | commands.executeCommand('setContext', 'projectCommandsRegistered', true); 51 | } catch (error) { 52 | console.error('Error registering project commands:', error); 53 | commands.executeCommand('setContext', 'projectCommandsRegistered', false); 54 | } 55 | }; 56 | -------------------------------------------------------------------------------- /src/commands/server/registry.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { ExtensionContext, commands } from 'vscode'; 14 | import { registerCommand } from '../../common/vscodeapi'; 15 | import { ZenExtension } from '../../services/ZenExtension'; 16 | import { serverCommands } from './cmds'; 17 | 18 | /** 19 | * Registers server-related commands for the extension. 20 | * 21 | * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. 22 | */ 23 | export const registerServerCommands = (context: ExtensionContext) => { 24 | try { 25 | const registeredCommands = [ 26 | registerCommand('zenml.connectServer', async () => serverCommands.connectServer()), 27 | registerCommand('zenml.disconnectServer', async () => serverCommands.disconnectServer()), 28 | registerCommand('zenml.refreshServerStatus', async () => 29 | serverCommands.refreshServerStatus() 30 | ), 31 | ]; 32 | 33 | registeredCommands.forEach(cmd => { 34 | context.subscriptions.push(cmd); 35 | ZenExtension.commandDisposables.push(cmd); 36 | }); 37 | 38 | commands.executeCommand('setContext', 'serverCommandsRegistered', true); 39 | } catch (error) { 40 | console.error('Error registering server commands:', error); 41 | commands.executeCommand('setContext', 'serverCommandsRegistered', false); 42 | } 43 | }; 44 | -------------------------------------------------------------------------------- /src/commands/stack/registry.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { ExtensionContext, commands } from 'vscode'; 14 | import { registerCommand } from '../../common/vscodeapi'; 15 | import { ZenExtension } from '../../services/ZenExtension'; 16 | import { StackDataProvider, StackTreeItem } from '../../views/activityBar'; 17 | import { stackCommands } from './cmds'; 18 | 19 | /** 20 | * Registers stack-related commands for the extension. 21 | * 22 | * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. 23 | */ 24 | export const registerStackCommands = (context: ExtensionContext) => { 25 | const stackDataProvider = StackDataProvider.getInstance(); 26 | try { 27 | const registeredCommands = [ 28 | registerCommand('zenml.setStackItemsPerPage', async () => 29 | stackDataProvider.updateItemsPerPage() 30 | ), 31 | registerCommand('zenml.refreshStackView', async () => stackCommands.refreshStackView()), 32 | registerCommand('zenml.registerStack', async () => stackCommands.registerStack()), 33 | registerCommand('zenml.updateStack', async (node: StackTreeItem) => 34 | stackCommands.updateStack(node) 35 | ), 36 | registerCommand('zenml.deleteStack', async (node: StackTreeItem) => 37 | stackCommands.deleteStack(node) 38 | ), 39 | registerCommand('zenml.renameStack', async (node: StackTreeItem) => 40 | stackCommands.renameStack(node) 41 | ), 42 | registerCommand('zenml.setActiveStack', async (node: StackTreeItem) => 43 | stackCommands.setActiveStack(node) 44 | ), 45 | registerCommand('zenml.goToStackUrl', (node: StackTreeItem) => 46 | stackCommands.goToStackUrl(node) 47 | ), 48 | registerCommand('zenml.copyStack', async (node: StackTreeItem) => 49 | stackCommands.copyStack(node) 50 | ), 51 | registerCommand('zenml.nextStackPage', async () => stackDataProvider.goToNextPage()), 52 | registerCommand('zenml.previousStackPage', async () => stackDataProvider.goToPreviousPage()), 53 | ]; 54 | 55 | registeredCommands.forEach(cmd => { 56 | context.subscriptions.push(cmd); 57 | ZenExtension.commandDisposables.push(cmd); 58 | }); 59 | 60 | commands.executeCommand('setContext', 'stackCommandsRegistered', true); 61 | } catch (error) { 62 | console.error('Error registering stack commands:', error); 63 | commands.executeCommand('setContext', 'stackCommandsRegistered', false); 64 | } 65 | }; 66 | -------------------------------------------------------------------------------- /src/common/WebviewBase.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as vscode from 'vscode'; 14 | 15 | /** 16 | * Provides functionality to share extension context among classes that inherit 17 | * from it. 18 | */ 19 | export default class WebviewBase { 20 | protected static context: vscode.ExtensionContext | null = null; 21 | 22 | /** 23 | * Sets the extension context so that descendant classes can correctly 24 | * path to their resources 25 | * @param {vscode.ExtensionContext} context ExtensionContext 26 | */ 27 | public static setContext(context: vscode.ExtensionContext) { 28 | WebviewBase.context = context; 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /src/common/api.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { LSClient } from '../services/LSClient'; 14 | import { 15 | ComponentsListResponse, 16 | Flavor, 17 | FlavorListResponse, 18 | StackComponent, 19 | } from '../types/StackTypes'; 20 | 21 | let flavors: Flavor[] = []; 22 | 23 | /** 24 | * Gets all component flavors and caches them 25 | * @returns {Flavor[]} List of flavors 26 | */ 27 | export const getAllFlavors = async (): Promise => { 28 | if (flavors.length > 0) { 29 | return flavors; 30 | } 31 | const lsClient = LSClient.getInstance(); 32 | 33 | let [page, maxPage] = [0, 1]; 34 | do { 35 | page++; 36 | const resp = await lsClient.sendLsClientRequest('listFlavors', [ 37 | page, 38 | 10000, 39 | ]); 40 | 41 | if ('error' in resp) { 42 | console.error(`Error retrieving flavors: ${resp.error}`); 43 | throw new Error(`Error retrieving flavors: ${resp.error}`); 44 | } 45 | 46 | maxPage = resp.total_pages; 47 | flavors = flavors.concat(resp.items); 48 | } while (page < maxPage); 49 | return flavors; 50 | }; 51 | 52 | /** 53 | * Gets all flavors of a specified component type 54 | * @param {string} type Type of component to filter by 55 | * @returns {Flavor[]} List of flavors that match the component type filter 56 | */ 57 | export const getFlavorsOfType = async (type: string): Promise => { 58 | const flavors = await getAllFlavors(); 59 | return flavors.filter(flavor => flavor.type === type); 60 | }; 61 | 62 | /** 63 | * Gets a specific flavor 64 | * @param {string} name The name of the flavor to get 65 | * @returns {Flavor} The specified flavor. 66 | */ 67 | export const getFlavor = async (name: string): Promise => { 68 | const flavors = await getAllFlavors(); 69 | const flavor = flavors.find(flavor => flavor.name === name); 70 | 71 | if (!flavor) { 72 | throw Error(`getFlavor: Flavor ${name} not found`); 73 | } 74 | 75 | return flavor; 76 | }; 77 | 78 | /** 79 | * Gets all stack components 80 | * @returns {object} Object containing all components keyed by each type. 81 | */ 82 | export const getAllStackComponents = async (): Promise<{ 83 | [type: string]: StackComponent[]; 84 | }> => { 85 | const lsClient = LSClient.getInstance(); 86 | let components: StackComponent[] = []; 87 | let [page, maxPage] = [0, 1]; 88 | 89 | do { 90 | page++; 91 | const resp = await lsClient.sendLsClientRequest('listComponents', [ 92 | page, 93 | 10000, 94 | ]); 95 | 96 | if ('error' in resp) { 97 | console.error(`Error retrieving components: ${resp.error}`); 98 | throw new Error(`Error retrieving components: ${resp.error}`); 99 | } 100 | 101 | maxPage = resp.total_pages; 102 | components = components.concat(resp.items); 103 | } while (page < maxPage); 104 | 105 | const out: { [type: string]: StackComponent[] } = {}; 106 | components.forEach(component => { 107 | if (!(component.type in out)) { 108 | out[component.type] = []; 109 | } 110 | out[component.type].push(component); 111 | }); 112 | 113 | return out; 114 | }; 115 | -------------------------------------------------------------------------------- /src/common/constants.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as path from 'path'; 14 | 15 | export const EXTENSION_ID = 'ZenML.zenml-vscode'; 16 | const folderName = path.basename(__dirname); 17 | export const EXTENSION_ROOT_DIR = 18 | folderName === 'common' ? path.dirname(path.dirname(__dirname)) : path.dirname(__dirname); 19 | export const BUNDLED_PYTHON_SCRIPTS_DIR = path.join(EXTENSION_ROOT_DIR, 'bundled'); 20 | export const SERVER_SCRIPT_PATH = path.join(BUNDLED_PYTHON_SCRIPTS_DIR, 'tool', `lsp_server.py`); 21 | export const DEBUG_SERVER_SCRIPT_PATH = path.join( 22 | BUNDLED_PYTHON_SCRIPTS_DIR, 23 | 'tool', 24 | `_debug_server.py` 25 | ); 26 | -------------------------------------------------------------------------------- /src/common/log/logging.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import * as util from 'util'; 15 | import { Disposable, LogOutputChannel } from 'vscode'; 16 | 17 | type Arguments = unknown[]; 18 | class OutputChannelLogger { 19 | constructor(private readonly channel: LogOutputChannel) {} 20 | 21 | public traceLog(...data: Arguments): void { 22 | this.channel.appendLine(util.format(...data)); 23 | } 24 | 25 | public traceError(...data: Arguments): void { 26 | this.channel.error(util.format(...data)); 27 | } 28 | 29 | public traceWarn(...data: Arguments): void { 30 | this.channel.warn(util.format(...data)); 31 | } 32 | 33 | public traceInfo(...data: Arguments): void { 34 | this.channel.info(util.format(...data)); 35 | } 36 | 37 | public traceVerbose(...data: Arguments): void { 38 | this.channel.debug(util.format(...data)); 39 | } 40 | } 41 | 42 | let channel: OutputChannelLogger | undefined; 43 | export function registerLogger(logChannel: LogOutputChannel): Disposable { 44 | channel = new OutputChannelLogger(logChannel); 45 | return { 46 | dispose: () => { 47 | channel = undefined; 48 | }, 49 | }; 50 | } 51 | 52 | export function traceLog(...args: Arguments): void { 53 | channel?.traceLog(...args); 54 | } 55 | 56 | export function traceError(...args: Arguments): void { 57 | channel?.traceError(...args); 58 | } 59 | 60 | export function traceWarn(...args: Arguments): void { 61 | channel?.traceWarn(...args); 62 | } 63 | 64 | export function traceInfo(...args: Arguments): void { 65 | channel?.traceInfo(...args); 66 | } 67 | 68 | export function traceVerbose(...args: Arguments): void { 69 | channel?.traceVerbose(...args); 70 | } 71 | -------------------------------------------------------------------------------- /src/common/panels.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as vscode from 'vscode'; 14 | 15 | /** 16 | * Handles creation and monitoring of webview panels. 17 | */ 18 | export default class Panels { 19 | private static instance: Panels | undefined; 20 | private openPanels: { [id: string]: vscode.WebviewPanel }; 21 | 22 | constructor() { 23 | this.openPanels = {}; 24 | } 25 | 26 | /** 27 | * Retrieves a singleton instance of Panels 28 | * @returns {Panels} The singleton instance 29 | */ 30 | public static getInstance(): Panels { 31 | if (Panels.instance === undefined) { 32 | Panels.instance = new Panels(); 33 | } 34 | return Panels.instance; 35 | } 36 | 37 | /** 38 | * Creates a webview panel 39 | * @param {string} id ID of the webview panel to create 40 | * @param {string} label Title of webview panel tab 41 | * @param {vscode.WebviewPanelOptions & vscode.WebviewOptions} options 42 | * Options applied to the webview panel 43 | * @returns {vscode.WebviewPanel} The webview panel created 44 | */ 45 | public createPanel( 46 | id: string, 47 | label: string, 48 | options?: vscode.WebviewPanelOptions & vscode.WebviewOptions 49 | ) { 50 | const panel = vscode.window.createWebviewPanel(id, label, vscode.ViewColumn.One, options); 51 | panel.webview.html = this.getLoadingContent(); 52 | 53 | this.openPanels[id] = panel; 54 | 55 | panel.onDidDispose(() => { 56 | this.deregisterPanel(id); 57 | }, null); 58 | 59 | return panel; 60 | } 61 | 62 | /** 63 | * Gets existing webview panel 64 | * @param {string} id ID of webview panel to retrieve. 65 | * @param {boolean} forceSpinner Whether to change the html content or not 66 | * @returns {vscode.WebviewPanel | undefined} The webview panel if it exists, 67 | * else undefined 68 | */ 69 | public getPanel(id: string, forceSpinner: boolean = false): vscode.WebviewPanel | undefined { 70 | const panel = this.openPanels[id]; 71 | 72 | if (panel && forceSpinner) { 73 | panel.webview.html = this.getLoadingContent(); 74 | } 75 | 76 | return panel; 77 | } 78 | 79 | private deregisterPanel(id: string) { 80 | delete this.openPanels[id]; 81 | } 82 | 83 | private getLoadingContent(): string { 84 | return ` 85 | 86 | 87 | 88 | 89 | 90 | 91 | Loading 92 | 107 | 108 | 109 | 110 | 111 | `; 112 | } 113 | } 114 | -------------------------------------------------------------------------------- /src/common/status.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { Disposable, l10n, LanguageStatusItem, LanguageStatusSeverity } from 'vscode'; 14 | import { Command } from 'vscode-languageclient'; 15 | import { getDocumentSelector } from './utilities'; 16 | import { createLanguageStatusItem } from './vscodeapi'; 17 | 18 | let _status: LanguageStatusItem | undefined; 19 | export function registerLanguageStatusItem(id: string, name: string, command: string): Disposable { 20 | _status = createLanguageStatusItem(id, getDocumentSelector()); 21 | _status.name = name; 22 | _status.text = name; 23 | _status.command = Command.create(l10n.t('Open logs'), command); 24 | 25 | return { 26 | dispose: () => { 27 | _status?.dispose(); 28 | _status = undefined; 29 | }, 30 | }; 31 | } 32 | 33 | export function updateStatus( 34 | status: string | undefined, 35 | severity: LanguageStatusSeverity, 36 | busy?: boolean, 37 | detail?: string 38 | ): void { 39 | if (_status) { 40 | _status.text = status && status.length > 0 ? `${_status.name}: ${status}` : `${_status.name}`; 41 | _status.severity = severity; 42 | _status.busy = busy ?? false; 43 | _status.detail = detail; 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /src/common/utilities.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as fs from 'fs-extra'; 14 | import * as path from 'path'; 15 | import { DocumentSelector, LogLevel, Uri, WorkspaceFolder } from 'vscode'; 16 | import { Trace } from 'vscode-jsonrpc/node'; 17 | import { getWorkspaceFolders, isVirtualWorkspace } from './vscodeapi'; 18 | 19 | function logLevelToTrace(logLevel: LogLevel): Trace { 20 | switch (logLevel) { 21 | case LogLevel.Error: 22 | case LogLevel.Warning: 23 | case LogLevel.Info: 24 | return Trace.Messages; 25 | 26 | case LogLevel.Debug: 27 | case LogLevel.Trace: 28 | return Trace.Verbose; 29 | 30 | case LogLevel.Off: 31 | default: 32 | return Trace.Off; 33 | } 34 | } 35 | 36 | export function getLSClientTraceLevel(channelLogLevel: LogLevel, globalLogLevel: LogLevel): Trace { 37 | if (channelLogLevel === LogLevel.Off) { 38 | return logLevelToTrace(globalLogLevel); 39 | } 40 | if (globalLogLevel === LogLevel.Off) { 41 | return logLevelToTrace(channelLogLevel); 42 | } 43 | const level = logLevelToTrace( 44 | channelLogLevel <= globalLogLevel ? channelLogLevel : globalLogLevel 45 | ); 46 | return level; 47 | } 48 | 49 | export async function getProjectRoot(): Promise { 50 | const workspaces: readonly WorkspaceFolder[] = getWorkspaceFolders(); 51 | if (workspaces.length === 0) { 52 | return { 53 | uri: Uri.file(process.cwd()), 54 | name: path.basename(process.cwd()), 55 | index: 0, 56 | }; 57 | } else if (workspaces.length === 1) { 58 | return workspaces[0]; 59 | } else { 60 | let rootWorkspace = workspaces[0]; 61 | let root = undefined; 62 | for (const w of workspaces) { 63 | if (await fs.pathExists(w.uri.fsPath)) { 64 | root = w.uri.fsPath; 65 | rootWorkspace = w; 66 | break; 67 | } 68 | } 69 | 70 | for (const w of workspaces) { 71 | if (root && root.length > w.uri.fsPath.length && (await fs.pathExists(w.uri.fsPath))) { 72 | root = w.uri.fsPath; 73 | rootWorkspace = w; 74 | } 75 | } 76 | return rootWorkspace; 77 | } 78 | } 79 | 80 | export function getDocumentSelector(): DocumentSelector { 81 | return isVirtualWorkspace() 82 | ? [{ language: 'python' }] 83 | : [ 84 | { scheme: 'file', language: 'python' }, 85 | { scheme: 'untitled', language: 'python' }, 86 | { scheme: 'vscode-notebook', language: 'python' }, 87 | { scheme: 'vscode-notebook-cell', language: 'python' }, 88 | ]; 89 | } 90 | -------------------------------------------------------------------------------- /src/common/vscodeapi.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import { 15 | commands, 16 | ConfigurationScope, 17 | Disposable, 18 | DocumentSelector, 19 | languages, 20 | LanguageStatusItem, 21 | LogOutputChannel, 22 | Uri, 23 | window, 24 | workspace, 25 | WorkspaceConfiguration, 26 | WorkspaceFolder, 27 | } from 'vscode'; 28 | 29 | export function createOutputChannel(name: string): LogOutputChannel { 30 | return window.createOutputChannel(name, { log: true }); 31 | } 32 | 33 | export function getConfiguration( 34 | config: string, 35 | scope?: ConfigurationScope 36 | ): WorkspaceConfiguration { 37 | return workspace.getConfiguration(config, scope); 38 | } 39 | 40 | export function registerCommand( 41 | command: string, 42 | callback: (...args: any[]) => any, 43 | thisArg?: any 44 | ): Disposable { 45 | return commands.registerCommand(command, callback, thisArg); 46 | } 47 | 48 | export const { onDidChangeConfiguration } = workspace; 49 | 50 | export function isVirtualWorkspace(): boolean { 51 | const isVirtual = 52 | workspace.workspaceFolders && workspace.workspaceFolders.every(f => f.uri.scheme !== 'file'); 53 | return !!isVirtual; 54 | } 55 | 56 | export function getWorkspaceFolders(): readonly WorkspaceFolder[] { 57 | return workspace.workspaceFolders ?? []; 58 | } 59 | 60 | export function getWorkspaceFolder(uri: Uri): WorkspaceFolder | undefined { 61 | return workspace.getWorkspaceFolder(uri); 62 | } 63 | 64 | export function createLanguageStatusItem( 65 | id: string, 66 | selector: DocumentSelector 67 | ): LanguageStatusItem { 68 | return languages.createLanguageStatusItem(id, selector); 69 | } 70 | -------------------------------------------------------------------------------- /src/dag/DagConfig.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | export interface DagConfig { 15 | paths: { 16 | rootPath: string[]; 17 | cssFile: string; 18 | jsFile: string; 19 | iconsDirectory: string; 20 | }; 21 | panZoom: { 22 | maxZoom: number; 23 | viewportSizeRatio: number; 24 | }; 25 | nodes: { 26 | step: { 27 | width: number; 28 | height: number; 29 | }; 30 | artifact: { 31 | width: number; 32 | height: number; 33 | }; 34 | }; 35 | layout: { 36 | rankdir: string; 37 | ranksep: number; 38 | nodesep: number; 39 | }; 40 | icons: { 41 | [key: string]: string; 42 | }; 43 | doubleClickTimeout: number; 44 | } 45 | 46 | export const DEFAULT_DAG_CONFIG: DagConfig = { 47 | paths: { 48 | rootPath: ['resources', 'dag-view'], 49 | cssFile: 'dag.css', 50 | jsFile: 'dag-packed.js', 51 | iconsDirectory: '/resources/dag-view/icons/', 52 | }, 53 | panZoom: { 54 | maxZoom: 40, 55 | viewportSizeRatio: 0.95, 56 | }, 57 | nodes: { 58 | step: { 59 | width: 300, 60 | height: 54, 61 | }, 62 | artifact: { 63 | width: 300, 64 | height: 48, 65 | }, 66 | }, 67 | layout: { 68 | rankdir: 'TB', 69 | ranksep: 35, 70 | nodesep: 5, 71 | }, 72 | icons: { 73 | failed: 'alert.svg', 74 | completed: 'check.svg', 75 | cached: 'cached.svg', 76 | initializing: 'initializing.svg', 77 | running: 'play.svg', 78 | database: 'database.svg', 79 | dataflow: 'dataflow.svg', 80 | }, 81 | doubleClickTimeout: 500, 82 | }; 83 | -------------------------------------------------------------------------------- /src/dag/index.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | export { DagConfig, DEFAULT_DAG_CONFIG } from './DagConfig'; 15 | export { default as DagRenderer } from './renderer/DagRenderer'; 16 | export { SvgRenderer } from './renderer/SvgRenderer'; 17 | export { HtmlTemplateBuilder } from './renderer/HtmlTemplateBuilder'; 18 | export { StatusUtils } from './utils/StatusUtils'; 19 | export { IconLoader } from './utils/IconLoader'; 20 | -------------------------------------------------------------------------------- /src/dag/utils/IconLoader.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import * as fs from 'fs/promises'; 15 | import * as path from 'path'; 16 | import { DagConfig } from '../DagConfig'; 17 | 18 | export class IconLoader { 19 | private iconSvgs: { [name: string]: string } = {}; 20 | private config: DagConfig; 21 | 22 | constructor(config: DagConfig) { 23 | this.config = config; 24 | } 25 | 26 | /** 27 | * Dynamically loads all SVG icons from the icons directory 28 | */ 29 | async loadIcons(basePath: string): Promise<{ [name: string]: string }> { 30 | // Always use the semantic mapping to ensure consistency 31 | await this.loadIconsManually(basePath); 32 | 33 | return this.iconSvgs; 34 | } 35 | 36 | /** 37 | * Fallback method using the original manual icon mapping 38 | */ 39 | private async loadIconsManually(basePath: string): Promise { 40 | const iconPath = path.join(basePath, this.config.paths.iconsDirectory.replace(/^\//, '')); 41 | 42 | const loadPromises = Object.entries(this.config.icons).map(async ([semanticName, fileName]) => { 43 | try { 44 | const filePath = path.join(iconPath, fileName); 45 | const content = await fs.readFile(filePath, 'utf-8'); 46 | this.iconSvgs[semanticName] = content; 47 | } catch (error) { 48 | this.iconSvgs[semanticName] = ''; 49 | console.error(`✗ Failed to load ${semanticName} from ${fileName}:`, error); 50 | } 51 | }); 52 | 53 | await Promise.all(loadPromises); 54 | } 55 | 56 | /** 57 | * Gets the loaded icons 58 | */ 59 | getIcons(): { [name: string]: string } { 60 | return this.iconSvgs; 61 | } 62 | } 63 | -------------------------------------------------------------------------------- /src/dag/utils/StatusUtils.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | export type StepStatus = 'initializing' | 'failed' | 'completed' | 'running' | 'cached'; 15 | 16 | export class StatusUtils { 17 | /** 18 | * Normalizes step status handling both string and object formats 19 | */ 20 | static normalizeStatus(status: string | { _value_: string } | undefined): string { 21 | if (!status) { 22 | return ''; 23 | } 24 | 25 | if (typeof status === 'string') { 26 | return status; 27 | } 28 | 29 | if (typeof status === 'object' && '_value_' in status) { 30 | return status._value_; 31 | } 32 | 33 | return ''; 34 | } 35 | 36 | /** 37 | * Checks if a pipeline run status indicates it needs updating 38 | */ 39 | static shouldShowUpdateButton(status: string): boolean { 40 | return status === 'running' || status === 'initializing'; 41 | } 42 | 43 | /** 44 | * Extracts and formats duration from step data 45 | */ 46 | static extractDuration(stepData: any): string | null { 47 | if (!stepData.start_time || !stepData.end_time) { 48 | return null; 49 | } 50 | 51 | try { 52 | const startTime = new Date(stepData.start_time); 53 | const endTime = new Date(stepData.end_time); 54 | 55 | const durationMs = endTime.getTime() - startTime.getTime(); 56 | const durationSeconds = Math.floor(durationMs / 1000); 57 | 58 | return StatusUtils.formatDuration(durationSeconds); 59 | } catch (error) { 60 | console.warn('Failed to parse step timestamps:', error); 61 | return null; 62 | } 63 | } 64 | 65 | /** 66 | * Formats duration in seconds to human-readable format 67 | */ 68 | static formatDuration(seconds: number): string { 69 | if (seconds < 1) { 70 | return '< 1s'; 71 | } else if (seconds < 60) { 72 | return `${seconds}s`; 73 | } else if (seconds < 3600) { 74 | const minutes = Math.floor(seconds / 60); 75 | const remainingSeconds = seconds % 60; 76 | if (remainingSeconds === 0) { 77 | return `${minutes}m`; 78 | } 79 | return `${minutes}m ${remainingSeconds}s`; 80 | } else { 81 | const hours = Math.floor(seconds / 3600); 82 | const minutes = Math.floor((seconds % 3600) / 60); 83 | if (minutes === 0) { 84 | return `${hours}h`; 85 | } 86 | return `${hours}h ${minutes}m`; 87 | } 88 | } 89 | 90 | /** 91 | * Determines the appropriate icon for a step based on its status 92 | */ 93 | static getStepIcon(status: string, iconSvgs: { [name: string]: string }): string { 94 | return iconSvgs[status] || ''; 95 | } 96 | 97 | /** 98 | * Determines the appropriate icon for an artifact based on its type 99 | */ 100 | static getArtifactIcon(artifactType: string, iconSvgs: { [name: string]: string }): string { 101 | if (artifactType === 'ModelArtifact') { 102 | return iconSvgs.dataflow || ''; 103 | } 104 | return iconSvgs.database || ''; 105 | } 106 | } 107 | -------------------------------------------------------------------------------- /src/extension.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as vscode from 'vscode'; 14 | import { registerEnvironmentCommands } from './commands/environment/registry'; 15 | import DagRenderer from './dag/renderer/DagRenderer'; 16 | import WebviewBase from './common/WebviewBase'; 17 | import { EventBus } from './services/EventBus'; 18 | import { LSClient } from './services/LSClient'; 19 | import { ZenExtension } from './services/ZenExtension'; 20 | import { LSP_ZENML_CLIENT_INITIALIZED } from './utils/constants'; 21 | import { toggleCommands } from './utils/global'; 22 | import { EnvironmentDataProvider } from './views/activityBar/environmentView/EnvironmentDataProvider'; 23 | 24 | export async function activate(context: vscode.ExtensionContext) { 25 | const eventBus = EventBus.getInstance(); 26 | const lsClient = LSClient.getInstance(); 27 | 28 | const handleZenMLClientInitialized = async (isInitialized: boolean) => { 29 | console.log('ZenML client initialized: ', isInitialized); 30 | if (isInitialized) { 31 | await toggleCommands(true); 32 | // await refreshUIComponents(); 33 | } 34 | }; 35 | 36 | eventBus.on(LSP_ZENML_CLIENT_INITIALIZED, handleZenMLClientInitialized); 37 | 38 | vscode.window.createTreeView('zenmlEnvironmentView', { 39 | treeDataProvider: EnvironmentDataProvider.getInstance(), 40 | }); 41 | registerEnvironmentCommands(context); 42 | 43 | await ZenExtension.activate(context, lsClient); 44 | 45 | context.subscriptions.push( 46 | new vscode.Disposable(() => { 47 | eventBus.off(LSP_ZENML_CLIENT_INITIALIZED, handleZenMLClientInitialized); 48 | }) 49 | ); 50 | 51 | WebviewBase.setContext(context); 52 | } 53 | 54 | /** 55 | * Deactivates the ZenML extension. 56 | * 57 | * @returns {Promise} A promise that resolves to void. 58 | */ 59 | export async function deactivate(): Promise { 60 | const lsClient = LSClient.getInstance().getLanguageClient(); 61 | 62 | if (lsClient) { 63 | await lsClient.stop(); 64 | EventBus.getInstance().emit('lsClientReady', false); 65 | } 66 | DagRenderer.getInstance()?.deactivate(); 67 | } 68 | -------------------------------------------------------------------------------- /src/services/EventBus.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { EventEmitter } from 'events'; 14 | 15 | export class EventBus extends EventEmitter { 16 | private static instance: EventBus; 17 | 18 | constructor() { 19 | super(); 20 | this.setMaxListeners(20); 21 | } 22 | 23 | /** 24 | * Retrieves the singleton instance of EventBus. 25 | * 26 | * @returns {EventBus} The singleton instance. 27 | */ 28 | public static getInstance(): EventBus { 29 | if (!EventBus.instance) { 30 | EventBus.instance = new EventBus(); 31 | } 32 | return EventBus.instance; 33 | } 34 | 35 | /** 36 | * Cleans up event listeners for a specific event and handler. 37 | * This is important to prevent memory leaks and MaxListenersExceededWarnings. 38 | * 39 | * @param {string} event - The event name to clean up 40 | * @param {Function} [handler] - Optional specific handler to remove 41 | */ 42 | public cleanupEventListener(event: string, handler?: (...args: any[]) => void): void { 43 | if (handler) { 44 | this.off(event, handler); 45 | } else { 46 | this.removeAllListeners(event); 47 | } 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /src/test/python_tests/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) ZenML GmbH 2024. All Rights Reserved. 2 | # Licensed under the Apache License, Version 2.0 (the "License"); 3 | # you may not use this file except in compliance with the License. 4 | # You may obtain a copy of the License at: 5 | # 6 | # http://www.apache.org/licenses/LICENSE-2.0 7 | # 8 | # Unless required by applicable law or agreed to in writing, software 9 | # distributed under the License is distributed on an "AS IS" BASIS, 10 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | # or implied. See the License for the specific language governing 12 | # permissions and limitations under the License. 13 | -------------------------------------------------------------------------------- /src/test/python_tests/lsp_test_client/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) ZenML GmbH 2024. All Rights Reserved. 2 | # Licensed under the Apache License, Version 2.0 (the "License"); 3 | # you may not use this file except in compliance with the License. 4 | # You may obtain a copy of the License at: 5 | # 6 | # http://www.apache.org/licenses/LICENSE-2.0 7 | # 8 | # Unless required by applicable law or agreed to in writing, software 9 | # distributed under the License is distributed on an "AS IS" BASIS, 10 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | # or implied. See the License for the specific language governing 12 | # permissions and limitations under the License. 13 | -------------------------------------------------------------------------------- /src/test/python_tests/lsp_test_client/constants.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) ZenML GmbH 2024. All Rights Reserved. 2 | # Licensed under the Apache License, Version 2.0 (the "License"); 3 | # you may not use this file except in compliance with the License. 4 | # You may obtain a copy of the License at: 5 | # 6 | # http://www.apache.org/licenses/LICENSE-2.0 7 | # 8 | # Unless required by applicable law or agreed to in writing, software 9 | # distributed under the License is distributed on an "AS IS" BASIS, 10 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | # or implied. See the License for the specific language governing 12 | # permissions and limitations under the License. 13 | 14 | """ 15 | Constants for use with tests. 16 | """ 17 | 18 | import pathlib 19 | 20 | TEST_ROOT = pathlib.Path(__file__).parent.parent 21 | PROJECT_ROOT = TEST_ROOT.parent.parent.parent 22 | TEST_DATA = TEST_ROOT / "test_data" 23 | -------------------------------------------------------------------------------- /src/test/python_tests/lsp_test_client/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) ZenML GmbH 2024. All Rights Reserved. 2 | # Licensed under the Apache License, Version 2.0 (the "License"); 3 | # you may not use this file except in compliance with the License. 4 | # You may obtain a copy of the License at: 5 | # 6 | # http://www.apache.org/licenses/LICENSE-2.0 7 | # 8 | # Unless required by applicable law or agreed to in writing, software 9 | # distributed under the License is distributed on an "AS IS" BASIS, 10 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | # or implied. See the License for the specific language governing 12 | # permissions and limitations under the License. 13 | 14 | """ 15 | Utility functions for use with tests. 16 | """ 17 | 18 | import json 19 | import os 20 | import pathlib 21 | import platform 22 | from random import choice 23 | 24 | from .constants import PROJECT_ROOT 25 | 26 | 27 | def normalizecase(path: str) -> str: 28 | """Fixes 'file' uri or path case for easier testing in windows.""" 29 | if platform.system() == "Windows": 30 | return path.lower() 31 | return path 32 | 33 | 34 | def as_uri(path: str) -> str: 35 | """Return 'file' uri as string.""" 36 | return normalizecase(pathlib.Path(path).as_uri()) 37 | 38 | 39 | class PythonFile: 40 | """Create python file on demand for testing.""" 41 | 42 | def __init__(self, contents, root): 43 | self.contents = contents 44 | self.basename = "".join( 45 | choice("abcdefghijklmnopqrstuvwxyz") if i < 8 else ".py" for i in range(9) 46 | ) 47 | self.fullpath = os.path.join(root, self.basename) 48 | 49 | def __enter__(self): 50 | """Creates a python file for testing.""" 51 | with open(self.fullpath, "w", encoding="utf8") as py_file: 52 | py_file.write(self.contents) 53 | return self 54 | 55 | def __exit__(self, exc_type, value, _tb): 56 | """Cleans up and deletes the python file.""" 57 | os.unlink(self.fullpath) 58 | 59 | 60 | def get_server_info_defaults(): 61 | """Returns server info from package.json""" 62 | package_json_path = PROJECT_ROOT / "package.json" 63 | package_json = json.loads(package_json_path.read_text()) 64 | return package_json["serverInfo"] 65 | 66 | 67 | def get_initialization_options(): 68 | """Returns initialization options from package.json""" 69 | package_json_path = PROJECT_ROOT / "package.json" 70 | package_json = json.loads(package_json_path.read_text()) 71 | 72 | server_info = package_json["serverInfo"] 73 | server_id = server_info["module"] 74 | 75 | properties = package_json["contributes"]["configuration"]["properties"] 76 | setting = {} 77 | for prop in properties: 78 | name = prop[len(server_id) + 1 :] 79 | value = properties[prop]["default"] 80 | setting[name] = value 81 | 82 | setting["workspace"] = as_uri(str(PROJECT_ROOT)) 83 | setting["interpreter"] = [] 84 | 85 | return {"settings": [setting]} 86 | -------------------------------------------------------------------------------- /src/test/python_tests/requirements.in: -------------------------------------------------------------------------------- 1 | # This file is used to generate ./src/test/python_tests/requirements.txt. 2 | # NOTE: 3 | # Use Python 3.8 or greater which ever is the minimum version of the python 4 | # you plan on supporting when creating the environment or using pip-tools. 5 | # Only run the commands below to manually upgrade packages in requirements.txt: 6 | # 1) python -m pip install pip-tools 7 | # 2) pip-compile --generate-hashes --upgrade ./src/test/python_tests/requirements.in 8 | # If you are using nox commands to setup or build package you don't need to 9 | # run the above commands manually. 10 | 11 | # Packages needed by the testing framework. 12 | pytest 13 | PyHamcrest 14 | python-jsonrpc-server 15 | -------------------------------------------------------------------------------- /src/test/python_tests/test_data/sample1/sample.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | print(x) 4 | -------------------------------------------------------------------------------- /src/test/python_tests/test_data/sample1/sample.unformatted: -------------------------------------------------------------------------------- 1 | import sys;print(x) -------------------------------------------------------------------------------- /src/test/ts_tests/__mocks__/MockEventBus.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | import { EventEmitter } from 'stream'; 13 | 14 | export class MockEventBus extends EventEmitter { 15 | public lsClientReady: boolean = false; 16 | private static instance: MockEventBus; 17 | public zenmlReady: boolean = false; 18 | 19 | constructor() { 20 | super(); 21 | this.setMaxListeners(20); 22 | 23 | this.on('lsClientReady', (isReady: boolean) => { 24 | this.lsClientReady = isReady; 25 | }); 26 | } 27 | 28 | /** 29 | * Retrieves the singleton instance of EventBus. 30 | * 31 | * @returns {MockEventBus} The singleton instance. 32 | */ 33 | public static getInstance(): MockEventBus { 34 | if (!MockEventBus.instance) { 35 | MockEventBus.instance = new MockEventBus(); 36 | } 37 | return MockEventBus.instance; 38 | } 39 | 40 | /** 41 | * Clears all event handlers. 42 | */ 43 | public clearAllHandlers() { 44 | this.removeAllListeners(); 45 | } 46 | 47 | /** 48 | * Cleans up event listeners for a specific event and handler. 49 | * This is important to prevent memory leaks and MaxListenersExceededWarnings. 50 | * 51 | * @param {string} event - The event name to clean up 52 | * @param {Function} [handler] - Optional specific handler to remove 53 | */ 54 | public cleanupEventListener(event: string, handler?: (...args: any[]) => void): void { 55 | if (handler) { 56 | this.off(event, handler); 57 | } else { 58 | this.removeAllListeners(event); 59 | } 60 | } 61 | 62 | /** 63 | * Simulates setting the LS Client readiness state. 64 | * 65 | * @param isReady A boolean indicating whether the LS Client is ready. 66 | * @returns void 67 | */ 68 | public setLsClientReady(isReady: boolean): void { 69 | this.lsClientReady = isReady; 70 | this.emit('lsClientReady', isReady); 71 | } 72 | 73 | /** 74 | * Gets event handlers for testing purposes. 75 | * 76 | * @returns {Map} A map of event names and whether they have handlers. 77 | */ 78 | public getEventHandlers(): Map { 79 | const handlers = new Map(); 80 | const eventNames = this.eventNames(); 81 | for (const eventName of eventNames) { 82 | handlers.set(eventName.toString(), this.listenerCount(eventName) > 0); 83 | } 84 | return handlers; 85 | } 86 | } 87 | -------------------------------------------------------------------------------- /src/test/ts_tests/__mocks__/MockViewProviders.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as sinon from 'sinon'; 14 | import { ServerStatus, ZenServerDetails } from '../../../types/ServerInfoTypes'; 15 | import { INITIAL_ZENML_SERVER_STATUS } from '../../../utils/constants'; 16 | import { 17 | ProjectDataProvider, 18 | ServerDataProvider, 19 | StackDataProvider, 20 | } from '../../../views/activityBar'; 21 | import ZenMLStatusBar from '../../../views/statusBar'; 22 | 23 | export class MockZenMLStatusBar extends ZenMLStatusBar { 24 | public refreshActiveStack = sinon.stub().resolves(); 25 | public refreshActiveProject = sinon.stub().resolves(); 26 | } 27 | 28 | export class MockStackDataProvider extends StackDataProvider { 29 | public refresh = sinon.stub().resolves(); 30 | public updateActiveStack = sinon.stub().resolves(); 31 | } 32 | 33 | export class MockProjectDataProvider extends ProjectDataProvider { 34 | public refresh = sinon.stub().resolves(); 35 | public updateActiveProject = sinon.stub().resolves(); 36 | } 37 | 38 | export class MockServerDataProvider extends ServerDataProvider { 39 | public refreshCalled: boolean = false; 40 | public currentServerStatus: ServerStatus = { 41 | ...INITIAL_ZENML_SERVER_STATUS, 42 | isConnected: true, 43 | url: 'http://mocked-server.com', 44 | dashboard_url: 'http://mocked-dashboard.zenml.io', 45 | deployment_type: 'cloud', 46 | active_workspace_id: 'mock-workspace-id', 47 | active_workspace_name: 'mock-workspace', 48 | }; 49 | 50 | public async refresh(updatedServerConfig?: ZenServerDetails): Promise { 51 | this.refreshCalled = true; 52 | if (updatedServerConfig) { 53 | this.currentServerStatus = { 54 | ...updatedServerConfig.storeInfo, 55 | isConnected: updatedServerConfig.storeConfig.type === 'rest', 56 | url: updatedServerConfig.storeConfig.url, 57 | store_type: updatedServerConfig.storeConfig.type, 58 | }; 59 | } 60 | } 61 | 62 | public getCurrentStatus(): ServerStatus | any[] { 63 | return this.currentServerStatus; 64 | } 65 | 66 | public resetMock(): void { 67 | this.refreshCalled = false; 68 | this.currentServerStatus = INITIAL_ZENML_SERVER_STATUS; 69 | } 70 | } 71 | -------------------------------------------------------------------------------- /src/test/ts_tests/__mocks__/constants.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as sinon from 'sinon'; 14 | import * as vscode from 'vscode'; 15 | import { ServerStatus, ZenServerDetails } from '../../../types/ServerInfoTypes'; 16 | 17 | export const MOCK_REST_SERVER_URL = 'https://zenml.example.com'; 18 | export const MOCK_SQL_SERVER_URL = 'sqlite:///path/to/sqlite.db'; 19 | export const MOCK_SERVER_ID = 'test-server'; 20 | export const MOCK_AUTH_SCHEME = 'OAUTH2_PASSWORD_BEARER'; 21 | export const MOCK_ZENML_VERSION = '0.63.0'; 22 | export const MOCK_ACCESS_TOKEN = 'valid_token'; 23 | 24 | export const MOCK_CONTEXT = { 25 | subscriptions: [], 26 | extensionUri: vscode.Uri.parse('file:///extension/path'), 27 | storagePath: '/path/to/storage', 28 | globalStoragePath: '/path/to/global/storage', 29 | workspaceState: { get: sinon.stub(), update: sinon.stub() }, 30 | globalState: { get: sinon.stub(), update: sinon.stub(), setKeysForSync: sinon.stub() }, 31 | logPath: '/path/to/log', 32 | asAbsolutePath: sinon.stub(), 33 | } as any; 34 | 35 | export const MOCK_REST_SERVER_STATUS: ServerStatus = { 36 | isConnected: true, 37 | id: MOCK_SERVER_ID, 38 | store_type: 'rest', 39 | url: MOCK_REST_SERVER_URL, 40 | version: MOCK_ZENML_VERSION, 41 | debug: false, 42 | deployment_type: 'kubernetes', 43 | database_type: 'sqlite', 44 | secrets_store_type: 'sql', 45 | auth_scheme: MOCK_AUTH_SCHEME, 46 | dashboard_url: '', 47 | }; 48 | 49 | export const MOCK_REST_SERVER_DETAILS: ZenServerDetails = { 50 | storeInfo: { 51 | id: MOCK_SERVER_ID, 52 | version: MOCK_ZENML_VERSION, 53 | debug: false, 54 | deployment_type: 'kubernetes', 55 | database_type: 'sqlite', 56 | secrets_store_type: 'sql', 57 | auth_scheme: MOCK_AUTH_SCHEME, 58 | dashboard_url: '', 59 | }, 60 | storeConfig: { 61 | type: 'rest', 62 | url: MOCK_REST_SERVER_URL, 63 | secrets_store: null, 64 | backup_secrets_store: null, 65 | username: null, 66 | password: null, 67 | api_key: 'api_key', 68 | verify_ssl: true, 69 | pool_pre_ping: true, 70 | http_timeout: 30, 71 | }, 72 | }; 73 | 74 | export const MOCK_SQL_SERVER_STATUS: ServerStatus = { 75 | isConnected: false, 76 | id: MOCK_SERVER_ID, 77 | store_type: 'sql', 78 | url: MOCK_SQL_SERVER_URL, 79 | version: MOCK_ZENML_VERSION, 80 | debug: false, 81 | deployment_type: 'local', 82 | database_type: 'sqlite', 83 | secrets_store_type: 'sql', 84 | auth_scheme: MOCK_AUTH_SCHEME, 85 | dashboard_url: '', 86 | }; 87 | 88 | export const MOCK_SQL_SERVER_DETAILS: ZenServerDetails = { 89 | storeInfo: { 90 | id: MOCK_SERVER_ID, 91 | version: MOCK_ZENML_VERSION, 92 | debug: false, 93 | deployment_type: 'local', 94 | database_type: 'sqlite', 95 | secrets_store_type: 'sql', 96 | auth_scheme: MOCK_AUTH_SCHEME, 97 | dashboard_url: '', 98 | }, 99 | storeConfig: { 100 | type: 'sql', 101 | url: MOCK_SQL_SERVER_URL, 102 | secrets_store: null, 103 | backup_secrets_store: null, 104 | username: null, 105 | password: null, 106 | verify_ssl: false, 107 | pool_pre_ping: true, 108 | http_timeout: 30, 109 | driver: '', 110 | database: '', 111 | ssl_ca: '', 112 | ssl_key: '', 113 | ssl_verify_server_cert: false, 114 | ssl_cert: '', 115 | pool_size: 0, 116 | max_overflow: 0, 117 | backup_strategy: '', 118 | backup_directory: '', 119 | backup_database: '', 120 | }, 121 | }; 122 | -------------------------------------------------------------------------------- /src/test/ts_tests/commands/dagRender.test.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import * as assert from 'assert'; 15 | import * as sinon from 'sinon'; 16 | import * as vscode from 'vscode'; 17 | import Panels from '../../../common/panels'; 18 | import WebviewBase from '../../../common/WebviewBase'; 19 | import DagRenderer from '../../../dag/renderer/DagRenderer'; 20 | import { IconLoader } from '../../../dag/utils/IconLoader'; 21 | import { LSClient } from '../../../services/LSClient'; 22 | import { ServerDataProvider } from '../../../views/activityBar/serverView/ServerDataProvider'; 23 | import { MockEventBus } from '../__mocks__/MockEventBus'; 24 | import { MockLSClient } from '../__mocks__/MockLSClient'; 25 | 26 | suite('DagRenderer Test Suite', () => { 27 | let sandbox: sinon.SinonSandbox; 28 | let mockLSClient: any; 29 | let mockEventBus: any; 30 | let mockWebviewPanel: any; 31 | 32 | setup(() => { 33 | sandbox = sinon.createSandbox(); 34 | mockEventBus = new MockEventBus(); 35 | mockLSClient = new MockLSClient(mockEventBus); 36 | 37 | // Mock extension context for WebviewBase 38 | const mockExtensionContext = { 39 | extensionUri: vscode.Uri.file('/mock/extension/path'), 40 | extensionPath: '/mock/extension/path', 41 | } as vscode.ExtensionContext; 42 | WebviewBase.setContext(mockExtensionContext); 43 | 44 | // Stub icon loading and SVG lib to prevent file system warnings 45 | sandbox.stub(IconLoader.prototype, 'loadIcons').resolves({}); 46 | sandbox.stub(DagRenderer.prototype, 'loadSvgWindowLib' as any).returns(undefined); 47 | 48 | // Mock webview panel 49 | mockWebviewPanel = { 50 | webview: { 51 | html: '', 52 | cspSource: 'vscode-webview:', 53 | asWebviewUri: sandbox.stub().returns(vscode.Uri.parse('vscode-webview://fake-uri')), 54 | onDidReceiveMessage: sandbox.stub(), 55 | postMessage: sandbox.stub().resolves(), 56 | }, 57 | onDidDispose: sandbox.stub(), 58 | onDidChangeViewState: sandbox.stub(), 59 | reveal: sandbox.stub(), 60 | dispose: sandbox.stub(), 61 | }; 62 | 63 | // Mock server status 64 | const mockServerStatus = { 65 | dashboard_url: 'https://test-dashboard.com', 66 | deployment_type: 'cloud', 67 | server_url: 'https://test-server.com', 68 | }; 69 | 70 | sandbox.stub(LSClient, 'getInstance').returns(mockLSClient); 71 | sandbox.stub(ServerDataProvider, 'getInstance').returns({ 72 | getCurrentStatus: sandbox.stub().returns(mockServerStatus), 73 | } as any); 74 | 75 | // Mock Panels singleton 76 | const mockPanels = { 77 | getPanel: sandbox.stub().returns(null), // No existing panel 78 | createPanel: sandbox.stub().returns(mockWebviewPanel), 79 | }; 80 | sandbox.stub(Panels, 'getInstance').returns(mockPanels as any); 81 | }); 82 | 83 | teardown(() => { 84 | sandbox.restore(); 85 | mockEventBus.clearAllHandlers(); 86 | }); 87 | 88 | test('should create DagRenderer instance', () => { 89 | const dagRenderer = new DagRenderer(); 90 | assert.ok(dagRenderer instanceof DagRenderer); 91 | }); 92 | 93 | test('should have working getInstance method', () => { 94 | const instance1 = DagRenderer.getInstance(); 95 | const instance2 = DagRenderer.getInstance(); 96 | assert.strictEqual(instance1, instance2); 97 | assert.ok(instance1 instanceof DagRenderer); 98 | }); 99 | }); 100 | -------------------------------------------------------------------------------- /src/test/ts_tests/commands/modelCommands.test.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import * as sinon from 'sinon'; 15 | import * as vscode from 'vscode'; 16 | import { modelCommands } from '../../../commands/models/cmds'; 17 | import { EventBus } from '../../../services/EventBus'; 18 | import { LSClient } from '../../../services/LSClient'; 19 | import { ModelDataProvider } from '../../../views/activityBar/modelView/ModelDataProvider'; 20 | import { MockEventBus } from '../__mocks__/MockEventBus'; 21 | import { MockLSClient } from '../__mocks__/MockLSClient'; 22 | 23 | interface MockModelDataProviderInterface { 24 | refresh: sinon.SinonStub; 25 | } 26 | 27 | suite('Model Commands Test Suite', () => { 28 | let sandbox: sinon.SinonSandbox; 29 | let mockLSClient: any; 30 | let mockEventBus: any; 31 | let mockModelDataProvider: MockModelDataProviderInterface; 32 | 33 | setup(() => { 34 | sandbox = sinon.createSandbox(); 35 | mockEventBus = new MockEventBus(); 36 | mockLSClient = new MockLSClient(mockEventBus); 37 | 38 | // Create a simple mock with just the methods we need 39 | mockModelDataProvider = { 40 | refresh: sinon.stub().resolves(), 41 | }; 42 | 43 | // Stub classes to return mock instances 44 | sandbox.stub(ModelDataProvider, 'getInstance').returns(mockModelDataProvider as any); 45 | sandbox.stub(LSClient, 'getInstance').returns(mockLSClient); 46 | sandbox.stub(EventBus, 'getInstance').returns(mockEventBus); 47 | 48 | sandbox.stub(vscode.window, 'withProgress').callsFake(async (options, task) => { 49 | const mockProgress = { 50 | report: sandbox.stub(), 51 | }; 52 | const mockCancellationToken = new vscode.CancellationTokenSource(); 53 | await task(mockProgress, mockCancellationToken.token); 54 | }); 55 | }); 56 | 57 | teardown(() => { 58 | sandbox.restore(); 59 | mockEventBus.clearAllHandlers(); 60 | 61 | const eventBus = EventBus.getInstance(); 62 | eventBus.cleanupEventListener('lsClientStateChanged'); 63 | eventBus.cleanupEventListener('zenml/clientInitialized'); 64 | }); 65 | 66 | test('refreshModelView calls refresh on the ModelDataProvider', async () => { 67 | await modelCommands.refreshModelView(); 68 | 69 | sinon.assert.calledOnce(mockModelDataProvider.refresh); 70 | }); 71 | 72 | test('ModelDataProvider.refresh can be called directly', async () => { 73 | await mockModelDataProvider.refresh(); 74 | sinon.assert.calledOnce(mockModelDataProvider.refresh); 75 | }); 76 | 77 | test('refreshModelView refreshes view on command execution', async () => { 78 | // Create a backup of the original executeCommand function 79 | const originalExecuteCommand = vscode.commands.executeCommand; 80 | 81 | // Instead of registering the command again, just stub executeCommand 82 | // to call our modelCommands.refreshModelView directly when it sees zenml.refreshModelView 83 | sandbox.stub(vscode.commands, 'executeCommand').callsFake(async (command, ...args) => { 84 | if (command === 'zenml.refreshModelView') { 85 | return modelCommands.refreshModelView(); 86 | } 87 | return originalExecuteCommand.call(vscode.commands, command, ...args); 88 | }); 89 | 90 | // Execute 91 | await vscode.commands.executeCommand('zenml.refreshModelView'); 92 | 93 | // Verify 94 | sinon.assert.calledOnce(mockModelDataProvider.refresh); 95 | }); 96 | }); 97 | -------------------------------------------------------------------------------- /src/test/ts_tests/commands/pipelineUtils.test.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import * as assert from 'assert'; 15 | import * as sinon from 'sinon'; 16 | import { getPipelineRunDashboardUrl } from '../../../commands/pipelines/utils'; 17 | import { ServerDataProvider } from '../../../views/activityBar/serverView/ServerDataProvider'; 18 | 19 | suite('Pipeline Utils Test Suite', () => { 20 | let sandbox: sinon.SinonSandbox; 21 | let mockServerDataProvider: any; 22 | 23 | setup(() => { 24 | sandbox = sinon.createSandbox(); 25 | 26 | mockServerDataProvider = { 27 | getCurrentStatus: sandbox.stub(), 28 | }; 29 | 30 | sandbox.stub(ServerDataProvider, 'getInstance').returns(mockServerDataProvider); 31 | }); 32 | 33 | teardown(() => { 34 | sandbox.restore(); 35 | }); 36 | 37 | test('getPipelineRunDashboardUrl should return empty string for empty id', () => { 38 | const result = getPipelineRunDashboardUrl(''); 39 | assert.strictEqual(result, ''); 40 | }); 41 | 42 | test('getPipelineRunDashboardUrl should return empty string for null server status', () => { 43 | mockServerDataProvider.getCurrentStatus.returns(null); 44 | 45 | const result = getPipelineRunDashboardUrl('test-run-id'); 46 | assert.strictEqual(result, ''); 47 | }); 48 | 49 | test('getPipelineRunDashboardUrl should return empty string for "other" deployment type', () => { 50 | mockServerDataProvider.getCurrentStatus.returns({ 51 | deployment_type: 'other', 52 | dashboard_url: 'https://test.com', 53 | }); 54 | 55 | const result = getPipelineRunDashboardUrl('test-run-id'); 56 | assert.strictEqual(result, ''); 57 | }); 58 | 59 | test('getPipelineRunDashboardUrl should handle missing dashboard_url', () => { 60 | mockServerDataProvider.getCurrentStatus.returns({ 61 | deployment_type: 'cloud', 62 | dashboard_url: null, 63 | server_url: 'https://server.zenml.cloud', 64 | }); 65 | 66 | const result = getPipelineRunDashboardUrl('test-run-id'); 67 | assert.strictEqual(result, ''); 68 | }); 69 | 70 | test('getPipelineRunDashboardUrl should handle undefined dashboard_url', () => { 71 | mockServerDataProvider.getCurrentStatus.returns({ 72 | deployment_type: 'cloud', 73 | server_url: 'https://server.zenml.cloud', 74 | }); 75 | 76 | const result = getPipelineRunDashboardUrl('test-run-id'); 77 | assert.strictEqual(result, ''); 78 | }); 79 | }); 80 | -------------------------------------------------------------------------------- /src/test/ts_tests/common/WebviewBase.test.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import * as assert from 'assert'; 15 | import * as sinon from 'sinon'; 16 | import WebviewBase from '../../../common/WebviewBase'; 17 | 18 | // Create a concrete test class since WebviewBase is abstract 19 | class TestWebviewBase extends WebviewBase { 20 | constructor() { 21 | super(); 22 | } 23 | 24 | public getHtmlContent(): string { 25 | return 'Test Content'; 26 | } 27 | } 28 | 29 | suite('WebviewBase Test Suite', () => { 30 | let sandbox: sinon.SinonSandbox; 31 | let testWebview: TestWebviewBase; 32 | 33 | setup(() => { 34 | sandbox = sinon.createSandbox(); 35 | testWebview = new TestWebviewBase(); 36 | }); 37 | 38 | teardown(() => { 39 | sandbox.restore(); 40 | }); 41 | 42 | test('should create WebviewBase instance', () => { 43 | assert.ok(testWebview instanceof WebviewBase); 44 | assert.ok(testWebview instanceof TestWebviewBase); 45 | }); 46 | 47 | test('getHtmlContent should return HTML string', () => { 48 | const content = testWebview.getHtmlContent(); 49 | assert.strictEqual(typeof content, 'string'); 50 | assert.match(content, /html/); 51 | assert.match(content, /Test Content/); 52 | }); 53 | 54 | test('concrete implementation should provide HTML content', () => { 55 | // Verify that our test implementation works 56 | const content = testWebview.getHtmlContent(); 57 | assert.strictEqual(typeof content, 'string'); 58 | assert.ok(content.length > 0); 59 | }); 60 | 61 | test('should be extendable for different webview types', () => { 62 | // Create another test class to verify extensibility 63 | class AnotherTestWebview extends WebviewBase { 64 | getHtmlContent(): string { 65 | return 'Another Test'; 66 | } 67 | } 68 | 69 | const anotherWebview = new AnotherTestWebview(); 70 | const content = anotherWebview.getHtmlContent(); 71 | assert.match(content, /Another Test/); 72 | }); 73 | 74 | test('should maintain proper inheritance chain', () => { 75 | assert.ok(testWebview instanceof WebviewBase); 76 | assert.ok(testWebview instanceof TestWebviewBase); 77 | }); 78 | }); 79 | -------------------------------------------------------------------------------- /src/test/ts_tests/extension.test.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as assert from 'assert'; 14 | import * as sinon from 'sinon'; 15 | import * as vscode from 'vscode'; 16 | import * as extension from '../../extension'; 17 | import { EventBus } from '../../services/EventBus'; 18 | import { LSClient } from '../../services/LSClient'; 19 | import { ZenExtension } from '../../services/ZenExtension'; 20 | import { MockEventBus } from './__mocks__/MockEventBus'; 21 | 22 | suite('Extension Activation Test Suite', () => { 23 | let sandbox: sinon.SinonSandbox; 24 | let contextMock: any; 25 | let initializeSpy: sinon.SinonSpy; 26 | let lsClient: LSClient; 27 | const mockEventBus = new MockEventBus(); 28 | 29 | setup(() => { 30 | sandbox = sinon.createSandbox(); 31 | contextMock = { 32 | subscriptions: [], 33 | extensionPath: '', 34 | extensionUri: vscode.Uri.file('/'), 35 | }; 36 | initializeSpy = sinon.spy(ZenExtension, 'activate'); 37 | lsClient = LSClient.getInstance(); 38 | sandbox.stub(EventBus, 'getInstance').returns(mockEventBus); 39 | }); 40 | 41 | teardown(() => { 42 | sandbox.restore(); 43 | initializeSpy.restore(); 44 | }); 45 | 46 | test('ZenML Extension should be present', () => { 47 | assert.ok(vscode.extensions.getExtension('ZenML.zenml-vscode')); 48 | }); 49 | 50 | test('activate function behaves as expected', async () => { 51 | await extension.activate(contextMock); 52 | sinon.assert.calledOnceWithExactly(initializeSpy, contextMock, lsClient); 53 | }); 54 | }); 55 | -------------------------------------------------------------------------------- /src/test/ts_tests/integration/serverConfigUpdate.test.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as assert from 'assert'; 14 | import * as sinon from 'sinon'; 15 | import * as vscode from 'vscode'; 16 | import { EventBus } from '../../../services/EventBus'; 17 | import { ZenServerDetails } from '../../../types/ServerInfoTypes'; 18 | import { LSCLIENT_READY, LSP_ZENML_SERVER_CHANGED } from '../../../utils/constants'; 19 | import { MOCK_REST_SERVER_DETAILS } from '../__mocks__/constants'; 20 | import { MockEventBus } from '../__mocks__/MockEventBus'; 21 | import { MockLSClient } from '../__mocks__/MockLSClient'; 22 | 23 | suite('Server Configuration Update Flow Tests', () => { 24 | let sandbox: sinon.SinonSandbox; 25 | let mockLSClientInstance: MockLSClient; 26 | let refreshUIComponentsStub: sinon.SinonStub; 27 | const mockEventBus = new MockEventBus(); 28 | 29 | setup(() => { 30 | sandbox = sinon.createSandbox(); 31 | sandbox.stub(vscode.window, 'showInformationMessage'); 32 | refreshUIComponentsStub = sandbox.stub().resolves(); 33 | 34 | // Mock LSClient 35 | mockLSClientInstance = new MockLSClient(mockEventBus); 36 | sandbox.stub(mockLSClientInstance, 'startLanguageClient').resolves(); 37 | 38 | // Mock EventBus 39 | sandbox.stub(EventBus, 'getInstance').returns(mockEventBus); 40 | mockEventBus.on(LSCLIENT_READY, async (isReady: boolean) => { 41 | if (isReady) { 42 | await refreshUIComponentsStub(); 43 | } 44 | }); 45 | mockEventBus.on(LSP_ZENML_SERVER_CHANGED, async (updatedServerConfig: ZenServerDetails) => { 46 | await refreshUIComponentsStub(updatedServerConfig); 47 | }); 48 | }); 49 | 50 | teardown(() => { 51 | sandbox.restore(); 52 | mockEventBus.clearAllHandlers(); 53 | }); 54 | 55 | test('LSClientReady event triggers UI refresh', async () => { 56 | mockEventBus.setLsClientReady(true); 57 | sinon.assert.calledOnce(refreshUIComponentsStub); 58 | }); 59 | 60 | test('MockLSClient triggerNotification works as expected', async () => { 61 | const mockNotificationType = 'testNotification'; 62 | const mockData = { key: 'value' }; 63 | mockLSClientInstance.mockLanguageClient.onNotification(mockNotificationType, (data: any) => { 64 | assert.deepStrictEqual(data, mockData); 65 | }); 66 | mockLSClientInstance.triggerNotification(mockNotificationType, mockData); 67 | }); 68 | 69 | test('zenml/serverChanged event updates global configuration and refreshes UI', async () => { 70 | mockLSClientInstance.mockLanguageClient.onNotification( 71 | LSP_ZENML_SERVER_CHANGED, 72 | (data: ZenServerDetails) => { 73 | assert.deepStrictEqual(data, MOCK_REST_SERVER_DETAILS); 74 | } 75 | ); 76 | 77 | mockLSClientInstance.triggerNotification(LSP_ZENML_SERVER_CHANGED, MOCK_REST_SERVER_DETAILS); 78 | 79 | await new Promise(resolve => setTimeout(resolve, 0)); 80 | 81 | sinon.assert.calledOnce(refreshUIComponentsStub); 82 | }); 83 | }); 84 | -------------------------------------------------------------------------------- /src/test/ts_tests/unit/ServerDataProvider.test.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as assert from 'assert'; 14 | import * as sinon from 'sinon'; 15 | import { serverUtils } from '../../../commands/server/utils'; 16 | import { EventBus } from '../../../services/EventBus'; 17 | import { LSClient } from '../../../services/LSClient'; 18 | import { ServerStatus } from '../../../types/ServerInfoTypes'; 19 | import { ServerDataProvider } from '../../../views/activityBar'; 20 | // Removed unused LOADING_TREE_ITEMS import 21 | import { MOCK_REST_SERVER_STATUS, MOCK_SQL_SERVER_STATUS } from '../__mocks__/constants'; 22 | import { MockEventBus } from '../__mocks__/MockEventBus'; 23 | import { MockLSClient } from '../__mocks__/MockLSClient'; 24 | 25 | suite('ServerDataProvider Tests', () => { 26 | let sandbox: sinon.SinonSandbox; 27 | let mockEventBus: MockEventBus; 28 | let serverDataProvider: ServerDataProvider; 29 | let mockLSClientInstance: any; 30 | 31 | setup(() => { 32 | sandbox = sinon.createSandbox(); 33 | serverDataProvider = ServerDataProvider.getInstance(); 34 | mockEventBus = new MockEventBus(); 35 | mockLSClientInstance = MockLSClient.getInstance(mockEventBus); 36 | sandbox.stub(LSClient, 'getInstance').returns(mockLSClientInstance); 37 | sandbox.stub(EventBus, 'getInstance').returns(mockEventBus); 38 | sandbox.stub(mockLSClientInstance, 'startLanguageClient').resolves(); 39 | serverDataProvider['zenmlClientReady'] = true; 40 | }); 41 | 42 | teardown(() => { 43 | sandbox.restore(); 44 | }); 45 | 46 | test('ServerDataProvider initializes correctly', async () => { 47 | assert.ok(serverDataProvider); 48 | }); 49 | 50 | test('ServerDataProvider should update server status correctly', async () => { 51 | sandbox.stub(serverUtils, 'checkServerStatus').callsFake(async () => { 52 | return Promise.resolve(MOCK_REST_SERVER_STATUS); 53 | }); 54 | 55 | await serverDataProvider.refresh(); 56 | const serverStatus = serverDataProvider.getCurrentStatus() as ServerStatus; 57 | 58 | assert.strictEqual( 59 | serverStatus.isConnected, 60 | true, 61 | 'Server should be reported as connected for REST config' 62 | ); 63 | }); 64 | 65 | test('ServerDataProvider should update server status to disconnected for non-REST type', async () => { 66 | sandbox.restore(); 67 | 68 | sandbox.stub(serverUtils, 'checkServerStatus').callsFake(async () => { 69 | return Promise.resolve(MOCK_SQL_SERVER_STATUS); 70 | }); 71 | 72 | await serverDataProvider.refresh(); 73 | const serverStatus = serverDataProvider.getCurrentStatus() as ServerStatus; 74 | 75 | assert.strictEqual( 76 | serverStatus.isConnected, 77 | false, 78 | 'Server should be reported as disconnected for SQL config' 79 | ); 80 | }); 81 | test('ServerDataProvider should handle zenmlClient not ready state', async () => { 82 | serverDataProvider['zenmlClientReady'] = false; 83 | 84 | await serverDataProvider.refresh(); 85 | const status = serverDataProvider.getCurrentStatus(); 86 | // Should return some kind of error/message item when client not ready 87 | assert.ok(Array.isArray(status), 'Should return an array'); 88 | assert.ok(status.length > 0, 'Should return at least one item'); 89 | 90 | serverDataProvider['zenmlClientReady'] = true; 91 | }); 92 | }); 93 | -------------------------------------------------------------------------------- /src/test/ts_tests/unit/eventBus.test.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as assert from 'assert'; 14 | import * as sinon from 'sinon'; 15 | import { LSCLIENT_READY } from '../../../utils/constants'; 16 | import { MockEventBus } from '../__mocks__/MockEventBus'; 17 | 18 | suite('MockEventBus and Event Handling', () => { 19 | let eventBus: MockEventBus; 20 | let spy: sinon.SinonSpy; 21 | 22 | setup(() => { 23 | eventBus = new MockEventBus(); 24 | spy = sinon.spy(); 25 | }); 26 | 27 | test('handles lsClientReady event correctly with mock', () => { 28 | eventBus.on(LSCLIENT_READY, spy); 29 | eventBus.emit(LSCLIENT_READY, true); 30 | assert.ok( 31 | spy.calledWith(true), 32 | 'lsClientReady event handler was not called with expected argument' 33 | ); 34 | }); 35 | 36 | test('can clear all event handlers and not trigger events', () => { 37 | eventBus.on(LSCLIENT_READY, spy); 38 | eventBus.clearAllHandlers(); 39 | 40 | // Try emitting the event after clearing all handlers 41 | eventBus.emit(LSCLIENT_READY, true); 42 | 43 | // Verify the spy was not called since all handlers were cleared 44 | assert.strictEqual( 45 | spy.called, 46 | false, 47 | 'lsClientReady event handler was called despite clearing all handlers' 48 | ); 49 | }); 50 | }); 51 | -------------------------------------------------------------------------------- /src/types/LSClientResponseTypes.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import { ZenServerDetails } from './ServerInfoTypes'; 15 | import { Components } from './StackTypes'; 16 | 17 | /***** Generic Response Types *****/ 18 | export interface SuccessMessageResponse { 19 | message: string; 20 | } 21 | 22 | export interface ErrorMessageResponse { 23 | error: string; 24 | message: string; 25 | } 26 | 27 | export interface VersionMismatchError { 28 | error: string; 29 | message: string; 30 | clientVersion: string; 31 | serverVersion: string; 32 | } 33 | 34 | export type GenericLSClientResponse = SuccessMessageResponse | ErrorMessageResponse; 35 | 36 | /***** Server Response Types *****/ 37 | export interface RestServerConnectionResponse { 38 | message: string; 39 | access_token: string; 40 | } 41 | 42 | export type ServerStatusInfoResponse = 43 | | ZenServerDetails 44 | | VersionMismatchError 45 | | ErrorMessageResponse; 46 | export type ConnectServerResponse = RestServerConnectionResponse | ErrorMessageResponse; 47 | 48 | /***** Stack Response Types *****/ 49 | export interface ActiveStackResponse { 50 | id: string; 51 | name: string; 52 | components: Components; 53 | } 54 | 55 | export type SetActiveStackResponse = ActiveStackResponse | ErrorMessageResponse; 56 | export type GetActiveStackResponse = ActiveStackResponse | ErrorMessageResponse; 57 | 58 | /***** Project Response Types *****/ 59 | export interface ActiveProjectResponse { 60 | id: string; 61 | name: string; 62 | display_name?: string; 63 | created?: string; 64 | updated?: string; 65 | } 66 | 67 | export type SetActiveProjectResponse = ActiveProjectResponse | ErrorMessageResponse; 68 | export type GetActiveProjectResponse = ActiveProjectResponse | ErrorMessageResponse; 69 | -------------------------------------------------------------------------------- /src/types/LSNotificationTypes.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | // create type for LSP server notification that returns {is_installed: boolean, version: string} 15 | 16 | export interface LSNotificationIsZenMLInstalled { 17 | is_installed: boolean; 18 | version?: string; 19 | } 20 | 21 | export interface LSNotification { 22 | is_ready: boolean; 23 | } 24 | -------------------------------------------------------------------------------- /src/types/ModelTypes.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import { ErrorMessageResponse, VersionMismatchError } from './LSClientResponseTypes'; 15 | 16 | export interface ModelUser { 17 | id?: string; 18 | name: string; 19 | is_service_account?: boolean; 20 | full_name?: string; 21 | email_opted_in?: boolean; 22 | is_admin?: boolean; 23 | } 24 | 25 | export interface ModelTag { 26 | name: string; 27 | } 28 | 29 | export interface Model { 30 | id?: string; 31 | name: string; 32 | latest_version_name?: string; 33 | latest_version_id?: string; 34 | user?: ModelUser; 35 | tags: string[]; 36 | } 37 | 38 | export interface ModelsData { 39 | index: number; 40 | max_size: number; 41 | total_pages: number; 42 | total: number; 43 | items: Model[]; 44 | } 45 | 46 | export interface ModelVersionModel { 47 | id: string; 48 | name: string; 49 | tags: string[]; 50 | user?: ModelUser; 51 | } 52 | 53 | export interface ModelVersion { 54 | id: string; 55 | name: string; 56 | created: string; 57 | updated: string; 58 | stage: string | null; 59 | number: number; 60 | model: ModelVersionModel; 61 | run_metadata?: Record; 62 | // These fields have been REMOVED in v0.82+ to reduce response size 63 | data_artifact_ids?: { 64 | [key: string]: { 65 | [key: string]: string; 66 | }; 67 | } | null; 68 | model_artifact_ids?: { 69 | [key: string]: { 70 | [key: string]: string; 71 | }; 72 | } | null; 73 | deployment_artifact_ids?: { 74 | [key: string]: { 75 | [key: string]: string; 76 | }; 77 | } | null; 78 | pipeline_run_ids?: { 79 | [key: string]: string; 80 | } | null; 81 | tags: ModelTag[]; 82 | // New resources field for optimized responses 83 | resources?: { 84 | data_artifact_ids?: { 85 | [key: string]: { 86 | [key: string]: string; 87 | }; 88 | }; 89 | model_artifact_ids?: { 90 | [key: string]: { 91 | [key: string]: string; 92 | }; 93 | }; 94 | pipeline_run_ids?: { 95 | [key: string]: string; 96 | }; 97 | user?: ModelUser; 98 | }; 99 | } 100 | 101 | export interface ModelVersionsData { 102 | index: number; 103 | max_size: number; 104 | total_pages: number; 105 | total: number; 106 | items: ModelVersion[]; 107 | } 108 | 109 | export type ModelsResponse = ModelsData | ErrorMessageResponse | VersionMismatchError; 110 | export type ModelVersionsResponse = ModelVersionsData | ErrorMessageResponse | VersionMismatchError; 111 | -------------------------------------------------------------------------------- /src/types/PipelineTypes.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import { ErrorMessageResponse, VersionMismatchError } from './LSClientResponseTypes'; 15 | 16 | export interface PipelineRunsData { 17 | runs: PipelineRun[]; 18 | total: number; 19 | total_pages: number; 20 | current_page: number; 21 | items_per_page: number; 22 | } 23 | 24 | export interface PipelineRunStep { 25 | status: string; 26 | start_time?: string; 27 | end_time?: string; 28 | id?: string; 29 | } 30 | 31 | export interface PipelineRunConfig { 32 | enable_cache?: boolean; 33 | enable_artifact_metadata?: boolean; 34 | enable_artifact_visualization?: boolean; 35 | enable_step_logs?: boolean; 36 | name?: string; 37 | model?: PipelineModel; 38 | } 39 | 40 | export interface PipelineModel { 41 | name: string; 42 | description?: string; 43 | tags?: string[]; 44 | version?: string; 45 | save_models_to_registry?: boolean; 46 | license?: string; 47 | } 48 | 49 | export interface PipelineRun { 50 | id: string; 51 | name: string; 52 | status: string; 53 | stackName: string; 54 | startTime: string; 55 | endTime: string; 56 | pipelineName: string; 57 | runMetadata?: Record; 58 | config?: PipelineRunConfig; 59 | steps?: { 60 | [stepName: string]: PipelineRunStep; 61 | } | null; // Steps may be null in optimized responses 62 | } 63 | 64 | export interface DagStep { 65 | id: string; 66 | type: 'step'; 67 | data: { 68 | execution_id: string; 69 | name: string; 70 | status: 'initializing' | 'failed' | 'completed' | 'running' | 'cached'; 71 | }; 72 | } 73 | 74 | export interface DagArtifact { 75 | id: string; 76 | type: 'artifact'; 77 | data: { 78 | execution_id: string; 79 | name: string; 80 | artifact_type: string; 81 | type?: string; 82 | }; 83 | } 84 | 85 | export type DagNode = DagStep | DagArtifact; 86 | 87 | export interface DagEdge { 88 | id: string; 89 | source: string; 90 | target: string; 91 | } 92 | 93 | export interface PipelineRunDag { 94 | nodes: Array; 95 | edges: Array; 96 | status: string; 97 | name: string; 98 | message?: string; // Optional message when step data is not available 99 | } 100 | 101 | export type PipelineRunsResponse = PipelineRunsData | ErrorMessageResponse | VersionMismatchError; 102 | -------------------------------------------------------------------------------- /src/types/ProjectTypes.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import { ErrorMessageResponse, VersionMismatchError } from './LSClientResponseTypes'; 15 | 16 | export interface MetadataType { 17 | [key: string]: string | number | boolean | object | string[] | number[] | boolean[] | object[]; 18 | } 19 | 20 | export interface Project { 21 | id: string; 22 | name: string; 23 | display_name?: string; 24 | created?: string; 25 | updated?: string; 26 | metadata?: MetadataType; 27 | } 28 | 29 | export interface ProjectsData { 30 | projects: Project[]; 31 | total: number; 32 | total_pages: number; 33 | current_page: number; 34 | items_per_page: number; 35 | } 36 | 37 | export interface GetProjectByNameResponse { 38 | id: string; 39 | name: string; 40 | display_name?: string; 41 | } 42 | 43 | export type ProjectsResponse = ProjectsData | ErrorMessageResponse | VersionMismatchError; 44 | -------------------------------------------------------------------------------- /src/types/QuickPickItemTypes.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | /** Custom interfaces for our QuickPick items */ 15 | 16 | import { QuickPickItem, ThemeIcon, Uri } from 'vscode'; 17 | export type IconPath = ThemeIcon | { light: Uri; dark: Uri }; 18 | 19 | export interface MainMenuQuickPickItem extends QuickPickItem { 20 | id: string; 21 | iconPath?: IconPath; 22 | buttons?: Array<{ 23 | iconPath: IconPath; 24 | tooltip: string; 25 | }>; 26 | } 27 | 28 | export interface StackQuickPickItem extends QuickPickItem { 29 | id?: string; 30 | iconPath?: IconPath; 31 | disabled?: boolean; 32 | } 33 | 34 | export interface ProjectQuickPickItem extends QuickPickItem { 35 | id?: string; 36 | name?: string; 37 | iconPath?: IconPath; 38 | disabled?: boolean; 39 | } 40 | -------------------------------------------------------------------------------- /src/types/StackTypes.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import { ErrorMessageResponse, VersionMismatchError } from './LSClientResponseTypes'; 15 | 16 | /************************************************************************************************ 17 | * LSClient parses the JSON response from the ZenML Client, and returns the following types. 18 | * Hydrated types are in the HydratedTypes.ts file. 19 | ************************************************************************************************/ 20 | interface StacksData { 21 | active_stack: Stack; 22 | stacks: Stack[]; 23 | total: number; 24 | total_pages: number; 25 | current_page: number; 26 | items_per_page: number; 27 | } 28 | 29 | interface Stack { 30 | id: string; 31 | name: string; 32 | components: Components; 33 | } 34 | 35 | interface Components { 36 | [componentType: string]: StackComponent[]; 37 | } 38 | 39 | interface StackComponent { 40 | id: string; 41 | name: string; 42 | flavor: Flavor; 43 | type: string; 44 | config: { [key: string]: any }; 45 | } 46 | 47 | export type StacksResponse = StacksData | ErrorMessageResponse | VersionMismatchError; 48 | 49 | interface ComponentsListData { 50 | index: number; 51 | max_size: number; 52 | total_pages: number; 53 | total: number; 54 | items: Array; 55 | } 56 | 57 | export type ComponentsListResponse = 58 | | ComponentsListData 59 | | ErrorMessageResponse 60 | | VersionMismatchError; 61 | 62 | interface Flavor { 63 | id: string; 64 | name: string; 65 | type: string; 66 | integration: string | null; 67 | source: string | null; 68 | logo_url: string; 69 | config_schema: { [key: string]: any }; 70 | docs_url: string | null; 71 | sdk_docs_url: string | null; 72 | connector_type: string | null; 73 | connector_resource_type: string | null; 74 | connector_resource_id_attr: string | null; 75 | created: string | null; 76 | updated: string | null; 77 | is_custom: boolean; 78 | } 79 | 80 | interface FlavorListData { 81 | index: number; 82 | max_size: number; 83 | total_pages: number; 84 | total: number; 85 | items: Flavor[]; 86 | } 87 | 88 | export type FlavorListResponse = FlavorListData | ErrorMessageResponse | VersionMismatchError; 89 | 90 | type ComponentTypes = string[]; 91 | 92 | export type ComponentTypesResponse = ComponentTypes | VersionMismatchError | ErrorMessageResponse; 93 | 94 | export { 95 | Components, 96 | ComponentsListData, 97 | ComponentTypes, 98 | Flavor, 99 | Stack, 100 | StackComponent, 101 | StacksData, 102 | }; 103 | -------------------------------------------------------------------------------- /src/utils/componentUtils.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import { Flavor } from '../types/StackTypes'; 15 | 16 | /** 17 | * Format flavor information into a nicely formatted markdown string for tooltips 18 | */ 19 | export function formatFlavorTooltip(flavor: Flavor | string): string { 20 | if (!flavor) { 21 | return ''; 22 | } 23 | 24 | if (typeof flavor === 'string') { 25 | return flavor; 26 | } 27 | 28 | const lines = []; 29 | 30 | if (flavor.name) { 31 | lines.push(`  name: ${flavor.name}`); 32 | } 33 | 34 | if (flavor.integration) { 35 | lines.push(`  integration: ${flavor.integration}`); 36 | } 37 | 38 | if (flavor.type) { 39 | lines.push(`  type: ${flavor.type}`); 40 | } 41 | 42 | const created = flavor.created; 43 | if (created) { 44 | lines.push(`  created: ${created}`); 45 | } 46 | 47 | const updated = flavor.updated; 48 | if (updated) { 49 | lines.push(`  updated: ${updated}`); 50 | } 51 | 52 | const sourceFromConfigSchema = flavor.config_schema?.title; 53 | const sourceFromBody = flavor.source?.split('.').pop() || flavor.source; 54 | const source = sourceFromConfigSchema || sourceFromBody; 55 | 56 | if (source) { 57 | lines.push(`  source: ${source}`); 58 | } 59 | 60 | // Join with line breaks + spaces for proper markdown rendering 61 | return lines.join(' \n'); 62 | } 63 | -------------------------------------------------------------------------------- /src/utils/constants.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { ServerStatus } from '../types/ServerInfoTypes'; 14 | 15 | export const PYTOOL_MODULE = 'zenml-python'; 16 | export const PYTOOL_DISPLAY_NAME = 'ZenML'; 17 | export const LANGUAGE_SERVER_NAME = 'zen-language-server'; 18 | export const MIN_ZENML_VERSION = '0.80.0'; 19 | export const ZENML_EMOJI = '⛩️'; 20 | 21 | export const ZENML_PYPI_URL = 'https://pypi.org/pypi/zenml/json'; 22 | export const DEFAULT_LOCAL_ZENML_SERVER_URL = 'http://127.0.0.1:8237'; 23 | export const BASE_CLOUD_ZENML_SERVER_URL = 'https://cloudapi.zenml.io'; 24 | 25 | // LSP server notifications 26 | export const LSP_IS_ZENML_INSTALLED = 'zenml/isInstalled'; 27 | export const LSP_ZENML_CLIENT_INITIALIZED = 'zenml/clientInitialized'; 28 | export const LSP_ZENML_SERVER_CHANGED = 'zenml/serverChanged'; 29 | export const LSP_ZENML_STACK_CHANGED = 'zenml/stackChanged'; 30 | export const LSP_ZENML_PROJECT_CHANGED = 'zenml/projectChanged'; 31 | export const LSP_ZENML_REQUIREMENTS_NOT_MET = 'zenml/requirementsNotMet'; 32 | 33 | // EventBus emitted events 34 | export const LSCLIENT_READY = 'lsClientReady'; 35 | export const LSCLIENT_STATE_CHANGED = 'lsClientStateChanged'; 36 | export const ZENML_CLIENT_STATE_CHANGED = 'zenmlClientStateChanged'; 37 | 38 | export const REFRESH_ENVIRONMENT_VIEW = 'refreshEnvironmentView'; 39 | 40 | export const SERVER_STATUS_UPDATED = 'serverStatusUpdated'; 41 | export const ITEMS_PER_PAGE_OPTIONS = ['5', '10', '15', '20', '25', '30', '35', '40', '45', '50']; 42 | 43 | export const INITIAL_ZENML_SERVER_STATUS: ServerStatus = { 44 | isConnected: false, 45 | url: '', 46 | dashboard_url: '', 47 | store_type: '', 48 | deployment_type: '', 49 | version: '', 50 | debug: false, 51 | database_type: '', 52 | secrets_store_type: '', 53 | username: null, 54 | }; 55 | -------------------------------------------------------------------------------- /src/utils/notifications.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as vscode from 'vscode'; 14 | 15 | /** 16 | * Shows an information message in the status bar for a specified duration. 17 | * 18 | * @param message The message to show. 19 | * @param duration Duration in milliseconds after which the message will disappear. 20 | */ 21 | export const showStatusBarInfoMessage = (message: string, duration: number = 5000): void => { 22 | const disposable = vscode.window.setStatusBarMessage(message); 23 | setTimeout(() => disposable.dispose(), duration); 24 | }; 25 | 26 | /** 27 | * Shows a warning message in the status bar for a specified duration. 28 | * 29 | * @param message The message to show. 30 | * @param duration Duration in milliseconds after which the message will disappear. 31 | */ 32 | export const showStatusBarWarningMessage = (message: string, duration: number = 5000): void => { 33 | const disposable = vscode.window.setStatusBarMessage(`$(alert) ${message}`); 34 | setTimeout(() => disposable.dispose(), duration); 35 | }; 36 | 37 | /** 38 | * Shows an error message in the status bar for a specified duration. 39 | * 40 | * @param message The message to show. 41 | * @param duration Duration in milliseconds after which the message will disappear. 42 | */ 43 | export const showStatusBarErrorMessage = (message: string, duration: number = 5000): void => { 44 | const disposable = vscode.window.setStatusBarMessage(`$(error) ${message}`); 45 | setTimeout(() => disposable.dispose(), duration); 46 | }; 47 | 48 | /** 49 | * Shows a modal pop up information message. 50 | * 51 | * @param message The message to display. 52 | */ 53 | export const showInformationMessage = (message: string): void => { 54 | vscode.window.showInformationMessage(message); 55 | }; 56 | 57 | /** 58 | * Shows a modal pop up error message, 59 | * 60 | * @param message The message to display. 61 | */ 62 | export const showErrorMessage = (message: string): void => { 63 | vscode.window.showErrorMessage(message); 64 | }; 65 | 66 | /** 67 | * Shows a warning message with actions (buttons) for the user to select. 68 | * 69 | * @param message The warning message to display. 70 | * @param actions An array of actions, each action being an object with a title and an action callback. 71 | */ 72 | export async function showWarningMessageWithActions( 73 | message: string, 74 | ...actions: { title: string; action: () => void | Promise }[] 75 | ): Promise { 76 | // Map actions to their titles to display as buttons. 77 | const items = actions.map(action => action.title); 78 | 79 | // Show warning message with buttons. 80 | const selection = await vscode.window.showWarningMessage(message, ...items); 81 | 82 | // Find the selected action based on the title and execute its callback. 83 | const selectedAction = actions.find(action => action.title === selection); 84 | if (selectedAction) { 85 | await selectedAction.action(); 86 | } 87 | } 88 | -------------------------------------------------------------------------------- /src/views/activityBar/common/LoadingTreeItem.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { ThemeIcon, TreeItem, TreeItemCollapsibleState } from 'vscode'; 14 | 15 | export class LoadingTreeItem extends TreeItem { 16 | constructor(message: string, description: string = 'Refreshing...') { 17 | super(message, TreeItemCollapsibleState.None); 18 | this.description = description; 19 | this.iconPath = new ThemeIcon('sync~spin'); 20 | } 21 | } 22 | 23 | export const LOADING_TREE_ITEMS = new Map([ 24 | ['server', new LoadingTreeItem('Refreshing Server View...')], 25 | ['stacks', new LoadingTreeItem('Refreshing Stacks View...')], 26 | ['components', new LoadingTreeItem('Refreshing Components View...')], 27 | ['pipelineRuns', new LoadingTreeItem('Refreshing Pipeline Runs...')], 28 | ['environment', new LoadingTreeItem('Refreshing Environments...')], 29 | ['projects', new LoadingTreeItem('Refreshing Projects View...')], 30 | ['models', new LoadingTreeItem('Refreshing Models...')], 31 | ['lsClient', new LoadingTreeItem('Waiting for Language Server to start...', '')], 32 | ['zenmlClient', new LoadingTreeItem('Waiting for ZenML Client to initialize...', '')], 33 | ]); 34 | -------------------------------------------------------------------------------- /src/views/activityBar/common/PaginationTreeItems.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { ThemeIcon, TreeItem, TreeItemCollapsibleState } from 'vscode'; 14 | 15 | /** 16 | * A TreeItem for displaying pagination in the VSCode TreeView. 17 | */ 18 | export class CommandTreeItem extends TreeItem { 19 | constructor( 20 | public readonly label: string, 21 | commandId: string, 22 | commandArguments?: any[], 23 | icon?: string 24 | ) { 25 | super(label); 26 | this.command = { 27 | title: label, 28 | command: commandId, 29 | arguments: commandArguments, 30 | }; 31 | if (icon) { 32 | this.iconPath = new ThemeIcon(icon); 33 | } 34 | } 35 | } 36 | 37 | export class SetItemsPerPageTreeItem extends TreeItem { 38 | constructor() { 39 | super('Set items per page', TreeItemCollapsibleState.None); 40 | this.tooltip = 'Click to set the number of items shown per page'; 41 | this.command = { 42 | command: 'zenml.setStacksPerPage', 43 | title: 'Set Stack Items Per Page', 44 | arguments: [], 45 | }; 46 | } 47 | } 48 | -------------------------------------------------------------------------------- /src/views/activityBar/common/TreeItemWithChildren.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import * as vscode from 'vscode'; 15 | 16 | /** 17 | * Interface for tree items that can have children 18 | */ 19 | export interface TreeItemWithChildren extends vscode.TreeItem { 20 | children?: vscode.TreeItem[]; 21 | } 22 | -------------------------------------------------------------------------------- /src/views/activityBar/common/index.ts: -------------------------------------------------------------------------------- 1 | export * from './ErrorTreeItem'; 2 | export * from './LoadingTreeItem'; 3 | export * from './PaginatedDataProvider'; 4 | export * from './TreeItemWithChildren'; 5 | -------------------------------------------------------------------------------- /src/views/activityBar/componentView/ComponentTreeItems.ts: -------------------------------------------------------------------------------- 1 | import * as vscode from 'vscode'; 2 | import { Flavor, StackComponent } from '../../../types/StackTypes'; 3 | import { formatFlavorTooltip } from '../../../utils/componentUtils'; 4 | import { CONTEXT_VALUES, TREE_ICONS } from '../../../utils/ui-constants'; 5 | import { TreeItemWithChildren } from '../common/TreeItemWithChildren'; 6 | 7 | /** 8 | * TreeItem for grouping components by type 9 | */ 10 | export class ComponentCategoryTreeItem extends vscode.TreeItem implements TreeItemWithChildren { 11 | public children?: vscode.TreeItem[]; 12 | 13 | constructor( 14 | public readonly type: string, 15 | public readonly childComponents: ComponentTreeItem[] 16 | ) { 17 | super(type, vscode.TreeItemCollapsibleState.Collapsed); 18 | this.contextValue = CONTEXT_VALUES.COMPONENT_CATEGORY; 19 | this.iconPath = TREE_ICONS.COMPONENT_CATEGORY; 20 | this.tooltip = `Component Type: ${type}`; 21 | this.children = childComponents; 22 | } 23 | } 24 | 25 | /** 26 | * Formats a component detail attribute for display 27 | */ 28 | export class ComponentDetailTreeItem extends vscode.TreeItem { 29 | constructor(label: string, value: string) { 30 | super(label, vscode.TreeItemCollapsibleState.None); 31 | this.description = value; 32 | this.contextValue = CONTEXT_VALUES.COMPONENT_DETAIL; 33 | this.tooltip = ''; 34 | } 35 | } 36 | 37 | /** 38 | * A TreeItem for displaying a component in the VSCode TreeView. 39 | */ 40 | export class ComponentTreeItem extends vscode.TreeItem implements TreeItemWithChildren { 41 | public children?: vscode.TreeItem[]; 42 | 43 | constructor( 44 | public component: StackComponent, 45 | public parentId?: string 46 | ) { 47 | super(component.name, vscode.TreeItemCollapsibleState.Collapsed); 48 | 49 | const flavorTooltip = component.flavor 50 | ? formatFlavorTooltip(component.flavor) 51 | : component.flavor || ''; 52 | 53 | this.tooltip = new vscode.MarkdownString(); 54 | this.tooltip.appendMarkdown(`**Component: ${component.name}**\n\n`); 55 | this.tooltip.appendMarkdown(`**Type:** ${component.type}\n\n`); 56 | this.tooltip.appendMarkdown(`**Flavor:** \n${flavorTooltip}\n\n`); 57 | this.tooltip.appendMarkdown(`**ID:** ${component.id}`); 58 | if (parentId) { 59 | this.tooltip.appendMarkdown(`\n\n**Stack ID:** ${parentId}`); 60 | } 61 | 62 | this.contextValue = CONTEXT_VALUES.COMPONENT; 63 | this.id = parentId ? `${parentId}-${component.id}` : `${component.id}`; 64 | this.iconPath = TREE_ICONS.COMPONENT; 65 | 66 | this.children = this.createDetailItems(parentId); 67 | } 68 | 69 | /** 70 | * Creates detail items for the component's attributes 71 | */ 72 | protected createDetailItems(parentId?: string): vscode.TreeItem[] { 73 | const details: vscode.TreeItem[] = []; 74 | 75 | details.push(new ComponentDetailTreeItem('name', this.component.name)); 76 | details.push(new ComponentDetailTreeItem('type', this.component.type)); 77 | const flavor = this.component.flavor as Flavor | string; 78 | if (flavor) { 79 | if (typeof flavor === 'string') { 80 | details.push(new ComponentDetailTreeItem('flavor', flavor)); 81 | } else { 82 | if (flavor.name) { 83 | details.push(new ComponentDetailTreeItem('flavor', flavor.name)); 84 | } 85 | if (flavor.integration) { 86 | details.push(new ComponentDetailTreeItem('integration', flavor.integration)); 87 | } 88 | } 89 | } 90 | 91 | details.push(new ComponentDetailTreeItem('id', this.component.id)); 92 | 93 | if (parentId) { 94 | details.push(new ComponentDetailTreeItem('stack_id', parentId)); 95 | } 96 | 97 | return details; 98 | } 99 | } 100 | 101 | /** 102 | * Specialized ComponentTreeItem for stack components 103 | */ 104 | export class StackComponentTreeItem extends ComponentTreeItem { 105 | constructor( 106 | public component: StackComponent, 107 | public stackId?: string 108 | ) { 109 | super(component, stackId); 110 | this.contextValue = CONTEXT_VALUES.STACK_COMPONENT; 111 | } 112 | } 113 | -------------------------------------------------------------------------------- /src/views/activityBar/environmentView/EnvironmentItem.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as path from 'path'; 14 | import { ThemeColor, ThemeIcon, TreeItem, TreeItemCollapsibleState, Uri } from 'vscode'; 15 | 16 | export class EnvironmentItem extends TreeItem { 17 | pythonLogo = path.join(__dirname, '..', 'resources', 'python.png'); 18 | zenmlLogo = path.join(__dirname, '..', 'resources', 'logo.png'); 19 | 20 | constructor( 21 | public readonly label: string, 22 | public readonly description?: string, 23 | public readonly collapsibleState: TreeItemCollapsibleState = TreeItemCollapsibleState.None, 24 | private readonly customIcon?: string, 25 | public readonly contextValue?: string, 26 | public readonly customTooltip?: string 27 | ) { 28 | super(label, collapsibleState); 29 | this.iconPath = this.determineIcon(label); 30 | this.contextValue = contextValue; 31 | 32 | // Use custom tooltip if provided, otherwise use description for tooltip if it's long enough to be truncated 33 | if (customTooltip) { 34 | this.tooltip = customTooltip; 35 | } else if (description && description.length > 60) { 36 | this.tooltip = description; 37 | } 38 | } 39 | 40 | /** 41 | * Determines the icon for the tree item based on the label. 42 | * 43 | * @param label The label of the tree item. 44 | * @returns The icon for the tree item. 45 | */ 46 | private determineIcon( 47 | label: string 48 | ): string | Uri | { light: Uri; dark: Uri } | ThemeIcon | undefined { 49 | if (this.customIcon) { 50 | switch (this.customIcon) { 51 | case 'check': 52 | return new ThemeIcon('check', new ThemeColor('gitDecoration.addedResourceForeground')); 53 | case 'close': 54 | return new ThemeIcon('close', new ThemeColor('gitDecoration.deletedResourceForeground')); 55 | case 'error': 56 | return new ThemeIcon('error', new ThemeColor('errorForeground')); 57 | case 'warning': 58 | return new ThemeIcon('warning', new ThemeColor('charts.yellow')); 59 | default: 60 | return new ThemeIcon(this.customIcon); 61 | } 62 | } 63 | switch (label) { 64 | case 'Workspace': 65 | case 'CWD': 66 | case 'File System': 67 | return new ThemeIcon('folder'); 68 | case 'Interpreter': 69 | case 'Name': 70 | case 'Python Version': 71 | case 'Path': 72 | case 'EnvType': 73 | return { 74 | light: Uri.file(this.pythonLogo), 75 | dark: Uri.file(this.pythonLogo), 76 | }; 77 | case 'ZenML Local': 78 | case 'ZenML Client': 79 | return { 80 | light: Uri.file(this.zenmlLogo), 81 | dark: Uri.file(this.zenmlLogo), 82 | }; 83 | default: 84 | return undefined; 85 | } 86 | } 87 | } 88 | -------------------------------------------------------------------------------- /src/views/activityBar/index.ts: -------------------------------------------------------------------------------- 1 | // /src/views/treeViews/index.ts 2 | 3 | // Component View 4 | export * from './componentView/ComponentDataProvider'; 5 | export * from './componentView/ComponentTreeItems'; 6 | 7 | // Pipeline View 8 | export * from './pipelineView/PipelineDataProvider'; 9 | export * from './pipelineView/PipelineTreeItems'; 10 | 11 | // Project View 12 | export * from './projectView/ProjectDataProvider'; 13 | export * from './projectView/ProjectTreeItems'; 14 | 15 | // Server View 16 | export * from './serverView/ServerDataProvider'; 17 | export * from './serverView/ServerTreeItems'; 18 | 19 | // Stack View 20 | export * from './stackView/StackDataProvider'; 21 | export * from './stackView/StackTreeItems'; 22 | 23 | // Model View 24 | export * from './modelView/ModelDataProvider'; 25 | export * from './modelView/ModelTreeItems'; 26 | 27 | // Environment View 28 | export * from './environmentView/EnvironmentDataProvider'; 29 | -------------------------------------------------------------------------------- /src/views/activityBar/modelView/ModelTreeItems.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | 14 | import * as vscode from 'vscode'; 15 | import { Model, ModelVersion } from '../../../types/ModelTypes'; 16 | import { 17 | CONTEXT_VALUES, 18 | MODEL_VERSION_SECTION_ICONS, 19 | MODEL_VERSION_STATUS_ICONS, 20 | TREE_ICONS, 21 | } from '../../../utils/ui-constants'; 22 | import { TreeItemWithChildren } from '../common/TreeItemWithChildren'; 23 | 24 | /** 25 | * Represents a Model Tree Item in the VS Code tree view. 26 | * Displays information about a model in the model registry. 27 | */ 28 | export class ModelTreeItem extends vscode.TreeItem implements TreeItemWithChildren { 29 | constructor( 30 | public readonly model: Model, 31 | public readonly id: string, 32 | public readonly description: string = '', 33 | public readonly contextValue: string = CONTEXT_VALUES.MODEL 34 | ) { 35 | super(model.name, vscode.TreeItemCollapsibleState.Collapsed); 36 | this.tooltip = `Model: ${model.name}${model.latest_version_name ? ` (Latest version: ${model.latest_version_name})` : ''}`; 37 | this.description = description; 38 | this.iconPath = TREE_ICONS.MODEL; 39 | } 40 | } 41 | 42 | /** 43 | * Represents a Model Version Tree Item in the VS Code tree view. 44 | * Displays information about a specific version of a model. 45 | */ 46 | export class ModelVersionTreeItem extends vscode.TreeItem implements TreeItemWithChildren { 47 | constructor( 48 | public readonly version: ModelVersion, 49 | public readonly description: string = '', 50 | public readonly contextValue: string = CONTEXT_VALUES.MODEL_VERSION 51 | ) { 52 | super(`${version.name}`, vscode.TreeItemCollapsibleState.Collapsed); 53 | 54 | // For description, show version number and stage if available 55 | const stage = (version.stage ?? 'Not set').toLowerCase(); 56 | this.description = stage !== 'not set' ? stage : ''; 57 | this.iconPath = MODEL_VERSION_STATUS_ICONS[stage] ?? TREE_ICONS.MODEL_VERSION; 58 | 59 | this.tooltip = `Model Version: ${version.name} 60 | Number: ${version.number} 61 | Stage: ${version.stage ?? 'Not set'} 62 | Created: ${new Date(version.created).toLocaleString()} 63 | Updated: ${new Date(version.updated).toLocaleString()}`; 64 | } 65 | } 66 | 67 | /** 68 | * Represents a detail item for a model or model version in the tree view. 69 | */ 70 | export class ModelDetailTreeItem extends vscode.TreeItem { 71 | public children?: ModelDetailTreeItem[]; 72 | 73 | constructor( 74 | public readonly label: string, 75 | public readonly description: string, 76 | collapsibleState: vscode.TreeItemCollapsibleState = vscode.TreeItemCollapsibleState.None, 77 | public readonly contextValue: string = CONTEXT_VALUES.MODEL_DETAIL 78 | ) { 79 | super(label, collapsibleState); 80 | this.description = description; 81 | this.contextValue = contextValue; 82 | this.tooltip = `${label}: ${description}`; 83 | } 84 | } 85 | 86 | /** 87 | * Represents a section for grouping model version details. 88 | */ 89 | export class ModelSectionTreeItem extends vscode.TreeItem implements TreeItemWithChildren { 90 | public children?: vscode.TreeItem[]; 91 | 92 | constructor( 93 | public readonly label: string, 94 | children: vscode.TreeItem[] = [] 95 | ) { 96 | super(label, vscode.TreeItemCollapsibleState.Collapsed); 97 | this.children = children; 98 | this.contextValue = CONTEXT_VALUES.MODEL_SECTION; 99 | this.iconPath = MODEL_VERSION_SECTION_ICONS[label]; 100 | } 101 | } 102 | -------------------------------------------------------------------------------- /src/views/activityBar/pipelineView/PipelineTreeItems.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as vscode from 'vscode'; 14 | import { PipelineRun } from '../../../types/PipelineTypes'; 15 | import { CONTEXT_VALUES, PIPELINE_RUN_STATUS_ICONS } from '../../../utils/ui-constants'; 16 | 17 | /** 18 | * Represents a Pipeline Run Tree Item in the VS Code tree view. 19 | * Displays its name, version and status. 20 | */ 21 | export class PipelineTreeItem extends vscode.TreeItem { 22 | public children: PipelineRunTreeItem[] | undefined; 23 | 24 | constructor( 25 | public readonly run: PipelineRun, 26 | public readonly id: string, 27 | children?: PipelineRunTreeItem[] 28 | ) { 29 | super( 30 | run.name, 31 | children === undefined 32 | ? vscode.TreeItemCollapsibleState.None 33 | : vscode.TreeItemCollapsibleState.Collapsed 34 | ); 35 | // Tooltip removed for demo recording (would contain ID) 36 | this.tooltip = ''; 37 | this.iconPath = PIPELINE_RUN_STATUS_ICONS[run.status]; 38 | this.children = children; 39 | } 40 | 41 | contextValue = CONTEXT_VALUES.PIPELINE_RUN; 42 | } 43 | 44 | /** 45 | * Represents details of a Pipeline Run Tree Item in the VS Code tree view. 46 | * Displays the stack name for the run, its start time, end time, machine details, Python version, and more. 47 | */ 48 | export class PipelineRunTreeItem extends vscode.TreeItem { 49 | public children?: PipelineRunTreeItem[]; 50 | 51 | constructor( 52 | public readonly label: string, 53 | public readonly description: string, 54 | collapsibleState: vscode.TreeItemCollapsibleState = vscode.TreeItemCollapsibleState.None 55 | ) { 56 | super(label, collapsibleState); 57 | this.tooltip = ''; 58 | } 59 | 60 | contextValue = CONTEXT_VALUES.PIPELINE_RUN_DETAIL; 61 | } 62 | -------------------------------------------------------------------------------- /src/views/activityBar/projectView/ProjectTreeItems.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import * as vscode from 'vscode'; 14 | import { Project } from '../../../types/ProjectTypes'; 15 | import { CONTEXT_VALUES, TREE_ICONS } from '../../../utils/ui-constants'; 16 | import { TreeItemWithChildren } from '../common/TreeItemWithChildren'; 17 | 18 | /** 19 | * A TreeItem for displaying a project detail in the VSCode TreeView. 20 | */ 21 | export class ProjectDetailItem extends vscode.TreeItem { 22 | constructor( 23 | public readonly label: string, 24 | public readonly detail: string 25 | ) { 26 | super(label, vscode.TreeItemCollapsibleState.None); 27 | this.description = detail; 28 | this.contextValue = CONTEXT_VALUES.PROJECT_DETAIL; 29 | } 30 | } 31 | 32 | /** 33 | * A TreeItem for displaying a project in the VSCode TreeView. 34 | */ 35 | export class ProjectTreeItem extends vscode.TreeItem implements TreeItemWithChildren { 36 | public isActive: boolean; 37 | public children?: vscode.TreeItem[]; 38 | 39 | constructor( 40 | public readonly project: Project, 41 | public readonly name: string, 42 | isActive?: boolean 43 | ) { 44 | super(project.name, vscode.TreeItemCollapsibleState.Collapsed); 45 | 46 | this.contextValue = CONTEXT_VALUES.PROJECT; 47 | this.isActive = isActive || false; 48 | 49 | if (isActive) { 50 | this.iconPath = TREE_ICONS.ACTIVE_PROJECT; 51 | this.description = 'Active'; 52 | } else { 53 | this.iconPath = TREE_ICONS.PROJECT; 54 | this.description = ''; 55 | } 56 | 57 | this.id = this.project.id; 58 | 59 | this.updateChildren(); 60 | } 61 | 62 | /** 63 | * Updates the children items when active status changes. 64 | */ 65 | public updateChildren(): void { 66 | const createdOn = this.project.created 67 | ? new Date(this.project.created).toLocaleString() 68 | : 'N/A'; 69 | const updatedOn = this.project.updated 70 | ? new Date(this.project.updated).toLocaleString() 71 | : 'N/A'; 72 | 73 | this.children = [ 74 | new ProjectDetailItem('id', this.project.id), 75 | new ProjectDetailItem('name', this.project.name), 76 | new ProjectDetailItem('display_name', this.project.display_name || 'N/A'), 77 | new ProjectDetailItem('created', createdOn), 78 | new ProjectDetailItem('updated', updatedOn), 79 | ]; 80 | } 81 | } 82 | -------------------------------------------------------------------------------- /src/views/panel/panelView/PanelDataProvider.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { EventEmitter, TreeDataProvider, TreeItem } from 'vscode'; 14 | import { 15 | JsonObject, 16 | PanelDetailTreeItem, 17 | PanelTreeItem, 18 | SourceCodeTreeItem, 19 | } from './PanelTreeItem'; 20 | 21 | import { LoadingTreeItem } from '../../activityBar/common/LoadingTreeItem'; 22 | 23 | export class PanelDataProvider implements TreeDataProvider { 24 | private _onDidChangeTreeData = new EventEmitter(); 25 | readonly onDidChangeTreeData = this._onDidChangeTreeData.event; 26 | 27 | private static instance: PanelDataProvider | null = null; 28 | private data: JsonObject | TreeItem = new TreeItem( 29 | 'No data has been requested for visualization yet' 30 | ); 31 | private dataType: string = ''; 32 | 33 | /** 34 | * Retrieves the singleton instance of PanelDataProvider 35 | * @returns {PanelDataProvider} The singleton instance 36 | */ 37 | public static getInstance(): PanelDataProvider { 38 | if (!PanelDataProvider.instance) { 39 | PanelDataProvider.instance = new PanelDataProvider(); 40 | } 41 | 42 | return PanelDataProvider.instance; 43 | } 44 | 45 | /** 46 | * Refreshes the ZenML Panel View 47 | */ 48 | public refresh(): void { 49 | this._onDidChangeTreeData.fire(undefined); 50 | } 51 | 52 | /** 53 | * Sets the data to be viewed in the ZenML Panel View 54 | * @param data Data to visualize 55 | * @param dataType Type of data being visualized 56 | */ 57 | public setData(data: JsonObject, dataType = 'data'): void { 58 | this.data = data; 59 | this.dataType = dataType; 60 | this.refresh(); 61 | } 62 | 63 | public setLoading(): void { 64 | this.data = new LoadingTreeItem('Retrieving data'); 65 | this.refresh(); 66 | } 67 | 68 | /** 69 | * Retrieves the tree item for a given data property 70 | * 71 | * @param element The data property. 72 | * @returns The corresponding VS Code tree item. 73 | */ 74 | public getTreeItem(element: TreeItem): TreeItem | Thenable { 75 | return element; 76 | } 77 | 78 | /** 79 | * Retrieves the children for a given tree item. 80 | * 81 | * @param element The parent tree item. If undefined, a PanelTreeItem is created 82 | * @returns An array of child tree items or undefined if there are no children. 83 | */ 84 | public getChildren(element?: TreeItem | undefined): TreeItem[] | undefined { 85 | if (element) { 86 | if ( 87 | element instanceof PanelTreeItem || 88 | element instanceof PanelDetailTreeItem || 89 | element instanceof SourceCodeTreeItem 90 | ) { 91 | return element.children; 92 | } 93 | 94 | return undefined; 95 | } 96 | 97 | if (this.data instanceof TreeItem) { 98 | return [this.data]; 99 | } 100 | 101 | return [new PanelTreeItem(this.dataType, this.data)]; 102 | } 103 | } 104 | -------------------------------------------------------------------------------- /src/views/panel/panelView/PanelTreeItem.ts: -------------------------------------------------------------------------------- 1 | // Copyright(c) ZenML GmbH 2024. All Rights Reserved. 2 | // Licensed under the Apache License, Version 2.0(the "License"); 3 | // you may not use this file except in compliance with the License. 4 | // You may obtain a copy of the License at: 5 | // 6 | // http://www.apache.org/licenses/LICENSE-2.0 7 | // 8 | // Unless required by applicable law or agreed to in writing, software 9 | // distributed under the License is distributed on an "AS IS" BASIS, 10 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express 11 | // or implied. See the License for the specific language governing 12 | // permissions and limitations under the License. 13 | import { ThemeColor, ThemeIcon, TreeItem, TreeItemCollapsibleState, Uri } from 'vscode'; 14 | 15 | type JsonType = string | number | Array | JsonObject; 16 | 17 | export interface JsonObject { 18 | [key: string]: JsonType; 19 | } 20 | 21 | export class PanelDetailTreeItem extends TreeItem { 22 | public children: PanelDetailTreeItem[] = []; 23 | 24 | /** 25 | * Constructs a PanelDetailTreeItem object 26 | * @param key Property key for TreeItem 27 | * @param value Property value for the TreeItem 28 | */ 29 | constructor(key: string, value: JsonType) { 30 | const simpleValue = typeof value === 'string' || typeof value === 'number'; 31 | super(key, simpleValue ? TreeItemCollapsibleState.None : TreeItemCollapsibleState.Collapsed); 32 | if (simpleValue) { 33 | this.description = String(value); 34 | } else if (value) { 35 | this.description = '...'; 36 | this.children = Object.entries(value).map( 37 | ([key, value]) => new PanelDetailTreeItem(key, value) 38 | ); 39 | } 40 | 41 | if (typeof value === 'string' && value.startsWith('http')) { 42 | this.command = { 43 | title: 'Open URL', 44 | command: 'vscode.open', 45 | arguments: [Uri.parse(value)], 46 | }; 47 | this.iconPath = new ThemeIcon('link', new ThemeColor('textLink.foreground')); 48 | this.tooltip = `Click to open ${value}`; 49 | } 50 | } 51 | } 52 | 53 | export class PanelTreeItem extends TreeItem { 54 | public children: Array = []; 55 | 56 | /** 57 | * Constructs a PanelTreeItem 58 | * @param label Data Type Label for the PanelTreeItem 59 | * @param data Object Data to build children 60 | */ 61 | constructor(label: string, data: JsonObject) { 62 | super(label, TreeItemCollapsibleState.Expanded); 63 | this.children = Object.entries(data).map(([key, value]) => { 64 | if (key === 'sourceCode' && typeof value === 'string') { 65 | return new SourceCodeTreeItem(key, value); 66 | } 67 | return new PanelDetailTreeItem(key, value); 68 | }); 69 | } 70 | } 71 | 72 | export class SourceCodeTreeItem extends TreeItem { 73 | public children: TreeItem[] = []; 74 | 75 | /** 76 | * Constructs a SourceCodeTreeItem that builds its childrens based on string passed to it 77 | * @param label Property Label for parent object 78 | * @param sourceCode Raw string of source code 79 | */ 80 | constructor(label: string, sourceCode: string) { 81 | super(label, TreeItemCollapsibleState.Collapsed); 82 | this.description = '...'; 83 | 84 | const lines = sourceCode.split('\n'); 85 | this.children = lines.map(line => new TreeItem(line)); 86 | } 87 | } 88 | -------------------------------------------------------------------------------- /tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "module": "Node16", 4 | "target": "ES2022", 5 | "lib": ["ES2022"], 6 | "sourceMap": true, 7 | "rootDir": "src", 8 | "strict": true, 9 | "skipLibCheck": true, 10 | "esModuleInterop": true 11 | }, 12 | "include": ["src/**/*"] 13 | } 14 | -------------------------------------------------------------------------------- /webpack.config.mjs: -------------------------------------------------------------------------------- 1 | //@ts-check 2 | 3 | import path from 'path'; 4 | import { fileURLToPath } from 'url'; 5 | 6 | // ES module equivalents for __dirname and __filename 7 | const __filename = fileURLToPath(import.meta.url); 8 | const __dirname = path.dirname(__filename); 9 | 10 | //@ts-check 11 | /** @typedef {import('webpack').Configuration} WebpackConfig **/ 12 | 13 | /** @type WebpackConfig */ 14 | const extensionConfig = { 15 | target: 'node', // VS Code extensions run in a Node.js-context 📖 -> https://webpack.js.org/configuration/node/ 16 | mode: 'none', // this leaves the source code as close as possible to the original (when packaging we set this to 'production') 17 | 18 | entry: './src/extension.ts', // the entry point of this extension, 📖 -> https://webpack.js.org/configuration/entry-context/ 19 | output: { 20 | // the bundle is stored in the 'dist' folder (check package.json), 📖 -> https://webpack.js.org/configuration/output/ 21 | path: path.resolve(__dirname, 'dist'), 22 | filename: 'extension.js', 23 | libraryTarget: 'commonjs2', 24 | }, 25 | externals: { 26 | 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/ 27 | // modules added here also need to be added in the .vscodeignore file 28 | }, 29 | resolve: { 30 | // support reading TypeScript and JavaScript files, 📖 -> https://github.com/TypeStrong/ts-loader 31 | extensions: ['.ts', '.js'], 32 | }, 33 | module: { 34 | rules: [ 35 | { 36 | test: /\.ts$/, 37 | exclude: /node_modules/, 38 | use: [ 39 | { 40 | loader: 'ts-loader', 41 | }, 42 | ], 43 | }, 44 | ], 45 | }, 46 | devtool: 'nosources-source-map', 47 | infrastructureLogging: { 48 | level: 'log', // enables logging required for problem matchers 49 | }, 50 | }; 51 | 52 | const dagWebviewConfig = { 53 | target: 'web', 54 | mode: 'none', 55 | entry: './resources/dag-view/dag.js', 56 | output: { 57 | path: path.resolve(__dirname, 'resources', 'dag-view'), 58 | filename: 'dag-packed.js', 59 | }, 60 | }; 61 | 62 | export default [extensionConfig, dagWebviewConfig]; 63 | --------------------------------------------------------------------------------