feat: configure auto-add script metadata to notebook file, when using…

… uv (#2102) * feat: configure auto-add script metadata to notebook file, when using uv * fix * snapshots * fix * improve * update * snapshots * add script metadat * fixes * 38com * name changes * Update package_management.md * Update package_management.md * add modules post installation * handle deletion * todo on using module watcher * deletion test --------- Co-authored-by: Akshay Agrawal <[email protected]>
marimo-team · Aug 27, 2024 · a9c7563 · a9c7563
1 parent 4c89075
commit a9c7563
Show file tree

Hide file tree

Showing 20 changed files with 465 additions and 38 deletions.
diff --git a/docs/guides/editor_features/configuration.md b/docs/guides/editor_features/configuration.md
diff --git a/docs/guides/editor_features/index.md b/docs/guides/editor_features/index.md
@@ -15,11 +15,12 @@ notebooks. We've taken a batteries-included approach to designing the editor:
 it comes _packed_ with features to make you productive when working
 with code and data.
 
-| Guide                | Description                                     |
-| :------------------- | :---------------------------------------------- |
-| {doc}`overview`      | An overview of editor features and configuration |
-| {doc}`ai_completion` | Code with the help of a language model          |
-| {doc}`hotkeys`       | Our hotkeys                                     |
+| Guide                     | Description                                      |
+| :------------------------ | :----------------------------------------------- |
+| {doc}`overview`           | An overview of editor features and configuration |
+| {doc}`package_management` | Using package managers in marimo                 |
+| {doc}`ai_completion`      | Code with the help of a language model           |
+| {doc}`hotkeys`            | Our hotkeys                                      |
 
 Highlights include:
 

diff --git a/docs/guides/editor_features/overview.md b/docs/guides/editor_features/overview.md
@@ -24,6 +24,7 @@ A non-exhaustive list of settings:
 
 - Outputs above or below code cells
 - [Disable/enable autorun](/guides/reactivity.md#runtime-configuration)
+- Package installation
 - Vim keybindings
 - Dark mode
 - Auto-save
@@ -67,14 +68,14 @@ marimo ships with the IDE panels that provide an overview of your notebook:
 
 1. **file explorer**: view the file tree, open other notebooks
 2. **variables**: explore variable values, see where they are defined and used, with go-to-definition
-2. **data explorer**: see dataframe and table schemas at a glance
-3. **dependency graph**: view dependencies between cells, drill-down on nodes and edges
-4. **table of contents**: corresponding to your markdown
-5. **documentation** - move your text cursor over a symbol to see its documentation
-6. **logs**: a continuous stream of stdout and stderr
-7. **scratchpad**: a scratchpad cell where you can execute throwaway code
-8. **snippets** - searchable snippets to copy directly into your notebook
-9. **feedback** - share feedback!
+3. **data explorer**: see dataframe and table schemas at a glance
+4. **dependency graph**: view dependencies between cells, drill-down on nodes and edges
+5. **table of contents**: corresponding to your markdown
+6. **documentation** - move your text cursor over a symbol to see its documentation
+7. **logs**: a continuous stream of stdout and stderr
+8. **scratchpad**: a scratchpad cell where you can execute throwaway code
+9. **snippets** - searchable snippets to copy directly into your notebook
+10. **feedback** - share feedback!
 :::
 
 ::::

diff --git a/docs/guides/editor_features/package_management.md b/docs/guides/editor_features/package_management.md
@@ -0,0 +1,27 @@
+# Package management
+
+marimo supports package management for `pip, rye, uv, poetry, pixi`. When marimo comes across a module that is not installed, you will be prompted to install it using your preferred package manager.
+
+Once the module is installed, all cells that depend on the module will be rerun.
+
+```{admonition} Package Installation
+:class: note
+
+We use some heuristic for guessing the package name in your registry (e.g. PyPI) from the module name. It is possible that the package name is different from the module name. If you encounter an error, please file an issue, so we can correct the heuristic.
+```
+
+## Auto-add inline script metadata (`uv` only)
+
+When using [uv](https://docs.astral.sh/uv), marimo can automatically add the package name metadata to your script, per [PEP 723](https://peps.python.org/pep-0723/). This metadata is used to manage the script's dependencies and Python version.
+
+For example, say you start marimo in a new virtual environment and spin up a new notebook. Whenever you add a new package, marimo will automatically add script metadata that looks like this:
+
+```python
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "pandas",
+#     "altair",
+# ]
+# ///
+```
diff --git a/frontend/src/components/app-config/user-config-form.tsx b/frontend/src/components/app-config/user-config-form.tsx
@@ -458,6 +458,47 @@ export const UserConfigForm: React.FC = () => {
               </FormItem>
             )}
           />
+          <FormField
+            control={form.control}
+            disabled={isWasmRuntime}
+            name="package_management.add_script_metadata"
+            render={({ field }) => {
+              if (form.getValues("package_management.manager") !== "uv") {
+                return <div />;
+              }
+
+              return (
+                <div className="flex flex-col gap-y-1">
+                  <FormItem className={formItemClasses}>
+                    <FormLabel className="font-normal">
+                      Auto-add script metadata
+                    </FormLabel>
+                    <FormControl>
+                      <Checkbox
+                        data-testid="auto-instantiate-checkbox"
+                        disabled={field.disabled}
+                        checked={field.value}
+                        onCheckedChange={field.onChange}
+                      />
+                    </FormControl>
+                  </FormItem>
+                  <FormDescription>
+                    Whether marimo should automatically add package metadata to
+                    scripts. See more about{" "}
+                    <a
+                      href="https://docs.marimo.io/guides/editor_features/package_management.html"
+                      target="_blank"
+                      rel="noreferrer"
+                      className="text-link hover:underline"
+                    >
+                      package metadata
+                    </a>
+                    .
+                  </FormDescription>
+                </div>
+              );
+            }}
+          />
         </SettingGroup>
         <SettingGroup title="Runtime">
           <FormField

diff --git a/frontend/src/core/config/__tests__/config-schema.test.ts b/frontend/src/core/config/__tests__/config-schema.test.ts
@@ -49,6 +49,7 @@ test("default UserConfig - empty", () => {
         "preset": "default",
       },
       "package_management": {
+        "add_script_metadata": false,
         "manager": "pip",
       },
       "runtime": {
@@ -99,6 +100,7 @@ test("default UserConfig - one level", () => {
         "preset": "default",
       },
       "package_management": {
+        "add_script_metadata": false,
         "manager": "pip",
       },
       "runtime": {

diff --git a/frontend/src/core/config/config-schema.ts b/frontend/src/core/config/config-schema.ts
@@ -90,6 +90,7 @@ export const UserConfigSchema = z
     package_management: z
       .object({
         manager: z.enum(PackageManagerNames).default("pip"),
+        add_script_metadata: z.boolean().default(false),
       })
       .default({ manager: "pip" }),
     ai: z

diff --git a/marimo/_config/config.py b/marimo/_config/config.py
@@ -159,9 +159,12 @@ class PackageManagementConfig(TypedDict):
     **Keys.**
 
     - `manager`: the package manager to use
+    - `add_script_metadata`: if true, add script metadata to the notebook
+        Currently only supports `uv`
     """
 
     manager: Literal["pip", "rye", "uv", "poetry", "pixi"]
+    add_script_metadata: bool
 
 
 @dataclass
@@ -230,7 +233,7 @@ class MarimoConfig(TypedDict):
         "autosave_delay": 1000,
         "format_on_save": False,
     },
-    "package_management": {"manager": "pip"},
+    "package_management": {"manager": "pip", "add_script_metadata": False},
     "server": {
         "browser": "default",
         "follow_symlink": False,

diff --git a/marimo/_runtime/packages/package_manager.py b/marimo/_runtime/packages/package_manager.py
@@ -4,6 +4,7 @@
 import abc
 import shutil
 import subprocess
+from typing import List
 
 
 class PackageManager(abc.ABC):
@@ -53,6 +54,21 @@ def run(self, command: list[str]) -> bool:
         proc = subprocess.run(command)  # noqa: ASYNC101
         return proc.returncode == 0
 
+    def update_notebook_script_metadata(
+        self,
+        filepath: str,
+        import_namespaces_to_add: List[str],
+        import_namespaces_to_remove: List[str],
+    ) -> None:
+        del filepath, import_namespaces_to_add, import_namespaces_to_remove
+        """
+        Add or remove inline script metadata metadata
+        in the marimo notebook.
+
+        This follows PEP 723 https://peps.python.org/pep-0723/
+        """
+        return
+
 
 class CanonicalizingPackageManager(PackageManager):
     """Base class for package managers.

diff --git a/marimo/_runtime/packages/pypi_package_manager.py b/marimo/_runtime/packages/pypi_package_manager.py
@@ -1,6 +1,8 @@
 # Copyright 2024 Marimo. All rights reserved.
 from __future__ import annotations
 
+from typing import List
+
 from marimo._runtime.packages.module_name_to_pypi_name import (
     module_name_to_pypi_name,
 )
@@ -48,6 +50,42 @@ class UvPackageManager(PypiPackageManager):
     async def _install(self, package: str) -> bool:
         return self.run(["uv", "pip", "install", package])
 
+    def update_notebook_script_metadata(
+        self,
+        filepath: str,
+        import_namespaces_to_add: List[str],
+        import_namespaces_to_remove: List[str],
+    ) -> None:
+        # Convert from module name to package name
+        packages_to_add = [
+            self.module_to_package(im) for im in import_namespaces_to_add
+        ]
+        packages_to_remove = [
+            self.module_to_package(im) for im in import_namespaces_to_remove
+        ]
+
+        # Filter to packages that are found by "uv pip show"
+        packages_to_add = [
+            im for im in packages_to_add if self._is_installed(im)
+        ]
+        packages_to_remove = [
+            im for im in packages_to_remove if self._is_installed(im)
+        ]
+
+        if packages_to_add:
+            self.run(
+                ["uv", "--quiet", "add", "--script", filepath]
+                + packages_to_add
+            )
+        if packages_to_remove:
+            self.run(
+                ["uv", "--quiet", "remove", "--script", filepath]
+                + packages_to_remove
+            )
+
+    def _is_installed(self, package: str) -> bool:
+        return self.run(["uv", "--quiet", "pip", "show", package])
+
 
 class RyePackageManager(PypiPackageManager):
     name = "rye"

diff --git a/marimo/_runtime/runtime.py b/marimo/_runtime/runtime.py
@@ -16,7 +16,7 @@
 import traceback
 from copy import copy
 from multiprocessing import connection
-from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional, cast
+from typing import TYPE_CHECKING, Any, Callable, Iterator, List, Optional, cast
 from uuid import uuid4
 
 from marimo import _loggers
@@ -639,6 +639,27 @@ def _maybe_register_cell(
                 cell_id, code, carried_imports=carried_imports
             )
 
+            # For any newly imported namespaces, add them to the metadata
+            #
+            # TODO(akshayka): Consider using the module watcher to discover
+            # packages used by a notebook; that would have the benefit of
+            # discovering transitive dependencies, ie if a notebook used a
+            # local module that in turn used packages available on PyPI.
+            if self._should_add_script_metadata():
+                cell = self.graph.cells.get(cell_id, None)
+                if cell:
+                    prev_imports: set[Name] = (
+                        set([im.namespace for im in previous_cell.imports])
+                        if previous_cell
+                        else set()
+                    )
+                    to_add = cell.imported_namespaces - prev_imports
+                    to_remove = prev_imports - cell.imported_namespaces
+                    self._add_script_metadata(
+                        import_namespaces_to_add=list(to_add),
+                        import_namespaces_to_remove=list(to_remove),
+                    )
+
         LOGGER.debug(
             "graph:\n\tcell id %s\n\tparents %s\n\tchildren %s\n\tsiblings %s",
             cell_id,
@@ -762,6 +783,14 @@ def _delete_cell(self, cell_id: CellId_t) -> set[CellId_t]:
         del self.cell_metadata[cell_id]
         cell = self.graph.cells[cell_id]
         cell.import_workspace.imported_defs = set()
+        if self._should_add_script_metadata():
+            self._add_script_metadata(
+                import_namespaces_to_add=[],
+                import_namespaces_to_remove=[
+                    im.namespace for im in cell.imports
+                ],
+            )
+
         return self._deactivate_cell(cell_id)
 
     def mutate_graph(
@@ -1164,6 +1193,7 @@ async def run(
     @kernel_tracer.start_as_current_span("rename_file")
     async def rename_file(self, filename: str) -> None:
         self.globals["__file__"] = filename
+        self.app_metadata.filename = filename
         roots = set()
         for cell in self.graph.cells.values():
             if "__file__" in cell.refs:
@@ -1615,6 +1645,12 @@ async def install_missing_packages(
             for pkg in package_statuses
             if package_statuses[pkg] == "installed"
         ]
+
+        # If a package was not installed at cell registration time, it won't
+        # yet be in the script metadata.
+        if self._should_add_script_metadata():
+            self._add_script_metadata(installed_modules, [])
+
         cells_to_run = set(
             cid
             for module in installed_modules
@@ -1623,6 +1659,39 @@ async def install_missing_packages(
         if cells_to_run:
             await self._if_autorun_then_run_cells(cells_to_run)
 
+    def _should_add_script_metadata(self) -> bool:
+        return (
+            self.user_config["package_management"]["add_script_metadata"]
+            and self.app_metadata.filename is not None
+            and self.package_manager is not None
+        )
+
+    def _add_script_metadata(
+        self,
+        import_namespaces_to_add: List[str],
+        import_namespaces_to_remove: List[str],
+    ) -> None:
+        filename = self.app_metadata.filename
+
+        if not filename or not self.package_manager:
+            return
+
+        try:
+            LOGGER.debug(
+                "Updating script metadata: %s. Adding namespaces: %s."
+                " Removing namespaces: %s",
+                filename,
+                import_namespaces_to_add,
+                import_namespaces_to_remove,
+            )
+            self.package_manager.update_notebook_script_metadata(
+                filepath=filename,
+                import_namespaces_to_add=import_namespaces_to_add,
+                import_namespaces_to_remove=import_namespaces_to_remove,
+            )
+        except Exception as e:
+            LOGGER.error("Failed to add script metadata to notebook: %s", e)
+
     @kernel_tracer.start_as_current_span("preview_dataset_column")
     async def preview_dataset_column(
         self, request: PreviewDatasetColumnRequest