Skip to content

Refactor verify-alpha-spec dependencies.yaml traversal #142

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

Merged
rapids-bot merged 6 commits into from
Jul 20, 2026
Merged

Refactor verify-alpha-spec dependencies.yaml traversal #142

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
4c3d5b8
Refactor `verify-alpha-spec` `dependencies.yaml` traversal
KyleFromNVIDIA Jul 17, 2026
869b473
Propagate parent context by default
KyleFromNVIDIA Jul 17, 2026
34688d3
Type annotation
KyleFromNVIDIA Jul 17, 2026
8edb81b
Traverse matrix
KyleFromNVIDIA Jul 20, 2026
65180d3
Refactor test
KyleFromNVIDIA Jul 20, 2026
99e8b2d
Merge branch 'main' into refactor-dependencies-yaml-traversal
KyleFromNVIDIA Jul 20, 2026
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
336 changes: 72 additions & 264 deletions src/rapids_pre_commit_hooks/alpha_spec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,17 +1,19 @@
# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import argparse
import os
import re
from functools import cache, total_ordering
from typing import Any, Optional

import yaml
from packaging.requirements import InvalidRequirement, Requirement
from rapids_metadata.metadata import RAPIDSMetadata, RAPIDSVersion
from rapids_metadata.remote import fetch_latest

from .lint import Linter, LintMain
from .utils.dependencies_yaml import Handler, traverse_dependencies_yaml

ALPHA_SPECIFIER: str = ">=0.0.0a0"

Expand All @@ -28,10 +30,6 @@ def all_metadata() -> "RAPIDSMetadata":
return fetch_latest()


def node_has_type(node: "yaml.Node", tag_type: str) -> bool:
return node.tag == f"tag:yaml.org,2002:{tag_type}"


def get_rapids_version(args: argparse.Namespace) -> "RAPIDSVersion":
md = all_metadata()
return (
Expand All @@ -49,277 +47,87 @@ def strip_cuda_suffix(args: argparse.Namespace, name: str) -> str:
return name


def check_and_mark_anchor(
anchors: dict[str, "yaml.Node"], used_anchors: set[str], node: "yaml.Node"
) -> tuple[bool, str | None]:
for key, value in anchors.items():
if value == node:
anchor = key
break
else:
anchor = None
if anchor in used_anchors:
return False, anchor
if anchor is not None:
used_anchors.add(anchor)
return True, anchor


def check_package_spec(
linter: Linter,
args: argparse.Namespace,
anchors: dict[str, "yaml.Node"],
used_anchors: set[str],
node: "yaml.Node",
) -> None:
@total_ordering
class SpecPriority:
def __init__(self, spec: str):
self.spec: str = spec

def __eq__(self, other: object) -> bool:
assert isinstance(other, SpecPriority)
return self.spec == other.spec
class AlphaSpecHandler(Handler):
def __init__(self, linter: Linter, args: argparse.Namespace):
self.linter = linter
self.args = args

def handle_package(
self,
packages_context: "Any", # noqa: ARG002
anchor: "Optional[str]",
node: "yaml.Node",
) -> None:
@total_ordering
class SpecPriority:
def __init__(self, spec: str):
self.spec: str = spec

def __eq__(self, other: object) -> bool:
assert isinstance(other, SpecPriority)
return self.spec == other.spec

def __lt__(self, other: object) -> bool:
assert isinstance(other, SpecPriority)
if self.spec == other.spec:
return False
if self.spec == ALPHA_SPECIFIER:
return False
if other.spec == ALPHA_SPECIFIER:
return True
return self.sort_str() < other.sort_str()

def sort_str(self) -> str:
return "".join(c for c in self.spec if c not in "<>=")

def create_specifier_string(specifiers: set[str]) -> str:
return ",".join(sorted(specifiers, key=SpecPriority))

def __lt__(self, other: object) -> bool:
assert isinstance(other, SpecPriority)
if self.spec == other.spec:
return False
if self.spec == ALPHA_SPECIFIER:
return False
if other.spec == ALPHA_SPECIFIER:
return True
return self.sort_str() < other.sort_str()

def sort_str(self) -> str:
return "".join(c for c in self.spec if c not in "<>=")

def create_specifier_string(specifiers: set[str]) -> str:
return ",".join(sorted(specifiers, key=SpecPriority))

if node_has_type(node, "str"):
try:
req = Requirement(node.value)
except InvalidRequirement:
return

if (
strip_cuda_suffix(args, req.name)
in get_rapids_version(args).prerelease_packages
strip_cuda_suffix(self.args, req.name)
not in get_rapids_version(self.args).prerelease_packages
):
descend, anchor = check_and_mark_anchor(
anchors, used_anchors, node
)
if descend:
has_alpha_spec = any(
str(s) == ALPHA_SPECIFIER for s in req.specifier
)
if args.mode == "development" and not has_alpha_spec:
linter.add_warning(
(node.start_mark.index, node.end_mark.index),
f"add alpha spec for RAPIDS package {req.name}",
).add_replacement(
(node.start_mark.index, node.end_mark.index),
str(
(f"&{anchor} " if anchor else "")
+ req.name
+ create_specifier_string(
{str(s) for s in req.specifier}
| {ALPHA_SPECIFIER},
)
),
return

has_alpha_spec = any(str(s) == ALPHA_SPECIFIER for s in req.specifier)
if self.args.mode == "development" and not has_alpha_spec:
self.linter.add_warning(
(node.start_mark.index, node.end_mark.index),
f"add alpha spec for RAPIDS package {req.name}",
).add_replacement(
(node.start_mark.index, node.end_mark.index),
str(
(f"&{anchor} " if anchor else "")
+ req.name
+ create_specifier_string(
{str(s) for s in req.specifier} | {ALPHA_SPECIFIER},
)
elif args.mode == "release" and has_alpha_spec:
linter.add_warning(
(node.start_mark.index, node.end_mark.index),
f"remove alpha spec for RAPIDS package {req.name}",
).add_replacement(
(node.start_mark.index, node.end_mark.index),
str(
(f"&{anchor} " if anchor else "")
+ req.name
+ create_specifier_string(
{str(s) for s in req.specifier}
- {ALPHA_SPECIFIER},
)
),
),
)
elif self.args.mode == "release" and has_alpha_spec:
self.linter.add_warning(
(node.start_mark.index, node.end_mark.index),
f"remove alpha spec for RAPIDS package {req.name}",
).add_replacement(
(node.start_mark.index, node.end_mark.index),
str(
(f"&{anchor} " if anchor else "")
+ req.name
+ create_specifier_string(
{str(s) for s in req.specifier} - {ALPHA_SPECIFIER},
)


def check_packages(
linter: Linter,
args: argparse.Namespace,
anchors: dict[str, "yaml.Node"],
used_anchors: set[str],
node: "yaml.Node",
) -> None:
if node_has_type(node, "seq"):
descend, _ = check_and_mark_anchor(anchors, used_anchors, node)
if descend:
for package_spec in node.value:
check_package_spec(
linter, args, anchors, used_anchors, package_spec
)


def check_common(
linter: Linter,
args: argparse.Namespace,
anchors: dict[str, "yaml.Node"],
used_anchors: set[str],
node: "yaml.Node",
) -> None:
if node_has_type(node, "seq"):
for dependency_set in node.value:
if node_has_type(dependency_set, "map"):
for (
dependency_set_key,
dependency_set_value,
) in dependency_set.value:
if (
node_has_type(dependency_set_key, "str")
and dependency_set_key.value == "packages"
):
check_packages(
linter,
args,
anchors,
used_anchors,
dependency_set_value,
)


def check_matrices(
linter: Linter,
args: argparse.Namespace,
anchors: dict[str, "yaml.Node"],
used_anchors: set[str],
node: "yaml.Node",
) -> None:
if node_has_type(node, "seq"):
for item in node.value:
if node_has_type(item, "map"):
for matrix_key, matrix_value in item.value:
if (
node_has_type(matrix_key, "str")
and matrix_key.value == "packages"
):
check_packages(
linter, args, anchors, used_anchors, matrix_value
)


def check_specific(
linter: Linter,
args: argparse.Namespace,
anchors: dict[str, "yaml.Node"],
used_anchors: set[str],
node: "yaml.Node",
) -> None:
if node_has_type(node, "seq"):
for matrix_matcher in node.value:
if node_has_type(matrix_matcher, "map"):
for (
matrix_matcher_key,
matrix_matcher_value,
) in matrix_matcher.value:
if (
node_has_type(matrix_matcher_key, "str")
and matrix_matcher_key.value == "matrices"
):
check_matrices(
linter,
args,
anchors,
used_anchors,
matrix_matcher_value,
)


def check_dependencies(
linter: Linter,
args: argparse.Namespace,
anchors: dict[str, "yaml.Node"],
used_anchors: set[str],
node: "yaml.Node",
) -> None:
if node_has_type(node, "map"):
for _, dependencies_value in node.value:
if node_has_type(dependencies_value, "map"):
for (
dependency_key,
dependency_value,
) in dependencies_value.value:
if node_has_type(dependency_key, "str"):
if dependency_key.value == "common":
check_common(
linter,
args,
anchors,
used_anchors,
dependency_value,
)
elif dependency_key.value == "specific":
check_specific(
linter,
args,
anchors,
used_anchors,
dependency_value,
)


def check_root(
linter: Linter,
args: argparse.Namespace,
anchors: dict[str, "yaml.Node"],
used_anchors: set[str],
node: "yaml.Node",
) -> None:
if node_has_type(node, "map"):
for root_key, root_value in node.value:
if (
node_has_type(root_key, "str")
and root_key.value == "dependencies"
):
check_dependencies(
linter, args, anchors, used_anchors, root_value
)


class AnchorPreservingLoader(yaml.SafeLoader):
"""A SafeLoader that preserves the anchors for later reference. The anchors
can be found in the document_anchors member, which is a list of
dictionaries, one dictionary for each parsed document.
"""

def __init__(self, stream) -> None:
super().__init__(stream)
self.document_anchors: list[dict[str, yaml.Node]] = []

def compose_document(self) -> "yaml.Node":
# Drop the DOCUMENT-START event.
self.get_event()

# Compose the root node.
node = self.compose_node(None, None) # type: ignore[arg-type]

# Drop the DOCUMENT-END event.
self.get_event()

self.document_anchors.append(self.anchors)
self.anchors = {}
assert node is not None
return node
),
)


def check_alpha_spec(linter: Linter, args: argparse.Namespace) -> None:
loader = AnchorPreservingLoader(linter.content)
try:
root = loader.get_single_node()
assert root is not None
finally:
loader.dispose()
check_root(linter, args, loader.document_anchors[0], set(), root)
handler = AlphaSpecHandler(linter, args)
traverse_dependencies_yaml(handler, linter.content)


def main() -> None:
Expand Down
Loading
Loading