From f53d00e8e2e803d9553099263855a268e6e1715a Mon Sep 17 00:00:00 2001
From: Daylin Morgan <daylinmorgan@gmail.com>
Date: Wed, 3 Jan 2024 13:49:27 -0600
Subject: [PATCH] feat: add viv.run() to accompany inline metatdata approach

---
 examples/pep723_run.py | 17 +++++++++++++
 src/viv/__init__.py    |  2 +-
 src/viv/viv.py         | 57 ++++++++++++++++++++++++++++++++++++------
 3 files changed, 67 insertions(+), 9 deletions(-)
 create mode 100755 examples/pep723_run.py

diff --git a/examples/pep723_run.py b/examples/pep723_run.py
new file mode 100755
index 0000000..9ce867a
--- /dev/null
+++ b/examples/pep723_run.py
@@ -0,0 +1,17 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#   "requests<3",
+#   "rich",
+# ]
+# ///
+
+__import__("viv").run()
+
+import requests
+from rich import print
+
+resp = requests.get("https://peps.python.org/api/peps.json")
+data = resp.json()
+print([(k, v["title"]) for k, v in data.items()][:10])
diff --git a/src/viv/__init__.py b/src/viv/__init__.py
index 4a399d1..44b3c04 100644
--- a/src/viv/__init__.py
+++ b/src/viv/__init__.py
@@ -1 +1 @@
-from .viv import __version__, use, main  # noqa
+from .viv import __version__, use, main, run  # noqa
diff --git a/src/viv/viv.py b/src/viv/viv.py
index 7be1f19..8849635 100755
--- a/src/viv/viv.py
+++ b/src/viv/viv.py
@@ -2775,6 +2775,16 @@ def use(*packages: str, track_exe: bool = False, name: str = "") -> Path:
     return vivenv.path
 
 
+def run() -> Path:
+    """create a vivenv and append to sys.path using embedded metadata"""
+    source = get_caller_path().read_text()
+    metadata = _read_metadata_block(source)
+    deps = metadata.get("dependencies", [])
+    if requires := metadata.get("requires-python", ""):
+        _check_python(requires)
+    return use(*deps)
+
+
 def combined_spec(reqs: List[str], requirements: Path) -> List[str]:
     if requirements:
         with requirements.open("r") as f:
@@ -2854,16 +2864,36 @@ def make_executable(path: Path) -> None:
     os.chmod(path, mode)
 
 
-def uses_viv(txt: str) -> bool:
+# TODO: abstract/deduplicate these functions
+def _uses_viv_use(txt: str) -> bool:
     return bool(
         re.search(
             r"""
             ^(?!\#)\s*
-            (?:__import__\(\s*["']viv["']\s*\))
+            (?:__import__\(\s*["']viv["']\s*\).use)
             |
             (?:from\ viv\ import\ use)
+            #|
+            #(?:import\ viv)
+            | 
+            (?:viv.use)
+        """,
+            txt,
+            re.VERBOSE | re.MULTILINE,
+        )
+    )
+
+
+def _uses_viv_run(txt: str) -> bool:
+    return bool(
+        re.search(
+            r"""
+            ^(?!\#)\s*
+            (?:__import__\(\s*["']viv["']\s*\).run)
             |
-            (?:import\ viv)
+            (?:from\ viv\ import\ run)
+            |
+            (?:viv.run)
         """,
             txt,
             re.VERBOSE | re.MULTILINE,
@@ -3405,21 +3435,23 @@ class Viv:
             else:
                 scriptpath = tmppath / name
                 script_text = fetch_script(script)
+                scriptpath.write_text(script_text)
 
-            viv_used = uses_viv(script_text)
+            has_use = _uses_viv_use(script_text)
+            has_run = _uses_viv_run(script_text)
             metadata = _read_metadata_block(script_text)
             deps = metadata.get("dependencies", [])
 
             if requires := metadata.get("requires-python", ""):
                 _check_python(requires)
 
-            if viv_used and deps:
+            if has_use and deps:
                 error(
                     "Script Dependencies block and "
                     "`viv.use` API can't be used in the same script"
                 )
 
-            if not self.local_source and viv_used:
+            if not self.local_source and (has_use or has_run):
                 log.debug("fetching remote copy to use for python api")
                 (tmppath / "viv.py").write_text(
                     fetch_script(
@@ -3427,10 +3459,9 @@ class Viv:
                     )
                 )
 
-            scriptpath.write_text(script_text)
             self._update_cache(env, keep, tmpdir)
 
-            if viv_used:
+            if has_use:
                 log.debug(f"script invokes viv.use passing along spec: \n  '{spec}'")
                 subprocess_run_quit(
                     [sys.executable, "-S", scriptpath, *rest],
@@ -3440,6 +3471,16 @@ class Viv:
                         PYTHONPATH=":".join((str(tmppath), env.get("PYTHONPATH", ""))),
                     ),
                 )
+            elif has_run:
+                log.debug("script invokes viv.run letting subprocess handle deps")
+                print(tmppath)
+                subprocess_run_quit(
+                    [sys.executable, "-S", scriptpath, *rest],
+                    env=dict(
+                        env,
+                        PYTHONPATH=":".join((str(tmppath), env.get("PYTHONPATH", ""))),
+                    ),
+                )
             elif not spec and not deps:
                 log.warning("using viv with empty spec, skipping vivenv creation")
                 subprocess_run_quit([sys.executable, "-S", scriptpath, *rest])