From 2d951c759a4c1e82fb13dc82b74225725dc41c1c Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Thu, 2 Jul 2026 19:28:44 +0200 Subject: [PATCH] Fail closed on class-body namespace aliases --- .../test_generate_popular_node_signatures.py | 128 ++++++++++++++++++ tools/generate_popular_node_signatures.py | 41 ++++-- 2 files changed, 161 insertions(+), 8 deletions(-) diff --git a/tests/test_generate_popular_node_signatures.py b/tests/test_generate_popular_node_signatures.py index 208b058..f7c0676 100644 --- a/tests/test_generate_popular_node_signatures.py +++ b/tests/test_generate_popular_node_signatures.py @@ -303,6 +303,50 @@ NODE_CLASS_MAPPINGS = { self.assertEqual({}, result["nodes"]) self.assertEqual("no_static_nodes", result["pack"]["status"]) + def test_duplicate_node_id_with_invalid_duplicate_mapping_literal_skips_static_node(self): + source_a = ''' +class StaticDupNode: + RETURN_TYPES = ("IMAGE",) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + }, + } + + +NODE_CLASS_MAPPINGS = { + "DupNode": StaticDupNode, +} +''' + source_b = ''' +def build_node(): + return object() + + +NODE_CLASS_MAPPINGS = { + "DupNode": build_node(), + "DupNode": build_node(), +} +''' + with tempfile.TemporaryDirectory() as tmp: + Path(tmp, "a.py").write_text(textwrap.dedent(source_a), encoding="utf-8") + Path(tmp, "b.py").write_text(textwrap.dedent(source_b), encoding="utf-8") + result = extract_repo_signatures( + Path(tmp), + { + "id": "invalid-duplicate-node-pack", + "title": "Invalid Duplicate Node Pack", + "repository": "https://github.com/example/invalid-duplicate-node-pack", + "rank": 1, + }, + ) + + self.assertEqual({}, result["nodes"]) + self.assertEqual("no_static_nodes", result["pack"]["status"]) + def test_unsupported_reassignment_invalidates_static_env_value(self): source = ''' def build_inputs(): @@ -3433,6 +3477,34 @@ H["NODE_CLASS_MAPPINGS"] = {} self.assertEqual({}, result["nodes"]) self.assertEqual("no_static_nodes", result["pack"]["status"]) + def test_class_body_globals_alias_subscript_assignment_invalidates_static_node_mapping(self): + source = ''' +class ClassBodyGlobalAliasSubscriptAssignmentNode: + RETURN_TYPES = ("IMAGE",) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + }, + } + + +NODE_CLASS_MAPPINGS = { + "ClassBodyGlobalAliasSubscriptAssignmentNode": ClassBodyGlobalAliasSubscriptAssignmentNode, +} + + +class MappingMutator: + ns = globals() + ns["NODE_CLASS_MAPPINGS"] = {} +''' + result = self._extract_source(source, "class-body-global-alias-subscript-assignment-pack") + + self.assertEqual({}, result["nodes"]) + self.assertEqual("no_static_nodes", result["pack"]["status"]) + def test_globals_alias_update_invalidates_static_node_mapping(self): source = ''' class GlobalAliasUpdateNode: @@ -3458,6 +3530,34 @@ G.update(NODE_CLASS_MAPPINGS={}) self.assertEqual({}, result["nodes"]) self.assertEqual("no_static_nodes", result["pack"]["status"]) + def test_class_body_globals_alias_update_invalidates_static_node_mapping(self): + source = ''' +class ClassBodyGlobalAliasUpdateNode: + RETURN_TYPES = ("IMAGE",) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + }, + } + + +NODE_CLASS_MAPPINGS = { + "ClassBodyGlobalAliasUpdateNode": ClassBodyGlobalAliasUpdateNode, +} + + +class MappingMutator: + ns = globals() + ns.update(NODE_CLASS_MAPPINGS={}) +''' + result = self._extract_source(source, "class-body-global-alias-update-pack") + + self.assertEqual({}, result["nodes"]) + self.assertEqual("no_static_nodes", result["pack"]["status"]) + def test_globals_alias_get_mutation_invalidates_static_node_mapping(self): source = ''' class GlobalAliasGetMutationNode: @@ -3483,6 +3583,34 @@ G.get("NODE_CLASS_MAPPINGS").clear() self.assertEqual({}, result["nodes"]) self.assertEqual("no_static_nodes", result["pack"]["status"]) + def test_class_body_globals_chained_alias_subscript_assignment_invalidates_static_node_mapping(self): + source = ''' +class ClassBodyGlobalChainedAliasSubscriptAssignmentNode: + RETURN_TYPES = ("IMAGE",) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + }, + } + + +NODE_CLASS_MAPPINGS = { + "ClassBodyGlobalChainedAliasSubscriptAssignmentNode": ClassBodyGlobalChainedAliasSubscriptAssignmentNode, +} + + +class MappingMutator: + ns = other = globals() + other["NODE_CLASS_MAPPINGS"] = {} +''' + result = self._extract_source(source, "class-body-global-chained-alias-subscript-assignment-pack") + + self.assertEqual({}, result["nodes"]) + self.assertEqual("no_static_nodes", result["pack"]["status"]) + def test_globals_dunder_setitem_invalidates_static_node_mapping(self): source = ''' class GlobalDunderSetitemNode: diff --git a/tools/generate_popular_node_signatures.py b/tools/generate_popular_node_signatures.py index bf07f22..31561f8 100644 --- a/tools/generate_popular_node_signatures.py +++ b/tools/generate_popular_node_signatures.py @@ -842,6 +842,7 @@ def _class_body_global_names(cls): def _class_body_module_mutation_names(cls): global_names = _class_body_global_names(cls) names = set() + namespace_aliases = set() def add_assignment_targets(stmt): names.update(_assignment_target_names(stmt).intersection(global_names)) @@ -926,7 +927,9 @@ def _class_body_module_mutation_names(cls): self.generic_visit(node) for stmt in cls.body: + names.update(_namespace_alias_mutation_target_names(stmt, namespace_aliases)) ClassBodyMutationVisitor().visit(stmt) + _update_namespace_aliases(stmt, namespace_aliases) return names @@ -2022,17 +2025,39 @@ def _node_class_mappings(tree): return {node_type: binding for node_type, (_class_name, binding) in mappings.items() if node_type} +def _literal_module_dict_string_keys(node, env): + if not isinstance(node, ast.Dict): + return set() + keys = set() + for key in node.keys: + if key is None: + continue + try: + key_value = _literal(key, env) + except UnsupportedStaticExpression: + continue + if isinstance(key_value, str) and key_value: + keys.add(key_value) + return keys + + def _node_class_mapping_keys(tree): if _has_module_wildcard_import(tree): return set() - mappings = _final_module_dict( - tree, - "NODE_CLASS_MAPPINGS", - lambda _value, _env, _class_bindings: True, - ) - if not all(isinstance(node_type, str) for node_type in mappings): - return set() - return {node_type for node_type in mappings if node_type} + keys = set() + env = {} + class_bindings = {} + for stmt in tree.body: + if isinstance(stmt, ast.Assign) and _name_is_assigned(stmt, "NODE_CLASS_MAPPINGS"): + keys.update(_literal_module_dict_string_keys(stmt.value, env)) + elif ( + isinstance(stmt, ast.AnnAssign) + and _name_is_assigned(stmt, "NODE_CLASS_MAPPINGS") + and stmt.value is not None + ): + keys.update(_literal_module_dict_string_keys(stmt.value, env)) + _apply_module_stmt_to_env(stmt, env, class_bindings) + return keys def _display_mappings(tree):