decorator

[koaning]

Add a decorator to a function to make it run inside of a uv env. Cool for versions.

[lmmx]

This can be achieved it just requires removing the decorator when inspecting the source of the function (IDK if there's a better way to remove the decorator, I suspect so...):

diff --git a/uvtrick/__init__.py b/uvtrick/__init__.py
index 8f57ddd..55d9286 100644
--- a/uvtrick/__init__.py
+++ b/uvtrick/__init__.py
@@ -137,7 +137,10 @@ class Env:
             # First pickle the inputs
             self.inputs.write_bytes(pickle.dumps((args, kwargs)))
             # Now write the contents of the script
-            contents = dedent(getsource(func)) + "\n\n" + self.maincall(func)
+            funcdef_code = dedent(getsource(func))
+            if funcdef_code.startswith("@with_package_versions"):
+                funcdef_code = funcdef_code[funcdef_code.find("\ndef "):]
+            contents = funcdef_code + "\n\n" + self.maincall(func)
             self.script.write_text(contents)
             # Then run the script with `uv run` in a subprocess
             if self.debug:

Then demo with:

from functools import wraps
from uvtrick import Env

def with_package_versions(package_name, versions, python_version="3.12"):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            results = []
            for version in versions:
                env = Env(f"{package_name}=={version}", python=python_version)
                result = env.run(func, *args, **kwargs)
                results.append((version, result))
            return results

        return wrapper

    return decorator

@with_package_versions("rich", versions=[10, 11, 12, 13])
def test_rich(a, b):
    from rich import print
    from importlib import metadata

    version = metadata.version("rich")
    print(f"Hello from rich=={version}")
    return a + b

if __name__ == "__main__":
    results = test_rich(1, 2)
    for version, result in results:
        print(f"Rich version {version}: add(1, 2) = {result}")

Works as expected. (The with_package_versions function would go in the uvtricks lib instead of the example of course)

[koaning]

The with_package_versions feels like it might be better if it was more done like this:

@runwith("rich==13", "scikit-learn==1.5", python="3.12")
def somefunc(a, b):
    import rich
    import sklearn
    # do stuff 
    return ...

A function might need to specify more than one dependency and it would be best to configure all of that with a single decorator.