4cfc63e74f
Every new .py this PR adds now carries the same two-line SPDX header the other new files in the branch already use (docker/jupyter/unsloth_branding.py), with the shebang kept first where present. Matches the licensing laid out in docker/NOTICE: the image bundles Studio (AGPL-3.0) while Unsloth Core stays Apache-2.0.
76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Regression tests for docker/unsloth_nb_pip_magic.py.
|
|
|
|
The input transformer rewrites explicit `!<python> -m pip|uv ...` shell lines
|
|
to `!pip|uv ...` so they resolve to the PATH shim. IPython input transformers
|
|
see the RAW cell text (brace expansion like `{sys.executable}` happens later,
|
|
in the system() execution path), so the braced and absolute-interpreter forms
|
|
notebooks use to target the running kernel must be rewritten too (item
|
|
3567875025); only matching literal `python`/`py` let module-pip bypass the
|
|
shim entirely.
|
|
"""
|
|
|
|
import importlib.util
|
|
import pathlib
|
|
|
|
_MOD_PATH = pathlib.Path(__file__).resolve().parents[2] / "docker" / "unsloth_nb_pip_magic.py"
|
|
_spec = importlib.util.spec_from_file_location("unsloth_nb_pip_magic", _MOD_PATH)
|
|
magic = importlib.util.module_from_spec(_spec)
|
|
_spec.loader.exec_module(magic)
|
|
|
|
|
|
def _rewrite(line):
|
|
return magic._rewrite_python_dash_m([line])[0]
|
|
|
|
|
|
def test_literal_python_rewritten():
|
|
assert _rewrite("!python -m pip install peft\n") == "!pip install peft\n"
|
|
|
|
|
|
def test_literal_python_version_rewritten():
|
|
assert _rewrite("!python3.12 -m pip install peft") == "!pip install peft"
|
|
|
|
|
|
def test_sys_executable_braces_rewritten():
|
|
assert _rewrite("!{sys.executable} -m pip install peft\n") == "!pip install peft\n"
|
|
|
|
|
|
def test_sys_executable_braces_quoted_rewritten():
|
|
assert _rewrite('!"{sys.executable}" -m pip install peft') == "!pip install peft"
|
|
|
|
|
|
def test_sys_executable_braces_spaced_rewritten():
|
|
assert _rewrite("!{ sys.executable } -m pip install peft") == "!pip install peft"
|
|
|
|
|
|
def test_absolute_interpreter_path_rewritten():
|
|
assert _rewrite("!/opt/unsloth-venv/bin/python -m pip install peft\n") == "!pip install peft\n"
|
|
|
|
|
|
def test_absolute_interpreter_versioned_path_rewritten():
|
|
assert _rewrite("!/usr/bin/python3.11 -m uv pip install peft") == "!uv pip install peft"
|
|
|
|
|
|
def test_quoted_interpreter_path_rewritten():
|
|
assert _rewrite('!"/opt/unsloth venv/bin/python" -m pip install peft') == "!pip install peft"
|
|
|
|
|
|
def test_indent_preserved():
|
|
assert _rewrite(" !{sys.executable} -m pip install peft") == " !pip install peft"
|
|
|
|
|
|
def test_python_script_not_rewritten():
|
|
line = "!python train.py --epochs 3"
|
|
assert _rewrite(line) == line
|
|
|
|
|
|
def test_module_other_than_pip_not_rewritten():
|
|
line = "!python -m venv .venv"
|
|
assert _rewrite(line) == line
|
|
|
|
|
|
def test_non_shell_line_not_rewritten():
|
|
line = "x = '{sys.executable} -m pip install peft'"
|
|
assert _rewrite(line) == line
|