diff --git a/src/rapids_pre_commit_hooks/alpha_spec.py b/src/rapids_pre_commit_hooks/alpha_spec.py index 6ea796e..ceddf13 100644 --- a/src/rapids_pre_commit_hooks/alpha_spec.py +++ b/src/rapids_pre_commit_hooks/alpha_spec.py @@ -79,7 +79,7 @@ def is_rapids_cuda_suffixed_package(name): ) -def check_package_spec(linter, args, node): +def check_package_spec(linter, args, anchors, used_anchors, node): @total_ordering class SpecPriority: def __init__(self, spec): @@ -111,42 +111,53 @@ def create_specifier_string(specifiers): if req.name in RAPIDS_ALPHA_SPEC_PACKAGES or is_rapids_cuda_suffixed_package( req.name ): - 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( - 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( - req.name - + create_specifier_string( - {str(s) for s in req.specifier} - {ALPHA_SPECIFIER} - ) - ), - ) - - -def check_packages(linter, args, node): + for key, value in anchors.items(): + if value == node: + anchor = key + break + else: + anchor = None + if anchor not in used_anchors: + if anchor is not None: + used_anchors.add(anchor) + 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}, + ) + ), + ) + 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}, + ) + ), + ) + + +def check_packages(linter, args, anchors, used_anchors, node): if node_has_type(node, "seq"): for package_spec in node.value: - check_package_spec(linter, args, package_spec) + check_package_spec(linter, args, anchors, used_anchors, package_spec) -def check_common(linter, args, node): +def check_common(linter, args, anchors, used_anchors, node): if node_has_type(node, "seq"): for dependency_set in node.value: if node_has_type(dependency_set, "map"): @@ -155,10 +166,12 @@ def check_common(linter, args, node): node_has_type(dependency_set_key, "str") and dependency_set_key.value == "packages" ): - check_packages(linter, args, dependency_set_value) + check_packages( + linter, args, anchors, used_anchors, dependency_set_value + ) -def check_matrices(linter, args, node): +def check_matrices(linter, args, anchors, used_anchors, node): if node_has_type(node, "seq"): for item in node.value: if node_has_type(item, "map"): @@ -167,10 +180,12 @@ def check_matrices(linter, args, node): node_has_type(matrix_key, "str") and matrix_key.value == "packages" ): - check_packages(linter, args, matrix_value) + check_packages( + linter, args, anchors, used_anchors, matrix_value + ) -def check_specific(linter, args, node): +def check_specific(linter, args, anchors, used_anchors, node): if node_has_type(node, "seq"): for matrix_matcher in node.value: if node_has_type(matrix_matcher, "map"): @@ -179,30 +194,66 @@ def check_specific(linter, args, node): node_has_type(matrix_matcher_key, "str") and matrix_matcher_key.value == "matrices" ): - check_matrices(linter, args, matrix_matcher_value) + check_matrices( + linter, args, anchors, used_anchors, matrix_matcher_value + ) -def check_dependencies(linter, args, node): +def check_dependencies(linter, args, anchors, used_anchors, node): 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, dependency_value) + check_common( + linter, args, anchors, used_anchors, dependency_value + ) elif dependency_key.value == "specific": - check_specific(linter, args, dependency_value) + check_specific( + linter, args, anchors, used_anchors, dependency_value + ) -def check_root(linter, args, node): +def check_root(linter, args, anchors, used_anchors, node): 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, root_value) + 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): + super().__init__(stream) + self.document_anchors = [] + + def compose_document(self): + # Drop the DOCUMENT-START event. + self.get_event() + + # Compose the root node. + node = self.compose_node(None, None) + + # Drop the DOCUMENT-END event. + self.get_event() + + self.document_anchors.append(self.anchors) + self.anchors = {} + return node def check_alpha_spec(linter, args): - check_root(linter, args, yaml.compose(linter.content)) + loader = AnchorPreservingLoader(linter.content) + try: + root = loader.get_single_node() + finally: + loader.dispose() + check_root(linter, args, loader.document_anchors[0], set(), root) def main(): diff --git a/test/rapids_pre_commit_hooks/test_alpha_spec.py b/test/rapids_pre_commit_hooks/test_alpha_spec.py index 8e75817..ddcd685 100644 --- a/test/rapids_pre_commit_hooks/test_alpha_spec.py +++ b/test/rapids_pre_commit_hooks/test_alpha_spec.py @@ -14,7 +14,7 @@ from itertools import chain from textwrap import dedent -from unittest.mock import Mock, call, patch +from unittest.mock import MagicMock, Mock, call, patch import pytest import yaml @@ -22,6 +22,15 @@ from rapids_pre_commit_hooks import alpha_spec, lint +def test_anchor_preserving_loader(): + loader = alpha_spec.AnchorPreservingLoader("- &a A\n- *a") + try: + root = loader.get_single_node() + finally: + loader.dispose() + assert loader.document_anchors == [{"a": root.value[0]}] + + @pytest.mark.parametrize( ["name", "is_suffixed"], [ @@ -91,7 +100,7 @@ def test_is_rapids_cuda_suffixed_package(name, is_suffixed): "cuml", "&cuml cuml>=24.04,<24.06,>=0.0.0a0", "release", - "cuml>=24.04,<24.06", + "&cuml cuml>=24.04,<24.06", ), ("packaging", "packaging", "development", None), (None, "--extra-index-url=https://pypi.nvidia.com", "development", None), @@ -103,8 +112,14 @@ def test_is_rapids_cuda_suffixed_package(name, is_suffixed): def test_check_package_spec(package, content, mode, replacement): args = Mock(mode=mode) linter = lint.Linter("dependencies.yaml", content) - composed = yaml.compose(content) - alpha_spec.check_package_spec(linter, args, composed) + loader = alpha_spec.AnchorPreservingLoader(content) + try: + composed = loader.get_single_node() + finally: + loader.dispose() + alpha_spec.check_package_spec( + linter, args, loader.document_anchors[0], set(), composed + ) if replacement is None: assert linter.warnings == [] else: @@ -113,10 +128,66 @@ def test_check_package_spec(package, content, mode, replacement): (composed.start_mark.index, composed.end_mark.index), f"{'add' if mode == 'development' else 'remove'} " f"alpha spec for RAPIDS package {package}", - ).add_replacement((0, len(content)), replacement) + ).add_replacement( + (composed.start_mark.index, composed.end_mark.index), replacement + ) assert linter.warnings == expected_linter.warnings +def test_check_package_spec_anchor(): + CONTENT = dedent( + """\ + - &cudf cudf>=24.04,<24.06 + - *cudf + - cuml>=24.04,<24.06 + - rmm>=24.04,<24.06 + """ + ) + args = Mock(mode="development") + linter = lint.Linter("dependencies.yaml", CONTENT) + loader = alpha_spec.AnchorPreservingLoader(CONTENT) + try: + composed = loader.get_single_node() + finally: + loader.dispose() + used_anchors = set() + + expected_linter = lint.Linter("dependencies.yaml", CONTENT) + expected_linter.add_warning( + (2, 26), "add alpha spec for RAPIDS package cudf" + ).add_replacement((2, 26), "&cudf cudf>=24.04,<24.06,>=0.0.0a0") + + alpha_spec.check_package_spec( + linter, args, loader.document_anchors[0], used_anchors, composed.value[0] + ) + assert linter.warnings == expected_linter.warnings + assert used_anchors == {"cudf"} + + alpha_spec.check_package_spec( + linter, args, loader.document_anchors[0], used_anchors, composed.value[1] + ) + assert linter.warnings == expected_linter.warnings + assert used_anchors == {"cudf"} + + expected_linter.add_warning( + (37, 55), "add alpha spec for RAPIDS package cuml" + ).add_replacement((37, 55), "cuml>=24.04,<24.06,>=0.0.0a0") + alpha_spec.check_package_spec( + linter, args, loader.document_anchors[0], used_anchors, composed.value[2] + ) + assert linter.warnings == expected_linter.warnings + assert used_anchors == {"cudf"} + + expected_linter.add_warning( + (58, 75), "add alpha spec for RAPIDS package rmm" + ).add_replacement((58, 75), "rmm>=24.04,<24.06,>=0.0.0a0") + alpha_spec.check_package_spec( + linter, args, loader.document_anchors[0], used_anchors, composed.value[3] + ) + assert linter.warnings == expected_linter.warnings + assert used_anchors == {"cudf"} + + @pytest.mark.parametrize( ["content", "indices"], [ @@ -141,10 +212,12 @@ def test_check_packages(content, indices): ) as mock_check_package_spec: args = Mock() linter = lint.Linter("dependencies.yaml", content) + anchors = Mock() + used_anchors = Mock() composed = yaml.compose(content) - alpha_spec.check_packages(linter, args, composed) + alpha_spec.check_packages(linter, args, anchors, used_anchors, composed) assert mock_check_package_spec.mock_calls == [ - call(linter, args, composed.value[i]) for i in indices + call(linter, args, anchors, used_anchors, composed.value[i]) for i in indices ] @@ -175,10 +248,13 @@ def test_check_common(content, indices): ) as mock_check_packages: args = Mock() linter = lint.Linter("dependencies.yaml", content) + anchors = Mock() + used_anchors = Mock() composed = yaml.compose(content) - alpha_spec.check_common(linter, args, composed) + alpha_spec.check_common(linter, args, anchors, used_anchors, composed) assert mock_check_packages.mock_calls == [ - call(linter, args, composed.value[i].value[j][1]) for i, j in indices + call(linter, args, anchors, used_anchors, composed.value[i].value[j][1]) + for i, j in indices ] @@ -207,10 +283,13 @@ def test_check_matrices(content, indices): ) as mock_check_packages: args = Mock() linter = lint.Linter("dependencies.yaml", content) + anchors = Mock() + used_anchors = Mock() composed = yaml.compose(content) - alpha_spec.check_matrices(linter, args, composed) + alpha_spec.check_matrices(linter, args, anchors, used_anchors, composed) assert mock_check_packages.mock_calls == [ - call(linter, args, composed.value[i].value[j][1]) for i, j in indices + call(linter, args, anchors, used_anchors, composed.value[i].value[j][1]) + for i, j in indices ] @@ -250,10 +329,13 @@ def test_check_specific(content, indices): ) as mock_check_matrices: args = Mock() linter = lint.Linter("dependencies.yaml", content) + anchors = Mock() + used_anchors = Mock() composed = yaml.compose(content) - alpha_spec.check_specific(linter, args, composed) + alpha_spec.check_specific(linter, args, anchors, used_anchors, composed) assert mock_check_matrices.mock_calls == [ - call(linter, args, composed.value[i].value[j][1]) for i, j in indices + call(linter, args, anchors, used_anchors, composed.value[i].value[j][1]) + for i, j in indices ] @@ -302,13 +384,16 @@ def test_check_dependencies(content, common_indices, specific_indices): ) as mock_check_specific: args = Mock() linter = lint.Linter("dependencies.yaml", content) + anchors = Mock() + used_anchors = Mock() composed = yaml.compose(content) - alpha_spec.check_dependencies(linter, args, composed) + alpha_spec.check_dependencies(linter, args, anchors, used_anchors, composed) assert mock_check_common.mock_calls == [ - call(linter, args, composed.value[i][1].value[j][1]) for i, j in common_indices + call(linter, args, anchors, used_anchors, composed.value[i][1].value[j][1]) + for i, j in common_indices ] assert mock_check_specific.mock_calls == [ - call(linter, args, composed.value[i][1].value[j][1]) + call(linter, args, anchors, used_anchors, composed.value[i][1].value[j][1]) for i, j in specific_indices ] @@ -334,10 +419,12 @@ def test_check_root(content, indices): ) as mock_check_dependencies: args = Mock() linter = lint.Linter("dependencies.yaml", content) + anchors = Mock() + used_anchors = Mock() composed = yaml.compose(content) - alpha_spec.check_root(linter, args, composed) + alpha_spec.check_root(linter, args, anchors, used_anchors, composed) assert mock_check_dependencies.mock_calls == [ - call(linter, args, composed.value[i][1]) for i in indices + call(linter, args, anchors, used_anchors, composed.value[i][1]) for i in indices ] @@ -345,9 +432,45 @@ def test_check_alpha_spec(): CONTENT = "dependencies: []" with patch( "rapids_pre_commit_hooks.alpha_spec.check_root", Mock() - ) as mock_check_root, patch("yaml.compose", Mock()) as mock_yaml_compose: + ) as mock_check_root, patch( + "rapids_pre_commit_hooks.alpha_spec.AnchorPreservingLoader", MagicMock() + ) as mock_anchor_preserving_loader: args = Mock() linter = lint.Linter("dependencies.yaml", CONTENT) alpha_spec.check_alpha_spec(linter, args) - mock_yaml_compose.assert_called_once_with(CONTENT) - mock_check_root.assert_called_once_with(linter, args, mock_yaml_compose()) + mock_anchor_preserving_loader.assert_called_once_with(CONTENT) + mock_check_root.assert_called_once_with( + linter, + args, + mock_anchor_preserving_loader().document_anchors[0], + set(), + mock_anchor_preserving_loader().get_single_node(), + ) + + +def test_check_alpha_spec_integration(): + CONTENT = dedent( + """\ + dependencies: + test: + common: + - output_types: pyproject + packages: + - cudf>=24.04,<24.06 + """ + ) + REPLACED = "cudf>=24.04,<24.06" + + args = Mock(mode="development") + linter = lint.Linter("dependencies.yaml", CONTENT) + alpha_spec.check_alpha_spec(linter, args) + + start = CONTENT.find(REPLACED) + end = start + len(REPLACED) + pos = (start, end) + + expected_linter = lint.Linter("dependencies.yaml", CONTENT) + expected_linter.add_warning( + pos, "add alpha spec for RAPIDS package cudf" + ).add_replacement(pos, "cudf>=24.04,<24.06,>=0.0.0a0") + assert linter.warnings == expected_linter.warnings