From 5dacf8ab6054a7752aa995561d8a96efd7c8cfb2 Mon Sep 17 00:00:00 2001
From: Daylin Morgan <daylinmorgan@gmail.com>
Date: Thu, 4 Jan 2024 15:53:04 -0600
Subject: [PATCH] refactor: abstract _uses_viv + add tests

---
 src/viv/viv.py        | 74 ++++++++++++++++++++++---------------------
 tests/test_readers.py | 52 ++++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 36 deletions(-)
 create mode 100644 tests/test_readers.py

diff --git a/src/viv/viv.py b/src/viv/viv.py
index 8a002a2..818bbfc 100755
--- a/src/viv/viv.py
+++ b/src/viv/viv.py
@@ -33,6 +33,7 @@ from argparse import (
 from argparse import ArgumentParser as StdArgParser
 from contextlib import contextmanager
 from datetime import datetime
+from enum import Enum
 from logging.handlers import RotatingFileHandler
 from pathlib import Path
 from textwrap import dedent, fill
@@ -2864,41 +2865,43 @@ def make_executable(path: Path) -> None:
     os.chmod(path, mode)
 
 
-# TODO: abstract/deduplicate these functions
-def _uses_viv_use(txt: str) -> bool:
-    return bool(
-        re.search(
+class _Viv_Mode(Enum):
+    NONE = 0
+    USE = 1
+    RUN = 2
+
+
+def _uses_viv(txt: str) -> _Viv_Mode:
+    matches = [
+        match.group("mode")
+        for match in re.finditer(
             r"""
-            ^(?!\#)\s*
-            (?:__import__\(\s*["']viv["']\s*\).use)
-            |
-            (?:from\ viv\ import\ use)
-            #|
-            #(?:import\ viv)
-            | 
-            (?:viv.use)
-        """,
+            ^(?!\#)\s*    # ignore comments/shebangs
+            (
+              (?:__import__\(\s*["']viv["']\s*\)\.)
+              |
+              (?:from\s+viv\s+import\s+)
+              |
+              (?:viv\.)
+            )
+            (?P<mode>(\w+))
+            """,
             txt,
             re.VERBOSE | re.MULTILINE,
         )
-    )
-
-
-def _uses_viv_run(txt: str) -> bool:
-    return bool(
-        re.search(
-            r"""
-            ^(?!\#)\s*
-            (?:__import__\(\s*["']viv["']\s*\).run)
-            |
-            (?:from\ viv\ import\ run)
-            |
-            (?:viv.run)
-        """,
-            txt,
-            re.VERBOSE | re.MULTILINE,
+    ]
+    if len(matches) == 0:
+        return _Viv_Mode.NONE
+    elif len(matches) > 1:
+        err_quit(
+            "Unexpected number of viv references in script.\n"
+            "Expected only 1, found: "
+            + ", ".join((a.style(match, "bold") for match in matches))
         )
-    )
+    elif (match := matches[0]) in {"run", "use"}:
+        return _Viv_Mode[match.upper()]
+    else:
+        err_quit(f"Unknown function {a.bold}{matches[0]}{a.end} associated with viv.")
 
 
 METADATA_BLOCK = (
@@ -3437,21 +3440,20 @@ class Viv:
                 script_text = fetch_script(script)
                 scriptpath.write_text(script_text)
 
-            has_use = _uses_viv_use(script_text)
-            has_run = _uses_viv_run(script_text)
+            mode = _uses_viv(script_text)
             metadata = _read_metadata_block(script_text)
             deps = metadata.get("dependencies", [])
 
             if requires := metadata.get("requires-python", ""):
                 _check_python(requires)
 
-            if has_use and deps:
+            if mode == _Viv_Mode.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 (has_use or has_run):
+            if not self.local_source and mode != _Viv_Mode.NONE:
                 log.debug("fetching remote copy to use for python api")
                 (tmppath / "viv.py").write_text(
                     fetch_script(
@@ -3461,7 +3463,7 @@ class Viv:
 
             self._update_cache(env, keep, tmpdir)
 
-            if has_use:
+            if mode == _Viv_Mode.USE:
                 log.debug(f"script invokes viv.use passing along spec: \n  '{spec}'")
                 subprocess_run_quit(
                     [sys.executable, "-S", scriptpath, *rest],
@@ -3471,7 +3473,7 @@ class Viv:
                         PYTHONPATH=":".join((str(tmppath), env.get("PYTHONPATH", ""))),
                     ),
                 )
-            elif has_run:
+            elif mode == _Viv_Mode.RUN:
                 log.debug("script invokes viv.run letting subprocess handle deps")
                 subprocess_run_quit(
                     [sys.executable, "-S", scriptpath, *rest],
diff --git a/tests/test_readers.py b/tests/test_readers.py
new file mode 100644
index 0000000..f4a3960
--- /dev/null
+++ b/tests/test_readers.py
@@ -0,0 +1,52 @@
+import pytest
+from viv.viv import _read_metadata_block, _uses_viv, _Viv_Mode
+
+RUN_METADATA_SCRIPT = """
+#!/usr/bin/env -S viv run -s
+# /// script
+# requires-python = ">3.10"
+# dependencies = [
+#   "rich"
+# ]
+# ///
+"""
+
+USE_SCRIPT = """
+#!/usr/bin/env python3
+__import__("viv").use("rich")
+
+from rich import print
+print("pretty!")
+"""
+
+
+def test_metadata():
+    metadata = _read_metadata_block(RUN_METADATA_SCRIPT)
+    assert metadata == {"requires-python": ">3.10", "dependencies": ["rich"]}
+
+
+def test_uses():
+    assert _uses_viv(RUN_METADATA_SCRIPT) == _Viv_Mode.NONE
+    assert (
+        _uses_viv(RUN_METADATA_SCRIPT + """\n__import__("viv").run()\n""")
+        == _Viv_Mode.RUN
+    )
+    assert _uses_viv(USE_SCRIPT) == _Viv_Mode.USE
+    assert _uses_viv("# from viv import use") == _Viv_Mode.NONE
+
+
+def test_uses_fail(caplog):
+    with pytest.raises(SystemExit):
+        _uses_viv("""__import__("viv").run()\n__import__("viv").use()""")
+    with pytest.raises(SystemExit):
+        _uses_viv("""__import__("viv").unknown()""")
+
+    assert [
+        (
+            "viv",
+            40,
+            "Unexpected number of viv references in script.\n"
+            "Expected only 1, found: run, use",
+        ),
+        ("viv", 40, "Unknown function unknown associated with viv."),
+    ] == caplog.record_tuples