From 8a44dabf39ea91765c87bdabe03e5f30460b954e Mon Sep 17 00:00:00 2001
From: Robin van der Noord <robinvandernoord@gmail.com>
Date: Mon, 4 Mar 2024 18:20:47 +0100
Subject: [PATCH] feat: implemented reinstall and list methods, save .metadata
 file in app-specific venv

---
 pyproject.toml        |   1 +
 src/uvx/_constants.py |   6 ++
 src/uvx/_python.py    |  41 +++++++++++
 src/uvx/_symlinks.py  |  14 ++++
 src/uvx/cli.py        |  84 +++++++++++++++++++++--
 src/uvx/core.py       | 153 +++++++++++++++++++++++++++++++-----------
 src/uvx/metadata.py   |  66 ++++++++++++++++++
 7 files changed, 322 insertions(+), 43 deletions(-)
 create mode 100644 src/uvx/_constants.py
 create mode 100644 src/uvx/_python.py
 create mode 100644 src/uvx/_symlinks.py
 create mode 100644 src/uvx/metadata.py

diff --git a/pyproject.toml b/pyproject.toml
index d8cebcb..6031925 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,6 +27,7 @@ dependencies = [
     "typer",
     "plumbum",
     "threadful>=0.2",
+    "quickle",
 ]
 
 [project.optional-dependencies]
diff --git a/src/uvx/_constants.py b/src/uvx/_constants.py
new file mode 100644
index 0000000..bfb96d3
--- /dev/null
+++ b/src/uvx/_constants.py
@@ -0,0 +1,6 @@
+from pathlib import Path
+
+BIN_DIR = Path.home() / ".local/bin"
+WORK_DIR = (
+    Path.home() / ".local/uvx"
+)  # use 'ensure_local_folder()' instead, whenever possible!
diff --git a/src/uvx/_python.py b/src/uvx/_python.py
new file mode 100644
index 0000000..36278f3
--- /dev/null
+++ b/src/uvx/_python.py
@@ -0,0 +1,41 @@
+import textwrap
+from pathlib import Path
+
+import plumbum
+from plumbum.cmd import grep, uv
+
+
+def _run_python_in_venv(*args: str, venv: Path) -> str:
+    python = venv / "bin" / "python"
+
+    return plumbum.local[python](*args)
+
+
+def run_python_code_in_venv(code: str, venv: Path) -> str:
+    """
+    Run Python code in a virtual environment.
+
+    Args:
+        code (str): The Python code to run.
+        venv (Path): The path of the virtual environment.
+
+    Returns:
+        str: The output of the Python code.
+    """
+    code = textwrap.dedent(code)
+    return _run_python_in_venv("-c", code, venv=venv)
+
+
+def get_python_version(venv: Path):
+    return _run_python_in_venv("--version", venv=venv).strip()
+
+
+def get_python_executable(venv: Path):
+    executable = venv / "bin" / "python"
+    return str(executable.resolve())  # /usr/bin/python3.xx
+
+
+def get_package_version(package: str) -> str:
+    # assumes `with virtualenv(venv)` block executing this function
+    # uv pip freeze | grep ^su6==
+    return (uv["pip", "freeze"] | grep[f"^{package}=="])().strip().split("==")[-1]
diff --git a/src/uvx/_symlinks.py b/src/uvx/_symlinks.py
new file mode 100644
index 0000000..92420de
--- /dev/null
+++ b/src/uvx/_symlinks.py
@@ -0,0 +1,14 @@
+import typing
+
+from ._constants import BIN_DIR, WORK_DIR
+
+
+def check_symlink(symlink: str, venv: str) -> bool:
+    symlink_path = BIN_DIR / symlink
+    target_path = WORK_DIR / "venvs" / venv
+
+    return symlink_path.is_symlink() and target_path in symlink_path.resolve().parents
+
+
+def check_symlinks(symlinks: typing.Iterable[str], venv: str) -> dict[str, bool]:
+    return {k: check_symlink(k, venv) for k in symlinks}
diff --git a/src/uvx/cli.py b/src/uvx/cli.py
index 3de7740..2818a26 100644
--- a/src/uvx/cli.py
+++ b/src/uvx/cli.py
@@ -1,17 +1,26 @@
 """This file builds the Typer cli."""
 
+from typing import Optional
+
+import rich
 import typer
 
-from .core import create_venv, install_package, uninstall_package
+from .core import (
+    format_bools,
+    install_package,
+    list_packages,
+    reinstall_package,
+    uninstall_package,
+)
+from .metadata import Metadata, read_metadata
 
 app = typer.Typer()
 
 
 @app.command()
-def install(package_name: str, force: bool = False):
+def install(package_name: str, force: bool = False, python: str = None):
     """Install a package (by pip name)."""
-    venv = create_venv(package_name)
-    install_package(package_name, venv, force=force)
+    install_package(package_name, python=python, force=force)
 
 
 @app.command(name="remove")
@@ -21,7 +30,70 @@ def uninstall(package_name: str, force: bool = False):
     uninstall_package(package_name, force=force)
 
 
+@app.command()
+def reinstall(package: str, python: Optional[str] = None, force: bool = False):
+    reinstall_package(package, python=python, force=force)
+
+
 # list
+def list_short(name: str, metadata: Optional[Metadata]):
+    rich.print("-", name, metadata.installed_version if metadata else "[red]?[/red]")
+
+
+TAB = " " * 3
+
+
+def list_normal(name: str, metadata: Optional[Metadata], verbose: bool = False):
+
+    if not metadata:
+        print("-", name)
+        rich.print(TAB, "[red]Missing metadata [/red]")
+        return
+    else:
+        extras = list(metadata.extras)
+        name_with_extras = name if not extras else f"{name}{extras}"
+        print("-", name_with_extras)
+
+    metadata.check_script_symlinks(name)
+
+    if verbose:
+        rich.print(TAB, metadata)
+    else:
+        rich.print(
+            TAB,
+            f"Installed Version: {metadata.installed_version} on {metadata.python}.",
+        )
+        rich.print(TAB, "Scripts:", format_bools(metadata.scripts))
+
+
+def list_venvs_json():
+    from json import dumps
+
+    print(
+        dumps(
+            {
+                name: metadata.check_script_symlinks(name).to_dict() if metadata else {}
+                for name, metadata in list_packages()
+            }
+        )
+    )
+
+
+@app.command(name="list")
+def list_venvs(short: bool = False, verbose: bool = False, json: bool = False):
+    """
+    List packages and apps installed with uvx.
+    """
+
+    if json:
+        return list_venvs_json()
+
+    for name, metadata in list_packages():
+        if short:
+            list_short(name, metadata)
+        else:
+            list_normal(name, metadata, verbose=verbose)
+
 
 # run
 
@@ -29,4 +101,8 @@ def uninstall(package_name: str, force: bool = False):
 
 # self-upgrade (uv and uvx)
 
+# inject
+
+# version or --version (incl. 'uv' version and Python version)
+
 # ...
diff --git a/src/uvx/core.py b/src/uvx/core.py
index 04822af..1393e33 100644
--- a/src/uvx/core.py
+++ b/src/uvx/core.py
@@ -2,19 +2,26 @@
 
 import shutil
 import sys
-import textwrap
 import typing
 from contextlib import contextmanager
 from pathlib import Path
 from typing import Optional
 
 import plumbum  # type: ignore
+import rich
 from plumbum import local  # type: ignore
 from plumbum.cmd import uv as _uv  # type: ignore
 from threadful import thread
 from threadful.bonus import animate
 
-BIN_DIR = Path.home() / ".local/bin"
+from ._constants import BIN_DIR, WORK_DIR
+from ._python import (
+    get_package_version,
+    get_python_executable,
+    get_python_version,
+    run_python_code_in_venv,
+)
+from .metadata import Metadata, collect_metadata, read_metadata, store_metadata
 
 
 @thread
@@ -84,23 +91,6 @@ def exit_on_pb_error() -> typing.Generator[None, None, None]:
         exit(e.retcode)
 
 
-def run_python_in_venv(code: str, venv: Path) -> str:
-    """
-    Run Python code in a virtual environment.
-
-    Args:
-        code (str): The Python code to run.
-        venv (Path): The path of the virtual environment.
-
-    Returns:
-        str: The output of the Python code.
-    """
-    python = venv / "bin" / "python"
-
-    code = textwrap.dedent(code)
-    return plumbum.local[python]("-c", code)
-
-
 def find_symlinks(library: str, venv: Path) -> list[str]:
     """
     Find the symlinks for a library in a virtual environment.
@@ -123,12 +113,18 @@ def find_symlinks(library: str, venv: Path) -> list[str]:
     """
 
     try:
-        raw = run_python_in_venv(code, venv)
+        raw = run_python_code_in_venv(code, venv)
         return [_ for _ in raw.split("\n") if _]
     except Exception:
         return []
 
 
+def format_bools(data: dict[str, bool], sep=" | ") -> str:
+    return sep.join(
+        [f"[green]{k}[/green]" if v else f"[red]{k}[/red]" for k, v in data.items()]
+    )
+
+
 def install_symlink(
     symlink: str, venv: Path, force: bool = False, binaries: tuple[str, ...] = ()
 ) -> bool:
@@ -172,7 +168,11 @@ def install_symlink(
 
 
 def install_symlinks(
-    library: str, venv: Path, force: bool = False, binaries: tuple[str, ...] = ()
+    library: str,
+    venv: Path,
+    force: bool = False,
+    binaries: tuple[str, ...] = (),
+    meta: Optional[Metadata] = None,
 ) -> bool:
     """
     Install symlinks for a library in a virtual environment.
@@ -182,21 +182,30 @@ def install_symlinks(
         venv (Path): The path of the virtual environment.
         force (bool, optional): If True, overwrites existing symlinks. Defaults to False.
         binaries (tuple[str, ...], optional): The binaries to install. Defaults to ().
+        meta: Optional metadata object to store results in
 
     Returns:
         bool: True if any symlink was installed, False otherwise.
     """
     symlinks = find_symlinks(library, venv)
 
-    results = []
+    results = {}
     for symlink in symlinks:
-        results.append(install_symlink(symlink, venv, force=force, binaries=binaries))
+        results[symlink] = install_symlink(
+            symlink, venv, force=force, binaries=binaries
+        )
+
+    if meta:
+        meta.scripts = results
 
-    return any(results)
+    return any(results.values())
 
 
 def install_package(
-    package_name: str, venv: Optional[Path] = None, force: bool = False
+    package_name: str,
+    venv: Optional[Path] = None,
+    python: Optional[str] = None,
+    force: bool = False,
 ):
     """
     Install a package in a virtual environment.
@@ -206,17 +215,58 @@ def install_package(
         venv (Optional[Path], optional): The path of the virtual environment. Defaults to None.
         force (bool, optional): If True, overwrites existing package. Defaults to False.
     """
+    meta = collect_metadata(package_name)
+
     if venv is None:
-        venv = create_venv(package_name)
+        venv = create_venv(meta.name, python=python, force=force)
 
     with virtualenv(venv), exit_on_pb_error():
-        animate(
-            uv("pip", "install", package_name),
-            # text=f"installing {package_name}"
+        try:
+            animate(uv("pip", "install", package_name), text=f"installing {meta.name}")
+
+            # must still be in the venv for these:
+            meta.installed_version = get_package_version(meta.name)
+            meta.python = get_python_version(venv)
+            meta.python_raw = get_python_executable(venv)
+
+        except plumbum.ProcessExecutionError as e:
+            remove_dir(venv)
+            raise e
+
+    if install_symlinks(meta.name, venv, force=force, meta=meta):
+        rich.print(f"📦 {meta.name} ({meta.installed_version}) installed!")  # :package:
+
+    store_metadata(meta, venv)
+
+
+def reinstall_package(
+    package_name: str, python: Optional[str] = None, force: bool = False
+):
+    new_metadata = collect_metadata(package_name)
+
+    workdir = ensure_local_folder()
+    venv = workdir / "venvs" / new_metadata.name
+
+    if not venv.exists():
+        rich.print(
+            f"'{new_metadata.name}' was not previously installed. "
+            f"Please run 'uvx install {package_name}' instead."
         )
+        exit(1)
+
+    existing_metadata = read_metadata(venv)
 
-    if install_symlinks(package_name, venv, force=force):
-        print(f"📦 {package_name} installed!")
+    # if a new version or extra is requested or no old metadata is available, install from cli arg package name.
+    # otherwise, install from old metadata spec
+    new_install_spec = bool(
+        new_metadata.requested_version or new_metadata.extras or not existing_metadata
+    )
+    install_spec = package_name if new_install_spec else existing_metadata.install_spec
+
+    python = python or (existing_metadata.python_raw if existing_metadata else None)
+
+    uninstall_package(new_metadata.name, force=force)
+    install_package(install_spec, python=python, force=force)
 
 
 def remove_symlink(symlink: str):
@@ -253,9 +303,11 @@ def uninstall_package(package_name: str, force: bool = False):
     workdir = ensure_local_folder()
     venv_path = workdir / "venvs" / package_name
 
+    meta = read_metadata(venv_path)
+
     if not venv_path.exists() and not force:
-        print(
-            f"No virtualenv for {package_name}, stopping. Use --force to remove an executable with that name anyway.",
+        rich.print(
+            f"No virtualenv for '{package_name}', stopping. Use '--force' to remove an executable with that name anyway.",
             file=sys.stderr,
         )
         exit(1)
@@ -266,7 +318,7 @@ def uninstall_package(package_name: str, force: bool = False):
         remove_symlink(symlink)
 
     remove_dir(venv_path)
-    print(f"🗑️ {package_name} removed!")
+    rich.print(f"🗑️ {package_name} ({meta.installed_version}) removed!")  # :trash:
 
 
 def ensure_local_folder() -> Path:
@@ -276,17 +328,18 @@ def ensure_local_folder() -> Path:
     Returns:
         Path: The path of the local folder.
     """
-    workdir = Path("~/.local/uvx/").expanduser()
-    (workdir / "venvs").mkdir(exist_ok=True, parents=True)
-    return workdir
+    (WORK_DIR / "venvs").mkdir(exist_ok=True, parents=True)
+    return WORK_DIR
 
 
-def create_venv(name: str) -> Path:
+def create_venv(name: str, python: Optional[str] = None, force: bool = False) -> Path:
     """
     Create a virtual environment.
 
     Args:
         name (str): The name of the virtual environment.
+        python (str): which version of Python to use (e.g. 3.11, python3.11)
+        force (bool): ignore existing venv
 
     Returns:
         Path: The path of the virtual environment.
@@ -295,6 +348,28 @@ def create_venv(name: str) -> Path:
 
     venv_path = workdir / "venvs" / name
 
-    uv("venv", venv_path).join()
+    if venv_path.exists() and not force:
+        rich.print(
+            f"'{name}' is already installed. "
+            f"Use 'uvx upgrade' to update existing tools or pass '--force' to this command to ignore this message.",
+            file=sys.stderr,
+        )
+        exit(1)
+
+    args = ["venv", venv_path]
+
+    if python:
+        args.extend(["--python", python])
+
+    uv(*args).join()
 
     return venv_path
+
+
+def list_packages() -> typing.Generator[tuple[str, Metadata | None], None, None]:
+    workdir = ensure_local_folder()
+
+    for subdir in workdir.glob("venvs/*"):
+        metadata = read_metadata(subdir)
+
+        yield subdir.name, metadata
diff --git a/src/uvx/metadata.py b/src/uvx/metadata.py
new file mode 100644
index 0000000..6e52586
--- /dev/null
+++ b/src/uvx/metadata.py
@@ -0,0 +1,66 @@
+import typing
+from pathlib import Path
+from typing import Optional
+
+import quickle
+from packaging.requirements import Requirement
+
+from ._symlinks import check_symlinks
+
+
+class Metadata(quickle.Struct):
+    name: str
+    scripts: dict[str, bool]  # {script: is_installed}
+    install_spec: str  # e.g. '2fas' or '2fas[gui]>=0.1.0'
+    extras: set[str]  # .e.g. {'gui'}
+    requested_version: Optional[str]  # e.g. ">=0.1.0"
+    installed_version: str
+    python: str = ""
+    python_raw: str = ""
+
+    def _convert_type(self, value):
+        if isinstance(value, set):
+            return list(value)
+        return value
+
+    def to_dict(self):
+        return {f: self._convert_type(getattr(self, f)) for f in self.__struct_fields__}
+
+    def check_script_symlinks(self, name: str):
+        self.scripts = check_symlinks(self.scripts.keys(), venv=name)
+        return self
+
+
+quickle_enc = quickle.Encoder(registry=[Metadata])
+quickle_dec = quickle.Decoder(registry=[Metadata])
+
+
+def collect_metadata(spec: str) -> Metadata:
+    parsed_spec = Requirement(spec)
+
+    return Metadata(
+        install_spec=spec,
+        name=parsed_spec.name,
+        scripts={},  # postponed
+        extras=parsed_spec.extras,
+        requested_version=str(parsed_spec.specifier),
+        installed_version="",  # postponed
+        python="",  # postponed
+    )
+
+
+def store_metadata(meta: Metadata, venv: Path):
+    with (venv / ".metadata").open("wb") as f:
+        f.write(quickle_enc.dumps(meta))
+
+
+def read_metadata(venv: Path) -> Metadata | None:
+    metafile = venv / ".metadata"
+    if not metafile.exists():
+        return None
+
+    with metafile.open("rb") as f:
+        return typing.cast(Metadata, quickle_dec.loads(f.read()))
+
+
+def get_metadata(): ...