Skip to content

Add --use-uv flag for faster venv creation using uv #409

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
budaesandrei wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add --use-uv flag for faster venv creation using uv #409

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,9 @@ python3 create_venvs.py --target <development | production>

# Or create venv for a specific version only
python3 create_venvs.py --target <development | production> --version 3.0.6

# Optionally use uv for faster package installation (requires uv: https://docs.astral.sh/uv/)
python3 create_venvs.py --target <development | production> --use-uv
```

3. Build a supported Airflow version Docker image
Expand Down
97 changes: 76 additions & 21 deletions create_venvs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,15 +32,24 @@ def verify_python_version():
sys.exit(1)


def create_venv(path: Path, development_build: bool, recreate: bool = False):
def is_uv_available() -> bool:
"""Check if uv is available on PATH."""
return shutil.which("uv") is not None


def create_venv(
path: Path, development_build: bool, recreate: bool = False, use_uv: bool = False
):
"""
Create a venv in the given directory and optionally recreate it if it already exists.

:param path: The path to create the venv in.
:param development_build: Is this a development build.
:param recreate: Whether to recreate the venv if it already exists.
:param use_uv: Whether to use uv instead of pip for package installation.
"""
venv_path = path / ".venv"
installer = uv_install if use_uv else pip_install
print(f">>> Creating a virtual environment under the path {venv_path}...")

if recreate and venv_path.exists():
Expand All @@ -49,28 +58,33 @@ def create_venv(path: Path, development_build: bool, recreate: bool = False):

if not venv_path.exists():
print(f"> Creating virtualenv in directory: {venv_path}")
venv.create(venv_path, with_pip=True, symlinks=True)
if use_uv:
subprocess.run(["uv", "venv", str(venv_path)], check=True)
else:
venv.create(venv_path, with_pip=True, symlinks=True)
else:
print(f"> Virtualenv already exists in {venv_path}")

print("> Upgrade pip...")
pip_install(venv_path, "-U", "pip")
print("")
if not use_uv:
print("> Upgrade pip...")
pip_install(venv_path, "-U", "pip")
print("")

requirements_path = generate_requirements(path, development_build)
print(f"> Install dependencies from {requirements_path}...")
pip_install(venv_path, "-r", str(requirements_path))
installer(venv_path, "-r", str(requirements_path))
print("")

dev_tools = ["pydocstyle", "pyright", "ruff"]
print(f"> Install/Upgrade development tools: {dev_tools}...")
pip_install(venv_path, "-U", *dev_tools)
installer(venv_path, "-U", *dev_tools)
print("")

print(f">>> Finished creating a virtual environment under the path {venv_path}.")
print("")
print("")


def generate_requirements(path: Path, development_build: bool) -> Path:
"""
If the requirements.txt file at the path needs to be updated for local development, generate
Expand All @@ -88,20 +102,23 @@ def generate_requirements(path: Path, development_build: bool) -> Path:
return requirements_path

if not re.search(r"images\/airflow\/[2-3]\.[0-9]+\.[0-9]+$", str(path.resolve())):
print(f"> No need to create dev requirements for {path.resolve()}. Using default.")
print(
f"> No need to create dev requirements for {path.resolve()}. Using default."
)
return requirements_path

with open(requirements_path.resolve(), 'r') as file:
with open(requirements_path.resolve(), "r") as file:
# psycopg2-binary is meant for development and removes the requirement to install pg_config
filedata = re.sub(r"\bpsycopg2\b", "psycopg2-binary", file.read())

dev_requirements_path = path.joinpath('requirements-dev.txt')
dev_requirements_path = path.joinpath("requirements-dev.txt")
print(f"> Creating {dev_requirements_path} from {requirements_path}")
with open(dev_requirements_path.resolve(), 'w') as file:
with open(dev_requirements_path.resolve(), "w") as file:
file.write(filedata)

return dev_requirements_path


def pip_install(venv_dir: Path, *args: str):
"""
Install dependencies from requirements.txt if it exists.
Expand All @@ -115,6 +132,26 @@ def pip_install(venv_dir: Path, *args: str):
)


def uv_install(venv_dir: Path, *args: str):
"""
Install dependencies using uv for faster resolution and installation.

:param venv_dir: The path to the venv directory.
:param args: Arguments to pass to uv pip install.
"""
subprocess.run(
[
"uv",
"pip",
"install",
"--python",
os.path.join(venv_dir, "bin", "python"),
*args,
],
check=True,
)


def main():
"""Start execution of the script."""
# Create the parser
Expand All @@ -129,33 +166,50 @@ def main():
parser.add_argument(
"--target", choices=build_targets, required=True, help="Sets the build target"
)

# Add version filter argument
parser.add_argument(
"--version", type=str, help="Only create venv for specific Airflow version (e.g., 3.0.6)"
"--version",
type=str,
help="Only create venv for specific Airflow version (e.g., 3.0.6)",
)

# Add uv flag for faster package installation
parser.add_argument(
"--use-uv",
action="store_true",
help="Use uv instead of pip for faster package installation (requires uv to be installed)",
)

# Parse the arguments
args = parser.parse_args()

verify_python_version()


if args.use_uv and not is_uv_available():
print("ERROR: --use-uv was specified but 'uv' is not installed or not on PATH.")
print("Install it with: pip install uv (or see https://docs.astral.sh/uv/)")
sys.exit(1)

# Filter directories based on version argument
if args.version:
# Validate that the version exists
version_path = Path(f"./images/airflow/{args.version}")
if not version_path.exists() or not version_path.is_dir():
# Get available versions
available_versions = sorted([
d.name for d in Path("./images/airflow").iterdir()
if d.is_dir() and not d.name.startswith('.')
])
available_versions = sorted(
[
d.name
for d in Path("./images/airflow").iterdir()
if d.is_dir() and not d.name.startswith(".")
]
)
print(f"ERROR: Version '{args.version}' not found in images/airflow/")
print("\nAvailable versions:")
for v in available_versions:
print(f" - {v}")
sys.exit(1)

project_dirs = [
Path("."),
version_path,
Expand All @@ -165,13 +219,14 @@ def main():
Path("."),
*Path("./images").glob("airflow/*"),
] # Include main project dir and each image dir

for dir_path in project_dirs:
if dir_path.is_dir() and (dir_path / "requirements.txt").exists():
create_venv(
dir_path,
development_build=args.target == development_target_choice,
recreate=args.recreate
recreate=args.recreate,
use_uv=args.use_uv,
)


Expand Down