Merge pull request #21 from marwan37/general-ui-enhancements

UI Improvements – Status Bar, Icons, and Cancellable Tasks
zenml-io · Apr 5, 2024 · dda2e4c · dda2e4c
2 parents 29dd99d + 277498c
commit dda2e4c
Show file tree

Hide file tree

Showing 36 changed files with 521 additions and 246 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -16,16 +16,18 @@ jobs:
           node-version: '18.x'
 
       - name: Install dependencies
-        run: npm install
+        run: |
+          npm install
+          pip install autoflake isort black ruff
 
       - name: Build VSCode Extension
         run: npm run compile
 
-      - name: Run VSCode Extension Formatter
-        run: npm run format
+      - name: Run Format Script for Python & TypeScript Files
+        run: ./scripts/format.sh
 
-      - name: Run VSCode Extension Linter
-        run: npm run lint
+      - name: Run Lint Script for Python & TypeScript Files
+        run: ./scripts/lint.sh
 
       - name: Run headless test
         uses: coactions/setup-xvfb@v1

diff --git a/.tool-versions b/.tool-versions
@@ -0,0 +1 @@
+nodejs 20.10.0
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -97,4 +97,4 @@ git push origin feature/your-feature-name
 
 - [ZenML VSCode Extension Repository](https://github.com/zenml-io/vscode-zenml)
 - [ZenML Documentation](https://docs.zenml.io)
-- [ZenML Slack Community](https://zenml.io/slack-invite)
+- [ZenML Slack Community](https://zenml.io/slack)
diff --git a/bundled/tool/lsp_jsonrpc.py b/bundled/tool/lsp_jsonrpc.py
@@ -110,11 +110,11 @@ def close(self):
         """Closes the underlying streams."""
         try:
             self._reader.close()
-        except:  # pylint: disable=bare-except
+        except:  # noqa: E722 # pylint: disable=bare-except
             pass
         try:
             self._writer.close()
-        except:  # pylint: disable=bare-except
+        except:  # noqa: E722 # pylint: disable=bare-except
             pass
 
     def send_data(self, data):
@@ -146,7 +146,7 @@ def stop_all_processes(self):
         for i in self._rpc.values():
             try:
                 i.send_data({"id": str(uuid.uuid4()), "method": "exit"})
-            except:  # pylint: disable=bare-except
+            except:  # noqa: E722 # pylint: disable=bare-except
                 pass
         self._thread_pool.shutdown(wait=False)
 
@@ -169,7 +169,7 @@ def _monitor_process():
                     del self._processes[workspace]
                     rpc = self._rpc.pop(workspace)
                     rpc.close()
-                except:  # pylint: disable=bare-except
+                except:  # noqa: E722 # pylint: disable=bare-except
                     pass
 
         self._thread_pool.submit(_monitor_process)

diff --git a/bundled/tool/lsp_runner.py b/bundled/tool/lsp_runner.py
@@ -41,8 +41,8 @@ def update_sys_path(path_to_add: str, strategy: str) -> None:
 
 
 # pylint: disable=wrong-import-position,import-error
-import lsp_jsonrpc as jsonrpc
-import lsp_utils as utils
+import lsp_jsonrpc as jsonrpc  # noqa: E402
+import lsp_utils as utils  # noqa: E402
 
 RPC = jsonrpc.create_json_rpc(sys.stdin.buffer, sys.stdout.buffer)
 

diff --git a/bundled/tool/lsp_server.py b/bundled/tool/lsp_server.py
@@ -45,20 +45,18 @@ def update_sys_path(path_to_add: str, strategy: str) -> None:
 # Imports needed for the language server goes below this.
 # **********************************************************
 # pylint: disable=wrong-import-position,import-error
-import lsp_jsonrpc as jsonrpc
-import lsprotocol.types as lsp
-from lsp_zenml import ZenLanguageServer
-from pygls import uris, workspace
+import lsp_jsonrpc as jsonrpc  # noqa: E402
+import lsprotocol.types as lsp  # noqa: E402
+from lsp_zenml import ZenLanguageServer  # noqa: E402
+from pygls import uris, workspace  # noqa: E402
 
 WORKSPACE_SETTINGS = {}
 GLOBAL_SETTINGS = {}
 RUNNER = pathlib.Path(__file__).parent / "lsp_runner.py"
 
 MAX_WORKERS = 5
 
-LSP_SERVER = ZenLanguageServer(
-    name="zen-language-server", version="0.0.1", max_workers=MAX_WORKERS
-)
+LSP_SERVER = ZenLanguageServer(name="zen-language-server", version="0.0.1", max_workers=MAX_WORKERS)
 
 # **********************************************************
 # Tool specific code goes below this.
@@ -176,9 +174,7 @@ def get_cwd(settings: Dict[str, Any], document: Optional[workspace.Document]) ->
 # *****************************************************
 # Logging and notification.
 # *****************************************************
-def log_to_output(
-    message: str, msg_type: lsp.MessageType = lsp.MessageType.Log
-) -> None:
+def log_to_output(message: str, msg_type: lsp.MessageType = lsp.MessageType.Log) -> None:
     """Log to output."""
     LSP_SERVER.show_message_log(message, msg_type)
 

diff --git a/bundled/tool/lsp_zenml.py b/bundled/tool/lsp_zenml.py
@@ -25,16 +25,14 @@
 from functools import wraps
 
 import lsprotocol.types as lsp
-from constants import MIN_ZENML_VERSION, TOOL_MODULE_NAME, IS_ZENML_INSTALLED
+from constants import IS_ZENML_INSTALLED, MIN_ZENML_VERSION, TOOL_MODULE_NAME
 from lazy_import import suppress_stdout_temporarily
 from packaging.version import parse as parse_version
 from pygls.server import LanguageServer
 from zen_watcher import ZenConfigWatcher
 from zenml_client import ZenMLClient
 
-zenml_init_error = {
-    "error": "ZenML is not initialized. Please check ZenML version requirements."
-}
+zenml_init_error = {"error": "ZenML is not initialized. Please check ZenML version requirements."}
 
 
 class ZenLanguageServer(LanguageServer):
@@ -60,9 +58,7 @@ async def is_zenml_installed(self) -> bool:
             if process.returncode == 0:
                 self.show_message_log("✅ ZenML installation check: Successful.")
                 return True
-            self.show_message_log(
-                "❌ ZenML installation check failed.", lsp.MessageType.Error
-            )
+            self.show_message_log("❌ ZenML installation check failed.", lsp.MessageType.Error)
             return False
         except Exception as e:
             self.show_message_log(
@@ -91,15 +87,13 @@ async def initialize_zenml_client(self):
         self.log_to_output("🚀 Initializing ZenML client...")
         try:
             self.zenml_client = ZenMLClient()
-            self.notify_user("✅ ZenML client initialized successfully.")
+            self.show_message_log("✅ ZenML client initialized successfully.")
             # register pytool module commands
             self.register_commands()
             # initialize watcher
             self.initialize_global_config_watcher()
         except Exception as e:
-            self.notify_user(
-                f"Failed to initialize ZenML client: {str(e)}", lsp.MessageType.Error
-            )
+            self.notify_user(f"Failed to initialize ZenML client: {str(e)}", lsp.MessageType.Error)
 
     def initialize_global_config_watcher(self):
         """Sets up and starts the Global Configuration Watcher."""
@@ -139,9 +133,7 @@ def wrapper(*args, **kwargs):
 
                 with suppress_stdout_temporarily():
                     if wrapper_name:
-                        wrapper_instance = getattr(
-                            self.zenml_client, wrapper_name, None
-                        )
+                        wrapper_instance = getattr(self.zenml_client, wrapper_name, None)
                         if not wrapper_instance:
                             return {"error": f"Wrapper '{wrapper_name}' not found."}
                         return func(wrapper_instance, *args, **kwargs)
@@ -185,33 +177,25 @@ def _construct_version_validation_response(self, meets_requirement, version_str)
 
     def send_custom_notification(self, method: str, args: dict):
         """Sends a custom notification to the LSP client."""
-        self.show_message_log(
-            f"Sending custom notification: {method} with args: {args}"
-        )
+        self.show_message_log(f"Sending custom notification: {method} with args: {args}")
         self.send_notification(method, args)
 
     def update_python_interpreter(self, interpreter_path):
         """Updates the Python interpreter path and handles errors."""
         try:
             self.python_interpreter = interpreter_path
-            self.show_message_log(
-                f"LSP_Python_Interpreter Updated: {self.python_interpreter}"
-            )
+            self.show_message_log(f"LSP_Python_Interpreter Updated: {self.python_interpreter}")
         # pylint: disable=broad-exception-caught
         except Exception as e:
             self.show_message_log(
                 f"Failed to update Python interpreter: {str(e)}", lsp.MessageType.Error
             )
 
-    def notify_user(
-        self, message: str, msg_type: lsp.MessageType = lsp.MessageType.Info
-    ):
+    def notify_user(self, message: str, msg_type: lsp.MessageType = lsp.MessageType.Info):
         """Logs a message and also notifies the user."""
         self.show_message(message, msg_type)
 
-    def log_to_output(
-        self, message: str, msg_type: lsp.MessageType = lsp.MessageType.Log
-    ) -> None:
+    def log_to_output(self, message: str, msg_type: lsp.MessageType = lsp.MessageType.Log) -> None:
         """Log to output."""
         self.show_message_log(message, msg_type)
 
@@ -264,7 +248,6 @@ def get_active_stack(wrapper_instance, *args, **kwargs):
         @self.zenml_command(wrapper_name="stacks_wrapper")
         def set_active_stack(wrapper_instance, args):
             """Sets the active ZenML stack to the specified stack."""
-            print(f"args received: {args}")
             return wrapper_instance.set_active_stack(args)
 
         @self.command(f"{TOOL_MODULE_NAME}.renameStack")

diff --git a/bundled/tool/zenml_wrappers.py b/bundled/tool/zenml_wrappers.py
@@ -47,9 +47,7 @@ def get_global_config_directory(self):
     def RestZenStoreConfiguration(self):
         """Returns the RestZenStoreConfiguration class for store configuration."""
         # pylint: disable=not-callable
-        return self.lazy_import(
-            "zenml.zen_stores.rest_zen_store", "RestZenStoreConfiguration"
-        )
+        return self.lazy_import("zenml.zen_stores.rest_zen_store", "RestZenStoreConfiguration")
 
     def get_global_config_directory_path(self) -> str:
         """Get the global configuration directory path.
@@ -176,9 +174,7 @@ def get_server_info(self) -> dict:
         store_info = json.loads(self.gc.zen_store.get_store_info().json(indent=2))
         # Handle both 'store' and 'store_configuration' depending on version
         store_attr_name = (
-            "store_configuration"
-            if hasattr(self.gc, "store_configuration")
-            else "store"
+            "store_configuration" if hasattr(self.gc, "store_configuration") else "store"
         )
         store_config = json.loads(getattr(self.gc, store_attr_name).json(indent=2))
         return {"storeInfo": store_info, "storeConfig": store_config}
@@ -200,9 +196,7 @@ def connect(self, args, **kwargs) -> dict:
         try:
             # pylint: disable=not-callable
             access_token = self.web_login(url=url, verify_ssl=verify_ssl)
-            self._config_wrapper.set_store_configuration(
-                remote_url=url, access_token=access_token
-            )
+            self._config_wrapper.set_store_configuration(remote_url=url, access_token=access_token)
             return {"message": "Connected successfully.", "access_token": access_token}
         except self.AuthorizationException as e:
             return {"error": f"Authorization failed: {str(e)}"}
@@ -218,9 +212,7 @@ def disconnect(self, args) -> dict:
         try:
             # Adjust for changes from 'store' to 'store_configuration'
             store_attr_name = (
-                "store_configuration"
-                if hasattr(self.gc, "store_configuration")
-                else "store"
+                "store_configuration" if hasattr(self.gc, "store_configuration") else "store"
             )
             url = getattr(self.gc, store_attr_name).url
             store_type = self.BaseZenStore.get_store_type(url)
@@ -289,21 +281,15 @@ def fetch_pipeline_runs(self, args):
                     "version": run.body.pipeline.body.version,
                     "stackName": run.body.stack.name,
                     "startTime": (
-                        run.metadata.start_time.isoformat()
-                        if run.metadata.start_time
-                        else None
+                        run.metadata.start_time.isoformat() if run.metadata.start_time else None
                     ),
                     "endTime": (
-                        run.metadata.end_time.isoformat()
-                        if run.metadata.end_time
-                        else None
+                        run.metadata.end_time.isoformat() if run.metadata.end_time else None
                     ),
                     "os": run.metadata.client_environment.get("os", "Unknown OS"),
                     "osVersion": run.metadata.client_environment.get(
                         "os_version",
-                        run.metadata.client_environment.get(
-                            "mac_version", "Unknown Version"
-                        ),
+                        run.metadata.client_environment.get("mac_version", "Unknown Version"),
                     ),
                     "pythonVersion": run.metadata.client_environment.get(
                         "python_version", "Unknown"
@@ -378,12 +364,11 @@ def ZenKeyError(self) -> Any:
 
     def fetch_stacks(self, args):
         """Fetches all ZenML stacks and components with pagination."""
-        page = args[0]
-        max_size = args[1]
+        if len(args) < 2:
+            return {"error": "Insufficient arguments provided."}
+        page, max_size = args
         try:
-            stacks_page = self.client.list_stacks(
-                page=page, size=max_size, hydrate=True
-            )
+            stacks_page = self.client.list_stacks(page=page, size=max_size, hydrate=True)
             stacks_data = self.process_stacks(stacks_page.items)
 
             return {
@@ -496,23 +481,17 @@ def copy_stack(self, args) -> dict:
         target_stack_name = args[1]
 
         if not source_stack_name_or_id or not target_stack_name:
-            return {
-                "error": "Both source stack name/id and target stack name are required"
-            }
+            return {"error": "Both source stack name/id and target stack name are required"}
 
         try:
-            stack_to_copy = self.client.get_stack(
-                name_id_or_prefix=source_stack_name_or_id
-            )
+            stack_to_copy = self.client.get_stack(name_id_or_prefix=source_stack_name_or_id)
             component_mapping = {
                 c_type: [c.id for c in components][0]
                 for c_type, components in stack_to_copy.components.items()
                 if components
             }
 
-            self.client.create_stack(
-                name=target_stack_name, components=component_mapping
-            )
+            self.client.create_stack(name=target_stack_name, components=component_mapping)
             return {
                 "message": (
                     f"Stack `{source_stack_name_or_id}` successfully copied "

diff --git a/package.json b/package.json
@@ -345,7 +345,6 @@
     "@types/fs-extra": "^11.0.4",
     "@types/mocha": "^10.0.6",
     "@types/node": "^18.19.18",
-    "@types/proxyquire": "^1.3.31",
     "@types/sinon": "^17.0.3",
     "@types/vscode": "^1.86.0",
     "@types/webpack": "^5.28.5",
@@ -357,14 +356,9 @@
     "eslint": "^8.56.0",
     "eslint-config-prettier": "^9.1.0",
     "prettier": "^3.2.5",
-    "proxyquire": "^2.1.3",
     "sinon": "^17.0.1",
     "ts-loader": "^9.5.1",
-    "ts-mockito": "^2.6.1",
     "ts-node": "^10.9.2",
-    "tsconfig-paths": "^4.2.0",
-    "tsconfig-paths-webpack-plugin": "^4.1.0",
-    "typemoq": "^2.1.0",
     "typescript": "^5.3.3",
     "webpack": "^5.90.0",
     "webpack-cli": "^5.1.4"

diff --git a/package.nls.json b/package.nls.json
@@ -1,6 +1,6 @@
 {
   "extension.description": "Integrates ZenML directly into VS Code, enhancing machine learning workflow with support for pipelines, stacks, and server management.",
-  "zenml.promptForInterpreter": "Prompt to select a Python interpreter from the command palette in VSCode.",
+  "command.promptForInterpreter": "Prompt to select a Python interpreter from the command palette in VSCode.",
   "command.connectServer": "Establishes a connection to a specified ZenML server.",
   "command.disconnectServer": "Disconnects from the currently connected ZenML server.",
   "command.refreshServerStatus": "Refreshes the status of the ZenML server to reflect the current state.",

diff --git a/requirements.in b/requirements.in
@@ -2,7 +2,7 @@
 # NOTE:
 # Use Python 3.8 or greater which ever is the minimum version of the python
 # you plan on supporting when creating the environment or using pip-tools.
-# Only run the commands below to manully upgrade packages in requirements.txt:
+# Only run the commands below to manually upgrade packages in requirements.txt:
 # 1) python -m pip install pip-tools
 # 2) pip-compile --generate-hashes --resolver=backtracking --upgrade ./requirements.in
 # If you are using nox commands to setup or build package you don't need to
@@ -13,4 +13,4 @@ pygls
 packaging
 # Tool-specific packages for ZenML extension
 watchdog
-PyYAML
+PyYAML
diff --git a/requirements.txt b/requirements.txt
@@ -24,9 +24,9 @@ packaging==24.0 \
     --hash=sha256:2ddfb553fdf02fb784c234c7ba6ccc288296ceabec964ad2eae3777778130bc5 \
     --hash=sha256:eb82c5e3e56209074766e6885bb04b8c38a0c015d0a30036ebe7ece34c9989e9
     # via -r ./requirements.in
-pygls==1.3.0 \
-    --hash=sha256:1b44ace89c9382437a717534f490eadc6fda7c0c6c16ac1eaaf5568e345e4fb8 \
-    --hash=sha256:d4a01414b6ed4e34e7e8fd29b77d3e88c29615df7d0bbff49bf019e15ec04b8f
+pygls==1.3.1 \
+    --hash=sha256:140edceefa0da0e9b3c533547c892a42a7d2fd9217ae848c330c53d266a55018 \
+    --hash=sha256:6e00f11efc56321bdeb6eac04f6d86131f654c7d49124344a9ebb968da3dd91e
     # via -r ./requirements.in
 pyyaml==6.0.1 \
     --hash=sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5 \