Require mapped classes to remain top-level bindings

tests/test_generate_popular_node_signatures.py

@@ -842,6 +842,82 @@ NODE_CLASS_MAPPINGS = build_mappings()
         self.assertEqual({}, result["nodes"])
         self.assertEqual("no_static_nodes", result["pack"]["status"])
 
+    def test_rebound_node_class_name_skips_static_mapping(self):
+        source = '''
+def build_node():
+    return object()
+
+
+class ReboundNode:
+    RETURN_TYPES = ("IMAGE",)
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+            },
+        }
+
+
+ReboundNode = build_node()
+
+NODE_CLASS_MAPPINGS = {
+    "ReboundNode": ReboundNode,
+}
+'''
+        result = self._extract_source(source, "rebound-node-class-pack")
+
+        self.assertEqual({}, result["nodes"])
+        self.assertEqual("no_static_nodes", result["pack"]["status"])
+
+    def test_conditional_class_mapping_skips_node(self):
+        source = '''
+if True:
+    class ConditionalNode:
+        RETURN_TYPES = ("IMAGE",)
+
+        @classmethod
+        def INPUT_TYPES(cls):
+            return {
+                "required": {
+                    "image": ("IMAGE",),
+                },
+            }
+
+
+NODE_CLASS_MAPPINGS = {
+    "ConditionalNode": ConditionalNode,
+}
+'''
+        result = self._extract_source(source, "conditional-node-class-pack")
+
+        self.assertEqual({}, result["nodes"])
+        self.assertEqual("no_static_nodes", result["pack"]["status"])
+
+    def test_top_level_class_mapping_still_extracts_node(self):
+        source = '''
+class TopLevelMappedNode:
+    RETURN_TYPES = ("IMAGE",)
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+            },
+        }
+
+
+NODE_CLASS_MAPPINGS = {
+    "TopLevelMappedNode": TopLevelMappedNode,
+}
+'''
+        result = self._extract_source(source, "top-level-node-class-pack")
+
+        self.assertIn("TopLevelMappedNode", result["nodes"])
+        self.assertEqual("ok", result["pack"]["status"])
+
     def test_mutated_node_class_mapping_skips_node(self):
         source = '''
 class MutatedMappingNode:

tools/generate_popular_node_signatures.py

@@ -283,13 +283,24 @@ def _has_module_wildcard_import(tree):
     return False
 
 
-def _collect_module_env(tree, class_envs=None):
+def _invalidate_class_bindings(class_bindings, names):
+    if class_bindings is None:
+        return
+    for name in names:
+        class_bindings.pop(name, None)
+
+
+def _collect_module_env(tree, class_bindings=None):
     env = {}
     for stmt in tree.body:
-        if class_envs is not None and isinstance(stmt, ast.ClassDef):
+        if isinstance(stmt, ast.ClassDef):
-            class_envs[stmt.name] = dict(env)
+            if class_bindings is not None:
+                class_bindings[stmt.name] = (stmt, dict(env))
+            env.pop(stmt.name, None)
+            continue
         if isinstance(stmt, ast.Assign):
             names = _assignment_target_names(stmt)
+            _invalidate_class_bindings(class_bindings, names)
             if len(stmt.targets) == 1 and isinstance(stmt.targets[0], ast.Name):
                 name = stmt.targets[0].id
                 if (
@@ -310,6 +321,7 @@ def _collect_module_env(tree, class_envs=None):
             continue
         if isinstance(stmt, ast.AnnAssign):
             names = _assignment_target_names(stmt)
+            _invalidate_class_bindings(class_bindings, names)
             if stmt.value is None:
                 continue
             if isinstance(stmt.target, ast.Name):
@@ -331,30 +343,46 @@ def _collect_module_env(tree, class_envs=None):
                     env.pop(name, None)
             continue
         if isinstance(stmt, ast.AugAssign):
-            for name in _assignment_target_names(stmt):
+            names = _assignment_target_names(stmt)
+            _invalidate_class_bindings(class_bindings, names)
+            for name in names:
                 env.pop(name, None)
             continue
         if isinstance(stmt, ast.Delete):
-            for name in _delete_target_names(stmt):
+            names = _delete_target_names(stmt)
+            _invalidate_class_bindings(class_bindings, names)
+            for name in names:
                 env.pop(name, None)
             continue
         if isinstance(stmt, ast.Expr):
-            for name in _mutating_call_target_names(stmt):
+            names = _mutating_call_target_names(stmt)
+            _invalidate_class_bindings(class_bindings, names)
+            for name in names:
                 env.pop(name, None)
-            for name in _bound_names(stmt):
+            names = _bound_names(stmt)
+            _invalidate_class_bindings(class_bindings, names)
+            for name in names:
                 env.pop(name, None)
             continue
         if isinstance(stmt, (ast.If, ast.For, ast.AsyncFor, ast.While, ast.Try, ast.With, ast.AsyncWith, ast.Match)):
             if _has_wildcard_import_in_control_flow(stmt):
                 env.clear()
+                if class_bindings is not None:
+                    class_bindings.clear()
                 continue
-            for name in _assigned_names_in_control_flow(stmt):
+            names = _assigned_names_in_control_flow(stmt)
+            _invalidate_class_bindings(class_bindings, names)
+            for name in names:
                 env.pop(name, None)
             continue
         if _has_wildcard_import(stmt):
             env.clear()
+            if class_bindings is not None:
+                class_bindings.clear()
             continue
-        for name in _bound_names(stmt):
+        names = _bound_names(stmt)
+        _invalidate_class_bindings(class_bindings, names)
+        for name in names:
             env.pop(name, None)
     return env
 
@@ -367,7 +395,7 @@ def normalise_input_spec(spec):
 
 
 def _class_defs(tree):
-    return {node.name: node for node in ast.walk(tree) if isinstance(node, ast.ClassDef)}
+    return {node.name: node for node in tree.body if isinstance(node, ast.ClassDef)}
 
 
 def _class_attr(cls, name, env):
@@ -602,16 +630,15 @@ def extract_repo_signatures(repo_dir, pack_meta):
         tree = _parse_python_file(path)
         if tree is None:
             continue
-        class_envs = {}
+        class_bindings = {}
-        env = _collect_module_env(tree, class_envs)
+        env = _collect_module_env(tree, class_bindings)
         mappings = _node_class_mappings(tree, env)
         displays = _display_mappings(tree, env)
-        classes = _class_defs(tree)
         for node_type, class_name in sorted(mappings.items()):
-            cls = classes.get(class_name)
+            binding = class_bindings.get(class_name)
-            if cls is None:
+            if binding is None:
                 continue
-            class_env = class_envs.get(class_name, env)
+            cls, class_env = binding
             sig = _signature_from_class(node_type, cls, displays.get(node_type), pack_meta, class_env)
             if sig is not None:
                 nodes[node_type] = sig